app.py

import gradio as gr
import torch
import clip
from PIL import Image
import numpy as np
import os
import cv2
import gc # Garbage collector
import logging
import random # For annotator colors

# --- YOLOv8 Imports ---
from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator # For drawing YOLO results

# --- Setup Logging ---
logging.getLogger("ultralytics").setLevel(logging.WARNING) # Reduce YOLO logging noise
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- Constants ---
# Damage segmentation classes (Order MUST match the training of 'model_best.pt')
DAMAGE_CLASSES = ['Cracked', 'Scratch', 'Flaking', 'Broken part', 'Corrosion', 'Dent', 'Paint chip', 'Missing part']
NUM_DAMAGE_CLASSES = len(DAMAGE_CLASSES)

# Part segmentation classes (Order MUST match the training of 'partdetection_yolobest.pt')
CAR_PART_CLASSES = [
    "Quarter-panel", "Front-wheel", "Back-window", "Trunk", "Front-door",
    "Rocker-panel", "Grille", "Windshield", "Front-window", "Back-door",
    "Headlight", "Back-wheel", "Back-windshield", "Hood", "Fender",
    "Tail-light", "License-plate", "Front-bumper", "Back-bumper", "Mirror",
    "Roof"
]
NUM_CAR_PART_CLASSES = len(CAR_PART_CLASSES)


# Paths within the Hugging Face Space repository
CLIP_TEXT_FEATURES_PATH = "./clip_text_features.pt"
DAMAGE_MODEL_WEIGHTS_PATH = "./best.pt" # <--- Your YOLOv8 damage model weights
PART_MODEL_WEIGHTS_PATH = "./partdetection_yolobest.pt" # <--- Your YOLOv8 part model weights

# Prediction Thresholds
DAMAGE_PRED_THRESHOLD = 0.4 # Threshold for showing damage masks
PART_PRED_THRESHOLD = 0.3 # Threshold for showing part masks

# --- Device Setup ---
if torch.cuda.is_available():
    DEVICE = "cuda"
    logger.info("CUDA available, using GPU.")
else:
    DEVICE = "cpu"
    logger.info("CUDA not available, using CPU.")

# --- MODEL LOADING (Load models globally ONCE on startup) ---
print("Loading models...")
clip_model = None
clip_preprocess = None
clip_text_features = None
damage_model = None
part_model = None

# --- Load CLIP Model (Model 1) ---
try:
    logger.info("Loading CLIP model...")
    clip_model, clip_preprocess = clip.load("ViT-B/16", device=DEVICE)
    clip_model.eval()
    logger.info("CLIP model loaded.")

    logger.info(f"Loading CLIP text features from {CLIP_TEXT_FEATURES_PATH}...")
    if not os.path.exists(CLIP_TEXT_FEATURES_PATH):
        raise FileNotFoundError(f"CLIP text features not found: {CLIP_TEXT_FEATURES_PATH}.")
    clip_text_features = torch.load(CLIP_TEXT_FEATURES_PATH, map_location=DEVICE)
    logger.info("CLIP text features loaded.")
except Exception as e:
    logger.error(f"Error loading CLIP model or features: {e}", exc_info=True)
    # Allow app to continue, functions will check for None

# --- Load Damage Segmentation Model (Model 2 - YOLOv8) ---
try:
    logger.info(f"Loading Damage Segmentation (YOLOv8) model from {DAMAGE_MODEL_WEIGHTS_PATH}...")
    if not os.path.exists(DAMAGE_MODEL_WEIGHTS_PATH):
        raise FileNotFoundError(f"Damage model weights not found: {DAMAGE_MODEL_WEIGHTS_PATH}.")
    damage_model = YOLO(DAMAGE_MODEL_WEIGHTS_PATH)
    damage_model.to(DEVICE) # Ensure model is on the correct device
    # Verify class names match
    loaded_damage_names = list(damage_model.names.values())
    if loaded_damage_names != DAMAGE_CLASSES:
         logger.warning(f"Mismatch between defined DAMAGE_CLASSES and names in {DAMAGE_MODEL_WEIGHTS_PATH}")
         logger.warning(f"  Model names: {loaded_damage_names}")
         DAMAGE_CLASSES = loaded_damage_names # Use names from model file
         logger.warning(f"  Updated DAMAGE_CLASSES to: {DAMAGE_CLASSES}")
    logger.info("Damage Segmentation (YOLOv8) model loaded.")
except Exception as e:
    logger.error(f"Error loading Damage Segmentation (YOLOv8) model: {e}", exc_info=True)
    damage_model = None

