Spaces:
Running
Running
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,162 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# --------------------------------------------------------
|
2 |
+
# YOLOv12 Streamlit App with Emoji-Powered UI ๐๐
|
3 |
+
# Based on yolov10: https://github.com/THU-MIG/yolov10/app.py
|
4 |
+
# --------------------------------------------------------
|
5 |
+
|
6 |
+
import streamlit as st
|
7 |
+
import cv2
|
8 |
+
import tempfile
|
9 |
+
from ultralytics import YOLO
|
10 |
+
from PIL import Image
|
11 |
+
import os
|
12 |
+
|
13 |
+
# Page config with a cool vibe
|
14 |
+
st.set_page_config(
|
15 |
+
page_title="YOLOv12 Detector ๐ต๏ธโโ๏ธ",
|
16 |
+
page_icon="๐",
|
17 |
+
layout="wide"
|
18 |
+
)
|
19 |
+
|
20 |
+
def yolov12_inference(uploaded_file, model_id, image_size, conf_threshold, input_type):
|
21 |
+
"""The magic happens here โจ"""
|
22 |
+
model = YOLO(model_id)
|
23 |
+
|
24 |
+
if input_type == "Image" and uploaded_file:
|
25 |
+
with st.spinner("๐ผ๏ธ Painting detections..."):
|
26 |
+
image = Image.open(uploaded_file)
|
27 |
+
results = model.predict(source=image, imgsz=image_size, conf=conf_threshold)
|
28 |
+
annotated_image = results[0].plot()
|
29 |
+
return annotated_image[:, :, ::-1], None
|
30 |
+
|
31 |
+
elif input_type == "Video" and uploaded_file:
|
32 |
+
with st.spinner("๐ฅ Cooking up a detected video..."):
|
33 |
+
video_path = tempfile.mktemp(suffix=".mp4")
|
34 |
+
with open(video_path, "wb") as f:
|
35 |
+
f.write(uploaded_file.read())
|
36 |
+
|
37 |
+
cap = cv2.VideoCapture(video_path)
|
38 |
+
fps = cap.get(cv2.CAP_PROP_FPS)
|
39 |
+
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
40 |
+
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
41 |
+
|
42 |
+
output_video_path = tempfile.mktemp(suffix=".mp4")
|
43 |
+
out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
|
44 |
+
|
45 |
+
frame_count = 0
|
46 |
+
while cap.isOpened():
|
47 |
+
ret, frame = cap.read()
|
48 |
+
if not ret:
|
49 |
+
break
|
50 |
+
results = model.predict(source=frame, imgsz=image_size, conf=conf_threshold)
|
51 |
+
annotated_frame = results[0].plot()
|
52 |
+
out.write(annotated_frame)
|
53 |
+
frame_count += 1
|
54 |
+
if frame_count % 30 == 0:
|
55 |
+
st.text(f"๐ณ Processed {frame_count} frames...")
|
56 |
+
|
57 |
+
cap.release()
|
58 |
+
out.release()
|
59 |
+
|
60 |
+
if os.path.exists(video_path):
|
61 |
+
os.remove(video_path)
|
62 |
+
|
63 |
+
return None, output_video_path
|
64 |
+
|
65 |
+
return None, None
|
66 |
+
|
67 |
+
def main():
|
68 |
+
# Header with flair
|
69 |
+
st.title("YOLOv12: Object Detection Superhero ๐ฆธโโ๏ธ")
|
70 |
+
st.markdown("Powered by xAI's cosmic tech ๐ | [arXiv ๐](https://arxiv.org/abs/2502.12524) | [GitHub ๐](https://github.com/sunsmarterjie/yolov12)")
|
71 |
+
|
72 |
+
# Layout in two columns
|
73 |
+
col1, col2 = st.columns([1, 1])
|
74 |
+
|
75 |
+
with col1:
|
76 |
+
st.subheader("๐ฎ Control Room")
|
77 |
+
|
78 |
+
# Upload section
|
79 |
+
uploaded_file = st.file_uploader(
|
80 |
+
"Drop your file here ๐ฅ - Images or Videos welcome!",
|
81 |
+
type=['jpg', 'jpeg', 'png', 'mp4'],
|
82 |
+
help="Upload an image or video to detect objects in!"
|
83 |
+
)
|
84 |
+
|
85 |
+
# Input type selector
|
86 |
+
input_type = st.radio(
|
87 |
+
"What's your flavor? ๐ฆ",
|
88 |
+
("Image", "Video"),
|
89 |
+
help="Tell me if it's a still or moving picture!"
|
90 |
+
)
|
91 |
+
|
92 |
+
# Model selection
|
93 |
+
model_id = st.selectbox(
|
94 |
+
"Pick your YOLO weapon โ๏ธ",
|
95 |
+
["yolov12n.pt", "yolov12s.pt", "yolov12m.pt", "yolov12l.pt", "yolov12x.pt"],
|
96 |
+
index=2,
|
97 |
+
help="Choose your model power level: n (nano) to x (extra spicy)!"
|
98 |
+
)
|
99 |
+
|
100 |
+
# Image size slider
|
101 |
+
image_size = st.slider(
|
102 |
+
"Zoom level ๐",
|
103 |
+
min_value=320,
|
104 |
+
max_value=1280,
|
105 |
+
value=640,
|
106 |
+
step=32,
|
107 |
+
help="Bigger numbers = sharper eyes (but slower)!"
|
108 |
+
)
|
109 |
+
|
110 |
+
# Confidence threshold
|
111 |
+
conf_threshold = st.slider(
|
112 |
+
"Certainty meter ๐ฏ",
|
113 |
+
min_value=0.0,
|
114 |
+
max_value=1.0,
|
115 |
+
value=0.25,
|
116 |
+
step=0.05,
|
117 |
+
help="How sure should I be? Higher = pickier!"
|
118 |
+
)
|
119 |
+
|
120 |
+
# The big red button
|
121 |
+
if st.button("Detect Objects! ๐", help="Click to unleash the detection magic!"):
|
122 |
+
if uploaded_file is None:
|
123 |
+
st.error("Yo! Upload something first ๐")
|
124 |
+
else:
|
125 |
+
annotated_image, annotated_video = yolov12_inference(
|
126 |
+
uploaded_file, model_id, image_size, conf_threshold, input_type
|
127 |
+
)
|
128 |
+
st.session_state['results'] = (annotated_image, annotated_video)
|
129 |
+
|
130 |
+
with col2:
|
131 |
+
st.subheader("๐ฅ๏ธ Detection HQ")
|
132 |
+
|
133 |
+
# Display results
|
134 |
+
if 'results' in st.session_state:
|
135 |
+
annotated_image, annotated_video = st.session_state['results']
|
136 |
+
|
137 |
+
if input_type == "Image" and annotated_image is not None:
|
138 |
+
st.image(
|
139 |
+
annotated_image,
|
140 |
+
caption="Your Detected Masterpiece ๐จ",
|
141 |
+
use_column_width=True
|
142 |
+
)
|
143 |
+
elif input_type == "Video" and annotated_video is not None:
|
144 |
+
st.video(
|
145 |
+
annotated_video,
|
146 |
+
format="video/mp4",
|
147 |
+
start_time=0
|
148 |
+
)
|
149 |
+
# Clean up temporary video file
|
150 |
+
if os.path.exists(annotated_video):
|
151 |
+
os.remove(annotated_video)
|
152 |
+
else:
|
153 |
+
st.warning("Nothing to show yet! Hit the button! โก")
|
154 |
+
else:
|
155 |
+
st.info("Awaiting your command, captain! ๐ Upload and detect to see results!")
|
156 |
+
|
157 |
+
# Footer with sass
|
158 |
+
st.markdown("---")
|
159 |
+
st.markdown("Built with ๐ by xAI's minions | Objects beware, YOLOv12 is here! ๐")
|
160 |
+
|
161 |
+
if __name__ == '__main__':
|
162 |
+
main()
|