File size: 1,835 Bytes
85b2ce2 9c22bfa 85b2ce2 9c22bfa 85b2ce2 9c22bfa 85b2ce2 9c22bfa 85b2ce2 9c22bfa 85b2ce2 9c22bfa 85b2ce2 9c22bfa 85b2ce2 9c22bfa 85b2ce2 9c22bfa 85b2ce2 9c22bfa 85b2ce2 9c22bfa 85b2ce2 9c22bfa 85b2ce2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
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 detect_objects_in_files(model_input, files):
"""
Processes uploaded images for object detection.
"""
if not files:
return "No files uploaded.", []
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)
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()
|