Update app.py

app.py

@@ -1,11 +1,10 @@
+
 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 
@@ -13,8 +12,7 @@ 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 pytorch_lightning.loggers import CSVLogger,TensorBoardLogger
 from tqdm import tqdm
 import torch
 import torch.optim as optim
@@ -139,83 +137,73 @@ class LitYolo(LightningModule):
                 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
+
+
+
+
+# gradio
+with gr.Blocks() as demo:
+
+    # colors for the bboxes
+    cmap = plt.get_cmap("tab20b")
+    class_labels = config.PASCAL_CLASSES
+    colors = [cmap(i) for i in np.linspace(0, 1, len(class_labels))]
+    colors_hex = {class_labels[i]:matplotlib.colors.rgb2hex(colors[i]) for i in range(0,len(class_labels))}
+
+    # app GUI
+    with gr.Row():
+        img_input  = gr.Image()
+        img_output = gr.AnnotatedImage().style(color_map = colors_hex)
+    section_btn = gr.Button("Identify Objects")
+
+    def yolo3_inference(input_img): # function for yolo inference
+
+      yololit = LitYolo()
+      inference_model = yololit.load_from_checkpoint("yolo3_model.ckpt")
+      anchors  = (torch.tensor(config.ANCHORS) * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2))
+      bboxes   = [[]]
+      sections = [] # to return image and annotations
+
+      # 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)
+      # infer the image
+      inference_model.eval()
+      test_img_out   = inference_model(pr_input_img)
+
+      # process the outputs to create bounding boxes
+      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_org.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
+      nms_boxes = utils_org.non_max_suppression(bboxes[0], iou_threshold=0.6, threshold=0.5, box_format="midpoint",)
+
+      # use gradio image annotations
+      height, width = 416, 416
+      for box in nms_boxes:
+        class_pred = box[0]
+        box = box[2:]
+        upper_left_x  = int((box[0] - box[2] / 2) * width)
+        upper_left_y  = max(int((box[1] - box[3] / 2) * height),0) # less than 0, box collapses
+        lower_right_x = int(upper_left_x + (box[2] * width))
+        lower_right_y = int(upper_left_y + (box[3] * height))
+        sections.append(((upper_left_x,upper_left_y,lower_right_x,lower_right_y), class_labels[int(class_pred)]))
+      return (np.array(pr_input_img.squeeze(0).permute(1,2,0)),sections)
+
+    section_btn.click(yolo3_inference, img_input, img_output)
+
+if __name__ == "__main__":
+    demo.launch()