Upload app.py

app.py

@@ -0,0 +1,407 @@
+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
+import time # For timing checks
+import traceback # For detailed error printing
+
+# --- 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, format='%(asctime)s - %(levelname)s - %(message)s')
+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
+clip_load_error_msg = None
+damage_load_error_msg = None
+part_load_error_msg = None
+
+
+# --- Load CLIP Model (Model 1) ---
+try:
+    logger.info("Loading CLIP model (ViT-B/16)...")
+    # jit=False might improve stability/compatibility in some environments
+    clip_model, clip_preprocess = clip.load("ViT-B/16", device=DEVICE, jit=False)
+    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}.")
+    # Load text features initially to the designated device
+    clip_text_features = torch.load(CLIP_TEXT_FEATURES_PATH, map_location=DEVICE)
+    logger.info(f"CLIP text features loaded (dtype: {clip_text_features.dtype}).")
+
+except Exception as e:
+    clip_load_error_msg = f"CLIP load error: {e}"
+    logger.error(clip_load_error_msg, exc_info=True)
+    clip_model = None # Set to None if loading fails
+
+# --- 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)
+    # Verify class names match
+    loaded_damage_names = list(damage_model.names.values())
+    if loaded_damage_names != DAMAGE_CLASSES:
+         logger.warning(f"Mismatch: Defined DAMAGE_CLASSES vs names in {DAMAGE_MODEL_WEIGHTS_PATH}")
+         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:
+    damage_load_error_msg = f"Damage YOLO load error: {e}"
+    logger.error(damage_load_error_msg, 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)
+    # Verify class names match
+    loaded_part_names = list(part_model.names.values())
+    if loaded_part_names != CAR_PART_CLASSES:
+         logger.warning(f"Mismatch: Defined CAR_PART_CLASSES vs names in {PART_MODEL_WEIGHTS_PATH}")
+         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:
+    part_load_error_msg = f"Part YOLO load error: {e}"
+    logger.error(part_load_error_msg, exc_info=True)
+    part_model = None
+
+print("Model loading process finished.")
+if clip_load_error_msg: print(f"WARNING: {clip_load_error_msg}")
+if damage_load_error_msg: print(f"WARNING: {damage_load_error_msg}")
+if part_load_error_msg: print(f"WARNING: {part_load_error_msg}")
+
+
+# --- Prediction Functions ---
+
+# --- Updated classify_image_clip (incorporating dtype handling) ---
+def classify_image_clip(image_pil):
+    """Classifies image using CLIP. Returns label and probability dictionary."""
+    if clip_model is None or clip_text_features is None:
+        logger.error(f"CLIP model or text features not loaded. Error: {clip_load_error_msg}")
+        return "Error: CLIP Model Not Loaded", {"Error": 1.0}
+
+    logger.info("Running CLIP classification...")
+    try:
+        # Ensure image is RGB PIL
+        if image_pil.mode != "RGB":
+             image_pil = image_pil.convert("RGB")
+
+        logger.info("  Preprocessing image for CLIP...")
+        image_input = clip_preprocess(image_pil).unsqueeze(0).to(DEVICE)
+        logger.info(f"  Image input tensor created (device: {image_input.device}, dtype: {image_input.dtype}).")
+
+        with torch.no_grad():
+            logger.info("  Encoding image with CLIP...")
+            image_features = clip_model.encode_image(image_input)
+            logger.info(f"  Image features encoded (dtype: {image_features.dtype}).")
+            image_features /= image_features.norm(dim=-1, keepdim=True)
+
+            # --- Ensure Text Features match Image Features dtype ---
+            text_features_matched = clip_text_features
+            if image_features.dtype != clip_text_features.dtype:
+                logger.warning(f"  Dtype mismatch! Image: {image_features.dtype}, Text: {clip_text_features.dtype}. Converting text features...")
+                text_features_matched = clip_text_features.to(image_features.dtype)
+            # -----------------------------------------------------
+
+            logit_scale = clip_model.logit_scale.exp()
+            logger.info("  Calculating similarity...")
+            similarity = (image_features @ text_features_matched.T) * logit_scale
+            probs = similarity.softmax(dim=-1).squeeze().cpu() # Move probabilities to CPU
+            logger.info("  Similarity calculated.")
+
+        # Indices based on your original feature creation: 0=Car, 1=Not Car
+        car_prob = probs[0].item()
+        not_car_prob = probs[1].item()
+
+        predicted_label = "Car" if car_prob > not_car_prob else "Not Car"
+        # Format probabilities for display
+        prob_dict = {"Car": f"{car_prob:.3f}", "Not Car": f"{not_car_prob:.3f}"}
+        logger.info(f"CLIP Result: {predicted_label}, Probs: {prob_dict}")
+
+        return predicted_label, prob_dict # Return dictionary
+
+    except Exception as e:
+        logger.error(f"Error during CLIP prediction: {e}", exc_info=True)
+        traceback.print_exc() # Print detailed traceback to logs
+        return "Error during CLIP processing", {"Error": 1.0}
+
+
+# --- Combined Processing and Overlap Logic (process_car_image) ---
+# (Keep the process_car_image function from the previous response, it should be fine)
+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), f"Error: Damage model not loaded ({damage_load_error_msg})"
+    if part_model is None:
+        logger.error("Part YOLOv8 model not available.")
+        return cv2.cvtColor(image_np_bgr, cv2.COLOR_BGR2RGB), f"Error: Part model not loaded ({part_load_error_msg})"
+
+    final_assignments = []
+    annotated_image_bgr = image_np_bgr.copy()
+    img_h, img_w = image_np_bgr.shape[:2]
+    logger.info("Starting combined image processing...")
+
+    try:
+        # --- 1. Predict Damages (YOLOv8) ---
+        logger.info(f"Running Damage Segmentation (Threshold: {DAMAGE_PRED_THRESHOLD})...")
+        damage_results = damage_model.predict(image_np_bgr, verbose=False, device=DEVICE, conf=DAMAGE_PRED_THRESHOLD)
+        damage_result = damage_results[0]
+        logger.info(f"Found {len(damage_result.boxes)} potential damages.")
+        damage_masks_raw = damage_result.masks.data.cpu() if damage_result.masks is not None else torch.empty((0,0,0))
+        damage_classes_ids = damage_result.boxes.cls.cpu().numpy().astype(int) if damage_result.boxes is not None else np.array([])
+        damage_boxes_xyxy = damage_result.boxes.xyxy.cpu() if damage_result.boxes is not None else torch.empty((0,4))
+
+        # --- 2. Predict Parts (YOLOv8) ---
+        logger.info(f"Running Part Segmentation (Threshold: {PART_PRED_THRESHOLD})...")
+        part_results = part_model.predict(image_np_bgr, verbose=False, device=DEVICE, conf=PART_PRED_THRESHOLD)
+        part_result = part_results[0]
+        logger.info(f"Found {len(part_result.boxes)} potential parts.")
+        part_masks_raw = part_result.masks.data.cpu() if part_result.masks is not None else torch.empty((0,0,0))
+        part_classes_ids = part_result.boxes.cls.cpu().numpy().astype(int) if part_result.boxes is not None else np.array([])
+        part_boxes_xyxy = part_result.boxes.xyxy.cpu() if part_result.boxes is not None else torch.empty((0,4))
+
+
+        # --- 3. Resize Masks if Necessary (Function definition) ---
+        def resize_masks(masks_tensor, target_h, target_w):
+            if masks_tensor.shape[0] == 0: return np.array([]) # Handle empty tensor
+            # Check if resize is needed
+            if masks_tensor.shape[1] == target_h and masks_tensor.shape[2] == target_w:
+                return masks_tensor.numpy().astype(bool) # No resize needed, convert to numpy bool
+
+            logger.info(f"Resizing {masks_tensor.shape[0]} masks from {masks_tensor.shape[1:]} to {(target_h, target_w)}")
+            masks_np = masks_tensor.numpy() # Convert to numpy first
+            resized_masks_list = []
+            for i in range(masks_np.shape[0]):
+                mask = masks_np[i]
+                mask_resized = cv2.resize(mask.astype(np.uint8), (target_w, target_h), interpolation=cv2.INTER_NEAREST)
+                resized_masks_list.append(mask_resized.astype(bool))
+            return np.array(resized_masks_list) # Return numpy array
+
+        # --- Perform resizing ---
+        damage_masks_np = resize_masks(damage_masks_raw, img_h, img_w)
+        part_masks_np = resize_masks(part_masks_raw, 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
+            # (Overlap calculation logic remains the same)
+            for i in range(len(damage_masks_np)):
+                damage_mask = damage_masks_np[i]
+                damage_class_id = damage_classes_ids[i]
+                try: damage_name = DAMAGE_CLASSES[damage_class_id]
+                except IndexError: logger.warning(f"Invalid damage ID {damage_class_id}"); continue
+
+                damage_area = np.sum(damage_mask)
+                if damage_area < 10: continue # Skip tiny masks
+
+                max_overlap = 0
+                assigned_part_name = "Unknown / Outside Parts"
+                for j in range(len(part_masks_np)):
+                    part_mask = part_masks_np[j]
+                    part_class_id = part_classes_ids[j]
+                    try: part_name = CAR_PART_CLASSES[part_class_id]
+                    except IndexError: logger.warning(f"Invalid part ID {part_class_id}"); continue
+
+                    intersection = np.logical_and(damage_mask, part_mask)
+                    overlap_ratio = np.sum(intersection) / 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)
+
+        # Handle cases with no damages or no parts found after thresholding
+        elif damage_masks_np.shape[0] > 0: final_assignments.append(f"{len(damage_masks_np)} damages found, but no parts detected/matched above threshold.")
+        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 above thresholds.")
+        logger.info(f"Assignment results: {final_assignments}")
+
+
+        # --- 5. Visualization using YOLO Annotator ---
+        logger.info("Visualizing results...")
+        annotator = Annotator(annotated_image_bgr, line_width=2, example=part_model.names) # Use names from part model
+
+        # Draw PART masks/boxes (Greenish) - Use original raw masks for drawing coordinates
+        if part_result.masks is not None:
+            colors_part = [(0, random.randint(100, 200), 0) for _ in part_classes_ids]
+            annotator.masks(part_result.masks.data, colors=colors_part, alpha=0.3)
+            for box, cls_id in zip(part_boxes_xyxy, part_classes_ids):
+                 try: label = f"{CAR_PART_CLASSES[cls_id]}"; annotator.box_label(box, label=label, color=(0, 200, 0))
+                 except IndexError: continue
+
+        # Draw DAMAGE masks/boxes (Reddish) - Use original raw masks for drawing coordinates
+        if damage_result.masks is not None:
+            colors_dmg = [(random.randint(100, 200), 0, 0) for _ in damage_classes_ids]
+            annotator.masks(damage_result.masks.data, colors=colors_dmg, alpha=0.4)
+            for box, cls_id in zip(damage_boxes_xyxy, damage_classes_ids):
+                try: label = f"{DAMAGE_CLASSES[cls_id]}"; annotator.box_label(box, label=label, color=(200, 0, 0))
+                except IndexError: continue
+
+        annotated_image_bgr = annotator.result() # Get final BGR image
+
+    except Exception as e:
+        logger.error(f"Error during combined processing: {e}", exc_info=True)
+        traceback.print_exc()
+        final_assignments.append("Error during segmentation/processing.")
+
+    # --- Prepare output ---
+    assignment_text = "\n".join(final_assignments) if final_assignments else "No specific damage assignments."
+    final_output_image_rgb = cv2.cvtColor(annotated_image_bgr, cv2.COLOR_BGR2RGB) # Convert final to RGB
+
+    return final_output_image_rgb, assignment_text
+
+
+# --- Main Gradio Function ---
+# (Keep the predict_pipeline function from the previous response, it calls the updated classify_image_clip and process_car_image)
+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("--- New Request ---")
+    start_time = time.time()
+    # Convert Gradio input (assumed RGB) to BGR for processing functions
+    image_np_bgr = cv2.cvtColor(image_np_input, cv2.COLOR_RGB2BGR)
+    image_pil = Image.fromarray(image_np_input) # PIL for CLIP (already RGB)
+
+    final_output_image = None
+    assignment_text = "Processing..."
+    classification_result = "Error"
+    probabilities = {}
+
+    # --- Stage 1: CLIP Classification ---
+    try:
+        classification_result, probabilities = classify_image_clip(image_pil)
+    except Exception as e:
+        logger.error(f"Error in CLIP stage: {e}", exc_info=True)
+        assignment_text = f"Error during classification: {e}"
+        # Show original image in case of classification error
+        final_output_image = cv2.cvtColor(image_np_bgr, 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:
+            final_output_image, assignment_text = process_car_image(image_np_bgr)
+        except Exception as e:
+             logger.error(f"Error in segmentation/assignment stage: {e}", exc_info=True)
+             assignment_text = f"Error during segmentation/assignment: {e}"
+             # Show original image in case of processing error
+             final_output_image = cv2.cvtColor(image_np_bgr, cv2.COLOR_BGR2RGB)
+
+    elif classification_result == "Not Car":
+        logger.info("Image classified as Not Car.")
+        final_output_image = cv2.cvtColor(image_np_bgr, 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_bgr, cv2.COLOR_BGR2RGB)
+
+
+    # --- Cleanup ---
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    end_time = time.time()
+    logger.info(f"Total processing time: {end_time - start_time:.2f} seconds.")
+    # Return all results
+    return classification_result, probabilities, final_output_image, assignment_text
+
+
+# --- Gradio Interface ---
+# (Keep the Gradio Interface definition from the previous response)
+logger.info("Setting up Gradio interface...")
+title = "🚗 Car Damage Analysis Pipeline (CLIP + YOLOv8 x2)"
+# ... (rest of Gradio interface setup: description, examples, inputs, outputs, iface.launch()) ...
+
+# Define Inputs and Outputs
+input_image = gr.Image(type="numpy", label="Upload Car Image") # Input numpy array (RGB from Gradio)
+output_classification = gr.Textbox(label="1. Classification Result")
+output_probabilities = gr.Label(label="Classification Probabilities") # Label is good for dicts
+output_image_display = gr.Image(type="numpy", label="3. Segmentation Visualization") # Output numpy array (RGB)
+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, # Add description back if needed
+    # examples=examples, # Add examples back if needed
+    allow_flagging="never"
+)
+
+if __name__ == "__main__":
+    logger.info("Launching Gradio app...")
+    iface.launch()