import cv2
import numpy as np
import gradio as gr

import torch
import torch.serialization
import ultralytics.nn.tasks as tasks
torch.serialization.add_safe_globals({tasks.DetectionModel})

from ultralyticsplus import YOLO

model = YOLO("keremberke/yolov8n-pothole-segmentation")
model.overrides['conf'] = 0.25



# Constants
SCALE = 0.005  # m² per pixel²
COST_PER_M2 = 1000  # ₹ per m²

def estimate_cost_from_image(image):
    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    h, w = img_bgr.shape[:2]

    # Save temporary image
    temp_path = "temp.jpg"
    cv2.imwrite(temp_path, img_bgr)

    result = model.predict(source=temp_path, verbose=False)[0]
    mask = result.masks.data.cpu().numpy()
    total_area_px = mask.sum()
    area_m2 = total_area_px * SCALE
    estimated_cost = area_m2 * COST_PER_M2

    # Overlay visualization
    overlay = img_bgr.copy()
    combined = np.max(mask, axis=0).astype(bool)
    overlay[combined] = [0, 0, 255]
    final = cv2.addWeighted(img_bgr, 0.7, overlay, 0.3, 0)
    final_rgb = cv2.cvtColor(final, cv2.COLOR_BGR2RGB)

    return final_rgb, f"{area_m2:.2f} m²", f"₹{estimated_cost:.0f}"

# Gradio Interface
demo = gr.Interface(
    fn=estimate_cost_from_image,
    inputs=gr.Image(type="numpy", label="Upload Pothole Image"),
    outputs=[
        gr.Image(type="numpy", label="Overlay"),
        gr.Text(label="Estimated Area"),
        gr.Text(label="Repair Cost"),
    ],
    title="Pothole Repair Estimator",
    description="Upload a pothole image to estimate repair cost using YOLOv8 segmentation model."
)

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