awacke1 commited on
Commit
4a339a2
ยท
verified ยท
1 Parent(s): 5cb8e23

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +162 -0
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()