Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from ultralytics import YOLO | |
| from PIL import Image | |
| import numpy as np | |
| import cv2 | |
| import os | |
| import time | |
| # Function to process video and save annotated results | |
| def process_video(video_in_filepath, video_out_filepath, model): | |
| video_reader = cv2.VideoCapture(video_in_filepath) | |
| # Get video properties | |
| nb_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| frame_h = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| frame_w = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| fps = video_reader.get(cv2.CAP_PROP_FPS) | |
| # Set up video writer | |
| video_writer = cv2.VideoWriter( | |
| video_out_filepath, | |
| cv2.VideoWriter_fourcc(*'mp4v'), | |
| fps, | |
| (frame_w, frame_h) | |
| ) | |
| progress_bar = st.progress(0) | |
| for frame_idx in range(nb_frames): | |
| success, frame = video_reader.read() | |
| if not success: | |
| break | |
| # YOLO inference | |
| results = model(frame) | |
| annotated_frame = results[0].plot() | |
| # Write annotated frame to output video | |
| video_writer.write(annotated_frame) | |
| # Update progress bar | |
| progress_bar.progress((frame_idx + 1) / nb_frames) | |
| video_reader.release() | |
| video_writer.release() | |
| progress_bar.empty() # Remove the progress bar after completion | |
| # Streamlit App Layout | |
| st.title("ITI107 Assignment: Taxi & License Plate Detection") | |
| uploaded_file = st.file_uploader("Image or video", type=["jpg", "jpeg", "png", "mp4"]) | |
| # Create two columns | |
| col1, col2 = st.columns(2) | |
| # Left Column: Upload functionality | |
| with col1: | |
| if uploaded_file: | |
| if uploaded_file.type.startswith("image"): | |
| image = Image.open(uploaded_file) | |
| st.image(image, caption="Uploaded Image", use_container_width=True) | |
| # YOLO inference on image | |
| model = YOLO("best.pt") | |
| results = model.predict(np.array(image)) | |
| # Annotate and display results | |
| result_image = Image.fromarray(results[0].plot()) | |
| with col2: # Show results in the right column | |
| st.image(result_image, caption="Detection Results", use_container_width=True) | |
| elif uploaded_file.type == "video/mp4": | |
| # Save uploaded video locally | |
| input_video_path = os.path.join(os.getcwd(), "uploaded_video.mp4") | |
| with open(input_video_path, "wb") as f: | |
| f.write(uploaded_file.read()) | |
| st.video(input_video_path) # Display uploaded video | |
| # Process video | |
| output_video_path = os.path.join(os.getcwd(), f"processed_video_{int(time.time())}.mp4") | |
| model = YOLO("best.pt") | |
| st.write("Processing video, please wait...") # Status message | |
| process_video(input_video_path, output_video_path, model) | |
| # Display output video | |
| with col2: # Show results in the right column | |
| st.video(output_video_path) # Display processed video | |
| st.write("Video processing complete!") # Completion message |