# --- Load Part Segmentation Model (Model 3 - YOLOv8) ---
try:
    logger.info(f"Loading Part Segmentation (YOLOv8) model from {PART_MODEL_WEIGHTS_PATH}...")
    if not os.path.exists(PART_MODEL_WEIGHTS_PATH):
        raise FileNotFoundError(f"Part model weights not found: {PART_MODEL_WEIGHTS_PATH}.")
    part_model = YOLO(PART_MODEL_WEIGHTS_PATH)
    part_model.to(DEVICE) # Ensure model is on the correct device
    # Verify class names match
    loaded_part_names = list(part_model.names.values())
    if loaded_part_names != CAR_PART_CLASSES:
         logger.warning(f"Mismatch between defined CAR_PART_CLASSES and names in {PART_MODEL_WEIGHTS_PATH}")
         logger.warning(f"  Model names: {loaded_part_names}")
         CAR_PART_CLASSES = loaded_part_names # Use names from model file
         logger.warning(f"  Updated CAR_PART_CLASSES to: {CAR_PART_CLASSES}")
    logger.info("Part Segmentation (YOLOv8) model loaded.")
except Exception as e:
    logger.error(f"Error loading Part Segmentation (YOLOv8) model: {e}", exc_info=True)
    part_model = None

print("Model loading complete.")

# --- Prediction Functions ---

def classify_image_clip(image_pil):
    """Classifies image using CLIP. Returns label and probabilities."""
    if clip_model is None or clip_text_features is None:
         return "Error: CLIP Model Not Loaded", {"Error": 1.0}
    try:
        # Ensure image is RGB PIL
        if image_pil.mode != "RGB":
             image_pil = image_pil.convert("RGB")
        image_input = clip_preprocess(image_pil).unsqueeze(0).to(DEVICE)
        with torch.no_grad():
            image_features = clip_model.encode_image(image_input)
            image_features /= image_features.norm(dim=-1, keepdim=True)
            logit_scale = clip_model.logit_scale.exp()
            similarity = (image_features @ clip_text_features.T) * logit_scale
            probs = similarity.softmax(dim=-1).squeeze().cpu()
        predicted_label = "Car" if probs[0] > probs[1] else "Not Car"
        prob_dict = {"Car": f"{probs[0]:.3f}", "Not Car": f"{probs[1]:.3f}"}
        return predicted_label, prob_dict
    except Exception as e:
        logger.error(f"Error during CLIP prediction: {e}", exc_info=True)
        return "Error during CLIP processing", {"Error": 1.0}

# --- Combined Processing and Overlap Logic ---

