Added few more features

app.py

@@ -10,6 +10,11 @@ import uuid
 import random
 from pathlib import Path
 import colorsys
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class PolygonAugmentation:
     def __init__(self, tolerance=0.2, area_threshold=0.01, debug=False):
@@ -25,7 +30,7 @@ class PolygonAugmentation:
         poly_np = np.array(points, dtype=np.float32)
         area = cv2.contourArea(poly_np)
         if self.debug:
-            print(f"[DEBUG] Calculating polygon area: {area:.2f}")
+            logger.info(f"[DEBUG] Calculating polygon area: {area:.2f}")
         return area
 
     def load_labelme_data(self, json_file: Any, image: np.ndarray) -> Tuple:
@@ -60,7 +65,7 @@ class PolygonAugmentation:
         for shape in shapes:
             if shape.get('shape_type') != 'polygon' or not shape.get('points') or len(shape['points']) < 3:
                 if self.debug:  
-                    print(f"[DEBUG] Skipping invalid shape: {shape}")
+                    logger.info(f"[DEBUG] Skipping invalid shape: {shape}")
                 continue
             try:
                 points = [[float(x), float(y)] for x, y in shape['points']]
@@ -69,11 +74,11 @@ class PolygonAugmentation:
                 original_areas.append(self.calculate_polygon_area(points))
             except (ValueError, TypeError) as e:
                 if self.debug:
-                    print(f"[DEBUG] Error processing points: {shape['points']}, error: {str(e)}")
+                    logger.info(f"[DEBUG] Error processing points: {shape['points']}, error: {str(e)}")
                 continue   
         
         if not polygons and self.debug:
-            print(f"[DEBUG] Warning: No valid polygons in JSON")
+            logger.info(f"[DEBUG] Warning: No valid polygons in JSON")
         return image, polygons, labels, original_areas, data, "input"
 
     def simplify_polygon(self, polygon: List[List[float]], tolerance: float = None, label: str = None) -> List[List[float]]:
@@ -81,25 +86,25 @@ class PolygonAugmentation:
         if label and label.lower() in ['background', 'bg', 'back']:
             tol = tol * 3
             if self.debug:
-                print(f"[DEBUG] Using increased tolerance {tol} for background label '{label}'")
+                logger.info(f"[DEBUG] Using increased tolerance {tol} for background label '{label}'")
         
         if len(polygon) < 3:
             if self.debug:
-                print(f"[DEBUG] Polygon has fewer than 3 points, skipping simplification.")
+                logger.info(f"[DEBUG] Polygon has fewer than 3 points, skipping simplification.")
             return polygon
         poly_np = np.array(polygon, dtype=np.float32)
         approx = cv2.approxPolyDP(poly_np, tol, closed=True)
         simplified = approx.reshape(-1, 2).tolist()
         
         if self.debug:
-            print(f"[DEBUG] Simplified polygon from {len(polygon)} to {len(simplified)} points with tolerance {tol}")
+            logger.info(f"[DEBUG] Simplified polygon from {len(polygon)} to {len(simplified)} points with tolerance {tol}")
         return simplified
 
     def create_donut_polygon(self, external_contour: np.ndarray, internal_contours: List[np.ndarray]) -> List[List[float]]:
         external_points = external_contour.reshape(-1, 2).tolist()
         if not internal_contours:
             if self.debug:
-                print("[DEBUG] No internal contours found, returning external points.")
+                logger.info("[DEBUG] No internal contours found, returning external points.")
             return external_points
         
         result_points = external_points.copy()
@@ -122,7 +127,7 @@ class PolygonAugmentation:
             bridge_from = internal_points[int_idx]
             
             if self.debug:
