import streamlit as st
import cv2
import numpy as np
import torch
import yolov5
from yolov5 import load

# Load YOLOv5 model
model = yolov5.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']})")