import streamlit as st
from ultralytics import YOLO
from PIL import Image
import json
import cv2
import numpy as np

model = YOLO("./models/best.pt")
classNames = ["license-plate", "vehicle"]

st.title("Number Plate and Vehicle Detection using YOLOv8")

st.write("This is a web app to detect number plates and vehicles in images using YOLOv8.")

image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

if image is not None:
    st.header("Original Image:")
    st.image(image, caption="Uploaded Image", use_column_width=True)
    st.info("Detecting...")
    img = Image.open(image)
    results = model.predict(img, conf=0.5)
    json_results = results[0].tojson()
    encoded_json_results = str(json_results).replace("\n", '').replace(" ", '')
    encoded_json_results = json.loads(encoded_json_results)
    for pred in encoded_json_results:
        x1 = int(pred['box']['x1'])
        y1 = int(pred['box']['y1'])
        x2 = int(pred['box']['x2'])
        y2 = int(pred['box']['y2'])

        img = np.array(img)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
        cv2.putText(img, pred['name'], (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
    st.header("Detected Image")
    st.image(img, caption="Detected Image", use_column_width=True)
    st.header("JSON Results")
    st.write(encoded_json_results)