Update app.py

app.py

@@ -12,10 +12,6 @@ import yolo9
 
 model = yolo9.load('best.pt')
 
-# Set model parameters
-model.conf = conf_threshold
-model.iou = iou_threshold
-
 classes = ('ball', 'goalkeeper', 'player', 'referee')
 
 model_layer_names = ["1", "2", "3"]
@@ -44,7 +40,10 @@ def resize_image_pil(image, new_width, new_height):
     return resized
 
 
-def inference(input_img, show_gradcam, layer_name, num_classes, transparancy = 0.5):
+def inference(input_img, num_classes, conf_threshold, show_gradcam, layer_name, num_classes, transparancy = 0.5):
+    # Set model parameters
+    model.conf = conf_threshold
+    model.iou = iou_threshold
     print(show_gradcam, layer_name, num_classes, transparancy)
     input_img = resize_image_pil(input_img,32,32)
     input_img = np.array(input_img)
@@ -55,6 +54,7 @@ def inference(input_img, show_gradcam, layer_name, num_classes, transparancy = 0
     input_img = transform(input_img)
     input_img = input_img.unsqueeze(0)
     outputs = model(input_img)
+    output_res = outputs.render()
     # print(outputs)
     softmax =  torch.nn.Softmax(dim=0)
     o = softmax(outputs.flatten())
@@ -62,12 +62,12 @@ def inference(input_img, show_gradcam, layer_name, num_classes, transparancy = 0
     output_numpy = np.squeeze(np.asarray(outputs.detach().numpy()))
     index_sort = np.argsort(output_numpy)[::-1]
     
-    confidences = {}
-    for i in range(int(num_classes)):
-        confidences[classes[index_sort[i]]] = float(o[index_sort[i]])
+    # confidences = {}
+    # for i in range(int(num_classes)):
+        # confidences[classes[index_sort[i]]] = float(o[index_sort[i]])
 
     
-    prediction= torch.max(outputs, 1)
+    # prediction= torch.max(outputs, 1)
 
     if show_gradcam:
         target_layers = get_layer(layer_name)
@@ -81,27 +81,28 @@ def inference(input_img, show_gradcam, layer_name, num_classes, transparancy = 0
         visualization = org_img
 
     
-    return classes[int(prediction[0].item())], visualization, confidences
+    return output_res[0], visualization
     
 demo = gr.Interface(
     inference,
     inputs = [
         gr.Image(width=256,height=256,label="Input image"),
-        gr.Number(value=3, maximum=10, minimum=1,step=1.0, precision=0,label="Number of classes to display"),
+        gr.Slider(label='Confidence threshold', minimum=0.1,maximum=1.0,step=0.1,value=0.4),
+        gr.Slider(label='IOU threshold', minimum=0.1,maximum=1.0,step=0.1,value=0.5),
         gr.Checkbox(True, label="Show GradCAM Image"),
         gr.Dropdown(model_layer_names, value=3, label="Which layer for Gradcam"),
         gr.Slider(0, 1, value=0.5,label="Overall opacity of the overlay"),
     ],
     outputs = [
-        gr.Label(label="Class", container=True, show_label= True),
         gr.Image(width= 256, height=256,label="Output Image"),
-        gr.Label(label="Confidences", container=True, show_label= True),
+        gr.Image(width= 256, height=256,label="GradCAM image")
     ],
-    title = "Yolo v9 on custom fooball dataset",
-    description = " A simple gradio inference, and detection results for custom trained yolov9 model",
-    examples = [["img1.jpg",1, True, 10, 0.4]]
+    title = "Yolo v9 on custom trained fooball dataset",
+    description = " A simple gradio inference, and detection results for custom trained yolov9 for 100 epochs",
+    examples = [["img1.jpg",1, True, 10, 0.4]],
+    cache_examples=True
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True)