Update app.py

app.py

@@ -33,7 +33,7 @@ CAR_PART_CLASSES = [
 NUM_CAR_PART_CLASSES = len(CAR_PART_CLASSES)
 
 CLIP_TEXT_FEATURES_PATH = "./clip_text_features.pt"
-DAMAGE_MODEL_WEIGHTS_PATH = "./best.pt"
+DAMAGE_MODEL_WEIGHTS_PATH = "./model_best.pt"
 PART_MODEL_WEIGHTS_PATH = "./partdetection_yolobest.pt"
 DEFAULT_DAMAGE_PRED_THRESHOLD = 0.4
 DEFAULT_PART_PRED_THRESHOLD = 0.3
@@ -107,9 +107,73 @@ except Exception as e:
     part_model = None
 
 print("--- Model loading process finished. ---")
-print(f"CLIP STATUS: {clip_load_error_msg}" if clip_load_error_msg else "CLIP STATUS: Loaded OK.")
-print(f"DAMAGE MODEL STATUS: {damage_load_error_msg}" if damage_load_error_msg else "DAMAGE MODEL STATUS: Loaded OK.")
-print(f"PART MODEL STATUS: {part_load_error_msg}" if part_load_error_msg else "PART MODEL STATUS: Loaded OK.")
+if clip_load_error_msg:
+    print(f"CLIP STATUS: {clip_load_error_msg}")
+else:
+    print("CLIP STATUS: Loaded OK.")
+    
+if damage_load_error_msg:
+    print(f"DAMAGE MODEL STATUS: {damage_load_error_msg}")
+else:
+    print("DAMAGE MODEL STATUS: Loaded OK.")
+    
+if part_load_error_msg:
+    print(f"PART MODEL STATUS: {part_load_error_msg}")
+else:
+    print("PART MODEL STATUS: Loaded OK.")
+
+# --- Add DirectVisualizer class for fallback visualization ---
+class DirectVisualizer:
+    """Fallback visualizer for when Ultralytics Annotator doesn't work"""
+    
+    def __init__(self, image):
+        self.image = image.copy()
+        
+    def draw_masks(self, masks_np, class_ids, class_names, color_type="damage"):
+        """Draw masks directly using OpenCV"""
+        if masks_np.shape[0] == 0:
+            return
+            
+        for i, (mask, class_id) in enumerate(zip(masks_np, class_ids)):
+            if not np.any(mask):  # Skip empty masks
+                continue
+                
+            # Set color based on type
+            if color_type == "damage":
+                color = (0, 0, 255)  # BGR Red
+                alpha = 0.4
+            else:
+                color = (0, 255, 0)  # BGR Green
+                alpha = 0.3
+                
+            # Create color overlay
+            overlay = self.image.copy()
+            overlay[mask] = color
+            
+            # Apply with transparency
+            cv2.addWeighted(overlay, alpha, self.image, 1-alpha, 0, self.image)
+            
+            # Draw contour
+            mask_8bit = mask.astype(np.uint8) * 255
+            contours, _ = cv2.findContours(mask_8bit, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+            cv2.drawContours(self.image, contours, -1, color, 2)
+            
+            # Add text label
+            try:
+                if 0 <= class_id < len(class_names):
+                    label = class_names[class_id]
+                    M = cv2.moments(mask_8bit)
+                    if M["m00"] > 0:
+                        cx = int(M["m10"] / M["m00"])
+                        cy = int(M["m01"] / M["m00"])
+                        cv2.putText(self.image, label, (cx, cy), 
+                                  cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
+            except Exception as e:
+                logger.warning(f"Error adding label: {e}")
+    
+    def result(self):
+        """Return the final image"""
+        return self.image
 
 # --- Prediction Functions ---
 def classify_image_clip(image_pil):
@@ -147,8 +211,17 @@ def process_car_image(image_np_bgr, damage_threshold, part_threshold):
     img_h, img_w = image_np_bgr.shape[:2]
     logger.info("Starting combined YOLO processing...")
     im_tensor_gpu_for_annotator = None
+    ultralytic_viz_success = False
 
     try:
+        # --- Check Ultralytics Version ---
+        try:
+            import pkg_resources
+            ultraytics_version = pkg_resources.get_distribution("ultralytics").version
+            logger.info(f"Ultralytics version: {ultraytics_version}")
+        except:
+            logger.warning("Could not determine Ultralytics version")
+
         # --- Prepare Image Tensor for Annotator ---
         logger.info("Preparing image tensor for annotator...")
         try:
@@ -169,11 +242,27 @@ def process_car_image(image_np_bgr, damage_threshold, part_threshold):
         damage_results = damage_model.predict(image_np_bgr, verbose=False, device=DEVICE, conf=damage_threshold)
         damage_result = damage_results[0]
         logger.info(f"Found {len(damage_result.boxes)} potential damages.")
-        damage_masks_raw = damage_result.masks.data if damage_result.masks is not None else torch.empty((0,0,0), device=DEVICE)
+        
+        # Check mask availability and format
         if damage_result.masks is None:
             logger.warning("No damage masks in result! Check if damage model is segmentation type.")
+            damage_masks_raw = torch.empty((0,0,0), device=DEVICE)
         else:
-            logger.info(f"Damage masks available: shape={damage_masks_raw.shape if damage_masks_raw.numel() > 0 else 'Empty'}")
+            logger.info(f"Damage masks type: {type(damage_result.masks)}")
+            try:
+                if hasattr(damage_result.masks, 'data'):
+                    damage_masks_raw = damage_result.masks.data
+                    logger.info(f"Damage masks via .data: shape={damage_masks_raw.shape}, dtype={damage_masks_raw.dtype}")
+                elif isinstance(damage_result.masks, torch.Tensor):
+                    damage_masks_raw = damage_result.masks
+                    logger.info(f"Damage masks as tensor: shape={damage_masks_raw.shape}, dtype={damage_masks_raw.dtype}")
+                else:
+                    logger.warning(f"Unknown mask type: {type(damage_result.masks)}")
+                    damage_masks_raw = torch.empty((0,0,0), device=DEVICE)
+            except Exception as e_mask:
+                logger.error(f"Error accessing damage masks: {e_mask}", exc_info=True)
+                damage_masks_raw = torch.empty((0,0,0), device=DEVICE)
+        
         damage_classes_ids_cpu = damage_result.boxes.cls.cpu().numpy().astype(int) if damage_result.boxes is not None else np.array([])
         damage_boxes_xyxy_cpu = damage_result.boxes.xyxy.cpu() if damage_result.boxes is not None else torch.empty((0,4))
 
@@ -182,25 +271,53 @@ def process_car_image(image_np_bgr, damage_threshold, part_threshold):
         part_results = part_model.predict(image_np_bgr, verbose=False, device=DEVICE, conf=part_threshold)
         part_result = part_results[0]
         logger.info(f"Found {len(part_result.boxes)} potential parts.")
-        part_masks_raw = part_result.masks.data if part_result.masks is not None else torch.empty((0,0,0), device=DEVICE)
+        
+        # Check mask availability and format
         if part_result.masks is None:
             logger.warning("No part masks in result! Check if part model is segmentation type.")
+            part_masks_raw = torch.empty((0,0,0), device=DEVICE)
         else:
-            logger.info(f"Part masks available: shape={part_masks_raw.shape if part_masks_raw.numel() > 0 else 'Empty'}")
+            logger.info(f"Part masks type: {type(part_result.masks)}")
+            try:
+                if hasattr(part_result.masks, 'data'):
+                    part_masks_raw = part_result.masks.data
+                    logger.info(f"Part masks via .data: shape={part_masks_raw.shape}, dtype={part_masks_raw.dtype}")
+                elif isinstance(part_result.masks, torch.Tensor):
+                    part_masks_raw = part_result.masks
+                    logger.info(f"Part masks as tensor: shape={part_masks_raw.shape}, dtype={part_masks_raw.dtype}")
+                else:
+                    logger.warning(f"Unknown mask type: {type(part_result.masks)}")
+                    part_masks_raw = torch.empty((0,0,0), device=DEVICE)
+            except Exception as e_mask:
+                logger.error(f"Error accessing part masks: {e_mask}", exc_info=True)
+                part_masks_raw = torch.empty((0,0,0), device=DEVICE)
+                
         part_classes_ids_cpu = part_result.boxes.cls.cpu().numpy().astype(int) if part_result.boxes is not None else np.array([])
         part_boxes_xyxy_cpu = part_result.boxes.xyxy.cpu() if part_result.boxes is not None else torch.empty((0,4))
 
         # --- 3. Resize Masks ---
         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)
+            if masks_tensor is None or masks_tensor.numel() == 0 or masks_tensor.shape[0] == 0:
+                logger.warning("Empty masks tensor passed to resize_masks")
+                return np.zeros((0, target_h, target_w), dtype=bool)
+            
+            try:
+                masks_np_bool = masks_tensor.cpu().numpy().astype(bool)
+                logger.info(f"Resizing masks from {masks_np_bool.shape} to ({target_h}, {target_w})")
+                
+                if 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)
+            except Exception as e_resize:
+                logger.error(f"Error resizing masks: {e_resize}", exc_info=True)
+                return np.zeros((0, target_h, target_w), dtype=bool)
             
         damage_masks_np = resize_masks(damage_masks_raw, img_h, img_w)
         part_masks_np = resize_masks(part_masks_raw, img_h, img_w)
@@ -246,64 +363,90 @@ def process_car_image(image_np_bgr, damage_threshold, part_threshold):
             final_assignments.append(f"No damages detected above threshold {damage_threshold}.")
         else:
             final_assignments.append(f"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=CAR_PART_CLASSES)
-
-        # Draw PART masks
-        if part_masks_raw.numel() > 0 and im_tensor_gpu_for_annotator is not None:
-            try:
-                logger.info("Attempting to draw part masks...")
-                colors_part = [(0, random.randint(100, 200), 0) for _ in part_classes_ids_cpu]
-                mask_data_part = part_masks_raw
-                if mask_data_part.device != im_tensor_gpu_for_annotator.device:
-                    mask_data_part = mask_data_part.to(im_tensor_gpu_for_annotator.device)
-                annotator.masks(mask_data_part, colors=colors_part, im_gpu=im_tensor_gpu_for_annotator, alpha=0.3)
-                logger.info("Successfully drew part masks.")
-                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))
-                    except IndexError:
-                        logger.warning(f"Invalid part ID {cls_id} during drawing")
-            except Exception as e_part_vis:
-                logger.error(f"Error drawing part masks/boxes: {e_part_vis}", exc_info=True)
-                traceback.print_exc()
-        elif part_masks_raw.numel() > 0:
-            logger.warning("Part masks exist but image tensor for annotator is None. Skipping part mask drawing.")
+        logger.info(f"Assignment results: {final_assignments}")
 
