Part5 obb

.gitignore

@@ -11,4 +11,6 @@ nn/cocodataset/images/*
 *checkpoint.pth.tar
 runs/
 **/video/*
-!**/video/.gitkeep
+!**/video/.gitkeep
+nn/dotadataset/*
+!nn/dotadataset/.gitkeep

nn/YOLO_VGG16/degug_notebooks/train_YOLO_VGG16_F_many_classes.ipynb

@@ -177,7 +177,7 @@
     "    load_checkpoint(model_path_base + f\"e108_b0_vgg16_{checkpoint_file}\", model, optimizer, leanring_rate, device) \n",
     "\n",
     "# Initialize TensorBoard writer\n",
-    "writer = SummaryWriter(log_dir='runs/YOLO_VGG16_F_2_classes_v1')\n",
+    "writer = SummaryWriter(log_dir='runs/DEBUG')\n",
     "\n"
    ]
   },

nn/YOLO_VGG16_OBB/model/YOLO_VGG16_OBB.py

@@ -0,0 +1,77 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+
+from ...YOLO_VGG16_OBB.model.layers import CNNBlock, ResidualBlock, ScalePrediction
+
+# Class for defining YOLOv3 model
+
+
+class YOLO_VGG16_OBB(nn.Module):
+    def __init__(self, in_channels=3, num_classes=20):
+        super().__init__()
+        self.num_classes = num_classes
+        self.in_channels = in_channels
+
+        vgg16 = models.vgg16_bn(weights=models.VGG16_BN_Weights.DEFAULT)
+        backbone = list(vgg16.features.children())
+
+        # Layers list for YOLOv3
+        self.layers = nn.ModuleList([
+            # layers until end channels = 256
+            # 23 equivalent to ResidualBlock(256, num_repeats=8),
+            *backbone[:24],
+            # layers until end channels = 512
+            # 33 equivalent to ResidualBlock(512, num_repeats=8),
+            *backbone[24:34],
+
+
+            CNNBlock(512, 1024, kernel_size=3, stride=2, padding=1),
+            ResidualBlock(1024, num_repeats=4),
+
+            # large objects detection
+            CNNBlock(1024, 512, kernel_size=1, stride=1, padding=0),
+            CNNBlock(512, 1024, kernel_size=3, stride=1, padding=1),
+            ResidualBlock(1024, use_residual=False, num_repeats=1),
+            CNNBlock(1024, 512, kernel_size=1, stride=1, padding=0),
+            ScalePrediction(512, num_classes=num_classes),
+
+            # medium objects detection
+            CNNBlock(512, 256, kernel_size=1, stride=1, padding=0),
+            nn.Upsample(scale_factor=2),
+            CNNBlock(768, 256, kernel_size=1, stride=1, padding=0),
+            CNNBlock(256, 512, kernel_size=3, stride=1, padding=1),
+            ResidualBlock(512, use_residual=False, num_repeats=1),
+            CNNBlock(512, 256, kernel_size=1, stride=1, padding=0),
+            ScalePrediction(256, num_classes=num_classes),
+
+            # small objects detection
+            CNNBlock(256, 128, kernel_size=1, stride=1, padding=0),
+            nn.Upsample(scale_factor=2),
+            CNNBlock(384, 128, kernel_size=1, stride=1, padding=0),
+            CNNBlock(128, 256, kernel_size=3, stride=1, padding=1),
+            ResidualBlock(256, use_residual=False, num_repeats=1),
+            CNNBlock(256, 128, kernel_size=1, stride=1, padding=0),
+            ScalePrediction(128, num_classes=num_classes)
+        ])
+
+    # Forward pass for YOLOv3 with route connections and scale predictions
+    def forward(self, x):
+        outputs = []
+        route_connections = []
+
+        layer_num = 0
+        for layer in self.layers:
+            if isinstance(layer, ScalePrediction):
+                outputs.append(layer(x))
+                continue
+            x = layer(x)
+            if layer_num == 23 or layer_num == 33:
+                route_connections.append(x)
+
+            elif isinstance(layer, nn.Upsample):
+                x = torch.cat([x, route_connections[-1]], dim=1)
+                route_connections.pop()
+            layer_num += 1
+
+        return outputs

nn/YOLO_VGG16_OBB/model/layers.py

@@ -0,0 +1,73 @@
+import torch.nn as nn 
+
+# Defining CNN Block 
+class CNNBlock(nn.Module): 
+    def __init__(self, in_channels, out_channels, use_batch_norm=True, **kwargs): 
+        super().__init__() 
+        self.conv = nn.Conv2d(in_channels, out_channels, bias=not use_batch_norm, **kwargs) 
+        self.bn = nn.BatchNorm2d(out_channels) 
+        self.activation = nn.LeakyReLU(0.1) 
+        self.use_batch_norm = use_batch_norm 
+
+    def forward(self, x): 
+        # Applying convolution 
+        x = self.conv(x) 
+        # Applying BatchNorm and activation if needed 
+        if self.use_batch_norm: 
+            x = self.bn(x) 
+            return self.activation(x) 
+        else: 
+            return x
+
+# Defining residual block 
+class ResidualBlock(nn.Module): 
+	def __init__(self, channels, use_residual=True, num_repeats=1): 
+		super().__init__() 
+
+		# Defining all the layers in a list and adding them based on number of 
+		# repeats mentioned in the design 
+		res_layers = [] 
+		for _ in range(num_repeats): 
+			res_layers += [ 
+				nn.Sequential( 
+					nn.Conv2d(channels, channels // 2, kernel_size=1), 
+					nn.BatchNorm2d(channels // 2), 
+					nn.LeakyReLU(0.1), 
+					nn.Conv2d(channels // 2, channels, kernel_size=3, padding=1), 
+					nn.BatchNorm2d(channels), 
+					nn.LeakyReLU(0.1) 
+				) 
+			] 
+		self.layers = nn.ModuleList(res_layers) 
+		self.use_residual = use_residual 
+		self.num_repeats = num_repeats 
+
+	# Defining forward pass 
+	def forward(self, x): 
+		for layer in self.layers: 
+			residual = x 
+			x = layer(x) 
+			if self.use_residual: 
+				x = x + residual 
+		return x
+
+# Defining scale prediction class 
+class ScalePrediction(nn.Module): 
+	def __init__(self, in_channels, num_classes): 
+		super().__init__() 
+		# Defining the layers in the network 
+		self.pred = nn.Sequential( 
+			nn.Conv2d(in_channels, 2*in_channels, kernel_size=3, padding=1), 
+			nn.BatchNorm2d(2*in_channels), 
+			nn.LeakyReLU(0.1), 
+			nn.Conv2d(2*in_channels, (num_classes + 6) * 3, kernel_size=1), 
+		) 
+		self.num_classes = num_classes 
+
+	# Defining the forward pass and reshaping the output to the desired output 
+	# format: (batch_size, 3, grid_size, grid_size, num_classes + 6) 
+	def forward(self, x): 
+		output = self.pred(x) 
+		output = output.view(x.size(0), 3, self.num_classes + 6, x.size(2), x.size(3)) 
+		output = output.permute(0, 1, 3, 4, 2) 
+		return output

nn/YOLO_VGG16_OBB/model/loss.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+
+from ..utils.helpers import iou
+
+# Defining YOLO loss class
+
+
+class YOLOLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mse = nn.MSELoss()
+        self.bce = nn.BCEWithLogitsLoss()
+        self.cross_entropy = nn.CrossEntropyLoss()
+        self.sigmoid = nn.Sigmoid()
+        self.than = nn.Tanh()
+
+    def forward(self, pred, target, anchors):
+        # Identifying which cells in target have objects
+        # and which have no objects
+        obj = target[..., 0] == 1
+        no_obj = target[..., 0] == 0
+
+        # Calculating No object loss
+        no_object_loss = self.bce(
+            (pred[..., 0:1][no_obj]), (target[..., 0:1][no_obj]),
+        )
+
+        # Reshaping anchors to match predictions
+        anchors = anchors.reshape(1, 3, 1, 1, 2)
+        pred_angle = torch.tanh(pred[..., 5])
+        # Box prediction confidence
+        box_preds = torch.cat([self.sigmoid(pred[..., 1:3]),
+                               torch.exp(pred[..., 3:5]) *
+                               anchors, pred_angle.unsqueeze(-1)
+                               ], dim=-1)
+        # Calculating intersection over union for prediction and target
+        ious = iou(box_preds[obj], target[..., 1:6][obj]).detach()
+        # Calculating Object loss
+        object_loss = self.mse(self.sigmoid(pred[..., 0:1][obj]),
+                               ious * target[..., 0:1][obj]).clamp(0, 100)
+
+        # Predicted box coordinates
+        pred[..., 1:3] = self.sigmoid(pred[..., 1:3])
+        # Target box coordinates
+        target[..., 3:5] = torch.log(1e-6 + target[..., 3:5] / anchors)
+        # Calculating box coordinate loss
+
+        box_loss = self.mse(pred[..., 1:5][obj],
+                            target[..., 1:5][obj])
+
+        # Calculate angle loss
+        angle_loss = self.mse(pred_angle[obj],
+                              target[..., 5][obj])
+        # Claculating class loss
+        class_loss = self.cross_entropy((pred[..., 6:][obj]),
+                                        target[..., 6][obj].long())
+
+        # print('~~~~~~~~~~~~~ FIND UNCAPPED ~~~~~~~~~~~~~')
+        # print('box_loss:', box_loss)
+        # print('angle_loss:', angle_loss)
+        # print('object_loss:', object_loss)
+        # print('no_object_loss:', no_object_loss)
+        # print('class_loss:', class_loss)
+
+        # Total loss
+        return (
+            box_loss
+            + angle_loss
+            + object_loss
+            + no_object_loss
+            + class_loss
+        )