Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py (Complete Final Version)
+# app.py (Complete Final Version - Fixed resize_masks SyntaxError)
 
 import gradio as gr
 import torch
@@ -40,7 +40,7 @@ NUM_CAR_PART_CLASSES = len(CAR_PART_CLASSES)
 
 # Paths within the Hugging Face Space repository (relative to app.py)
 CLIP_TEXT_FEATURES_PATH = "./clip_text_features.pt"
-DAMAGE_MODEL_WEIGHTS_PATH = "./best.pt" # Your YOLOv8 damage model weights
+DAMAGE_MODEL_WEIGHTS_PATH = "./model_best.pt" # Your YOLOv8 damage model weights
 PART_MODEL_WEIGHTS_PATH = "./partdetection_yolobest.pt" # Your YOLOv8 part model weights
 
 # Default Prediction Thresholds (can be overridden by sliders)
@@ -167,7 +167,6 @@ def classify_image_clip(image_pil):
         traceback.print_exc()
         return "Error during CLIP processing", {"Error": 1.0}
 
-# --- CORRECTED process_car_image Function ---
 def process_car_image(image_np_bgr, damage_threshold, part_threshold):
     """
     Runs damage and part segmentation (YOLOv8), calculates overlap, visualizes.
@@ -186,7 +185,6 @@ def process_car_image(image_np_bgr, damage_threshold, part_threshold):
 
     try:
         # --- Create the image tensor ONCE for the annotator ---
-        # Needs to be HWC format on the correct device
         try:
              im_tensor_gpu_for_annotator = torch.from_numpy(image_np_bgr).to(DEVICE) # Keep HWC
              if not isinstance(im_tensor_gpu_for_annotator, torch.Tensor) or im_tensor_gpu_for_annotator.ndim != 3:
@@ -217,9 +215,38 @@ def process_car_image(image_np_bgr, damage_threshold, part_threshold):
         yolo_end_time = time.time()
         logger.info(f" YOLO predictions took {yolo_end_time - yolo_start_time:.2f}s")
 
-        # --- 3. Resize Masks ---
+        # --- 3. Resize Masks (Corrected Function Definition) ---
         def resize_masks(masks_tensor, target_h, target_w):
-             masks_np_bool = masks_tensor.cpu().numpy().astype(bool); if masks_np_bool.shape[0] == 0 or (masks_np_bool.shape[1] == target_h and masks_np_bool.shape[2] == target_w): return masks_np_bool; resized_masks_list = []; for i in range(masks_np_bool.shape[0]): mask = masks_np_bool[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)
+            """Resizes masks tensor to target H, W using CPU numpy and OpenCV."""
+            masks_np_bool = masks_tensor.cpu().numpy().astype(bool) # Move to CPU *before* resizing
+
+            # Handle empty tensor or already correct size
+            if masks_np_bool.shape[0] == 0:
+                return np.array([]) # Return empty numpy array
+            if masks_np_bool.ndim == 3 and masks_np_bool.shape[1] == target_h and masks_np_bool.shape[2] == target_w:
+                return masks_np_bool # Return if already correct size
+
+            # Ensure masks_np_bool is 3D [N, H, W] even if only one mask
+            if masks_np_bool.ndim == 2: # Handle case of single mask output [H, W]
+                 masks_np_bool = np.expand_dims(masks_np_bool, axis=0)
+                 logger.warning("Detected 2D mask input, expanding to 3D for resize loop.")
+
+            # Check dimensions *before* logging resize message
+            if masks_np_bool.ndim != 3:
+                 logger.error(f"Unexpected mask dimension: {masks_np_bool.ndim}. Expected 3D [N, H, W]. Cannot resize.")
+                 return np.array([]) # Return empty if shape is wrong
+
+            # Proceed with resizing if necessary
+            # logger.info(f"Resizing {masks_np_bool.shape[0]} masks from {masks_np_bool.shape[1:]} to {(target_h, target_w)}") # Optional verbose log
+            resized_masks_list = []
+            for i in range(masks_np_bool.shape[0]):
+                mask = masks_np_bool[i] # Get the single [H, W] mask
+                # Resize needs uint8
+                mask_resized = cv2.resize(mask.astype(np.uint8), (target_w, target_h), interpolation=cv2.INTER_NEAREST)
+                resized_masks_list.append(mask_resized.astype(bool)) # Append boolean mask
+            return np.array(resized_masks_list) # Return numpy array [N, target_h, target_w]
+
+        # --- Perform resizing ---
         resize_start_time = time.time()
         damage_masks_np = resize_masks(damage_masks_raw, img_h, img_w)
         part_masks_np = resize_masks(part_masks_raw, img_h, img_w)
@@ -251,7 +278,6 @@ def process_car_image(image_np_bgr, damage_threshold, part_threshold):
         if part_result.masks is not None and im_tensor_gpu_for_annotator is not None:
             try:
                 colors_part = [(0, random.randint(100, 200), 0) for _ in part_classes_ids_cpu]
-                # Pass CORRECT image tensor (HWC on GPU) to annotator.masks
                 annotator.masks(part_masks_raw, colors=colors_part, im_gpu=im_tensor_gpu_for_annotator, alpha=0.3)
                 for box, cls_id in zip(part_boxes_xyxy_cpu, part_classes_ids_cpu):
                      try: label = f"{CAR_PART_CLASSES[cls_id]}"; annotator.box_label(box, label=label, color=(0, 200, 0))
@@ -264,7 +290,6 @@ def process_car_image(image_np_bgr, damage_threshold, part_threshold):
         if damage_result.masks is not None and im_tensor_gpu_for_annotator is not None:
             try:
                 colors_dmg = [(random.randint(100, 200), 0, 0) for _ in damage_classes_ids_cpu]
-                # Pass CORRECT image tensor (HWC on GPU) to annotator.masks
                 annotator.masks(damage_masks_raw, colors=colors_dmg, im_gpu=im_tensor_gpu_for_annotator, alpha=0.4)
                 for box, cls_id in zip(damage_boxes_xyxy_cpu, damage_classes_ids_cpu):
                     try: label = f"{DAMAGE_CLASSES[cls_id]}"; annotator.box_label(box, label=label, color=(200, 0, 0))
@@ -281,7 +306,6 @@ def process_car_image(image_np_bgr, damage_threshold, part_threshold):
         logger.error(f"Error during combined processing: {e}", exc_info=True)
         traceback.print_exc()
         final_assignments.append("Error during segmentation/processing.")
-        # annotated_image_bgr remains the original copy in case of error
 
     # --- Prepare output ---
     assignment_text = "\n".join(final_assignments) if final_assignments else "No damage assignments generated."