def process_car_image(image_np_bgr):
    """
    Runs damage and part segmentation (both YOLOv8), calculates overlap, and returns results.
    Returns:
        - combined_image_rgb: Image with both part and damage masks drawn.
        - assignment_text: String describing damage-part assignments.
    """
    if damage_model is None:
        logger.error("Damage YOLOv8 model not available.")
        return cv2.cvtColor(image_np_bgr, cv2.COLOR_BGR2RGB), "Error: Damage model not loaded."
    if part_model is None:
        logger.error("Part YOLOv8 model not available.")
        return cv2.cvtColor(image_np_bgr, cv2.COLOR_BGR2RGB), "Error: Part model not loaded."

    final_assignments = []
    # Use original BGR image for drawing, convert to RGB only for final display
    annotated_image_bgr = image_np_bgr.copy()
    img_h, img_w = image_np_bgr.shape[:2]

    try:
        # --- 1. Predict Damages (YOLOv8) ---
        logger.info("Running Damage Segmentation (YOLOv8)...")
        # Use conf threshold here
        damage_results = damage_model.predict(image_np_bgr, verbose=False, device=DEVICE, conf=DAMAGE_PRED_THRESHOLD)
        damage_result = damage_results[0] # Result for the first image
        logger.info(f"Found {len(damage_result.boxes)} potential damages.")

        damage_masks_np = damage_result.masks.data.cpu().numpy().astype(bool) if damage_result.masks is not None else np.array([])
        damage_classes_ids = damage_result.boxes.cls.cpu().numpy().astype(int) if damage_result.boxes is not None else np.array([])
        damage_boxes = damage_result.boxes.xyxy.cpu().numpy() if damage_result.boxes is not None else np.array([]) # For drawing

        # --- 2. Predict Parts (YOLOv8) ---
        logger.info("Running Part Segmentation (YOLOv8)...")
        part_results = part_model.predict(image_np_bgr, verbose=False, device=DEVICE, conf=PART_PRED_THRESHOLD)
        part_result = part_results[0] # Result for the first image
        logger.info(f"Found {len(part_result.boxes)} potential parts.")

        part_masks_np = part_result.masks.data.cpu().numpy().astype(bool) if part_result.masks is not None else np.array([]) # [N_part, H, W]
        part_classes_ids = part_result.boxes.cls.cpu().numpy().astype(int) if part_result.boxes is not None else np.array([])
        part_boxes = part_result.boxes.xyxy.cpu().numpy() if part_result.boxes is not None else np.array([]) # For drawing

        # --- 3. Resize Masks if Necessary ---
        # YOLO segmentation masks might be smaller than original image, resize them
        def resize_masks(masks_np, target_h, target_w):
            if masks_np.shape[0] == 0 or (masks_np.shape[1] == target_h and masks_np.shape[2] == target_w):
                return masks_np # Return if empty or already correct size
            logger.info(f"Resizing {masks_np.shape[0]} masks from {masks_np.shape[1:]} to {(target_h, target_w)}")
            resized_masks = []
            for mask in masks_np:
                mask_resized = cv2.resize(mask.astype(np.uint8), (target_w, target_h), interpolation=cv2.INTER_NEAREST)
                resized_masks.append(mask_resized.astype(bool))
            return np.array(resized_masks)

        damage_masks_np = resize_masks(damage_masks_np, img_h, img_w)
        part_masks_np = resize_masks(part_masks_np, img_h, img_w)


        # --- 4. Calculate Overlap ---
        logger.info("Calculating overlap...")
        if damage_masks_np.shape[0] > 0 and part_masks_np.shape[0] > 0:
            overlap_threshold = 0.4 # Minimum overlap ratio

            for i in range(len(damage_masks_np)): # Iterate through each detected damage
                damage_mask = damage_masks_np[i]
                damage_class_id = damage_classes_ids[i]
                try:
                    damage_name = DAMAGE_CLASSES[damage_class_id]
                except IndexError: continue # Skip if invalid class ID

                damage_area = np.sum(damage_mask)
                if damage_area < 10: continue # Skip tiny damage masks

                max_overlap = 0
                assigned_part_name = "Unknown / Outside Parts"

                for j in range(len(part_masks_np)): # Iterate through each detected part
                    part_mask = part_masks_np[j]
                    part_class_id = part_classes_ids[j]
                    try:
                        part_name = CAR_PART_CLASSES[part_class_id]
                    except IndexError: continue # Skip if invalid class ID

                    intersection = np.logical_and(damage_mask, part_mask)
                    intersection_area = np.sum(intersection)
                    overlap_ratio = intersection_area / damage_area if damage_area > 0 else 0

                    if overlap_ratio > max_overlap:
                        max_overlap = overlap_ratio
                        if max_overlap >= overlap_threshold:
                            assigned_part_name = part_name

                assignment_desc = f"{damage_name} in {assigned_part_name}"
                if assigned_part_name == "Unknown / Outside Parts":
                     assignment_desc += f" (Overlap < {overlap_threshold*100:.0f}%)"
                final_assignments.append(assignment_desc)
                logger.info(f"Overlap result: {assignment_desc}")

        elif damage_masks_np.shape[0] > 0: final_assignments.append(f"{len(damage_masks_np)} damages found, but no parts detected/matched.")
        elif part_masks_np.shape[0] > 0: final_assignments.append(f"No damages detected (above threshold).")
        else: final_assignments.append("No damages or parts detected.")


        # --- 5. Visualization using YOLO Annotator ---
        logger.info("Visualizing results...")
        # Create annotator ONCE on the BGR image copy
        annotator = Annotator(annotated_image_bgr, line_width=2, example=part_model.names) # Use part model names

        # Draw PART masks first (Green boxes, light semi-transparent masks)
        if part_result.masks is not None:
            colors_part = [(0, random.randint(100, 200), 0) for _ in part_classes_ids] # Shades of green
            annotator.masks(part_result.masks.data, colors=colors_part, alpha=0.3)
            for box, cls_id in zip(part_boxes, part_classes_ids):
                 try:
                     label = f"{CAR_PART_CLASSES[cls_id]}"
                     annotator.box_label(box, label=label, color=(0, 200, 0)) # Darker green box/text
                 except IndexError: continue

        # Draw DAMAGE masks second (Red boxes, light semi-transparent masks)
        if damage_result.masks is not None:
            colors_dmg = [(random.randint(100, 200), 0, 0) for _ in damage_classes_ids] # Shades of red
            annotator.masks(damage_result.masks.data, colors=colors_dmg, alpha=0.4) # Slightly more opaque
            for box, cls_id in zip(damage_boxes, damage_classes_ids):
                try:
                    label = f"{DAMAGE_CLASSES[cls_id]}"
                    annotator.box_label(box, label=label, color=(200, 0, 0)) # Darker red box/text
                except IndexError: continue

        # Get the final annotated image (still BGR)
        annotated_image_bgr = annotator.result()

    except Exception as e:
        logger.error(f"Error during combined processing: {e}", exc_info=True)
        final_assignments.append("Error during processing.")
        # Return original image in case of error (but annotated_image_bgr might exist)

    # --- Prepare output ---
    assignment_text = "\n".join(final_assignments) if final_assignments else "No specific damage assignments."
    # Convert final annotated image to RGB for Gradio display
    final_output_image_rgb = cv2.cvtColor(annotated_image_bgr, cv2.COLOR_BGR2RGB)

    return final_output_image_rgb, assignment_text