-        # Draw DAMAGE masks
-        if damage_masks_raw.numel() > 0 and im_tensor_gpu_for_annotator is not None:
+        # --- 5. Try BOTH visualization approaches ---
+        # First attempt: Ultralytics annotator
+        try:
+            logger.info("Trying Ultralytics Annotator for visualization...")
+            annotator = Annotator(annotated_image_bgr.copy(), line_width=2, example=CAR_PART_CLASSES)
+            
+            # Draw part masks with Ultralytics
+            if part_masks_raw.numel() > 0 and im_tensor_gpu_for_annotator is not None:
+                try:
+                    colors_part = [(0, random.randint(100, 200), 0) for _ in part_classes_ids_cpu]
+                    mask_data_part = part_masks_raw
+                    if mask_data_part.device != im_tensor_gpu_for_annotator.device:
+                        mask_data_part = mask_data_part.to(im_tensor_gpu_for_annotator.device)
+                    annotator.masks(mask_data_part, 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))
+                        except IndexError:
+                            logger.warning(f"Invalid part ID {cls_id}")
+                    logger.info("Successfully drew part masks with annotator")
+                    ultralytic_viz_success = True
+                except Exception as e_part:
+                    logger.error(f"Error drawing part masks with annotator: {e_part}", exc_info=True)
+            
+            # Draw damage masks with Ultralytics
+            if damage_masks_raw.numel() > 0 and im_tensor_gpu_for_annotator is not None:
+                try:
+                    colors_dmg = [(random.randint(100, 200), 0, 0) for _ in damage_classes_ids_cpu]
+                    mask_data_dmg = damage_masks_raw
+                    if mask_data_dmg.device != im_tensor_gpu_for_annotator.device:
+                        mask_data_dmg = mask_data_dmg.to(im_tensor_gpu_for_annotator.device)
+                    annotator.masks(mask_data_dmg, 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))
+                        except IndexError:
+                            logger.warning(f"Invalid damage ID {cls_id}")
+                    logger.info("Successfully drew damage masks with annotator")
+                    ultralytic_viz_success = True
+                except Exception as e_dmg:
+                    logger.error(f"Error drawing damage masks with annotator: {e_dmg}", exc_info=True)
+            
+            # Get result from annotator if successful
+            if ultralytic_viz_success:
+                annotated_image_bgr = annotator.result()
+                logger.info("Using Ultralytics annotator visualization")
+            else:
+                logger.warning("Ultralytics annotator visualization failed, will try direct approach")
+        except Exception as e_anno:
+            logger.error(f"Error with Ultralytics annotator: {e_anno}", exc_info=True)
+            ultralytic_viz_success = False
+        
+        # Second attempt: Direct visualization with OpenCV if Ultralytics failed
+        if not ultralytic_viz_success:
             try:
