Commit
·
5f178a4
1
Parent(s):
52f9205
app
Browse files
app.py
CHANGED
@@ -5,21 +5,19 @@ import os
|
|
5 |
import cv2
|
6 |
import torch
|
7 |
|
8 |
-
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"
|
9 |
|
10 |
-
|
|
|
|
|
11 |
"""
|
12 |
Processes uploaded images for object detection.
|
13 |
"""
|
14 |
if not files:
|
15 |
return "No files uploaded.", []
|
16 |
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
print("Using GPU for inference")
|
21 |
-
else:
|
22 |
-
print("Using CPU for inference")
|
23 |
|
24 |
results_images = []
|
25 |
for file in files:
|
@@ -43,10 +41,7 @@ def detect_objects_in_files(model_input, files):
|
|
43 |
|
44 |
interface = gr.Interface(
|
45 |
fn=detect_objects_in_files,
|
46 |
-
inputs=[
|
47 |
-
gr.Textbox(value=DEFAULT_MODEL_URL, label="Model URL", placeholder="Enter the model URL"),
|
48 |
-
gr.Files(file_types=["image"], label="Select Images"),
|
49 |
-
],
|
50 |
outputs=[
|
51 |
gr.Textbox(label="Status"),
|
52 |
gr.Gallery(label="Results")
|
@@ -57,3 +52,4 @@ interface = gr.Interface(
|
|
57 |
|
58 |
if __name__ == "__main__":
|
59 |
interface.launch()
|
|
|
|
5 |
import cv2
|
6 |
import torch
|
7 |
|
|
|
8 |
|
9 |
+
|
10 |
+
|
11 |
+
def detect_objects_in_files(files):
|
12 |
"""
|
13 |
Processes uploaded images for object detection.
|
14 |
"""
|
15 |
if not files:
|
16 |
return "No files uploaded.", []
|
17 |
|
18 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
19 |
+
model = YOLO("https://github.com/luisarizmendi/ai-apps/raw/refs/heads/main/models/luisarizmendi/object-detector-safety/object-detector-safety-v1.pt")
|
20 |
+
model.to(device)
|
|
|
|
|
|
|
21 |
|
22 |
results_images = []
|
23 |
for file in files:
|
|
|
41 |
|
42 |
interface = gr.Interface(
|
43 |
fn=detect_objects_in_files,
|
44 |
+
inputs=gr.Files(file_types=["image"], label="Select Images"),
|
|
|
|
|
|
|
45 |
outputs=[
|
46 |
gr.Textbox(label="Status"),
|
47 |
gr.Gallery(label="Results")
|
|
|
52 |
|
53 |
if __name__ == "__main__":
|
54 |
interface.launch()
|
55 |
+
|