Spaces:
Running
Running
| import streamlit as st | |
| import cv2 | |
| import numpy as np | |
| from ultralytics import YOLO | |
| # Load the YOLO model | |
| model = YOLO('yolov8_Medium.pt') # Ensure the model file is in the root directory of your Space | |
| def run_yolo(image): | |
| # Run the model on the image and get results | |
| results = model(image) | |
| return results | |
| def process_results(results, image): | |
| # Draw bounding boxes and labels on the image | |
| boxes = results[0].boxes # Get boxes from results | |
| for box in boxes: | |
| # Get the box coordinates and label | |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) # Convert to integer coordinates | |
| conf = box.conf[0] # Confidence score | |
| cls = int(box.cls[0]) # Class index | |
| label = model.names[cls] # Get class name from index | |
| # Draw rectangle and label on the image | |
| cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2) # Blue box | |
| cv2.putText(image, f"{label} {conf:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2) | |
| return image | |
| def process_video(uploaded_file): | |
| # Read the video file | |
| video = cv2.VideoCapture(uploaded_file) | |
| frames = [] | |
| while True: | |
| ret, frame = video.read() | |
| if not ret: | |
| break # Break the loop if there are no frames left | |
| # Run YOLO model on the current frame | |
| results = run_yolo(frame) | |
| # Process the results and draw boxes on the current frame | |
| processed_frame = process_results(results, frame) | |
| frames.append(processed_frame) # Save the processed frame | |
| video.release() | |
| # Create a video writer to save the processed frames | |
| height, width, _ = frames[0].shape | |
| out = cv2.VideoWriter('processed_video.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 30, (width, height)) | |
| for frame in frames: | |
| out.write(frame) # Write each processed frame to the video | |
| out.release() | |
| def main(): | |
| st.title("Motorbike Violation Detection") | |
| # Upload file | |
| uploaded_file = st.file_uploader("Choose an image or video...", type=["jpg", "jpeg", "png", "mp4"]) | |
| if uploaded_file is not None: | |
| if uploaded_file.type in ["image/jpeg", "image/png", "image/jpg"]: | |
| # Process the image | |
| image = np.array(cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)) | |
| results = run_yolo(image) | |
| # Process the results and draw boxes on the image | |
| processed_image = process_results(results, image) | |
| # Display the processed image | |
| st.image(processed_image, caption='Detected Image', use_column_width=True) | |
| elif uploaded_file.type == "video/mp4": | |
| # Process the video | |
| process_video(uploaded_file) # Process the video and save the output | |
| st.video('processed_video.mp4') # Display the processed video | |
| if __name__ == "__main__": | |
| main() | |