Updated inferencing code

app.py

@@ -1,27 +1,24 @@
 import os
 import cv2
+import glob
 import math
 import torch
 import numpy as np
 import gradio as gr
-import albumentations
 import matplotlib.pyplot as plt
-from glob import glob
+
 from PIL import Image
-from pytorch_grad_cam import EigenCAM
-from models.common import DetectMultiBackend
-from albumentations.pytorch import ToTensorV2
-from utils.augmentations import letterbox
 from utils.plots import Annotator, colors
-from pytorch_grad_cam.utils.image import show_cam_on_image, scale_cam_image
+from utils.augmentations import letterbox
+from models.common import DetectMultiBackend
+from utils.general import non_max_suppression, scale_boxes
 from utils.torch_utils import select_device, smart_inference_mode
-from utils.general import check_img_size, Profile, non_max_suppression, scale_boxes
 
 weights = "runs/train/best_striped.pt"
 data = "data.yaml"
 # Load model
 device = select_device('cpu')
-model = DetectMultiBackend(weights, device=device, dnn=False, data=data, fp16=False)
+model = DetectMultiBackend(weights=weights, device=device, fp16=False, data=data)
 #target_layers = [model.model.model[-1]]
 
 false_detection_data = glob(os.path.join("false_detection", '*.jpg'))
@@ -65,41 +62,32 @@ def display_false_detection_data(false_detection_data, number_of_samples):
     return fig
 
 def inference(input_img, conf_thres, iou_thres, is_false_detection_images=True, num_false_detection_images=10):
-    im0 = input_img.copy()
-    rgb_img = cv2.resize(im0, (640, 640))
     stride, names, pt = model.stride, model.names, model.pt
-    imgsz = check_img_size((640, 640), s=stride)  # check image size
-
-    bs = 1
-    # Run inference
-    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz)) 
-    seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
-
-    with dt[0]:
-        im = letterbox(input_img, imgsz, stride=stride, auto=True)[0]  # padded resize
-        im = im.transpose((2, 0, 1))[::-1]
-        im = np.ascontiguousarray(im)
-        im = torch.from_numpy(im).to(model.device)
-        im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
-        im /= 255  # 0 - 255 to 0.0 - 1.0
-        if len(im.shape) == 3:
-            im = im[None]  # expand for batch dim
-    
+
+    # Load image
+    img0 = input_img.copy()
+    img = letterbox(img0, 640, stride=stride, auto=True)[0]
+    img = img[:, :, ::-1].transpose(2, 0, 1)
+    img = np.ascontiguousarray(img)
+    img = torch.from_numpy(img).to(device).float()
+    img /= 255.0
+    if img.ndimension() == 3:
+        img = img.unsqueeze(0)
+
     # Inference
-    with dt[1]:
-        pred = model(im, augment=False, visualize=False)
+    pred = model(img, augment=False, visualize=False)
 
-    # NMS
-    with dt[2]:
-        pred = non_max_suppression(pred, conf_thres, iou_thres, None, False, max_det=10)
+    # Apply NMS
+    pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, max_det=1000)
 
     # Process predictions
+    seen = 0
     for i, det in enumerate(pred):  # per image
         seen += 1
-        annotator = Annotator(im0, line_width=2, example=str(model.names))
+        annotator = Annotator(img0, line_width=2, example=str(model.names))
         if len(det):
             # Rescale boxes from img_size to im0 size
-            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
+            det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], img0.shape).round()
 
             # Write results
             for *xyxy, conf, cls in reversed(det):
@@ -117,7 +105,7 @@ def inference(input_img, conf_thres, iou_thres, is_false_detection_images=True,
     # grayscale_cam = cam(im)[0, :]
     # cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)
 
-    return im0, misclassified_images
+    return img0, misclassified_images
 
 title = "YOLOv9 model to detect shirt/tshirt"
 description = "A simple Gradio interface to infer on YOLOv9 model and detect tshirt in image"