-                print(f"[DEBUG] Creating bridge between external index {ext_idx} and internal index {int_idx}, distance {min_dist:.2f}")
+                logger.info(f"[DEBUG] Creating bridge between external index {ext_idx} and internal index {int_idx}, distance {min_dist:.2f}")
             
             new_points = (
                 result_points[:ext_idx+1] +
@@ -180,7 +185,7 @@ class PolygonAugmentation:
                 poly_np = np.array(poly, dtype=np.int32)
                 if len(poly_np) < 3:
                     if self.debug:
-                        print(f"[DEBUG] Skipping polygon {poly_idx}: fewer than 3 points")
+                        logger.info(f"[DEBUG] Skipping polygon {poly_idx}: fewer than 3 points")
                     continue
                 mask = np.zeros((height, width), dtype=np.uint8)
                 cv2.fillPoly(mask, [poly_np], 1)
@@ -188,7 +193,7 @@ class PolygonAugmentation:
                 all_labels.append(label)
             except Exception as e:
                 if self.debug:
-                    print(f"[DEBUG] Error processing polygon {poly_idx}: {str(e)}")
+                    logger.info(f"[DEBUG] Error processing polygon {poly_idx}: {str(e)}")
         
         if not all_masks:
             return np.zeros((0, height, width), dtype=np.uint8), []
@@ -217,7 +222,7 @@ class PolygonAugmentation:
             all_polygons.append(poly_labelme)
             all_labels.append(label)
             if self.debug:
-                print(f"[DEBUG] Added simplified external polygon with {len(poly_labelme)} points.")
+                logger.info(f"[DEBUG] Added simplified external polygon with {len(poly_labelme)} points.")
 
         for internal_contour in internal_contours:
             internal_points = internal_contour.reshape(-1, 2).tolist()
@@ -230,7 +235,7 @@ class PolygonAugmentation:
                 all_polygons.append(poly_labelme)
                 all_labels.append(label)
                 if self.debug:
-                    print(f"[DEBUG] Added simplified internal polygon with {len(poly_labelme)} points.")
+                    logger.info(f"[DEBUG] Added simplified internal polygon with {len(poly_labelme)} points.")
 
     def masks_to_labelme_polygons(
         self,
@@ -249,12 +254,12 @@ class PolygonAugmentation:
         for mask_idx, (mask, label) in enumerate(zip(masks, labels)):
             if mask.sum() < 10:
                 if self.debug:
-                    print(f"[DEBUG] Skipping mask {mask_idx}: very small or empty.")
+                    logger.info(f"[DEBUG] Skipping mask {mask_idx}: very small or empty.")
                 continue
             contours, hierarchy = cv2.findContours(mask, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
             if hierarchy is None or len(contours) == 0:
                 if self.debug:
-                    print(f"[DEBUG] No contours found in mask {mask_idx}.")
+                    logger.info(f"[DEBUG] No contours found in mask {mask_idx}.")
                 continue
             
             hierarchy = hierarchy[0]
@@ -274,7 +279,7 @@ class PolygonAugmentation:
             
             if not external_contours:
                 if self.debug:
-                    print(f"[DEBUG] No external contours found in mask {mask_idx}.")
+                    logger.info(f"[DEBUG] No external contours found in mask {mask_idx}.")
                 continue
             
             for ext_idx, external_contour in enumerate(external_contours):
@@ -286,7 +291,7 @@ class PolygonAugmentation:
                     relative_area = ext_area / original_areas[mask_idx]
                     if relative_area < area_thresh:
                         if self.debug:
-                            print(f"[DEBUG] Skipping contour {ext_idx} (area too small: {relative_area:.4f})")
+                            logger.info(f"[DEBUG] Skipping contour {ext_idx} (area too small: {relative_area:.4f})")
                         continue
                 
                 is_background = label.lower() in ['background', 'bg', 'back']
@@ -301,10 +306,10 @@ class PolygonAugmentation:
                             all_polygons.append(poly_labelme)
                             all_labels.append(label)
                             if self.debug:
-                                print(f"[DEBUG] Added donut polygon with {len(poly_labelme)} points.")
+                                logger.info(f"[DEBUG] Added donut polygon with {len(poly_labelme)} points.")
                     except Exception as e:
                         if self.debug:
-                            print(f"[DEBUG] Error creating donut: {str(e)}, fallback to separate polygons.")
+                            logger.info(f"[DEBUG] Error creating donut: {str(e)}, fallback to separate polygons.")
                         self.process_contours(
                             external_contour, internal_contours, width, height,
                             label, all_polygons, all_labels, tol
@@ -327,6 +332,7 @@ class PolygonAugmentation:
         aug_type: str,
         aug_param: float
     ) -> 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)
@@ -351,22 +357,28 @@ class PolygonAugmentation:
         transform = A.Compose([
             aug_dict[aug_type],
             A.RandomCrop(width=crop_width, height=crop_height, p=0.8)
-        ])
+        ], additional_targets={'mask': 'mask'})
         
         masks, mask_labels = self.polygons_to_masks(image, polygons, labels)
         if masks.shape[0] == 0:
             raise ValueError("No valid masks created from polygons")
         
+        # Ensure masks are processed correctly
         aug_result = transform(image=image, masks=masks)
         aug_image = aug_result['image']
         aug_masks = aug_result['masks']
         
+        # Validate augmented image
+        if aug_image is None or aug_image.size == 0:
+            raise ValueError("Augmented image is empty or invalid")
+        
         aug_polygons, aug_labels = self.masks_to_labelme_polygons(
             aug_masks, mask_labels, original_areas, self.area_threshold, self.tolerance
         )
         
         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")
         return aug_image, aug_data
 
 def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param: float):
