Sirapatrwan commited on
Commit
08d2f61
·
verified ·
1 Parent(s): 1725a15

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +64 -0
app.py ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app.py
2
+
3
+ from transformers import pipeline
4
+ import gradio as gr
5
+ from PIL import Image, ImageDraw, ImageFont
6
+ import random
7
+
8
+ # Load the YOLO-based object detection pipeline
9
+ detector = pipeline("object-detection", model="hustvl/yolos-tiny")
10
+
11
+ # Generate a random color for each label
12
+ label_colors = {}
13
+
14
+ def get_color(label):
15
+ if label not in label_colors:
16
+ label_colors[label] = (
17
+ random.randint(0, 255),
18
+ random.randint(0, 255),
19
+ random.randint(0, 255)
20
+ )
21
+ return label_colors[label]
22
+
23
+ # Detection function
24
+ def detect_objects(img):
25
+ results = detector(img)
26
+ draw = ImageDraw.Draw(img)
27
+ font = ImageFont.load_default()
28
+
29
+ for obj in results:
30
+ label = obj["label"]
31
+ score = obj["score"]
32
+ box = obj["box"]
33
+ color = get_color(label)
34
+
35
+ # Draw bounding box
36
+ draw.rectangle(
37
+ [box["xmin"], box["ymin"], box["xmax"], box["ymax"]],
38
+ outline=color,
39
+ width=3
40
+ )
41
+
42
+ # Prepare label with confidence
43
+ label_text = f"{label} ({score:.2f})"
44
+ text_bbox = draw.textbbox((box["xmin"], box["ymin"]), label_text, font=font)
45
+ text_background = [text_bbox[0], text_bbox[1], text_bbox[2], text_bbox[3]]
46
+
47
+ # Draw background for text
48
+ draw.rectangle(text_background, fill=color)
49
+ draw.text((text_bbox[0], text_bbox[1]), label_text, fill="black", font=font)
50
+
51
+ return img
52
+
53
+ # Gradio interface
54
+ interface = gr.Interface(
55
+ fn=detect_objects,
56
+ inputs=gr.Image(type="pil"),
57
+ outputs=gr.Image(type="pil"),
58
+ title="YOLO Object Detection with Color-coded Labels",
59
+ description="Upload an image. Detected objects are shown with bounding boxes and color-coded labels using YOLOS-Tiny."
60
+ )
61
+
62
+ if __name__ == "__main__":
63
+ interface.launch()
64
+