import gradio as gr
import numpy as np
import torch
from PIL import Image, ImageDraw
from transformers import SegformerForSemanticSegmentation
from torchvision.transforms import Resize, ToTensor, Normalize

# Load models from Hugging Face
part_seg_model = SegformerForSemanticSegmentation.from_pretrained("Mohaddz/huggingCars")
damage_seg_model = SegformerForSemanticSegmentation.from_pretrained("Mohaddz/DamageSegMohaddz/DamageSeg")

# Define your labels
part_labels = ["front-bumper", "fender", "hood", "door", "trunk", "cars-8I1q"]  # Add all your part labels
damage_labels = ["dent", "scratch", "misalignment", "crack", "etc"]  # Add all your damage labels

def preprocess_image(image):
    # Resize and normalize the image
    transform = Resize((512, 512))
    image = transform(Image.fromarray(image))
    image = ToTensor()(image)
    image = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(image)
    return image.unsqueeze(0)  # Add batch dimension

def inference_seg(model, image):
    with torch.no_grad():
        outputs = model(image)
    logits = outputs.logits
    mask = torch.argmax(logits, dim=1).squeeze().numpy()
    return mask

def inference_part_seg(image):
    preprocessed_image = preprocess_image(image)
    return inference_seg(part_seg_model, preprocessed_image)

def inference_damage_seg(image):
    preprocessed_image = preprocess_image(image)
    return inference_seg(damage_seg_model, preprocessed_image)

def combine_masks(part_mask, damage_mask):
    part_damage_pairs = []
    for part_id, part_name in enumerate(part_labels):
        if part_name == "cars-8I1q":
            continue
        for damage_id, damage_name in enumerate(damage_labels):
            if damage_name == "etc":
                continue
            part_binary = (part_mask == part_id)
            damage_binary = (damage_mask == damage_id)
            intersection = np.logical_and(part_binary, damage_binary)
            if np.any(intersection):
                part_damage_pairs.append((part_name, damage_name))
    return part_damage_pairs

def create_one_hot_vector(part_damage_pairs):
    vector = np.zeros(len(part_labels) * len(damage_labels))
    for part, damage in part_damage_pairs:
        if part in part_labels and damage in damage_labels:
            part_index = part_labels.index(part)
            damage_index = damage_labels.index(damage)
            vector_index = part_index * len(damage_labels) + damage_index
            vector[vector_index] = 1
    return vector

def visualize_results(image, part_mask, damage_mask):
    img = Image.fromarray(image)
    draw = ImageDraw.Draw(img)
    
    for i in range(img.width):
        for j in range(img.height):
            part = part_labels[part_mask[j, i]]
            damage = damage_labels[damage_mask[j, i]]
            if part != "cars-8I1q" and damage != "etc":
                draw.point((i, j), fill="red")
    
    return img

def process_image(image):
    # Perform inference
    part_mask = inference_part_seg(image)
    damage_mask = inference_damage_seg(image)
    
    # Combine masks
    part_damage_pairs = combine_masks(part_mask, damage_mask)
    
    # Create one-hot encoded vector
    one_hot_vector = create_one_hot_vector(part_damage_pairs)
    
    # Visualize results
    result_image = visualize_results(image, part_mask, damage_mask)
    
    return result_image, part_damage_pairs, one_hot_vector.tolist()

def gradio_interface(input_image):
    result_image, part_damage_pairs, one_hot_vector = process_image(input_image)
    
    # Convert part_damage_pairs to a string for display
    damage_description = "\n".join([f"{part} : {damage}" for part, damage in part_damage_pairs])
    
    return result_image, damage_description, str(one_hot_vector)

iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Image(type="numpy"),
    outputs=[
        gr.Image(type="pil", label="Detected Damage"),
        gr.Textbox(label="Damage Description"),
        gr.Textbox(label="One-hot Encoded Vector")
    ],
    title="Car Damage Assessment",
    description="Upload an image of a damaged car to get an assessment of the damage."
)

iface.launch()