luisarizmendi commited on
Commit
5f178a4
·
1 Parent(s): 52f9205
Files changed (1) hide show
  1. app.py +8 -12
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
- def detect_objects_in_files(model_input, files):
 
 
11
  """
12
  Processes uploaded images for object detection.
13
  """
14
  if not files:
15
  return "No files uploaded.", []
16
 
17
- model = YOLO(str(model_input))
18
- if torch.cuda.is_available():
19
- model.to('cuda')
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
+