Spaces:
Build error
Build error
File size: 2,779 Bytes
cab5c05 d5ce280 cab5c05 d5ce280 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import streamlit as st
from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
from ultralytics import YOLO
import torch
import cv2
import numpy as np
from PIL import Image
st.set_page_config(page_title="🚗 YOLOv8 Vehicle Detector", layout="wide")
# Load YOLOv8 model once
@st.cache_resource
def load_model():
model = YOLO("runs/detect/train/weights/best.pt")
model.conf = 0.4 # Confidence threshold
return model
model = load_model()
class_names = ['car', 'emv', 'htv']
st.title("🚗 YOLOv8 Vehicle Detector")
st.markdown("Detect vehicles (`car`, `emv`, `htv`) from uploaded images or live webcam.")
# Tabs: One for Upload, One for Webcam
tab1, tab2 = st.tabs(["📸 Image Upload", "🎥 Live Detection"])
# --------------- 📸 Upload Tab ---------------
with tab1:
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
if uploaded_file:
image = Image.open(uploaded_file).convert("RGB")
img_array = np.array(image)
st.image(image, caption="Uploaded Image", use_column_width=True)
st.subheader("🔍 Running Detection...")
results = model(img_array)[0]
for box in results.boxes:
cls_id = int(box.cls.item())
conf = float(box.conf.item())
x1, y1, x2, y2 = map(int, box.xyxy[0])
label = f"{class_names[cls_id]} {conf:.2f}"
cv2.rectangle(img_array, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(img_array, label, (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
st.image(img_array, caption="🖼️ Detection Result", use_column_width=True)
# --------------- 🎥 Real-Time Detection Tab ---------------
with tab2:
st.markdown("Use your webcam for real-time detection. Press `Q` in the live feed to quit.")
class YOLOTransformer(VideoTransformerBase):
def transform(self, frame):
image = frame.to_ndarray(format="bgr24")
results = model(image)[0]
for box in results.boxes:
cls_id = int(box.cls.item())
conf = float(box.conf.item())
if cls_id >= len(class_names): continue
x1, y1, x2, y2 = map(int, box.xyxy[0])
label = f"{class_names[cls_id]} {conf:.2f}"
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 255), 2)
cv2.putText(image, label, (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
return image
webrtc_streamer(
key="realtime",
video_transformer_factory=YOLOTransformer,
media_stream_constraints={"video": True, "audio": False},
async_transform=True,
)
|