app

app.py

@@ -5,19 +5,21 @@ import os
 import cv2
 import torch
 
+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(files):
+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-safety/object-detector-safety-v1.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:
@@ -41,7 +43,10 @@ def detect_objects_in_files(files):
 
 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")
@@ -52,4 +57,3 @@ interface = gr.Interface(
 
 if __name__ == "__main__":
     interface.launch()
-