Update app.py

app.py

@@ -109,41 +109,55 @@ class PolygonAugmentation:
         return simplified
 
     def create_donut_polygon(self, external_contour: np.ndarray, internal_contours: List[np.ndarray]) -> List[List[float]]:
+        """Create a donut/ring polygon by connecting external and internal contours with bridges"""
         external_points = external_contour.reshape(-1, 2).tolist()
         if not internal_contours:
             if self.debug:
                 logger.info("[DEBUG] No internal contours found, returning external points.")
             return external_points
         
+        # Start with external contour points
         result_points = external_points.copy()
         
-        for internal_contour in internal_contours:
+        # Process each internal contour (hole)
+        for hole_idx, internal_contour in enumerate(internal_contours):
             internal_points = internal_contour.reshape(-1, 2).tolist()
+            
+            # Find the closest point between external and internal contours
             min_dist = float('inf')
-            ext_idx = 0
-            int_idx = 0
+            best_ext_idx = 0
+            best_int_idx = 0
             
-            for i, p1 in enumerate(external_points):
-                for j, p2 in enumerate(internal_points):
-                    dist = np.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)
+            # Check all combinations to find minimum distance
+            for i, ext_point in enumerate(result_points):
+                for j, int_point in enumerate(internal_points):
+                    dist = np.sqrt((ext_point[0] - int_point[0])**2 + (ext_point[1] - int_point[1])**2)
                     if dist < min_dist:
                         min_dist = dist
-                        ext_idx = i
-                        int_idx = j
+                        best_ext_idx = i
+                        best_int_idx = j
             
-            bridge_to = external_points[ext_idx]
-            bridge_from = internal_points[int_idx]
+            # Create bridge points
+            bridge_start = result_points[best_ext_idx]
+            connect_point = internal_points[best_int_idx]
             
             if self.debug:
-                logger.info(f"[DEBUG] Creating bridge between external index {ext_idx} and internal index {int_idx}, distance {min_dist:.2f}")
+                logger.info(f"[DEBUG] Creating bridge for hole {hole_idx}: ext_idx={best_ext_idx}, int_idx={best_int_idx}, distance={min_dist:.2f}")
             
-            new_points = (
-                result_points[:ext_idx+1] +
-                internal_points[int_idx:] + internal_points[:int_idx+1] +
-                [bridge_to] +
-                external_points[ext_idx+1:]
+            # Insert the internal contour into the result
+            # Order: external_points[:best_ext_idx+1] + internal_hole + back_to_external + external_points[best_ext_idx+1:]
+            new_result = (
+                result_points[:best_ext_idx+1] +  # External points up to bridge
+                internal_points[best_int_idx:] +  # Internal points from connection point to end
+                internal_points[:best_int_idx+1] +  # Internal points from start to connection point
+                [bridge_start] +  # Bridge back to external
+                result_points[best_ext_idx+1:]  # Remaining external points
             )
-            result_points = new_points
+            
+            result_points = new_result
+        
+        if self.debug:
+            logger.info(f"[DEBUG] Created donut polygon with {len(result_points)} total points")
         
         return result_points
 
@@ -171,16 +185,22 @@ class PolygonAugmentation:
                 "confidence": 1.0
             })
         
+        # Get actual dimensions from augmented image
+        aug_height, aug_width = aug_image.shape[:2]
+        
         aug_data = {
             "version": original_data.get("version", "5.0.1"),
             "flags": original_data.get("flags", {}),
             "shapes": new_shapes,
             "imagePath": aug_img_name,
             "imageData": None,
-            "imageHeight": aug_image.shape[0],
-            "imageWidth": aug_image.shape[1]
+            "imageHeight": aug_height,  # Use actual augmented image height
+            "imageWidth": aug_width     # Use actual augmented image width
         }
         
+        if self.debug:
+            logger.info(f"[DEBUG] Created augmented data: {len(new_shapes)} shapes, size: {aug_width}x{aug_height}")
+        
         return aug_data
 
     def polygons_to_masks(self, image: np.ndarray, polygons: List[List[List[float]]], labels: List[str]) -> Tuple[np.ndarray, List[str]]:
@@ -302,27 +322,49 @@ class PolygonAugmentation:
                             logger.info(f"[DEBUG] Skipping contour {ext_idx} (area too small: {relative_area:.4f})")
                         continue
                 
+                # Check if this is a ring/donut shape or complex polygon
+                is_ring_shape = label.lower() in ['ring', 'donut', 'annulus', 'circle', 'round'] or len(internal_contours) > 0
                 is_background = label.lower() in ['background', 'bg', 'back']
