app update

app.py

@@ -1 +1,92 @@
-print("hello")
+import torch
+import torch.optim as optim
+import lightning.pytorch as pl
+from torchvision import transforms
+from custom_yolo.custom_library.utils import cells_to_bboxes, non_max_suppression
+from custom_yolo.custom_library import config
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from custom_yolo.custom_library.lightning_model import YOLOv3Lightning
+import cv2
+import numpy as np
+from pytorch_grad_cam.utils.image import show_cam_on_image
+from custom_yolo.custom_library.gradio_utils import  draw_predictions, YoloCAM
+import gradio as gr
+import os
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+model = YOLOv3Lightning(config=config)
+model.load_state_dict(torch.load("custom_yolo_model.pth", map_location=torch.device('cpu')), strict=False)
+model.setup(stage="test")
+classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
+
+scaled_anchors = (torch.tensor(config.ANCHORS)* torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)).to(config.DEVICE)
+
+transforms = A.Compose(
+    [
+        A.LongestMaxSize(max_size=config.IMAGE_SIZE),
+        A.PadIfNeeded(
+            min_height=config.IMAGE_SIZE, min_width=config.IMAGE_SIZE, border_mode=cv2.BORDER_CONSTANT
+        ),
+        A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),
+        ToTensorV2(),
+    ],
+)
+
+def model_inference(image, iou_threshold=0.5, threshold=0.4, show_cam="No", transparency=0.5, target_layer=-2):
+    # Transforming image
+    transformed_image = transforms(image=image)["image"].unsqueeze(0)
+    output = model(transformed_image)
+    # Selecting layer for gradCAM
+    if target_layer == -2:
+      layer = [model.model.layers[-2]]
+    else:
+      layer = [model.model.layers[-1]]
+
+    cam = YoloCAM(model=model, target_layers=layer, use_cuda=False)
+
+    bboxes = [[] for _ in range(1)]
+    for i in range(3):
+        batch_size, A, S, _, _ = output[i].shape
+        anchor = scaled_anchors[i]
+        boxes_scale_i = cells_to_bboxes(output[i], anchor, S=S, is_preds=True)
+        for idx, (box) in enumerate(boxes_scale_i):
+            bboxes[idx] += box
+
+    nms_boxes = non_max_suppression(bboxes[0], iou_threshold=iou_threshold, threshold=threshold, box_format="midpoint")
+    plot_img = draw_predictions(image.copy(), nms_boxes, class_labels=config.PASCAL_CLASSES)
+    if  show_cam == "No":
+      return [plot_img]
+    else:
+      grayscale_cam = cam(transformed_image, scaled_anchors)[0, :, :]
+      img = cv2.resize(image, (416, 416))
+      img = np.float32(img) / 255
+      cam_image = show_cam_on_image(img, grayscale_cam, use_rgb=True, image_weight=transparency)
+      return [plot_img, cam_image]
+
+
+title = "Custom YOLOv3"
+description = "Pytorch Lightning implemetation of YOLOv3 on Pascal VOC dataset.\
+                        Supported classes are aeroplane, bicycle, bird, boat, bottle, bus, car, cat, chair, cow, dining table, dog, horse, motorbike, person, potted plant, sheep, sofa, train, and TV/monitor."
+
+# examples = [
+#     ["images/000014.jpg", 0.5, 0.4, True, 0.5],
+#     ["images/000017.jpg", 0.6, 0.5, True, 0.5],
+#     ["images/000018.jpg", 0.55, 0.45, True, 0.5],
+#     ["images/000030.jpg", 0.5, 0.4, True, 0.5],
+#     ["images/Puppies.jpg", 0.6, 0.7, True, 0.5],
+# ]
+
+demo = gr.Interface(model_inference, inputs=[gr.Image(shape=(416, 416), label="Input an image"),
+                                                                          gr.Slider(0, 1, value=0.5, label="IOU Threshold"),
+                                                                          gr.Slider(0, 1, value=0.4, label="Threshold"),
+                                                                          gr.Radio(["Yes", "No"], value="No"  , label="Show GradCAM outputs"),
+                                                                          gr.Slider(0, 1, value=0.5, label="Opacity of GradCAM"),
+                                                                          gr.Slider(-2, -1, value=-2, step=1, label="Which Layer?")],
+                                  outputs=[gr.Gallery(rows=2, columns=1)],
+                                  title=title, description=description)
+demo.launch() 
+
+