Added few more features

app.py

@@ -1,34 +1,446 @@
-# -*- coding: utf-8 -*-
 import gradio as gr
 from PIL import Image
 import numpy as np
 import cv2
+import json
+import albumentations as A
+from typing import List, Tuple, Dict, Any
+import supervision as sv
+import uuid
+import random
+from pathlib import Path
 
-def dummy_polygon_augment(image: Image.Image):
-    # Convert to NumPy
-    img = np.array(image)
-    
-    # Dummy augmentation: Draw a donut-like polygon
-    overlay = img.copy()
-    height, width = img.shape[:2]
-    center = (width // 2, height // 2)
-    outer_radius = min(width, height) // 4
-    inner_radius = outer_radius // 2
-    
-    cv2.circle(overlay, center, outer_radius, (0, 255, 0), -1)
-    cv2.circle(overlay, center, inner_radius, (0, 0, 0), -1)
-    
-    alpha = 0.4
-    cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0, img)
-    
-    return Image.fromarray(img)
+class PolygonAugmentation:
+    def __init__(self, tolerance=0.2, area_threshold=0.01, debug=False):
+        self.tolerance = tolerance
+        self.area_threshold = area_threshold
+        self.debug = debug
+        self.supported_extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.PNG', '.JPEG']
+        
+    def __getattr__(self, name: str) -> Any:
+        raise AttributeError(f"'PolygonAugmentation' object has no attribute '{name}'")    
+
+    def calculate_polygon_area(self, points: List[List[float]]) -> float:
+        poly_np = np.array(points, dtype=np.float32)
+        area = cv2.contourArea(poly_np)
+        if self.debug:
+            print(f"[DEBUG] Calculating polygon area: {area:.2f}")
+        return area
+
+    def load_labelme_data(self, json_file: Any, image: np.ndarray) -> Tuple:
+        if isinstance(json_file, str):
+            with open(json_file, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+        else:
+            data = json.load(json_file)
+        if 'shapes' not in data or not isinstance(data['shapes'], list):
+            raise ValueError("Invalid JSON: 'shapes' key missing or not a list")
+        
+        polygons = []
+        labels = []
+        original_areas = []
+        
+        for shape in data['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")
+                continue
+            try:
+                points = [[float(x), float(y)] for x, y in shape['points']]
+                polygons.append(points)
+                labels.append(shape['label'])
+                original_areas.append(self.calculate_polygon_area(points))
+            except (ValueError, TypeError):
+                if self.debug:
+                    print(f"[DEBUG] Error processing points: {shape['points']}")
+                continue   
+        
+        if not polygons and self.debug:
+            print(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]]:
+        tol = tolerance if tolerance is not None else self.tolerance
+        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}'")
+        
+        if len(polygon) < 3:
+            if self.debug:
+                print(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}")
+        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.")
+            return external_points
+        
+        result_points = external_points.copy()
+        
+        for internal_contour in internal_contours:
+            internal_points = internal_contour.reshape(-1, 2).tolist()
+            min_dist = float('inf')
+            ext_idx = 0
+            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)
+                    if dist < min_dist:
+                        min_dist = dist
+                        ext_idx = i
+                        int_idx = j
+            
+            bridge_to = external_points[ext_idx]
+            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}")
+            
+            new_points = (
+                result_points[:ext_idx+1] +
+                internal_points[int_idx:] + internal_points[:int_idx+1] +
+                [bridge_to] +
+                external_points[ext_idx+1:]
+            )
+            result_points = new_points
+        
+        return result_points
+
+    def save_augmented_data(
+        self,
+        aug_image: np.ndarray,
+        aug_polygons: List[List[List[float]]],
+        aug_labels: List[str],
+        original_data: Dict[str, Any],
+        base_name: str
+    ) -> Dict[str, Any]:
+        aug_id = uuid.uuid4().hex[:4]
+        aug_img_name = f"{base_name}_{aug_id}_aug.png"
+        
+        new_shapes = []
+        for poly, label in zip(aug_polygons, aug_labels):
+            if not poly or len(poly) < 3:
+                continue
+            new_shapes.append({
+                "label": label,
+                "points": poly,
+                "group_id": None,
+                "shape_type": "polygon",
+                "flags": {}
+            })
+        
+        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]
+        }
+        
+        return aug_data
 
