Harika12323's picture
init
7fd9342 verified
raw
history blame
2.24 kB
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']})")