Adding images

app.py

@@ -1,18 +1,12 @@
 import gradio as gr
-# import spaces
-from huggingface_hub import hf_hub_download
+import numpy as np
 import os
+import spaces
+
+from huggingface_hub import hf_hub_download
+from ultralytics import YOLO
 
-# make sure you have the following dependencies
-# import torch
-import numpy as np
-# 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.augmentations import letterbox
-# import PIL.Image
 
-#@smart_inference_mode()
 @spaces.GPU
 def yolov9_inference(img_path, model_id='YOLOv9-S_X_LOCO-converted.pt', img_size=640, conf_thres=0.1, iou_thres=0.4):
     """
@@ -27,36 +21,13 @@ def yolov9_inference(img_path, model_id='YOLOv9-S_X_LOCO-converted.pt', img_size
     :return: A tuple containing the detections (boxes, scores, categories) and the results object for further actions like displaying.
     """
     
-    return np.array(PIL.Image.open(img_path))
-
-    # # Load the model
-    # model_path = download_models(model_id)
-
-    # # Initialize
-    # device = select_device('0')
-    # model = DetectMultiBackend(model_path, device="0", fp16=False, data='data/coco.yaml')
-    # stride, names, pt = model.stride, model.names, model.pt
-
-    # # Load image
-    # img = np.array(PIL.Image.open(img_path))
-    # img = letterbox(img0, img_size, 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
-    # results = model(img, augment=False, visualize=False)
+    model_path = download_models(model_id)
 
-    # # Apply NMS
-    # results = non_max_suppression(results[0][0], conf_thres, iou_thres, classes=None, max_det=1000)
+    model = YOLO(model_path)
 
-    # output = results.render()
-
-    # return output[0]
+    results = model.predict(img_path, conf=conf_thres, iou=iou_thres, imgsz=img_size)
 
+    return np.array(PIL.Image.open(img_path))
 
 
 def download_models(model_id):
@@ -64,38 +35,6 @@ def download_models(model_id):
                     token=os.getenv("HF_TOKEN"))
     return f"./{model_id}"
 
-# @spaces.GPU
-# def yolov9_inference(img_path, model_id, image_size, conf_threshold, iou_threshold):
-#     """
-#     Load a YOLOv9 model, configure it, perform inference on an image, and optionally adjust 
-#     the input size and apply test time augmentation.
-    
-#     :param model_path: Path to the YOLOv9 model file.
-#     :param conf_threshold: Confidence threshold for NMS.
-#     :param iou_threshold: IoU threshold for NMS.
-#     :param img_path: Path to the image file.
-#     :param size: Optional, input size for inference.
-#     :return: A tuple containing the detections (boxes, scores, categories) and the results object for further actions like displaying.
-#     """
-#     # Import YOLOv9
-#     import yolov9
-    
-#     # Load the model
-#     model_path = download_models(model_id)
-#     model = yolov9.load(model_path, device="cuda:0")
-    
-#     # Set model parameters
-#     model.conf = conf_threshold
-#     model.iou = iou_threshold
-    
-#     # Perform inference
-#     results = model(img_path, size=image_size)
-
-#     # Optionally, show detection bounding boxes on image
-#     output = results.render()
-    
-#     return output[0]
-
 
 def app():
     with gr.Blocks():
@@ -106,9 +45,9 @@ def app():
                     label="Model",
                     choices=[
                         "YOLOv9-S_X_LOCO-converted.pt",
-                        "YOLOv9_S_X_LOCO.pt",
+                        "YOLOv9-S_X_LOCO.pt",
                         "YOLOv9-E_X_LOCO-converted.pt",
-                        "YOLOv9_E_X_LOCO.pt",
+                        "YOLOv9-E_X_LOCO.pt",
                     ],
                     value="YOLOv9-S_X_LOCO-converted.pt",
                 )

loco.yaml

@@ -0,0 +1,7 @@
+nc: 5
+names:
+  0: small_load_carrier
+  1: forklift
+  2: pallet
+  3: stillage
+  4: pallet_truck