@@ -387,11 +399,13 @@ def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param:
             raise ValueError(f"Parameter {aug_param} for {aug_type} is out of range [{min_val}, {max_val}]")
 
         # Convert PIL image to NumPy
+        if image is None:
+            raise ValueError("Input image is None")
         img_np = np.array(image)
         img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
         
         # Initialize augmenter
-        augmenter = PolygonAugmentation(tolerance=2.0, area_threshold=0.01, debug=False)
+        augmenter = PolygonAugmentation(tolerance=2.0, area_threshold=0.01, debug=True)
         
         # Load data
         img_np, polygons, labels, original_areas, original_data, _ = augmenter.load_labelme_data(json_file, img_np)
@@ -401,6 +415,10 @@ def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param:
             img_np, polygons, labels, original_areas, original_data, aug_type, aug_param
         )
         
+        # Validate augmented image
+        if aug_image is None or aug_image.size == 0:
+            raise ValueError("Augmented image is empty or invalid")
+        
         # Create a dynamic color map for unique labels
         unique_labels = list(set(shape['label'] for shape in aug_data['shapes']))
         if not unique_labels:
@@ -416,22 +434,29 @@ def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param:
         
         # Convert augmented image to RGB for visualization
         aug_image_rgb = cv2.cvtColor(aug_image, cv2.COLOR_BGR2RGB)
+        if aug_image_rgb is None or aug_image_rgb.size == 0:
+            raise ValueError("Failed to convert augmented image to RGB")
+        
+        # Create a clean copy of the augmented image for visualization
         overlay = aug_image_rgb.copy()
         
-        # Create masks for visualization
+        # Create masks and outlines for visualization
         height, width = aug_image.shape[:2]
         for shape in aug_data['shapes']:
             label = shape['label']
-            color = label_color_map.get(label, (0, 255, 0))
+            color = label_color_map.get(label, (0, 255, 0))  # Fallback to green
             points = np.array(shape['points'], dtype=np.int32)
+            if len(points) < 3:
+                logger.warning(f"Skipping invalid polygon for label {label}: fewer than 3 points")
+                continue
             
-            # Draw filled mask with transparency
+            # Draw semi-transparent mask
             mask = np.zeros((height, width), dtype=np.uint8)
             cv2.fillPoly(mask, [points], 1)
             colored_mask = np.zeros_like(aug_image_rgb)
             colored_mask[mask == 1] = color
-            alpha = 0.3
-            cv2.addWeighted(colored_mask, alpha, overlay, 1 - alpha, 0, overlay)
+            alpha = 0.3  # Transparency for mask
+            overlay = cv2.addWeighted(overlay, 1.0, colored_mask, alpha, 0.0)
             
             # Draw polygon outline
             cv2.polylines(overlay, [points], isClosed=True, color=color, thickness=2)
@@ -442,8 +467,10 @@ def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param:
         # Format JSON for display
         aug_json_str = json.dumps(aug_data, indent=2)
         
+        logger.info("Visualization completed successfully")
         return aug_image_pil, aug_json_str
     except Exception as e:
+        logger.error(f"Error in augment_image: {str(e)}")
         return None, f"Error: {str(e)}"
 
 # Define augmentation types and parameter ranges
@@ -459,7 +486,7 @@ aug_options = {
 def create_interface():
     with gr.Blocks(title="Donut Polygon Augmentation") as demo:
         gr.Markdown("# Donut Polygon Augmentation 🌀")
-        gr.Markdown("Upload an image and a LabelMe JSON file to apply topology-preserving augmentation to donut-shaped polygons. Each class is visualized with a unique color and semi-transparent mask.")
+        gr.Markdown("Upload an image and a LabelMe JSON file to apply topology-preserving augmentation to donut-shaped polygons. Each class is visualized with a unique color and semi-transparent mask over the augmented image.")
         
         with gr.Row():
             with gr.Column():
@@ -479,7 +506,6 @@ def create_interface():
                 )
                 
                 def update_slider(display_name):
-                    # Map display_name back to internal key
                     aug_key = next(k for k, v in aug_options.items() if v["display_name"] == display_name)
                     return {
                         aug_param: gr.update(
@@ -500,7 +526,6 @@ def create_interface():
                 output_json = gr.Textbox(label="Augmented LabelMe JSON", lines=10, max_lines=20)
         
         def submit(image, json_file, display_name, aug_param):
-            # Map display_name to internal key
             aug_key = next(k for k, v in aug_options.items() if v["display_name"] == display_name)
             return augment_image(image, json_file, aug_key, aug_param)