Create app.py

app.py

@@ -0,0 +1,221 @@
+import gradio as gr
+from torchvision import datasets, transforms
+import cv2
+import albumentations as Al
+from albumentations.pytorch import ToTensorV2
+from PIL import Image
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import io
+import numpy as np
+import pandas as pd 
+from torch.optim.lr_scheduler import OneCycleLR
+from pytorch_lightning import LightningModule, Trainer, seed_everything
+from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
+from pytorch_lightning.callbacks.progress import TQDMProgressBar
+from pytorch_lightning.loggers import CSVLogger
+from pytorch_lightning.loggers import TensorBoardLogger
+from tqdm import tqdm
+import torch
+import torch.optim as optim
+
+# my files
+import utils
+import config
+from model import YOLOv3
+from utils import (
+    mean_average_precision,
+    cells_to_bboxes,
+    get_evaluation_bboxes,
+    save_checkpoint,
+    load_checkpoint,
+    check_class_accuracy,
+    plot_couple_examples,
+    accuracy_fn,
+    get_loaders
+)
+from loss import YoloLoss
+
+
+# custom functions for yolo
+# loss function for yolov3
+loss_fn = YoloLoss()
+
+def model_criterion(out, y,anchors):
+    loss = (  loss_fn(out[0], y[0], anchors[0])
+            + loss_fn(out[1], y[1], anchors[1])
+            + loss_fn(out[2], y[2], anchors[2])
+            )
+    return loss
+
+
+# accuracy function for yolov3
+def accuracy_fn(y, out, threshold,correct_class, correct_obj,correct_noobj, tot_class_preds,tot_obj, tot_noobj):
+
+    for i in range(3):
+
+      obj = y[i][..., 0] == 1 # in paper this is Iobj_i
+      noobj = y[i][..., 0] == 0  # in paper this is Iobj_i
+
+      correct_class += torch.sum(
+          torch.argmax(out[i][..., 5:][obj], dim=-1) == y[i][..., 5][obj]
+      )
+      tot_class_preds += torch.sum(obj)
+
+      obj_preds = torch.sigmoid(out[i][..., 0]) > threshold
+      correct_obj += torch.sum(obj_preds[obj] == y[i][..., 0][obj])
+      tot_obj += torch.sum(obj)
+      correct_noobj += torch.sum(obj_preds[noobj] == y[i][..., 0][noobj])
+      tot_noobj += torch.sum(noobj)
+
+    return((correct_class/(tot_class_preds+1e-16))*100,
+         (correct_noobj/(tot_noobj+1e-16))*100,
+         (correct_obj/(tot_obj+1e-16))*100)
+
+# pytorch lightning
+class LitYolo(LightningModule):
+    def __init__(self, num_classes=config.NUM_CLASSES, lr=1E-3,weight_decay=config.WEIGHT_DECAY,threshold=config.CONF_THRESHOLD):
+        super().__init__()
+
+        self.save_hyperparameters()
+        self.model = YOLOv3(num_classes=self.hparams.num_classes)
+        self.criterion = model_criterion
+        self.accuracy_fn = accuracy_fn
+        self.scaled_anchors = (torch.tensor(config.ANCHORS) * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2))
+        self.tot_class_preds, self.correct_class = 0, 0
+        self.tot_noobj, self.correct_noobj = 0, 0
+        self.tot_obj, self.correct_obj = 0, 0
+
+    def forward(self, x):
+        out = self.model(x)
+        return out
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        out  = self(x)
+        loss = self.criterion(out,y,self.scaled_anchors)
+        acc  = self.accuracy_fn(y,out,self.hparams.threshold,self.correct_class,
+                                                                     self.correct_obj,
+                                                                     self.correct_noobj,
+                                                                     self.tot_class_preds,
+                                                                     self.tot_obj,
+                                                                     self.tot_noobj)
+
+        self.log('train_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
+        self.log_dict({"class_accuracy": acc[0], "no_object_accuracy": acc[1], "object_accuracy":acc[2]},prog_bar=True,on_step=False, on_epoch=True)
+        return loss
+
+
+    def evaluate(self, batch, stage=None):
+        x, y = batch
+        out = self(x)
+        loss = self.criterion(out,y,self.scaled_anchors)
+        acc  = self.accuracy_fn(y,out,self.hparams.threshold,self.correct_class,
+                                                                     self.correct_obj,
+                                                                     self.correct_noobj,
+                                                                     self.tot_class_preds,
+                                                                     self.tot_obj,
+                                                                     self.tot_noobj)
+
+        if stage:
+            self.log(f"{stage}_loss", loss, prog_bar=True)
+            self.log_dict({"class_accuracy": acc[0], "no_object_accuracy": acc[1], "object_accuracy":acc[2]},prog_bar=True)
+
+    def test_step(self, batch, batch_idx):
+        self.evaluate(batch, "test")
+
+    def validation_step(self, batch, batch_idx):
+        self.evaluate(batch, "val")
+
+    def configure_optimizers(self):
+        optimizer = optim.Adam(self.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay)
+        scheduler = OneCycleLR(
+                optimizer,
+                max_lr= 1E-3,
+                pct_start = 5/self.trainer.max_epochs,
+                epochs=self.trainer.max_epochs,
+                steps_per_epoch=len(train_loader),
+                div_factor=100,verbose=True,
+                three_phase=False
+            )
+        return ([optimizer],[scheduler])
+yololit = LitYolo()
+inference_model = yololit.load_from_checkpoint("yolo3_model.ckpt")
+
+def yolo3_inference(input_img):
+
+  anchors = (torch.tensor(config.ANCHORS) * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2))
+  bboxes = [[]]
+
+  # color of the boxes
+  cmap = plt.get_cmap("tab20b")
+  class_labels = config.PASCAL_CLASSES
+  colors = [cmap(i) for i in np.linspace(0, 1, len(class_labels))]
+
+
+  # image transformation
+  test_transforms = Al.Compose(
+    [
+        Al.LongestMaxSize(max_size=416),
+        Al.PadIfNeeded(
+            min_height=416, min_width=416, border_mode=cv2.BORDER_CONSTANT
+        ),
+        Al.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),
+        ToTensorV2(),
+    ]
+  )
+  pr_input_img = test_transforms(image=input_img)
+  pr_input_img = pr_input_img['image'].unsqueeze(0)
+  test_img_out   = inference_model(pr_input_img)
+
+  # process the outputs
+  for i in range(3):
+      batch_size, A, S, _, _ = test_img_out[i].shape # 1, anchors = 3, scaling = 13/26/52
+      anchor = anchors[i]
+      boxes_scale_i = utils.cells_to_bboxes(test_img_out[i], anchor, S=S, is_preds=True)
+      for idx, (box) in enumerate(boxes_scale_i):
+          bboxes[idx] += box
+  # nms
+  boxes = utils.non_max_suppression(bboxes[0], iou_threshold=0.6, threshold=0.5, box_format="midpoint",)
+
+  # create matplotlib plot
+  fig, ax = plt.subplots(1)
+  # Display the image
+  ax.imshow(input_img)
+  height, width, _ = input_img.shape
+
+  # add boxes to the image
+  for box in boxes:
+    assert len(box) == 6, "box should contain class pred, confidence, x, y, width, height"
+    class_pred = box[0]
+    box = box[2:]
+    upper_left_x = box[0] - box[2] / 2
+    upper_left_y = box[1] - box[3] / 2
+    rect = patches.Rectangle(
+        (upper_left_x * width, upper_left_y * height),
+        box[2] * width,
+        box[3] * height,
+        linewidth=2,
+        edgecolor=colors[int(class_pred)],
+        facecolor="none",
+    )
+    # Add the patch to the Axes
+    ax.add_patch(rect)
+    plt.text(
+        upper_left_x * width,
+        upper_left_y * height,
+        s=class_labels[int(class_pred)],
+        color="white",
+        verticalalignment="top",
+        bbox={"color": colors[int(class_pred)], "pad": 0},
+    )
+  #plt.show()
+  img_buf = io.BytesIO()
+  fig.savefig(img_buf, format='png')
+  img_buf.seek(0)
+  img_arr = np.frombuffer(img_buf.getvalue(), dtype=np.uint8)
+  img_buf.close()
+  output_img = cv2.imdecode(img_arr, 1)
+  output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2RGB)
+
+  return output_img