import os
# Force Keras 2 logic to prevent recursion/quantization errors from Kaggle .h5 files
os.environ["TF_USE_LEGACY_KERAS"] = "1" 

import gradio as gr
import tensorflow as tf
import tf_keras as keras 
import numpy as np
import cv2
from PIL import Image
from huggingface_hub import hf_hub_download

# --- CONFIGURATION ---
REPO_ID = "mediaportal/Roadsegmentation" 
MODEL_FILENAME = "trained_model_33_cpu.h5"

# BDD100K Color Dictionary from your notebook
COLOR_DICT = {
    0: (128, 128, 128),  # road - gray
    1: (230, 230, 50),   # sidewalk - yellow
    8: (50, 150, 50),    # vegetation - green
    10: (128, 180, 255), # sky - blue
    11: (255, 0, 0),     # person - red
    13: (0, 0, 255),     # car - blue
    19: (0, 0, 0)        # unknown - black
}

hf_token = os.getenv("HF_TOKEN")
model = None

def load_model():
    global model
    try:
        path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME, token=hf_token)
        # compile=False is used because the notebook uses SparseCategoricalCrossentropy
        model = keras.models.load_model(path, compile=False)
        return "✅ Road Segmentation Model Loaded"
    except Exception as e:
        return f"❌ Error: {str(e)}"

def segment_road(img):
    if model is None:
        return None, None
    
    # 1. Store original size for scaling back
    h_orig, w_orig = img.shape[:2]
    
    # 2. Preprocessing (Notebook uses 192 height, 256 width)
    img_resized = cv2.resize(img, (256, 192)) 
    img_array = img_resized.astype('float32') / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    
    # 3. Predict (Returns 20 channels for 20 classes)
    prediction = model.predict(img_array)[0]
    
    # Get the class index with the highest probability for each pixel
    mask = np.argmax(prediction, axis=-1).astype(np.uint8)
    
    # 4. Create Outputs
    # A. Full Semantic Map (Colorizing all classes)
    full_mask_color = np.zeros((192, 256, 3), dtype=np.uint8)
    for class_idx, color in COLOR_DICT.items():
        full_mask_color[mask == class_idx] = color
    
    # B. Road Highlight Overlay (Class 0 is Road)
    road_mask = (mask == 0).astype(np.uint8) * 255
    road_mask_resized = cv2.resize(road_mask, (w_orig, h_orig), interpolation=cv2.INTER_NEAREST)
    
    overlay = img.copy()
    overlay[road_mask_resized > 0] = [0, 255, 0] # Highlight road in green
    
    # Blend: 70% original image, 30% green highlight
    highlighted_road = cv2.addWeighted(img, 0.7, overlay, 0.3, 0)
    
    # Resize full mask back to original aspect ratio for display
    full_mask_resized = cv2.resize(full_mask_color, (w_orig, h_orig), interpolation=cv2.INTER_NEAREST)
    
    return highlighted_road, full_mask_resized

# --- GRADIO INTERFACE ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚗 ADAS Road & Scene Segmentation")
    gr.Markdown("Upload a dashboard image to identify the drivable road surface and other objects.")
    
    status = gr.Markdown("⏳ Initializing system...")
    
    with gr.Row():
        input_img = gr.Image(label="Input Dashboard Image", type="numpy")
        output_overlay = gr.Image(label="Drivable Road (Green Highlight)")
        output_full = gr.Image(label="Full Semantic Map")
    
    btn = gr.Button("Analyze Scene", variant="primary")
    
    demo.load(load_model, outputs=status)
    btn.click(fn=segment_road, inputs=input_img, outputs=[output_overlay, output_full])

if __name__ == "__main__":
    demo.queue().launch()