import streamlit as st
import numpy as np
import json
import os
import cv2
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier

# Load COCO JSON dataset
def load_coco_dataset(json_file, image_dir):
    with open(json_file, 'r') as f:
        data = json.load(f)
    
    images, labels = [], []
    label_dict = {}
    
    for annotation in data['annotations']:
        image_id = annotation['image_id']
        category_id = annotation['category_id']
        
        image_info = next((img for img in data['images'] if img['id'] == image_id), None)
        if image_info is None:
            continue
        
        image_path = os.path.join(image_dir, image_info['file_name'])
        if not os.path.exists(image_path):
            continue
        
        image = cv2.imread(image_path)
        image = cv2.resize(image, (64, 64))
        image = image / 255.0  # Normalize to [0,1]
        
        images.append(image)
        labels.append(category_id)
        
        if category_id not in label_dict:
            label_dict[category_id] = next(cat['name'] for cat in data['categories'] if cat['id'] == category_id)
    
    return np.array(images), np.array(labels), label_dict

# Load dataset paths
train_json_file = "Dataset/Train/_annotations.coco.json"
train_image_directory = "Dataset/Train/Images"

test_json_file = "Dataset/Test/_annotations.coco.json"
test_image_directory = "Dataset/Test/Images"

# Load datasets
X_train, y_train, _ = load_coco_dataset(train_json_file, train_image_directory)
X_test, y_test, label_dict = load_coco_dataset(test_json_file, test_image_directory)

# Flatten images for KNN
X_train = X_train.reshape(X_train.shape[0], -1)
X_test = X_test.reshape(X_test.shape[0], -1)

# Scale the data
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

# Train KNN model
model = KNeighborsClassifier(n_neighbors=5)
model.fit(X_train, y_train)

# Streamlit UI
st.set_page_config(page_title="Road Object Classifier", layout="wide")
st.title("🚗🚶🐾 KNN Road Object Classifier App")

# Tabs for UI
tabs = st.tabs(["Classifier", "How It Works"])

with tabs[0]:
    st.header("Upload an Image for Classification")
    uploaded_file = st.file_uploader("Upload an image (Car, People, or Animal)", type=["jpg", "png", "jpeg"])
    if uploaded_file:
        image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), cv2.IMREAD_COLOR)
        image = cv2.resize(image, (64, 64))
        image_display = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = image / 255.0
        image = image.reshape(1, -1)
        image = scaler.transform(image)
        prediction = model.predict(image)[0]
        predicted_label = label_dict.get(prediction, "Unknown")
        
        col1, col2 = st.columns(2)
        with col1:
            st.image(image_display, caption="Uploaded Image", use_column_width=True)
        with col2:
            st.markdown(f"## **Predicted Object:** {predicted_label}")

with tabs[1]:
    st.header("How the App Works")
    st.write("This app uses a K-Nearest Neighbors (KNN) model to classify objects on roads. The model is trained on a dataset labeled using COCO annotations and processes uploaded images to predict whether they contain a car, a person, or an animal.")