Update app.py

app.py

@@ -63,7 +63,7 @@ scaled_anchors = (
     * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
 )
 
-def process_image_and_plot(image,iou_threshold=0.5, threshold=0.4,target_layer=-2):
+def process_image_and_plot(image,iou_threshold=0.5, threshold=0.4):
 
     transformed_image = transforms(image=image)["image"].unsqueeze(0)
     output = model(transformed_image)
@@ -80,13 +80,7 @@ def process_image_and_plot(image,iou_threshold=0.5, threshold=0.4,target_layer=-
         bboxes[0], iou_threshold=iou_threshold, threshold=threshold, box_format="midpoint",
     )
     fig = plot_image(transformed_image[0].permute(1, 2, 0), nms_boxes)
-
-    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)
+    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))
@@ -129,8 +123,7 @@ description_with_icon = f"""
 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.Slider(-2, -1, value=-1, step=1, label="Which Layer?"),],
+           gr.Slider(0, 1, value=0.4, label="Threshold"),],
     outputs=[
         gr.Plot(label="Output with Classes",), 
         gr.Image(shape=(32, 32), label="GradCAM Output"),