-                if is_background and internal_contours:
+                
+                # Handle different polygon types
+                if (is_background or is_ring_shape) and internal_contours:
                     try:
+                        # Create donut polygon for rings, backgrounds, or shapes with holes
                         donut_points = self.create_donut_polygon(external_contour, internal_contours)
                         simplified_donut = self.simplify_polygon(donut_points, tolerance=tol, label=label)
+                        
                         if len(simplified_donut) >= 3:
-                            poly_labelme = [[round(max(0, min(float(x), width - 1)), 2),
-                                            round(max(0, min(float(y), height - 1)), 2)]
-                                            for x, y in simplified_donut]
+                            # Ensure all points are within image boundaries
+                            poly_labelme = []
+                            for x, y in simplified_donut:
+                                clipped_x = round(max(0, min(float(x), width - 1)), 2)
+                                clipped_y = round(max(0, min(float(y), height - 1)), 2)
+                                poly_labelme.append([clipped_x, clipped_y])
+                            
                             all_polygons.append(poly_labelme)
                             all_labels.append(label)
+                            
                             if self.debug:
-                                logger.info(f"[DEBUG] Added donut polygon with {len(poly_labelme)} points.")
+                                logger.info(f"[DEBUG] Added {'ring' if is_ring_shape else 'background'} donut polygon with {len(poly_labelme)} points, {len(internal_contours)} holes")
+                        else:
+                            if self.debug:
+                                logger.info(f"[DEBUG] Donut polygon too small after simplification, falling back to separate contours")
+                            # Fallback to separate contours
+                            self.process_contours(
+                                external_contour, internal_contours, width, height,
+                                label, all_polygons, all_labels, tol
+                            )
+                            
                     except Exception as e:
                         if self.debug:
-                            logger.info(f"[DEBUG] Error creating donut: {str(e)}, fallback to separate polygons.")
+                            logger.info(f"[DEBUG] Error creating donut for {label}: {str(e)}, fallback to separate polygons.")
+                        # Fallback to processing contours separately
                         self.process_contours(
                             external_contour, internal_contours, width, height,
                             label, all_polygons, all_labels, tol
                         )
                 else:
+                    # Handle regular polygons (no holes or simple shapes)
                     self.process_contours(
                         external_contour, internal_contours, width, height,
                         label, all_polygons, all_labels, tol
@@ -342,31 +384,36 @@ class PolygonAugmentation:
     ) -> Tuple[np.ndarray, Dict[str, Any]]:
         logger.info(f"Applying augmentation: {aug_type} with parameter {aug_param}")
         height, width = image.shape[:2]
-        crop_scale = random.uniform(0.8, 0.9)
-        crop_height = int(height * crop_scale)
-        crop_width = int(width * crop_scale)
-        
-        aug_dict = {
-            "rotate": A.Rotate(limit=aug_param, p=1.0),
-            "horizontal_flip": A.HorizontalFlip(p=1.0 if aug_param == 1 else 0.0),
-            "vertical_flip": A.VerticalFlip(p=1.0 if aug_param == 1 else 0.0),
-            "scale": A.Affine(scale=aug_param, p=1.0),
-            "brightness_contrast": A.RandomBrightnessContrast(
-                brightness_limit=aug_param,
-                contrast_limit=aug_param,
-                p=1.0
-            ),
-            "pixel_dropout": A.PixelDropout(dropout_prob=min(max(aug_param, 0.0), 1.0), p=1.0)
-        }
         
-        if aug_type not in aug_dict:
+        # Setup augmentation based on type with proper parameters
+        if aug_type == "rotate":
+            # For rotation, use the parameter as degrees and make it more visible
+            rotation_angle = aug_param if abs(aug_param) >= 5 else (15 if aug_param >= 0 else -15)
+            aug_transform = A.Rotate(limit=abs(rotation_angle), p=1.0, border_mode=cv2.BORDER_CONSTANT, value=0)
+            logger.info(f"Applying rotation: {rotation_angle} degrees")
+        elif aug_type == "horizontal_flip":
+            aug_transform = A.HorizontalFlip(p=1.0 if aug_param == 1 else 0.0)
+        elif aug_type == "vertical_flip":
+            aug_transform = A.VerticalFlip(p=1.0 if aug_param == 1 else 0.0)
+        elif aug_type == "scale":
+            # Ensure scale parameter is reasonable
+            scale_factor = max(0.5, min(2.0, aug_param))
+            aug_transform = A.Affine(scale=scale_factor, p=1.0, keep_ratio=True)
+            logger.info(f"Applying scale: {scale_factor}")
+        elif aug_type == "brightness_contrast":
+            brightness_factor = max(-0.5, min(0.5, aug_param))
+            aug_transform = A.RandomBrightnessContrast(
+                brightness_limit=abs(brightness_factor),
+                contrast_limit=abs(brightness_factor),
+                p=1.0
+            )
+        elif aug_type == "pixel_dropout":
+            dropout_prob = min(max(aug_param, 0.0), 0.2)
+            aug_transform = A.PixelDropout(dropout_prob=dropout_prob, p=1.0)
+        else:
             raise ValueError(f"Unsupported augmentation type: {aug_type}")
         
