ethanrom/helm

1import cv22import torch3import torch.backends.cudnn as cudnn4from models.experimental import attempt_load5from utils.general import non_max_suppression6from torchvision import models7from torchvision import transforms8from PIL import Image9import time10import streamlit as st11import IPython12import numpy as np13 14yolov5_weight_file = 'model100e.pt'15 16device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')17yolov5_model = attempt_load(yolov5_weight_file, device=device, inplace=True, fuse=True)18cudnn.benchmark = True 19names = yolov5_model.module.names if hasattr(yolov5_model, 'module') else yolov5_model.names20 21conf_set = 0.122frame_size = (800, 480)23 24colors = {25    'helmet': (255, 0, 0),26    'rider': (0, 255, 0),27    'number': (0, 0, 255),28    'no_helmet': (0, 100, 255),29}30 31def detect_objects(frame):32    img = torch.from_numpy(frame)33    img = img.permute(2, 0, 1).float().to(device)34    img /= 255.035    if img.ndimension() == 3:36        img = img.unsqueeze(0)37    with torch.no_grad():38        pred = yolov5_model(img, augment=False)[0]39        pred = non_max_suppression(pred, conf_set, 0.30)40        detections = []41        for det in pred:42            if len(det):43                for d in det:  # d = (x1, y1, x2, y2, conf, cls)44                    x1 = int(d[0].item())45                    y1 = int(d[1].item())46                    x2 = int(d[2].item())47                    y2 = int(d[3].item())48                    conf = round(d[4].item(), 2)49                    c = int(d[5].item())50                    detected_name = names[c]51                    detections.append((x1, y1, x2, y2, conf, detected_name))52            53                    color = colors.get(detected_name, (255, 255, 255))54                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)55                    cv2.putText(frame, detected_name, (x1, y1), cv2.FONT_HERSHEY_DUPLEX, 1, color, 2)56                    57        return detections58 59def display_detections(input_image, output_image, detections):60    for det in detections:61        x1, y1, x2, y2, conf, detected_name = det62        color = colors.get(detected_name, (255, 255, 255))63        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)64        cv2.putText(output_image, f"{detected_name} ({conf:.2f})", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)65    return output_image66    67def app():68    st.title("Helmet Detection App")69    st.write("This app uses YOLOv5 to detect helmets and riders in images and videos.")70    71    # Select input type72    input_type = st.radio("Select input type:", options=["Image", "Video"])73    74    # Upload file or use webcam75    if input_type == "Image":76        uploaded_file = st.file_uploader("Upload image", type=["jpg", "jpeg", "png"])77        if uploaded_file is not None:78            image = Image.open(uploaded_file)79            st.image(image, caption="Uploaded Image", use_column_width=True)80            detections = detect_objects(np.array(image))81            output_image = display_detections(np.array(image), np.array(image), detections)82            st.image(output_image, caption="Output Image", use_column_width=True)83 84            85    elif input_type == "Video":86        st.write("Select an option to get the input video:")87        video_option = st.radio("", options=["Webcam", "Upload video"])88        89        if video_option == "Webcam":90            cap = cv2.VideoCapture(0)91        elif video_option == "Upload video":92            uploaded_file = st.file_uploader("Upload video", type=["mp4"])93            if uploaded_file is not None:94                temp_file = NamedTemporaryFile(delete=False)95                temp_file.write(uploaded_file.read())96                st.write("Video uploaded successfully!")97                cap = cv2.VideoCapture(temp_file.name)98        99        if 'cap' in locals():100            frame_size = (800, 480)101            show_video = st.checkbox("Show video", value=True)102            save_video = st.checkbox("Save video", value=False)103            font = cv2.FONT_HERSHEY_DUPLEX104            105            while True:106                ret, frame = cap.read()107                if ret:108                    frame = cv2.resize(frame, frame_size)109                    detections = detect_objects(frame)110                    display_frame = display_detections(frame, detections)111                    fps = 1 / (time.time() - start_time)112                    start_time = time.time()113                    cv2.putText(display_frame, f'FPS: {fps:.2f}', (10, 30), font, 1, (0, 255, 0), 2, cv2.LINE_AA)114                    if show_video:115                        stframe.image(display_frame, channels="BGR")116                    if save_video:117                        out.write(display_frame)118                    if cv2.waitKey(1) & 0xFF == ord('q'):119                        break120            cap.release()121            if save_video:122                out.release()123 124if __name__ == "__main__":125    app()