Upload 21 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+sample/street.jpeg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,57 @@
+import os
+
+import gradio as gr
+
+from src.run.yolov3.inference import YoloInfer
+
+
+infer = YoloInfer(model_path="./checkpoint/model.pt")
+
+demo = gr.Interface(
+    fn=infer.infer,
+    inputs=[
+        gr.Image(
+            shape=(416, 416),
+            label="Input Image",
+            value="./sample/bird_plane.jpeg",
+        ),
+        gr.Slider(
+            minimum=0,
+            maximum=1,
+            value=0.2,
+            label="IOU Threshold",
+            info="Permissible overlap for the same class bounding boxes",
+        ),
+        gr.Slider(
+            minimum=0,
+            maximum=1,
+            value=0.95,
+            label="Objectness Threshold",
+            info="Confidence for each pixel to predict an object",
+        ),
+        gr.Slider(
+            minimum=0,
+            maximum=1,
+            value=0.5,
+            label="Class Threshold",
+            info="Confidence for each pixel to predict a class",
+        ),
+        gr.Slider(
+            minimum=0,
+            maximum=10,
+            value=1,
+            label="Font Size",
+            info="Bounding box text size",
+        ),
+    ],
+    outputs=[
+        gr.Image(),
+    ],
+    examples=[
+        [os.path.join("./sample/", f)]
+        for f in os.listdir("./sample/")
+    ],
+)
+
+
+demo.launch()

requirements.txt

@@ -0,0 +1,10 @@
+torch
+gradio
+torchvision
+numpy
+grad-cam
+Pillow
+albumentations
+tqdm
+cv2
+matplotlib

src/datasets/pascal_voc/__init__.py

@@ -0,0 +1 @@
+from .dataset import YOLODataset

src/datasets/pascal_voc/dataset.py

@@ -0,0 +1,117 @@
+"""
+Creates a Pytorch dataset to load the Pascal VOC & MS COCO datasets
+"""
+
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+from PIL import Image, ImageFile
+from torch.utils.data import Dataset
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+class YOLODataset(Dataset):
+    def __init__(
+        self,
+        csv_file,
+        img_dir,
+        label_dir,
+        anchors,
+        image_size=416,
+        S=[13, 26, 52],
+        transform=None,
+        load_mosaic=True,
+    ):
+        self.annotations = pd.read_csv(csv_file)
+        self.img_dir = img_dir
+        self.label_dir = label_dir
+        self.image_size = image_size
+        self.transform = transform
+        self.S = S
+        self.load_mosaic = load_mosaic
+
+        # shape: [number of states, number of anchors, 2]
+        self.anchors = torch.tensor(anchors)
+        self.num_anchors_per_scale = self.anchors.shape[1]
+
+    def __len__(self):
+        return len(self.annotations)
+
+    @staticmethod
+    def iou(box, anchors):
+        """
+        box:
+            tensor shape: [2]
+        anchors:
+            tensor shape: [number of states, number of anchors, 2]
+
+        * 2 above is for width and height
+        """
+
+        intersection = torch.prod(torch.min(box, anchors), dim=-1)
+        union = torch.prod(box) + torch.prod(anchors, dim=-1) - intersection
+        return intersection / union
+
+    def __getitem__(self, index):
+        label_path = os.path.join(self.label_dir, self.annotations.iloc[index, 1])
+
+        # bboxes = np.loadtxt(fname=label_path, delimiter=" ", ndmin=2)
+        bboxes = np.roll(
+            np.loadtxt(fname=label_path, delimiter=" ", ndmin=2), 4, axis=1
+        )
+        img_path = os.path.join(self.img_dir, self.annotations.iloc[index, 0])
+        image = np.array(Image.open(img_path).convert("RGB"))
+
+        if self.transform:
+            augmentations = self.transform(image=image, bboxes=bboxes)
+            image = augmentations["image"]
+            bboxes = augmentations["bboxes"]
+
+        """
+        Below assumes 3 scale predictions (as paper) and same num of anchors per scale
+        6 = [objectness, cx, cy, w, h, class]
+        """
+        targets = [torch.zeros((self.num_anchors_per_scale, S, S, 6)) for S in self.S]
+
+        for bbox in bboxes:
+            iou = self.iou(torch.tensor(bbox[2:4]), self.anchors)
+
+            idx = torch.argsort(iou, descending=True, dim=-1)
+            idx = idx[:, 0].tolist()
+
+            dimensions, class_ = np.array(bbox[:-1]), bbox[-1]  # +1
+
+            for scale_idx, anchor_id in enumerate(idx):
+                scale_dim = self.S[scale_idx]
+                scale_cx, scale_cy, scale_w, scale_h = dimensions * scale_dim
+
+                row, col = int(scale_cy), int(scale_cx)
+
+                # fill values
+                scale_cx = scale_cx - col
+                scale_cy = scale_cy - row
+
+                box_target = torch.tensor(
+                    [1, scale_cx, scale_cy, scale_w, scale_h, class_]
+                )
+
+                targets[scale_idx][anchor_id, row, col] = box_target
+
+        return image, targets
+
+
+if __name__ == "__main__":
+    from src.run.yolov3 import config
+
+    IMAGE_SIZE = config.IMAGE_SIZE
+    train_dataset = YOLODataset(
+        config.DATASET + "/2examples.csv",
+        transform=config.train_transforms,
+        S=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
+        img_dir=config.IMG_DIR,
+        label_dir=config.LABEL_DIR,
+        anchors=config.ANCHORS,
+    )