-        transform = A.Compose([
-            aug_dict[aug_type],
-            A.RandomCrop(width=crop_width, height=crop_height, p=0.8)
-        ])
-        
+        # Create masks from polygons
         masks, mask_labels = self.polygons_to_masks(image, polygons, labels)
         if masks.shape[0] == 0:
             raise ValueError("No valid masks created from polygons")
@@ -377,10 +424,9 @@ class PolygonAugmentation:
         # Create additional targets for each mask
         additional_targets = {f'mask{i}': 'mask' for i in range(len(masks_list))}
         
-        # Update transform with additional targets
+        # Create transform with proper mask handling
         transform = A.Compose([
-            aug_dict[aug_type],
-            A.RandomCrop(width=crop_width, height=crop_height, p=0.8)
+            aug_transform
         ], additional_targets=additional_targets)
         
         # Prepare input dictionary
@@ -392,10 +438,16 @@ class PolygonAugmentation:
         aug_result = transform(**input_dict)
         aug_image = aug_result['image']
         
-        # Collect augmented masks
+        # Collect augmented masks and ensure they match image dimensions
         aug_masks_list = []
+        aug_height, aug_width = aug_image.shape[:2]
+        
         for i in range(len(masks_list)):
-            aug_masks_list.append(aug_result[f'mask{i}'])
+            aug_mask = aug_result[f'mask{i}']
+            # Ensure mask dimensions match augmented image
+            if aug_mask.shape[:2] != (aug_height, aug_width):
+                aug_mask = cv2.resize(aug_mask, (aug_width, aug_height), interpolation=cv2.INTER_NEAREST)
+            aug_masks_list.append(aug_mask)
         
         aug_masks = np.array(aug_masks_list, dtype=np.uint8)
         
@@ -403,13 +455,46 @@ class PolygonAugmentation:
         if aug_image is None or aug_image.size == 0:
             raise ValueError("Augmented image is empty or invalid")
         
+        # Convert augmented masks back to polygons
         aug_polygons, aug_labels = self.masks_to_labelme_polygons(
             aug_masks, mask_labels, original_areas, self.area_threshold, self.tolerance
         )
         
+        # Apply random crop as post-processing to add variety
+        if random.random() < 0.3:  # 30% chance of cropping
+            crop_scale = random.uniform(0.85, 0.95)
+            crop_height = int(aug_height * crop_scale)
+            crop_width = int(aug_width * crop_scale)
+            
+            # Create crop transform
+            crop_transform = A.Compose([
+                A.RandomCrop(width=crop_width, height=crop_height, p=1.0)
+            ], additional_targets={f'mask{i}': 'mask' for i in range(len(aug_masks_list))})
+            
+            # Apply crop
+            crop_input = {'image': aug_image}
+            for i, mask in enumerate(aug_masks_list):
+                crop_input[f'mask{i}'] = mask
+            
+            crop_result = crop_transform(**crop_input)
+            aug_image = crop_result['image']
+            
+            # Update masks after crop
+            cropped_masks = []
+            for i in range(len(aug_masks_list)):
+                cropped_masks.append(crop_result[f'mask{i}'])
+            
+            aug_masks = np.array(cropped_masks, dtype=np.uint8)
+            
+            # Re-convert masks to polygons after crop
+            aug_polygons, aug_labels = self.masks_to_labelme_polygons(
+                aug_masks, mask_labels, original_areas, self.area_threshold, self.tolerance
+            )
+        
+        # Create augmented data with correct dimensions
         aug_data = self.save_augmented_data(aug_image, aug_polygons, aug_labels, original_data, "input")
         
-        logger.info(f"Augmentation completed: {len(aug_polygons)} polygons generated")
+        logger.info(f"Augmentation completed: {len(aug_polygons)} polygons generated, final size: {aug_image.shape[:2]}")
         return aug_image, aug_data
 
     def batch_augment_images(self, image_json_pairs, aug_configs, num_augmentations):