Spaces:
Runtime error
Runtime error
File size: 3,666 Bytes
7e5c89f 45d532d 7e5c89f 45d532d 7e5c89f 45d532d 7e5c89f 45d532d 7e5c89f 45d532d |
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 |
import streamlit as st
from ultralytics import YOLO
from supervision.draw.color import ColorPalette
from supervision.tracker.byte_tracker.core import ByteTrack
from supervision.detection.annotate import BoxAnnotator
from supervision.detection.line_counter import LineZone, LineZoneAnnotator
from supervision.geometry.core import Point
from streamlit_webrtc import webrtc_streamer, RTCConfiguration
from video_processor import VideoProcessor
from utils import process_image, process_video_realtime
import tempfile
import numpy as np
import cv2
RTC_CONFIGURATION = RTCConfiguration({"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]})
def main():
st.title('Vehicle Detection and Counting')
model = YOLO("yolov8n.pt")
model.fuse()
label_map = model.model.names
byte_tracker = ByteTrack()
box_annotator = BoxAnnotator(color=ColorPalette.default(), thickness=4, text_thickness=4, text_scale=2)
line_start = Point(50, 1500)
line_end = Point(3840 - 50, 1500)
line_counter = LineZone(start=line_start, end=line_end)
line_annotator = LineZoneAnnotator(thickness=4, text_thickness=4, text_scale=2)
source_option = st.selectbox("Select Source Data", ("Webcam", "Image file", "Video file"))
if source_option == "Webcam":
ctx = webrtc_streamer(
key="example",
video_processor_factory=lambda: VideoProcessor(model, label_map, byte_tracker, box_annotator),
rtc_configuration=RTC_CONFIGURATION,
media_stream_constraints={"video": True, "audio": False}
)
if st.button('Take Screenshot and Process'):
if ctx.video_processor and ctx.video_processor.frame is not None:
image = ctx.video_processor.frame
processed_image, object_counts = process_image(image, model, label_map, byte_tracker, box_annotator)
st.image(processed_image, caption="Processed Image", use_column_width=True)
st.write("Results:", object_counts)
else:
st.warning("No frame available yet.")
elif source_option == "Image file":
uploaded_image = st.file_uploader("Upload an image file", type=["jpg", "jpeg", "png"])
if uploaded_image is not None:
image = cv2.imdecode(np.frombuffer(uploaded_image.read(), np.uint8), cv2.IMREAD_COLOR)
st.image(image, caption="Uploaded Image", use_column_width=True)
if st.button('Process Image'):
processed_image, object_counts = process_image(image, model, label_map, byte_tracker, box_annotator)
st.image(processed_image, caption="Processed Image", use_column_width=True)
st.write("Results:", object_counts)
elif source_option == "Video file":
uploaded_video = st.file_uploader("Upload a video file", type=["mp4"])
demo_video = "vehicle-counting.mp4"
if uploaded_video is not None:
temp_video_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
with open(temp_video_path, 'wb') as f:
f.write(uploaded_video.read())
st.video(temp_video_path)
if st.button('Process Uploaded Video'):
process_video_realtime(temp_video_path, model, label_map, byte_tracker, box_annotator, line_counter, line_annotator)
else:
st.header("Demo Video Preview")
st.video(demo_video)
if st.button('Process Demo Video'):
process_video_realtime(demo_video, model, label_map, byte_tracker, box_annotator, line_counter, line_annotator)
if __name__ == "__main__":
main() |