Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import yolov5 | |
| from yolov5 import YOLOv5 | |
| from models.yolo import Model | |
| # Initialize YOLOv5 model | |
| model = load('best.pt') # Replace with your model path | |
| def detect_number_plate(frame): | |
| results = model(frame) | |
| detections = results.pandas().xyxy[0] | |
| plates = [] | |
| for _, row in detections.iterrows(): | |
| if row['name'] == 'number_plate': # Adjust class name | |
| plates.append({ | |
| 'class': row['name'], | |
| 'confidence': row['confidence'], | |
| 'x_min': row['xmin'], | |
| 'y_min': row['ymin'], | |
| 'x_max': row['xmax'], | |
| 'y_max': row['ymax'] | |
| }) | |
| return plates | |
| def detect_smoke(frame): | |
| gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) | |
| blur = cv2.GaussianBlur(gray, (21, 21), 0) | |
| _, thresh = cv2.threshold(blur, 200, 255, cv2.THRESH_BINARY) | |
| smoke_intensity = np.sum(thresh) / (thresh.shape[0] * thresh.shape[1]) | |
| smoke_detected = smoke_intensity > 0.1 # Adjust this threshold | |
| return smoke_detected, smoke_intensity | |
| def process_frame(frame): | |
| plates = detect_number_plate(frame) | |
| smoke_detected, smoke_intensity = detect_smoke(frame) | |
| return { | |
| 'smoke_detected': smoke_detected, | |
| 'smoke_intensity': smoke_intensity, | |
| 'number_plates': plates | |
| } | |
| # Streamlit app | |
| st.title("Vehicle Number Plate and Smoke Detection") | |
| uploaded_file = st.file_uploader("Choose an image...", type="jpg") | |
| if uploaded_file is not None: | |
| # Convert file to image | |
| in_memory_file = uploaded_file.read() | |
| np_arr = np.frombuffer(in_memory_file, np.uint8) | |
| frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR) | |
| # Process the frame | |
| results = process_frame(frame) | |
| st.subheader("Results") | |
| st.write(f"Smoke Detected: {results['smoke_detected']}") | |
| st.write(f"Smoke Intensity: {results['smoke_intensity']:.2f}") | |
| st.subheader("Number Plates Detected") | |
| for plate in results['number_plates']: | |
| st.write(f"Class: {plate['class']}, Confidence: {plate['confidence']:.2f}") | |
| st.write(f"Bounding Box: ({plate['x_min']}, {plate['y_min']}) to ({plate['x_max']}, {plate['y_max']})") | |