src/loss/yolov3/__init__.py

@@ -0,0 +1 @@
+from .loss import YoloLoss

src/loss/yolov3/loss.py

@@ -0,0 +1,191 @@
+"""
+Implementation of Yolo Loss Function similar to the one in Yolov3 paper,
+the difference from what I can tell is I use CrossEntropy for the classes
+instead of BinaryCrossEntropy.
+"""
+import random
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class YoloLoss(nn.Module):
+    def __init__(self, nclasses):
+        super().__init__()
+        self.mse = nn.MSELoss()
+        self.bce = nn.BCEWithLogitsLoss()
+        self.entropy = nn.CrossEntropyLoss()
+        self.sigmoid = nn.Sigmoid()
+
+        # Constants signifying how much to pay for each respective part of the loss
+        self.lambda_class = 5  # 1.5
+        self.lambda_noobj = 2
+        self.lambda_obj = 1
+        self.lambda_box = 2
+
+        self.nclasses = nclasses
+
+    # intersection over union
+    @staticmethod
+    def iou(box1, box2):
+        """
+        boxi shape = [any shape, 4] i.e [4] or [3,4] or [2,3,4] etc.
+
+        * 4 = [x, y, w, h]
+
+        output shape = [batch]
+        """
+        # box1 x1, x2
+        box1_x1 = box1[..., 0] - box1[..., 2] / 2
+        box1_x2 = box1[..., 0] + box1[..., 2] / 2
+
+        # box2 x1, x2
+        box2_x1 = box2[..., 0] - box2[..., 2] / 2
+        box2_x2 = box2[..., 0] + box2[..., 2] / 2
+
+        # the width of intersection (x)
+        x1 = torch.max(box1_x1, box2_x1)
+        x2 = torch.min(box1_x2, box2_x2)
+
+        x = (x2 - x1).clamp(0)
+
+        # box1 y1, y2
+        box1_y1 = box1[..., 1] - box1[..., 3] / 2
+        box1_y2 = box1[..., 1] + box1[..., 3] / 2
+
+        # box2 y1, y2
+        box2_y1 = box2[..., 1] - box2[..., 3] / 2
+        box2_y2 = box2[..., 1] + box2[..., 3] / 2
+
+        # the height of intersection (y)
+        y1 = torch.max(box1_y1, box2_y1)
+        y2 = torch.min(box1_y2, box2_y2)
+
+        y = (y2 - y1).clamp(0)
+
+        # intersection
+        intersection = x * y
+
+        # union
+        area_box1 = (box1_x2 - box1_x1) * (box1_y2 - box1_y1)
+        area_box2 = (box2_x2 - box2_x1) * (box2_y2 - box2_y1)
+
+        union = area_box1 + area_box2 - intersection + 1e-6
+
+        return intersection / union
+
+    def forward(self, predictions, target, anchors):
+        """
+        predictions: [batch, 3, 13, 13, 25] where 25 = [objectness, cx, cy, w, h] + 20 classes
+        target: [batch, 3, 13, 13, 6] where 6 = [objectness, cx, cy, w, h, true class]
+        anchors: [3, 2]
+
+        * 13 is S
+        * 3 is number of anchors
+        """
+        obj = target[..., 0] == 1
+        noobj = target[..., 0] == 0
+
+        """
+        both - no object and object loss
+        uncomment noobj above and below for only no_object_loss)
+        
+        predictions shape: [batch, 3, 13, 13]
+        target shape: [batch, 3, 13, 13]
+        """
+        no_object_loss = self.bce(predictions[..., 0][noobj], target[..., 0][noobj])
+
+        """
+        object loss
+        
+        predictions[..., 0][obj] shape: [total_object_in_batch]
+        target[..., 0][obj] shape: [total_object_in_batch]
+        """
+        object_loss = self.bce(predictions[..., 0][obj], target[..., 0][obj])
+
+        anchors = anchors.reshape(1, 3, 1, 1, 2)
+        box_preds = torch.cat(
+            [
+                self.sigmoid(predictions[..., 1:3]),
+                torch.exp(predictions[..., 3:5]) * anchors,
+            ],
+            dim=-1,
+        )
+
+        ious = self.iou(box_preds[obj], target[..., 1:5][obj]).detach()
+        object_loss += self.mse(
+            self.sigmoid(predictions[..., 0][obj]), ious * target[..., 0][obj]
+        )
+
+        """
+        coordinate loss or box loss
+        
+        predictions[..., 1:5][obj] shape: [total_obj_in_batch, 4]
+        target[..., 1:5][obj] shape: [total_obj_in_batch, 4]
+        """
+        # x, y coordinates
+        predictions[..., 1:3] = self.sigmoid(predictions[..., 1:3])
+
+        # width, height coordinates
+        target[..., 3:5] = torch.log((1e-16 + target[..., 3:5] / anchors))
+        box_loss = self.mse(predictions[..., 1:5][obj], target[..., 1:5][obj])
+
+        """
+        classification loss : cross entropy
+        
+        predictions[..., 5:][obj] shape: [total_obj_in_batch, nclasses]
+        target[..., 5][obj].long() shape: [total_obj_in_batch]
+        """
+        class_loss = self.entropy(
+            (predictions[..., 5:][obj]),
+            (target[..., 5][obj].long()),
+        )
+
+        """
+        classification loss : binary cross entropy
+        
+        This is my innovation: could be wrong
+        Train and test without it as well.        
+        """
+        binary_class_loss = self.bce(
+            predictions[..., 5:][obj],
+            F.one_hot(target[..., 5][obj].long(), num_classes=self.nclasses).float(),
+        )
+
+        return (
+            self.lambda_box * box_loss
+            + self.lambda_obj * object_loss
+            + self.lambda_noobj * no_object_loss
+            + self.lambda_class * class_loss
+            + self.lambda_class * binary_class_loss
+        )
+
+
+if __name__ == "__main__":
+    from src.run.yolov3 import config
+    from src.datasets.pascal_voc import YOLODataset
+
+    S = 13
+    yl = YoloLoss(nclasses=20)
+
+    predictions = torch.rand((20, 3, S, S, 25))
+
+    # build target
+    IMAGE_SIZE = config.IMAGE_SIZE
+
+    train_dataset = YOLODataset(
+        config.DATASET + "/train.csv",
+        transform=None,  # config.train_transforms,
+        S=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
+        img_dir=config.IMG_DIR,
+        label_dir=config.LABEL_DIR,
+        anchors=config.ANCHORS,
+    )
+    _, target = train_dataset[3]
+    target = target[0].unsqueeze(0)  # target[0] if S=13
+    target = torch.cat([target, target] * 10)
+
+    # anchor
+    anchor = S * train_dataset.anchors[0]
+
+    print(yl(predictions, target, anchor))

