Amandeep01 commited on
Commit
573dab7
·
verified ·
1 Parent(s): b0a82cf

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +26 -33
app.py CHANGED
@@ -14,45 +14,39 @@ def get_color_for_label(label):
14
 
15
  # Main function: detect, draw, and return outputs
16
  def detect_and_draw(image, threshold):
17
- try:
18
- # Perform object detection
19
- results = detector(image)
20
- image = image.convert("RGB")
21
- draw = ImageDraw.Draw(image)
22
-
23
- # Try to load a font for annotations, else use default
24
- try:
25
- font = ImageFont.truetype("arial.ttf", 16)
26
- except:
27
- font = ImageFont.load_default()
28
 
29
- annotations = []
 
 
 
30
 
31
- for obj in results:
32
- score = obj["score"]
33
- if score < threshold:
34
- continue
35
 
36
- label = f"{obj['label']} ({score:.2f})"
37
- box = obj["box"]
38
- color = get_color_for_label(obj["label"])
 
39
 
40
- # Draw the bounding box and label
41
- draw.rectangle(
42
- [(box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])],
43
- outline=color,
44
- width=3,
45
- )
46
 
47
- draw.text((box["xmin"] + 5, box["ymin"] + 5), label, fill=color, font=font)
 
 
 
 
48
 
49
- box_coords = (box["xmin"], box["ymin"], box["xmax"], box["ymax"])
50
- annotations.append((box_coords, label))
51
 
52
- return image, annotations
 
53
 
54
- except Exception as e:
55
- return f"Error during detection: {e}", None
56
 
57
  # Gradio UI setup
58
  demo = gr.Interface(
@@ -66,7 +60,6 @@ demo = gr.Interface(
66
  ],
67
  title="YOLOS Object Detection",
68
  description="Upload an image to detect objects using the YOLOS-small model. Adjust the confidence threshold using the slider.",
69
- live=True
70
  )
71
 
72
- demo.launch()
 
14
 
15
  # Main function: detect, draw, and return outputs
16
  def detect_and_draw(image, threshold):
17
+ results = detector(image)
18
+ image = image.convert("RGB")
19
+ draw = ImageDraw.Draw(image)
 
 
 
 
 
 
 
 
20
 
21
+ try:
22
+ font = ImageFont.truetype("arial.ttf", 16)
23
+ except:
24
+ font = ImageFont.load_default()
25
 
26
+ annotations = []
 
 
 
27
 
28
+ for obj in results:
29
+ score = obj["score"]
30
+ if score < threshold:
31
+ continue
32
 
33
+ label = f"{obj['label']} ({score:.2f})"
34
+ box = obj["box"]
35
+ color = get_color_for_label(obj["label"])
 
 
 
36
 
37
+ draw.rectangle(
38
+ [(box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])],
39
+ outline=color,
40
+ width=3,
41
+ )
42
 
43
+ draw.text((box["xmin"] + 5, box["ymin"] + 5), label, fill=color, font=font)
 
44
 
45
+ box_coords = (box["xmin"], box["ymin"], box["xmax"], box["ymax"])
46
+ annotations.append((box_coords, label))
47
 
48
+ # Return the annotated image and annotations (no download option)
49
+ return image, annotations
50
 
51
  # Gradio UI setup
52
  demo = gr.Interface(
 
60
  ],
61
  title="YOLOS Object Detection",
62
  description="Upload an image to detect objects using the YOLOS-small model. Adjust the confidence threshold using the slider.",
 
63
  )
64
 
65
+ demo.launch()