-                logger.info("Attempting to draw damage masks...")
-                colors_dmg = [(random.randint(100, 200), 0, 0) for _ in damage_classes_ids_cpu]
-                mask_data_dmg = damage_masks_raw
-                logger.info(f"Damage mask data: shape={mask_data_dmg.shape}, dtype={mask_data_dmg.dtype}, device={mask_data_dmg.device}")
-                if mask_data_dmg.device != im_tensor_gpu_for_annotator.device:
-                    logger.warning(f"Moving damage masks to match image tensor device ({im_tensor_gpu_for_annotator.device})")
-                    mask_data_dmg = mask_data_dmg.to(im_tensor_gpu_for_annotator.device)
-                annotator.masks(mask_data_dmg, colors=colors_dmg, im_gpu=im_tensor_gpu_for_annotator, alpha=0.4)
-                logger.info("Successfully drew damage masks.")
-                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))
-                    except IndexError:
-                        logger.warning(f"Invalid damage ID {cls_id} during drawing")
-            except Exception as e_dmg_vis:
-                logger.error(f"Error drawing damage masks/boxes: {e_dmg_vis}", exc_info=True)
-                traceback.print_exc()
-        elif damage_masks_raw.numel() > 0:
-            logger.warning("Damage masks exist but image tensor for annotator is None. Skipping damage mask drawing.")
-
-        annotated_image_bgr = annotator.result()
+                logger.info("Using DirectVisualizer as fallback...")
+                direct_viz = DirectVisualizer(image_np_bgr.copy())
+                
+                # Draw part masks first (background layer)
+                if part_masks_np.shape[0] > 0:
+                    logger.info(f"Drawing {part_masks_np.shape[0]} part masks directly")
+                    direct_viz.draw_masks(part_masks_np, part_classes_ids_cpu, CAR_PART_CLASSES, "part")
+                
+                # Draw damage masks on top
+                if damage_masks_np.shape[0] > 0:
+                    logger.info(f"Drawing {damage_masks_np.shape[0]} damage masks directly")
+                    direct_viz.draw_masks(damage_masks_np, damage_classes_ids_cpu, DAMAGE_CLASSES, "damage")
+                
+                annotated_image_bgr = direct_viz.result()
+                logger.info("Direct visualization successful")
+            except Exception as e_direct:
+                logger.error(f"Error with direct visualization: {e_direct}", exc_info=True)
+                # If direct visualization also fails, use the original image
+                annotated_image_bgr = image_np_bgr.copy()
 
     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.")
+        final_assignments.append(f"Error during processing: {str(e)}")
+        annotated_image_bgr = image_np_bgr.copy()
 
     assignment_text = "\n".join(final_assignments) if final_assignments else "No damage assignments generated."
     final_output_image_rgb = cv2.cvtColor(annotated_image_bgr, cv2.COLOR_BGR2RGB)