-demo = gr.Interface(
-    fn=dummy_polygon_augment,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Donut Polygon Augmentation",
-    description="Upload an image to visualize a sample donut-shaped polygon augmentation."
-)
+    def polygons_to_masks(self, image: np.ndarray, polygons: List[List[List[float]]], labels: List[str]) -> Tuple[np.ndarray, List[str]]:
+        height, width = image.shape[:2]
+        all_masks = []
+        all_labels = []
+        
+        for poly_idx, (poly, label) in enumerate(zip(polygons, labels)):
+            try:
+                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")
+                    continue
+                mask = np.zeros((height, width), dtype=np.uint8)
+                cv2.fillPoly(mask, [poly_np], 1)
+                all_masks.append(mask)
+                all_labels.append(label)
+            except Exception as e:
+                if self.debug:
+                    print(f"[DEBUG] Error processing polygon {poly_idx}: {str(e)}")
+        
+        if not all_masks:
+            return np.zeros((0, height, width), dtype=np.uint8), []
+        
+        return np.array(all_masks, dtype=np.uint8), all_labels
+
+    def process_contours(
+        self,
+        external_contour: np.ndarray,
+        internal_contours: List[np.ndarray],
+        width: int,
+        height: int,
+        label: str,
+        all_polygons: List[List[List[float]]],
+        all_labels: List[str],
+        tolerance: float = None
+    ) -> None:
+        tol = tolerance if tolerance is not None else self.tolerance
+        external_points = external_contour.reshape(-1, 2).tolist()
+        simplified_external = self.simplify_polygon(external_points, tolerance=tol, label=label)
+        
+        if len(simplified_external) >= 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_external]
+            all_polygons.append(poly_labelme)
+            all_labels.append(label)
+            if self.debug:
+                print(f"[DEBUG] Added simplified external Nghia with {len(poly_labelme)} points.")
+
+        for internal_contour in internal_contours:
+            internal_points = internal_contour.reshape(-1, 2).tolist()
+            simplified_internal = self.simplify_polygon(internal_points, tolerance=tol, label=label)
+            
+            if len(simplified_internal) >= 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_internal]
+                all_polygons.append(poly_labelme)
+                all_labels.append(label)
+                if self.debug:
+                    print(f"[DEBUG] Added simplified internal polygon with {len(poly_labelme)} points.")
+
+    def masks_to_labelme_polygons(
+        self,
+        masks: np.ndarray,
+        labels: List[str],
+        original_areas: List[float],
+        area_threshold: float = None,
+        tolerance: float = None
+    ) -> Tuple[List[List[List[float]]], List[str]]:
+        tol = tolerance if tolerance is not None else self.tolerance
+        area_thresh = area_threshold if area_threshold is not None else self.area_threshold
+        height, width = masks[0].shape if len(masks) > 0 else (0, 0)
+        all_polygons = []
+        all_labels = []
+        
+        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.")
+                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}.")
+                continue
+            
+            hierarchy = hierarchy[0]
+            external_contours = []
+            internal_contours_map = {}
+            
+            for i, (contour, h) in enumerate(zip(contours, hierarchy)):
+                if h[3] == -1:
+                    external_contours.append(contour)
+                    internal_contours_map[len(external_contours)-1] = []
+                else:
+                    parent_idx = h[3]
+                    for j, _ in enumerate(external_contours):
+                        if parent_idx == j:
+                            internal_contours_map[j].append(contour)
+                            break
+            
+            if not external_contours:
+                if self.debug:
+                    print(f"[DEBUG] No external contours found in mask {mask_idx}.")
+                continue
+            
+            for ext_idx, external_contour in enumerate(external_contours):
+                internal_contours = internal_contours_map.get(ext_idx, [])
+                ext_area = cv2.contourArea(external_contour)
+                if ext_area <= 0:
+                    continue
+                if mask_idx < len(original_areas) and original_areas[mask_idx] > 0:
+                    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})")
+                        continue
+                
+                is_background = label.lower() in ['background', 'bg', 'back']
+                if is_background and internal_contours:
+                    try:
+                        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]
+                            all_polygons.append(poly_labelme)
+                            all_labels.append(label)
+                            if self.debug:
+                                print(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.")
+                        self.process_contours(
+                            external_contour, internal_contours, width, height,
+                            label, all_polygons, all_labels, tol
+                        )
+                else:
+                    self.process_contours(
+                        external_contour, internal_contours, width, height,
+                        label, all_polygons, all_labels, tol
+                    )
+        
+        return all_polygons, all_labels
+
+    def augment_single_image(
+        self,
+        image: np.ndarray,
+        polygons: List[List[List[float]]],
+        labels: List[str],
+        original_areas: List[float],
+        original_data: Dict[str, Any],
+        aug_type: str,
+        aug_param: float
+    ) -> Tuple[np.ndarray, Dict[str, Any]]:
+        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 > 0 else 0.0),
+            "vertical_flip": A.VerticalFlip(p=1.0 if aug_param > 0 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=aug_param, p=1.0)
+        }
+        
+        if aug_type not in aug_dict:
+            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)
+        ])
+        
+        masks, mask_labels = self.polygons_to_masks(image, polygons, labels)
+        if masks.shape[0] == 0:
+            raise ValueError("No valid masks created from polygons")
+        
+        aug_result = transform(image=image, masks=masks)
+        aug_image = aug_result['image']
+        aug_masks = aug_result['masks']
+        
+        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")
+        
+        return aug_image, aug_data
+
+def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param: float):
+    try:
+        # Convert PIL image to NumPy
+        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)
+        
+        # Load data
+        img_np, polygons, labels, original_areas, original_data, _ = augmenter.load_labelme_data(json_file, img_np)
+        
+        # Perform augmentation
+        aug_image, aug_data = augmenter.augment_single_image(
+            img_np, polygons, labels, original_areas, original_data, aug_type, aug_param
+        )
+        
+        # Overlay polygons on augmented image for visualization
+        aug_image_rgb = cv2.cvtColor(aug_image, cv2.COLOR_BGR2RGB)
+        for poly in aug_data['shapes']:
+            points = np.array(poly['points'], dtype=np.int32)
+            cv2.polylines(aug_image_rgb, [points], isClosed=True, color=(0, 255, 0), thickness=2)
+        
+        # Convert augmented image back to PIL
+        aug_image_pil = Image.fromarray(aug_image_rgb)
+        
+        # Format JSON for display
+        aug_json_str = json.dumps(aug_data, indent=2)
+        
+        return aug_image_pil, aug_json_str
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+# Define augmentation types and parameter ranges
+aug_options = {
+    "Rotate": {"param_name": "Angle (degrees)", "range": (-30, 30), "default": 0},
+    "Horizontal Flip": {"param_name": "Apply Flip (0 or 1)", "range": (0, 1), "default": 0},
+    "Vertical Flip": {"param_name": "Apply Flip (0 or 1)", "range": (0, 1), "default": 0},
+    "Scale": {"param_name": "Scale Factor", "range": (0.5, 1.5), "default": 1.0},
+    "Brightness/Contrast": {"param_name": "Brightness/Contrast Limit", "range": (-0.3, 0.3), "default": 0},
+    "Pixel Dropout": {"param_name": "Dropout Probability", "range": (0.01, 0.1), "default": 0.05}
+}
+
+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.")
+        
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(type="pil", label="Input Image")
+                json_input = gr.File(label="LabelMe JSON File", file_types=[".json"])
+                aug_type = gr.Dropdown(
+                    choices=list(aug_options.keys()),
+                    label="Augmentation Type",
+                    value="Rotate"
+                )
+                aug_param = gr.Slider(
+                    minimum=aug_options["Rotate"]["range"][0],
+                    maximum=aug_options["Rotate"]["range"][1],
+                    value=aug_options["Rotate"]["default"],
+                    label=aug_options["Rotate"]["param_name"]
+                )
+                
+                def update_slider(aug_type):
+                    return {
+                        aug_param: gr.update(
+                            minimum=aug_options[aug_type]["range"][0],
+                            maximum=aug_options[aug_type]["range"][1],
+                            value=aug_options[aug_type]["default"],
+                            label=aug_options[aug_type]["param_name"]
+                        )
+                    }
+                
+                aug_type.change(fn=update_slider, inputs=aug_type, outputs=[aug_param])
+                
+                submit_btn = gr.Button("Apply Augmentation")
+            
+            with gr.Column():
+                output_image = gr.Image(type="pil", label="Augmented Image")
+                output_json = gr.Textbox(label="Augmented LabelMe JSON", lines=10, max_lines=20)
+        
+        submit_btn.click(
+            fn=augment_image,
+            inputs=[image_input, json_input, aug_type, aug_param],
+            outputs=[output_image, output_json]
+        )
+    
+    return demo
 
-demo.launch()
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()

requirements.txt

@@ -2,3 +2,5 @@ opencv-python-headless
 gradio==5.30.0
 Pillow
 numpy
+albumentations 
+supervision