# --- Main Gradio Function ---
def predict_pipeline(image_np_input):
    """
    Main pipeline: Classify -> Segment -> Assign -> Visualize
    """
    if image_np_input is None:
        return "Please upload an image.", {}, None, "N/A"

    logger.info("Received image for processing...")
    final_output_image = None
    assignment_text = "Processing..."
    classification_result = "Error"
    probabilities = {}

    # --- Stage 1: CLIP Classification ---
    try:
        image_pil = Image.fromarray(cv2.cvtColor(image_np_input, cv2.COLOR_BGR2RGB))
        classification_result, probabilities = classify_image_clip(image_pil)
        logger.info(f"CLIP Result: {classification_result}, Probs: {probabilities}")
    except Exception as e:
        logger.error(f"Error in CLIP stage: {e}", exc_info=True)
        assignment_text = "Error during classification."
        # Show original image in case of classification error
        final_output_image = cv2.cvtColor(image_np_input, cv2.COLOR_BGR2RGB)

    # --- Stage 2 & 3: Segmentation and Assignment (if 'Car') ---
    if classification_result == "Car":
        logger.info("Image classified as Car. Running segmentation and assignment...")
        try:
            # Pass the original BGR numpy array
            final_output_image, assignment_text = process_car_image(image_np_input)
        except Exception as e:
             logger.error(f"Error in segmentation/assignment stage: {e}", exc_info=True)
             assignment_text = "Error during segmentation/assignment."
             # Show original image in case of processing error
             final_output_image = cv2.cvtColor(image_np_input, cv2.COLOR_BGR2RGB)

    elif classification_result == "Not Car":
        logger.info("Image classified as Not Car.")
        final_output_image = cv2.cvtColor(image_np_input, cv2.COLOR_BGR2RGB) # Show original
        assignment_text = "Image classified as Not Car."
    # Else: Handle CLIP error case (already logged, show original image)
    elif final_output_image is None: # Ensure image is set if CLIP error occurred
         final_output_image = cv2.cvtColor(image_np_input, cv2.COLOR_BGR2RGB)


    # --- Cleanup ---
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # Return all results
    return classification_result, probabilities, final_output_image, assignment_text


# --- Gradio Interface ---
logger.info("Setting up Gradio interface...")

title = "🚗 Car Damage Analysis Pipeline (YOLOv8)"
description = """
1. **Upload** an image of a vehicle.
2. **Classification:** Determines if the image contains a car (using CLIP).
3. **Segmentation:** If it's a car, detects car parts and damages (using YOLOv8 for both).
4. **Assignment:** Assigns detected damages to the corresponding car part based on mask overlap.
5. **Output:** Shows the image with overlaid masks (Green=Part, Red=Damage) and lists the damage assignments.
"""
examples = [] # Add example image paths if uploaded

# Define Inputs and Outputs
input_image = gr.Image(type="numpy", label="Upload Car Image")
output_classification = gr.Textbox(label="1. Classification Result")
output_probabilities = gr.Label(label="Classification Probabilities")
output_image_display = gr.Image(type="numpy", label="3. Segmentation Visualization")
output_assignment = gr.Textbox(label="2. Damage Assignments", lines=5, interactive=False)


# Launch the interface
iface = gr.Interface(
    fn=predict_pipeline,
    inputs=input_image,
    outputs=[output_classification, output_probabilities, output_image_display, output_assignment],
    title=title,
    description=description,
    examples=examples,
    allow_flagging="never"
)

if __name__ == "__main__":
    logger.info("Launching Gradio app...")
    iface.launch()