src/model/yolov3/__init__.py

@@ -0,0 +1 @@
+from .model import YOLOv3

src/model/yolov3/model.py

@@ -0,0 +1,173 @@
+"""
+Implementation of YOLOv3 architecture
+"""
+
+import torch
+import torch.nn as nn
+
+""" 
+Information about architecture config:
+Tuple is structured by (filters, kernel_size, stride) 
+Every conv is a same convolution. 
+List is structured by "B" indicating a residual block followed by the number of repeats
+"S" is for scale prediction block and computing the yolo loss
+"U" is for upsampling the feature map and concatenating with a previous layer
+"""
+model_config = [
+    (32, 3, 1),
+    (64, 3, 2),
+    ["B", 1],
+    (128, 3, 2),
+    ["B", 2],
+    (256, 3, 2),
+    ["B", 8],
+    (512, 3, 2),
+    ["B", 8],
+    (1024, 3, 2),
+    ["B", 4],  # To this point is Darknet-53
+    (512, 1, 1),
+    (1024, 3, 1),
+    "S",
+    (256, 1, 1),
+    "U",
+    (256, 1, 1),
+    (512, 3, 1),
+    "S",
+    (128, 1, 1),
+    "U",
+    (128, 1, 1),
+    (256, 3, 1),
+    "S",
+]
+
+
+class CNNBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, bn_act=True, **kwargs):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, bias=not bn_act, **kwargs)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.leaky = nn.LeakyReLU(0.1)
+        self.use_bn_act = bn_act
+
+    def forward(self, x):
+        if self.use_bn_act:
+            return self.leaky(self.bn(self.conv(x)))
+        else:
+            return self.conv(x)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels, use_residual=True, num_repeats=1):
+        super().__init__()
+        self.layers = nn.ModuleList()
+        for repeat in range(num_repeats):
+            self.layers += [
+                nn.Sequential(
+                    CNNBlock(channels, channels // 2, kernel_size=1),
+                    CNNBlock(channels // 2, channels, kernel_size=3, padding=1),
+                )
+            ]
+
+        self.use_residual = use_residual
+        self.num_repeats = num_repeats
+
+    def forward(self, x):
+        for layer in self.layers:
+            if self.use_residual:
+                x = x + layer(x)
+            else:
+                x = layer(x)
+
+        return x
+
+
+class ScalePrediction(nn.Module):
+    def __init__(self, in_channels, num_classes):
+        super().__init__()
+        self.pred = nn.Sequential(
+            CNNBlock(in_channels, 2 * in_channels, kernel_size=3, padding=1),
+            CNNBlock(
+                2 * in_channels, (num_classes + 5) * 3, bn_act=False, kernel_size=1
+            ),
+        )
+        self.num_classes = num_classes
+
+    def forward(self, x):
+        x = (
+            self.pred(x)
+            .reshape(x.shape[0], 3, self.num_classes + 5, x.shape[2], x.shape[3])
+            .permute(0, 1, 3, 4, 2)
+        )
+        return x
+
+
+class YOLOv3(nn.Module):
+    def __init__(self, in_channels=3, num_classes=80):
+        super().__init__()
+        self.num_classes = num_classes
+        self.in_channels = in_channels
+        self.layers = self._create_conv_layers()
+
+    def forward(self, x):
+        outputs = []  # for each scale
+        route_connections = []
+
+        for layer in self.layers:
+            if isinstance(layer, ScalePrediction):
+                outputs.append(layer(x))
+                continue
+
+            x = layer(x)
+
+            if isinstance(layer, ResidualBlock) and layer.num_repeats == 8:
+                route_connections.append(x)
+
+            elif isinstance(layer, nn.Upsample):
+                x = torch.cat([x, route_connections[-1]], dim=1)
+                route_connections.pop()
+
+        return outputs
+
+    def _create_conv_layers(self):
+        layers = nn.ModuleList()
+        in_channels = self.in_channels
+
+        for module in model_config:
+            if isinstance(module, tuple):
+                out_channels, kernel_size, stride = module
+                layers.append(
+                    CNNBlock(
+                        in_channels,
+                        out_channels,
+                        kernel_size=kernel_size,
+                        stride=stride,
+                        padding=1 if kernel_size == 3 else 0,
+                    )
+                )
+                in_channels = out_channels
+
+            elif isinstance(module, list):
+                num_repeats = module[1]
+                layers.append(
+                    ResidualBlock(
+                        in_channels,
+                        num_repeats=num_repeats,
+                    )
+                )
+
+            elif isinstance(module, str):
+                if module == "S":
+                    layers += [
+                        ResidualBlock(in_channels, use_residual=False, num_repeats=1),
+                        CNNBlock(in_channels, in_channels // 2, kernel_size=1),
+                        ScalePrediction(in_channels // 2, num_classes=self.num_classes),
+                    ]
+                    in_channels = in_channels // 2
+
+                elif module == "U":
+                    layers.append(
+                        nn.Upsample(scale_factor=2),
+                    )
+                    in_channels = in_channels * 3
+
+        return layers

src/run/yolov3/config.py

@@ -0,0 +1,196 @@
+import albumentations as A
+import cv2
+import torch
+from albumentations.pytorch import ToTensorV2
+
+
+DATASET = "PASCAL_VOC"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+NUM_WORKERS = 0
+BATCH_SIZE = 32
+IMAGE_SIZE = 416
+NUM_CLASSES = 20
+LEARNING_RATE = 1e-5
+WEIGHT_DECAY = 1e-4
+NUM_EPOCHS = 100
+CONF_THRESHOLD = 0.05
+MAP_IOU_THRESH = 0.5
+NMS_IOU_THRESH = 0.45
+S = [IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8]
+PIN_MEMORY = True
+LOAD_MODEL = False
+SAVE_MODEL = True
+CHECKPOINT_FILE = "checkpoint.pth.tar"
+IMG_DIR = DATASET + "/images/"
+LABEL_DIR = DATASET + "/labels/"
+
+ANCHORS = [
+    [(0.28, 0.22), (0.38, 0.48), (0.9, 0.78)],
+    [(0.07, 0.15), (0.15, 0.11), (0.14, 0.29)],
+    [(0.02, 0.03), (0.04, 0.07), (0.08, 0.06)],
+]  # Note these have been rescaled to be between [0, 1]
+
+means = [0.485, 0.456, 0.406]
+
+scale = 1.1
+train_transforms = A.Compose(
+    [
+        A.LongestMaxSize(max_size=int(IMAGE_SIZE * scale)),
+        A.PadIfNeeded(
+            min_height=int(IMAGE_SIZE * scale),
+            min_width=int(IMAGE_SIZE * scale),
+            border_mode=cv2.BORDER_CONSTANT,
+        ),
+        A.Rotate(limit=10, interpolation=1, border_mode=4),
+        A.RandomCrop(width=IMAGE_SIZE, height=IMAGE_SIZE),
+        A.ColorJitter(brightness=0.6, contrast=0.6, saturation=0.6, hue=0.6, p=0.4),
+        A.OneOf(
+            [
+                A.ShiftScaleRotate(
+                    rotate_limit=20, p=0.5, border_mode=cv2.BORDER_CONSTANT
+                ),
+                # A.Affine(shear=15, p=0.5, mode="constant"),
+            ],
+            p=1.0,
+        ),
+        A.HorizontalFlip(p=0.5),
+        A.Blur(p=0.1),
+        A.CLAHE(p=0.1),
+        A.Posterize(p=0.1),
+        A.ToGray(p=0.1),
+        A.ChannelShuffle(p=0.05),
+        A.Normalize(
+            mean=[0, 0, 0],
+            std=[1, 1, 1],
+            max_pixel_value=255,
+        ),
+        ToTensorV2(),
+    ],
+    bbox_params=A.BboxParams(
+        format="yolo",
+        min_visibility=0.4,
+        label_fields=[],
+    ),
+)
+test_transforms = A.Compose(
+    [
+        A.LongestMaxSize(max_size=IMAGE_SIZE),
+        A.PadIfNeeded(
+            min_height=IMAGE_SIZE, min_width=IMAGE_SIZE, border_mode=cv2.BORDER_CONSTANT
+        ),
+        A.Normalize(
+            mean=[0, 0, 0],
+            std=[1, 1, 1],
+            max_pixel_value=255,
+        ),
+        ToTensorV2(),
+    ],
+    bbox_params=A.BboxParams(format="yolo", min_visibility=0.4, label_fields=[]),
+)
+
+PASCAL_CLASSES = [
+    "aeroplane",
+    "bicycle",
+    "bird",
+    "boat",
+    "bottle",
+    "bus",
+    "car",
+    "cat",
+    "chair",
+    "cow",
+    "diningtable",
+    "dog",
+    "horse",
+    "motorbike",
+    "person",
+    "pottedplant",
+    "sheep",
+    "sofa",
+    "train",
+    "tvmonitor",
+]
+
+COCO_LABELS = [
+    "person",
+    "bicycle",
+    "car",
+    "motorcycle",
+    "airplane",
+    "bus",
+    "train",
+    "truck",
+    "boat",
+    "traffic light",
+    "fire hydrant",
+    "stop sign",
+    "parking meter",
+    "bench",
+    "bird",
+    "cat",
+    "dog",
+    "horse",
+    "sheep",
+    "cow",
+    "elephant",
+    "bear",
+    "zebra",
+    "giraffe",
+    "backpack",
+    "umbrella",
+    "handbag",
+    "tie",
+    "suitcase",
+    "frisbee",
+    "skis",
+    "snowboard",
+    "sports ball",
+    "kite",
+    "baseball bat",
+    "baseball glove",
+    "skateboard",
+    "surfboard",
+    "tennis racket",
+    "bottle",
+    "wine glass",
+    "cup",
+    "fork",
+    "knife",
+    "spoon",
+    "bowl",
+    "banana",
+    "apple",
+    "sandwich",
+    "orange",
+    "broccoli",
+    "carrot",
+    "hot dog",
+    "pizza",
+    "donut",
+    "cake",
+    "chair",
+    "couch",
+    "potted plant",
+    "bed",
+    "dining table",
+    "toilet",
+    "tv",
+    "laptop",
+    "mouse",
+    "remote",
+    "keyboard",
+    "cell phone",
+    "microwave",
+    "oven",
+    "toaster",
+    "sink",
+    "refrigerator",
+    "book",
+    "clock",
+    "vase",
+    "scissors",
+    "teddy bear",
+    "hair drier",
+    "toothbrush",
+]

src/run/yolov3/dataloader.py

@@ -0,0 +1,58 @@
+from torch.utils.data import DataLoader
+from src.run.yolov3 import config
+from src.datasets.pascal_voc import YOLODataset
+
+
+def get_loaders(train_csv_path, test_csv_path):
+    IMAGE_SIZE = config.IMAGE_SIZE
+    train_dataset = YOLODataset(
+        train_csv_path,
+        transform=config.train_transforms,
+        S=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
+        img_dir=config.IMG_DIR,
+        label_dir=config.LABEL_DIR,
+        anchors=config.ANCHORS,
+    )
+    test_dataset = YOLODataset(
+        test_csv_path,
+        transform=config.test_transforms,
+        S=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
+        img_dir=config.IMG_DIR,
+        label_dir=config.LABEL_DIR,
+        anchors=config.ANCHORS,
+    )
+    train_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=config.BATCH_SIZE,
+        num_workers=config.NUM_WORKERS,
+        pin_memory=config.PIN_MEMORY,
+        shuffle=True,
+        drop_last=False,
+    )
+    test_loader = DataLoader(
+        dataset=test_dataset,
+        batch_size=config.BATCH_SIZE,
+        num_workers=config.NUM_WORKERS,
+        pin_memory=config.PIN_MEMORY,
+        shuffle=False,
+        drop_last=False,
+    )
+
+    train_eval_dataset = YOLODataset(
+        train_csv_path,
+        transform=config.test_transforms,
+        S=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
+        img_dir=config.IMG_DIR,
+        label_dir=config.LABEL_DIR,
+        anchors=config.ANCHORS,
+    )
+    train_eval_loader = DataLoader(
+        dataset=train_eval_dataset,
+        batch_size=config.BATCH_SIZE,
+        num_workers=config.NUM_WORKERS,
+        pin_memory=config.PIN_MEMORY,
+        shuffle=False,
+        drop_last=False,
+    )
+
+    return train_loader, test_loader, train_eval_loader

src/run/yolov3/inference.py

@@ -0,0 +1,333 @@
+import os
+
+import albumentations as A
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from albumentations.pytorch import ToTensorV2
+from PIL import Image
+from tqdm import tqdm
+
+from src.loss.yolov3 import YoloLoss
+from src.model.yolov3 import YOLOv3 as Model
+from src.run.yolov3 import config
+
+
+class YoloInfer:
+    def __init__(self, model_path):
+        self.model = self.load_model(model_path=model_path)
+        self.transform = A.Compose(
+            [
+                A.Resize(config.IMAGE_SIZE, config.IMAGE_SIZE),
+                A.Normalize(
+                    mean=[0, 0, 0],
+                    std=[1, 1, 1],
+                    max_pixel_value=255,
+                ),
+                ToTensorV2(),
+            ]
+        )
+
+        self.scaled_anchors = (
+            torch.tensor(config.ANCHORS) * torch.tensor(config.S).reshape(-1, 1, 1)
+        ).to(config.DEVICE)
+
+    def load_model(self, model_path):
+        model = Model(num_classes=config.NUM_CLASSES).to(config.DEVICE)
+
+        if os.path.isfile(model_path):
+            checkpoint = torch.load(model_path, map_location=config.DEVICE)
+            model.load_state_dict(checkpoint["model_state_dict"], strict=False)
+
+        return model
+
+    @staticmethod
+    def pred_to_boxes(prediction, anchors):
+        """
+        prediction tensor = [batch, num_anchors_per_scale, scale, scale, 5 + num_classes]
+        5 = [objness, cx, cy, w, h]
+        anchors tensor = [num_anchors_per_scale, 2]
+
+        Note: The below operation could been done entirely inplace.
+            Slightly unoptimsed implementation below to maintain readability
+
+        Output shape: [batch, num_anchors_per_scale, scale, scale, 7]
+        7: [predicted_class's_idx, obj score, cx, cy, width, height, predicted class probability score]
+        """
+        scale = prediction.shape[2]
+
+        # reversing the equations of box loss and obj in the loss function
+        anchors = anchors.reshape(1, len(anchors), 1, 1, 2)
+        cx_cy = torch.sigmoid(prediction[..., 1:3])
+        width_height = anchors * torch.exp(prediction[..., 3:5])
+
+        # reversing the equations we wrote while making training data
+        arange = torch.arange(scale, device=config.DEVICE)
+        cx = (arange.reshape(1, 1, 1, scale, 1) + cx_cy[..., 0:1]) / scale
+        cy = (arange.reshape(1, 1, scale, 1, 1) + cx_cy[..., 1:2]) / scale
+        width_height = width_height / scale
+
+        # class prediction
+        class_predictions = torch.softmax(prediction[..., 5:], dim=-1)
+        class_score, class_idx = torch.max(class_predictions, dim=-1)
+        class_score, class_idx = class_score.unsqueeze(-1), class_idx.unsqueeze(-1)
+
+        # objectness score
+        obj_score = torch.sigmoid(prediction[..., 0:1])
+        return torch.cat(
+            [class_idx, obj_score, cx, cy, width_height, class_score], dim=-1
+        )
+
+    @staticmethod
+    def sort_3Dtensor_rows_on_two_columns(
+        tensor, index1, index2, descending1=True, descending2=True
+    ):
+        """
+        tensor = tensor([[[1, 2, 3],
+                        [1, 3, 4],
+                        [0, 2, 1]],
+
+                        [[0, 2, 3],
+                        [1, 4, 5],
+                        [0, 1, 2]]])
+
+        sort_tensor_rows_on_two_columns(tensor,
+                                        index1=0,
+                                        index2=1,
+                                        descending1=False,
+                                        descending2=True)
+
+        output = tensor([[[0, 2, 1],
+                        [1, 3, 4],
+                        [1, 2, 3]],
+
+                        [[0, 2, 3],
+                        [0, 1, 2],
+                        [1, 4, 5]]])
+        """
+        inner_sorting = torch.argsort(tensor[..., index2], descending=descending1)
+        inner_sorted = torch.gather(
+            tensor, 1, inner_sorting.unsqueeze(-1).expand(-1, -1, tensor.size(2))
+        )
+
+        outer_sorting = torch.argsort(
+            inner_sorted[:, :, index1], stable=True, descending=descending2
+        )
+        outer_sorted = torch.gather(
+            inner_sorted,
+            1,
+            outer_sorting.unsqueeze(-1).expand(-1, -1, inner_sorted.size(2)),
+        )
+        return outer_sorted
+
+    @staticmethod
+    def non_max_supression(
+        prediction, iou_threshold, object_threshold, class_threshold
+    ):
+        """
+        prediction = [batch, summation(num_anchors_per_scale * scale * scale), 7]
+        i.e. [batch, (3 * 13 * 13 + 3 * 26 * 26 + 3 * 52 * 52), 7]
+
+        7: [class_pred, obj_score, cx, cy, width, height, class_score]
+        """
+        """
+        inside each batch output,
+        first sort by class prediction, 
+        and inside each class sort objectness in descending 
+        """
+        prediction = YoloInfer.sort_3Dtensor_rows_on_two_columns(
+            tensor=prediction, index1=0, index2=1, descending1=True, descending2=True
+        )
+
+        """
+        remove predictions with object threshold below the given threshold
+        and split prediction to get a list of tensors
+        
+        length of list = batch size
+        each element in the list = results/output of 1 image
+        """
+        # objectness condition [threshold]
+        objectness = (prediction[..., 1] > object_threshold) & (
+            prediction[..., 6] > class_threshold
+        )
+        indices = torch.nonzero(objectness)
+        batch_boxes = torch.split(
+            tensor=prediction[objectness],
+            split_size_or_sections=torch.bincount(indices[:, 0]).tolist(),
+            dim=0,
+        )
+
+        # iterate for output
+        output = []
+
+        for boxes in tqdm(batch_boxes, disable=True):
+            # boxes shape = [-1, 7]
+            boxes = boxes.tolist()
+            final_boxes = []
+
+            while boxes:
+                top_box = boxes.pop(0)
+
+                idx = 0
+
+                while idx < len(boxes):
+                    box = boxes[idx]
+
+                    # class match
+                    if box[0] != top_box[0]:
+                        break
+
+                    # iou match
+                    if (
+                        YoloLoss.iou(torch.tensor(top_box[2:6]), torch.tensor(box[2:6]))
+                        > iou_threshold
+                    ):
+                        del boxes[idx]
+
+                        idx -= 1
+
+                    idx += 1
+
+                final_boxes.append(top_box)
+
+            output.append(final_boxes)
+
+        return output
+
+    @staticmethod
+    def draw_bounding_boxes(image, boxes, font_size=1):
+        """Draws bounding boxes on the image using OpenCV"""
+        cmap = plt.get_cmap("tab20b")
+        class_labels = (
+            config.COCO_LABELS if config.DATASET == "COCO" else config.PASCAL_CLASSES
+        )
+        colors = [cmap(i) for i in np.linspace(0, 1, len(class_labels))]
+        im = np.array(image)
+        height, width, _ = im.shape
+
+        # font = ImageFont.truetype("DejaVuSans.ttf", 20)  # Load the DejaVuSans font
+
+        for box in boxes:
+            assert (
+                len(box) == 7
+            ), "box should contain class pred, confidence, x, y, width, height, class score"
+            class_pred = box[0]
+            class_score = round(box[-1], 2)
+
+            upper_left_x = int((box[2] - box[4] / 2) * width)
+            upper_left_y = int((box[3] - box[5] / 2) * height)
+            lower_right_x = int((box[2] + box[4] / 2) * width)
+            lower_right_y = int((box[3] + box[5] / 2) * height)
+
+            color = colors[int(class_pred)]
+            color_rgb = (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255))
+            thickness = max(
+                int((0.0005 * (image.shape[0] + image.shape[1]) / 2) + 1), 1
+            )
+
+            cv2.rectangle(
+                im,
+                (upper_left_x, upper_left_y),
+                (lower_right_x, lower_right_y),
+                color_rgb,
+                thickness=thickness,
+            )
+
+            # label
+            font_scale = font_size
+            label = f"{class_labels[int(class_pred)]} {class_score}"
+            text_size = cv2.getTextSize(
+                label,
+                fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                fontScale=font_scale,
+                thickness=1,
+            )[0]
+
+            # Draw rectangle background
+            cv2.rectangle(
+                im,
+                (upper_left_x, upper_left_y),
+                (upper_left_x + text_size[0], upper_left_y - text_size[1]),
+                color_rgb,
+                thickness=-1,
+            )
+            cv2.putText(
+                im,
+                label,
+                (upper_left_x, upper_left_y),
+                fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                fontScale=font_scale,
+                color=[0, 0, 0],
+                thickness=1,
+                lineType=cv2.LINE_AA,
+            )
+
+        return im
+
+    def infer(
+        self,
+        image: np.array,
+        iou_threshold=0.75,
+        object_threshold=0.75,
+        class_threshold=0.5,
+        font_size=1,
+    ):
+        self.model.eval()
+        input_tensor = self.transform(image=image)["image"].unsqueeze(0)
+
+        with torch.no_grad():
+            """
+            output = list of tensors
+            tensor shape=[batch, num_anchors_per_scale, scale, scale, 5 + num_classes]
+            """
+            output = self.model(input_tensor.to(config.DEVICE))
+
+            # convert model prediction to actual box prediction
+            output = torch.cat(
+                [
+                    self.pred_to_boxes(out, self.scaled_anchors[idx]).reshape(
+                        out.shape[0], -1, 7
+                    )
+                    for idx, out in enumerate(output)
+                ],
+                dim=1,
+            )
+
+            # non max suppression
+            output = self.non_max_supression(
+                prediction=output,
+                iou_threshold=iou_threshold,
+                object_threshold=object_threshold,
+                class_threshold=class_threshold,
+            )
+
+        return self.draw_bounding_boxes(image, output[0], font_size=font_size)
+
+    @staticmethod
+    def load_image_as_array(image_path):
+        # Load a PIL image
+        pil_image = Image.open(image_path)
+
+        # Convert PIL image to NumPy array
+        return np.array(pil_image.convert("RGB"))
+
+    @staticmethod
+    def plot_array(array: np.array, figsize=(10, 10)):
+        plt.figure(figsize=figsize)
+        plt.imshow(array)
+        plt.show()
+
+    @staticmethod
+    def save_numpy_as_image(numpy_array, image_path):
+        """
+        Saves a NumPy array as an image.
+        Args:
+            numpy_array (numpy.ndarray): The NumPy array to be saved as an image.
+            image_path (str): The path where the image will be saved.
+        """
+        # Convert the NumPy array to a PIL image
+        image = Image.fromarray(numpy_array)
+
+        # Save the PIL image to the specified path
+        image.save(image_path)