File size: 2,917 Bytes
eaa4e30
 
7d35b07
 
eaa4e30
7d35b07
9ec5726
eaa4e30
7d35b07
9ec5726
eaa4e30
 
 
9ec5726
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8be0123
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7d35b07
 
eaa4e30
7d35b07
 
 
 
 
 
 
 
 
9ec5726
 
 
 
 
7d35b07
 
 
8be0123
 
7d35b07
 
 
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
81
82
83
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()