Spaces:
Sleeping
Sleeping
import streamlit as st | |
import cv2 | |
import numpy as np | |
import torch | |
import yolov5 | |
from yolov5 import load | |
# Load YOLOv5 model | |
model = load('best.pt') # Replace with your model path | |
def detect_number_plate(frame): | |
# Convert frame to RGB | |
img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
# Perform inference | |
results = model(img) | |
# Parse results | |
detections = results.pandas().xyxy[0] | |
plates = [] | |
for _, row in detections.iterrows(): | |
if row['name'] == 'number_plate': # Adjust based on your model�s class names | |
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']})") | |