v1

app.py

@@ -3,23 +3,21 @@ from ultralytics import YOLO
 from PIL import Image
 import os
 import cv2 
+import torch 
+
 
 
-def load_model(model_input):
-    """
-    Load the model either from a file or from Hugging Face.
-    """
-    model = YOLO(model_input)
-    return model
 
-def detect_objects_in_files(model_input, files):
+def detect_objects_in_files(files):
     """
     Processes uploaded images for object detection.
     """
     if not files:
         return "No files uploaded.", []
 
-    model = load_model(model_input)
+    device = "cuda" if torch.cuda.is_available() else "cpu"  
+    model = YOLO("https://github.com/luisarizmendi/ai-apps/raw/refs/heads/main/models/luisarizmendi/safety-hat/safety-hat-v1.pt")
+    model.to(device)
     
     results_images = []
     for file in files:
@@ -36,14 +34,14 @@ def detect_objects_in_files(model_input, files):
         except Exception as e:
             return f"Error processing file: {file}. Exception: {str(e)}", []
 
+    del model  
+    torch.cuda.empty_cache()
+    
     return "Processing completed.", results_images
 
-iface = gr.Interface(
+interface = gr.Interface(
     fn=detect_objects_in_files,
-    inputs=[
-        gr.File(label="Upload Model file"),
-        gr.Files(file_types=["image"], label="Select Images"),
-    ],
+    inputs=gr.Files(file_types=["image"], label="Select Images"),
     outputs=[
         gr.Textbox(label="Status"),
         gr.Gallery(label="Results")
@@ -53,5 +51,5 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    interface.launch()
 

requirements.txt

@@ -2,3 +2,4 @@ gradio
 ultralytics
 Pillow
 opencv-python
+torch