Update app.py

app.py

@@ -63,12 +63,12 @@ scaled_anchors = (
     * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
 )
 
-def process_image_and_plot(image, model, scaled_anchors):
+def process_image_and_plot(image,iou_threshold=0.5, threshold=0.4,show_cam="No"):
 
     transformed_image = transforms(image=image)["image"].unsqueeze(0)
     output = model(transformed_image)
     bboxes = [[] for _ in range(1)]
-
+    
     for i in range(3):
         batch_size, A, S, _, _ = output[i].shape
         anchor = scaled_anchors[i]
@@ -77,17 +77,25 @@ def process_image_and_plot(image, model, scaled_anchors):
             bboxes[idx] += box
 
     nms_boxes = non_max_suppression(
-        bboxes[0], iou_threshold=0.5, threshold=0.4, box_format="midpoint",
+        bboxes[0], iou_threshold=iou_threshold, threshold=threshold, box_format="midpoint",
     )
     fig = plot_image(transformed_image[0].permute(1, 2, 0), nms_boxes)
 
-    cam = YoloCAM(model=model, target_layers=[model.model.layers[-2]], use_cuda=False)
-    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)
-
-    return fig,cam_image
+    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)
+    if  show_cam == "No":
+        return fig
+    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)
+    
+        return fig,cam_image
 
 
 examples = [
@@ -96,9 +104,6 @@ examples = [
     ["images/automobile.jpg"],
 ]
 
-def processed_image(image):
-    figure,gradcam = process_image_and_plot(image, model, scaled_anchors)
-    return figure,gradcam
 
 icon_html = '<i class="fas fa-chart-bar"></i>'
 title_with_icon = f"""
@@ -123,8 +128,12 @@ description_with_icon = f"""
 """
 
 
-demo = gr.Interface(processed_image,
-    inputs=[gr.Image(label="Input Image"),],
+demo = gr.Interface(process_image_and_plot,
+    inputs=[gr.Image(label="Input Image"),
+           gr.Slider(0, 1, value=0.5, label="Intersection over Union (IOU) Threshold",info="Determines how much overlap between two boxes, Set it low to filter out weaker predicts"),
+           gr.Slider(0, 1, value=0.4, label="Threshold"),
+           gr.Radio(["Yes", "No"], value="No"  , label="Show GradCAM outputs"),
+           gr.Slider(-2, -1, value=-1, step=1, label="Which Layer?"),],
     outputs=[
         gr.Plot(label="Output with Classes",), 
         gr.Image(shape=(32, 32), label="GradCAM Output"),