Update app.py

app.py

@@ -1,11 +1,11 @@
 import torch
 import torchvision
-from torchvision import transforms
+# from torchvision import transforms
 import numpy as np
 import gradio as gr
 from PIL import Image
-from pytorch_grad_cam import GradCAM
-from pytorch_grad_cam.utils.image import show_cam_on_image
+# from pytorch_grad_cam import GradCAM
+# from pytorch_grad_cam.utils.image import show_cam_on_image
 import gradio as gr
 from models.yolo import Model
 
@@ -16,97 +16,81 @@ import yolov9
                 
 model = yolov9.load('best.pt')
 
-classes = ('ball', 'goalkeeper', 'player', 'referee')
+# classes = ('ball', 'goalkeeper', 'player', 'referee')
 
-model_layer_names = ["1", "2", "3"]
-
-def get_layer(layer_name):
-    print("layer name:", layer_name)
-    if layer_name == 1:
-        return [model.layer1[-1]]
-    elif layer_name == 2:
-        return [model.layer2[-1]]
-    elif layer_name == 3:
-        return [model.layer3[-1]]
-    else:
-        return None
-
-def resize_image_pil(image, new_width, new_height):
-    img =  Image.fromarray(np.array(image))
-    width, height = img.size
-
-    width_scale = new_width/width
-    height_scale = new_height/height
-    scale = min(width_scale, height_scale)
-    resized = img.resize((int(width*scale), int(height*scale)), Image.NEAREST)
-    resized = resized.crop((0,0,new_width, new_height))
-
-    return resized
-
-
-def inference(input_img, conf_threshold, iou_threshold, show_gradcam, layer_name, transparancy = 0.5):
+def inference(input_img, conf_threshold, iou_threshold):
     # 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)
-    org_img = input_img
-    
-    input_img= input_img.reshape((32,32,3))
-    transform = transforms.ToTensor()
-    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())
+    # Perform inference
+    results = model(img_path, size=image_size)
+
+    # Optionally, show detection bounding boxes on image
+    output = results.render()
     
-    output_numpy = np.squeeze(np.asarray(outputs.detach().numpy()))
-    index_sort = np.argsort(output_numpy)[::-1]
+    return output[0]
     
-    # confidences = {}
-    # for i in range(int(num_classes)):
-        # confidences[classes[index_sort[i]]] = float(o[index_sort[i]])
 
-    
-    # prediction= torch.max(outputs, 1)
+def app():
+    with gr.Blocks():
+        with gr.Row():
+            with gr.Column():
+                input_img = gr.Image(width= 256, height=256,label="Image")
+                conf_threshold = gr.Slider(
+                    label="Confidence threshold",
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.1,
+                    value=0.4
+                )
+                iou_threshold = gr.Slider(
+                    label="IoU Threshold",
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.1,
+                    value=0.5,
+                )
+                yolo_inf = gr.Button(value="Inference")
 
-    if show_gradcam:
-        target_layers = get_layer(layer_name)
-        print("target layer",target_layers)
-        cam = GradCAM(model=model, target_layers=target_layers)
-        grayscale_cam = cam(input_tensor= input_img)
-        grayscale_cam = grayscale_cam[0, :]
-        visualization = show_cam_on_image(org_img/255,grayscale_cam,use_rgb=True,
-                                     image_weight=transparancy)
-    else:
-        visualization = org_img
+            with gr.Column():
+                output_val = gr.Image(width= 256, height=256,label="Output Image")
+        yolo_inf.click(
+            fn= inference,
+            inputs = [
+                input_img,
+                conf_threshold,
+                iou_threshold
+            ],
+            outputs = [output_val],
+        )
+        gr.Examples(["img1.jpg",0.4, 0.6, 0.4],
+                   fn= inference,
+                   inputs = [
+                       input_img,
+                       conf_threshold,
+                       iou_threshold
+            ],
+            outputs = [output_val],
+            cache_examples=True,
+                   )
 
-    
-    return output_res[0], visualization
-    
-demo = gr.Interface(
-    inference,
-    inputs = [
-        gr.Image(width=256,height=256,label="Input image"),
-        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.Image(width= 256, height=256,label="Output Image"),
-        gr.Image(width= 256, height=256,label="GradCAM image")
-    ],
-    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",0.4, 0.6, True, 3, 0.4]],
-    cache_examples=True
-)
+demo = gr.Blocks()
+with demo:
+    gr.HTML(
+        """
+    <h1 style='text-align: center'>
+    YOLOv9
+    </h1>
+    """)
+    gr.HTML(
+        """
+        <h3 style='text-align: center'>
+        Inferencing yolov9 with custom dataset - football players dataset 
+        </h3>
+        """)
+    with gr.Row():
+        with gr.Column():
+            app()
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch(debug=True)