Update README.md

README.md

@@ -153,43 +153,51 @@ import gradio as gr
 from ultralytics import YOLO
 from PIL import Image
 import os
-import cv2 
-import torch 
+import cv2
+import torch
 
-def detect_objects_in_files(files):
+DEFAULT_MODEL_URL = "https://github.com/luisarizmendi/ai-apps/raw/refs/heads/main/models/luisarizmendi/object-detector-hardhat-or-hat/object-detector-hardhat-or-hat.pt"
+
+def detect_objects_in_files(model_input, files):
     """
     Processes uploaded images for object detection.
     """
     if not files:
         return "No files uploaded.", []
 
-    device = "cuda" if torch.cuda.is_available() else "cpu"  
-    model = YOLO("https://github.com/luisarizmendi/ai-apps/raw/refs/heads/main/models/luisarizmendi/object-detector-hardhat-or-hat/object-detector-hardhat-or-hat.pt")
-    model.to(device)
-    
+    model = YOLO(str(model_input))
+    if torch.cuda.is_available():
+        model.to('cuda')
+        print("Using GPU for inference")
+    else:
+        print("Using CPU for inference")
+
     results_images = []
     for file in files:
         try:
             image = Image.open(file).convert("RGB")
-            results = model(image) 
+            results = model(image)
             result_img_bgr = results[0].plot()
             result_img_rgb = cv2.cvtColor(result_img_bgr, cv2.COLOR_BGR2RGB)
-            results_images.append(result_img_rgb)   
-         
+            results_images.append(result_img_rgb)
+
             # If you want that images appear one by one (slower)
-            #yield "Processing image...", results_images  
-                
+            #yield "Processing image...", results_images
+
         except Exception as e:
             return f"Error processing file: {file}. Exception: {str(e)}", []
 
-    del model  
+    del model
     torch.cuda.empty_cache()
-    
+
     return "Processing completed.", results_images
 
 interface = gr.Interface(
     fn=detect_objects_in_files,
-    inputs=gr.Files(file_types=["image"], label="Select Images"),
+    inputs=[
+        gr.Textbox(value=DEFAULT_MODEL_URL, label="Model URL", placeholder="Enter the model URL"),
+        gr.Files(file_types=["image"], label="Select Images"),
+    ],
     outputs=[
         gr.Textbox(label="Status"),
         gr.Gallery(label="Results")