first commit

app.py

@@ -0,0 +1,64 @@
+import gradio as gr
+from ultralytics import YOLO
+from PIL import Image
+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 load_model(model_input):
+    model = YOLO(model_input)
+    if torch.cuda.is_available():
+        model.to('cuda')
+        print("Using GPU for inference")
+    else:
+        print("Using CPU for inference")
+
+    return model
+
+def detect_objects_in_files(model_input, files):
+    """
+    Processes uploaded images for object detection.
+    """
+    if not files:
+        return "No files uploaded.", []
+
+    model = load_model(model_input)
+
+    results_images = []
+    for file in files:
+        try:
+            image = Image.open(file).convert("RGB")
+            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)
+
+            # If you want that images appear one by one (slower)
+            #yield "Processing image...", results_images
+
+        except Exception as e:
+            return f"Error processing file: {file}. Exception: {str(e)}", []
+
+    del model
+    torch.cuda.empty_cache()
+
+    return "Processing completed.", results_images
+
+interface = gr.Interface(
+    fn=detect_objects_in_files,
+    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")
+    ],
+    title="Object Detection on Images",
+    description="Upload images to perform object detection. The model will process each image and display the results."
+)
+
+if __name__ == "__main__":
+    interface.launch()