Update app.py

app.py

@@ -11,6 +11,9 @@ import random
 from pathlib import Path
 import colorsys
 import logging
+import zipfile
+import io
+from datetime import datetime
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -22,6 +25,7 @@ class PolygonAugmentation:
         self.area_threshold = area_threshold
         self.debug = debug
         self.supported_extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.PNG', '.JPEG']
+        self.augmented_results = []  # Store all augmentation results
         
     def __getattr__(self, name: str) -> Any:
         raise AttributeError(f"'PolygonAugmentation' object has no attribute '{name}'")    
@@ -381,44 +385,72 @@ class PolygonAugmentation:
         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):
-    try:
-        # Validate aug_param based on aug_type
-        aug_ranges = {
-            "rotate": (-30, 30),
-            "horizontal_flip": (0, 1),
-            "vertical_flip": (0, 1),
-            "scale": (0.5, 1.5),
-            "brightness_contrast": (-0.3, 0.3),
-            "pixel_dropout": (0.01, 0.1)
-        }
-        if aug_type not in aug_ranges:
-            raise ValueError(f"Invalid augmentation type: {aug_type}")
-        min_val, max_val = aug_ranges[aug_type]
-        if not (min_val <= aug_param <= max_val):
-            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=True)
+    def batch_augment_images(self, image_json_pairs, aug_configs, num_augmentations):
+        """Batch process multiple images with multiple augmentation configurations"""
+        self.augmented_results = []
+        results = []
         
-        # 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
-        )
-        
-        # Validate augmented image
-        if aug_image is None or aug_image.size == 0:
-            raise ValueError("Augmented image is empty or invalid")
+        for pair_idx, (image, json_file) in enumerate(image_json_pairs):
+            if image is None or json_file is None:
+                continue
+                
+            try:
+                # Convert PIL image to NumPy
+                img_np = np.array(image)
+                img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+                
+                # Load data
+                img_np, polygons, labels, original_areas, original_data, _ = self.load_labelme_data(json_file, img_np)
+                
+                # Apply each augmentation configuration
+                for config in aug_configs:
+                    for aug_idx in range(num_augmentations):
+                        # Generate random parameter within range
+                        min_val, max_val = config['param_range']
+                        if config['aug_type'] in ['horizontal_flip', 'vertical_flip']:
+                            aug_param = random.choice([0, 1])
+                        else:
+                            aug_param = random.uniform(min_val, max_val)
+                        
+                        try:
+                            aug_image, aug_data = self.augment_single_image(
+                                img_np, polygons, labels, original_areas, 
+                                original_data, config['aug_type'], aug_param
+                            )
+                            
+                            # Create visualization
+                            aug_image_vis = self.create_visualization(aug_image, aug_data)
+                            
+                            # Store result
+                            result_data = {
+                                'image': aug_image_vis,
+                                'json_data': aug_data,
+                                'metadata': {
+                                    'original_image_index': pair_idx,
+                                    'augmentation_index': aug_idx,
+                                    'augmentation_type': config['aug_type'],
+                                    'parameter_value': aug_param,
+                                    'parameter_range': config['param_range'],
+                                    'timestamp': datetime.now().isoformat(),
+                                    'filename': f'aug_{pair_idx}_{config["aug_type"]}_{aug_idx}.png'
+                                }
+                            }
+                            
+                            self.augmented_results.append(result_data)
+                            results.append(aug_image_vis)
+                            
+                        except Exception as e:
+                            logger.error(f"Error augmenting image {pair_idx} with {config['aug_type']}: {str(e)}")
+                            continue
+                            
+            except Exception as e:
+                logger.error(f"Error processing image pair {pair_idx}: {str(e)}")
+                continue
         
+        return results
+
+    def create_visualization(self, aug_image, aug_data):
+        """Create visualization with colored polygons and masks"""
         # Create a dynamic color map for unique labels
         unique_labels = list(set(shape['label'] for shape in aug_data['shapes']))
         if not unique_labels:
@@ -434,20 +466,15 @@ 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 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))  # Fallback to green
+            color = label_color_map.get(label, (0, 255, 0))
             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 semi-transparent mask
@@ -455,84 +482,254 @@ def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param:
             cv2.fillPoly(mask, [points], 1)
             colored_mask = np.zeros_like(aug_image_rgb)
             colored_mask[mask == 1] = color
-            alpha = 0.3  # Transparency for mask
+            alpha = 0.3
             overlay = cv2.addWeighted(overlay, 1.0, colored_mask, alpha, 0.0)
             
             # Draw polygon outline
             cv2.polylines(overlay, [points], isClosed=True, color=color, thickness=2)
         
-        # Convert overlay back to PIL
-        aug_image_pil = Image.fromarray(overlay)
-        
-        # 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)}"
+        return Image.fromarray(overlay)
 
-# Define augmentation types and parameter ranges
-aug_options = {
-    "rotate": {"display_name": "Rotate", "param_name": "Angle (degrees)", "range": (-30, 30), "default": 0},
-    "horizontal_flip": {"display_name": "Horizontal Flip", "param_name": "Apply Flip (0 or 1)", "range": (0, 1), "default": 0},
-    "vertical_flip": {"display_name": "Vertical Flip", "param_name": "Apply Flip (0 or 1)", "range": (0, 1), "default": 0},
-    "scale": {"display_name": "Scale", "param_name": "Scale Factor", "range": (0.5, 1.5), "default": 1.0},
-    "brightness_contrast": {"display_name": "Brightness/Contrast", "param_name": "Brightness/Contrast Limit", "range": (-0.3, 0.3), "default": 0},
-    "pixel_dropout": {"display_name": "Pixel Dropout", "param_name": "Dropout Probability", "range": (0.01, 0.1), "default": 0.05}
-}
+    def create_download_package(self):
+        """Create a zip file with all augmented images and JSON files"""
+        if not self.augmented_results:
+            return None
+            
+        zip_buffer = io.BytesIO()
+        
+        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
+            # Add all augmented images and their JSON files
+            for idx, result in enumerate(self.augmented_results):
+                # Save image
+                img_buffer = io.BytesIO()
+                result['image'].save(img_buffer, format='PNG')
+                img_filename = result['metadata']['filename']
+                zip_file.writestr(img_filename, img_buffer.getvalue())
+                
+                # Save JSON data
+                json_filename = f"metadata_{idx}.json"
+                json_str = json.dumps(result['json_data'], indent=2)
+                zip_file.writestr(json_filename, json_str)
+            
+            # Add summary metadata
+            summary = {
+                'total_augmentations': len(self.augmented_results),
+                'generation_timestamp': datetime.now().isoformat(),
+                'augmentation_summary': [result['metadata'] for result in self.augmented_results]
+            }
+            zip_file.writestr('augmentation_summary.json', json.dumps(summary, indent=2))
+        
+        zip_buffer.seek(0)
+        return zip_buffer.getvalue()
 
 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 over the augmented image.")
+    augmenter = PolygonAugmentation(tolerance=2.0, area_threshold=0.01, debug=True)
+    
+    def process_batch_augmentation(
+        images, json_files, num_augmentations,
+        rotate_enabled, rotate_min, rotate_max,
+        hflip_enabled, vflip_enabled,
+        scale_enabled, scale_min, scale_max,
+        brightness_enabled, brightness_min, brightness_max,
+        dropout_enabled, dropout_min, dropout_max
+    ):
+        if not images or not json_files:
+            return [], "No images or JSON files uploaded", None
+        
+        # Pair images with JSON files
+        image_json_pairs = []
+        min_length = min(len(images), len(json_files))
+        
+        for i in range(min_length):
+            if images[i] is not None and json_files[i] is not None:
+                try:
+                    image = Image.open(images[i].name)
+                    image_json_pairs.append((image, images[i].name))
+                except Exception as e:
+                    logger.error(f"Error loading image {i}: {str(e)}")
+                    continue
+        
+        if not image_json_pairs:
+            return [], "No valid image-JSON pairs found", None
+        
+        # Configure augmentations based on user selections
+        aug_configs = []
+        
+        if rotate_enabled:
+            aug_configs.append({
+                'aug_type': 'rotate',
+                'param_range': (rotate_min, rotate_max)
+            })
+        
+        if hflip_enabled:
+            aug_configs.append({
+                'aug_type': 'horizontal_flip',
+                'param_range': (0, 1)
+            })
+        
+        if vflip_enabled:
+            aug_configs.append({
+                'aug_type': 'vertical_flip',
+                'param_range': (0, 1)
+            })
+        
+        if scale_enabled:
+            aug_configs.append({
+                'aug_type': 'scale',
+                'param_range': (scale_min, scale_max)
+            })
+        
+        if brightness_enabled:
+            aug_configs.append({
+                'aug_type': 'brightness_contrast',
+                'param_range': (brightness_min, brightness_max)
+            })
+        
+        if dropout_enabled:
+            aug_configs.append({
+                'aug_type': 'pixel_dropout',
+                'param_range': (dropout_min, dropout_max)
+            })
+        
+        if not aug_configs:
+            return [], "No augmentation types selected", None
+        
+        # Process augmentations
+        try:
+            augmented_images = augmenter.batch_augment_images(
+                image_json_pairs, aug_configs, num_augmentations
+            )
+            
+            # Create JSON summary
+            json_summary = json.dumps([result['metadata'] for result in augmenter.augmented_results], indent=2)
+            
+            status = f"Generated {len(augmented_images)} augmented images from {len(image_json_pairs)} input pairs"
+            return augmented_images, json_summary, status
+            
+        except Exception as e:
+            logger.error(f"Batch augmentation error: {str(e)}")
+            return [], f"Error: {str(e)}", None
+    
+    def download_package():
+        return augmenter.create_download_package()
+    
+    def show_mask_overlay(evt: gr.SelectData):
+        if evt.index < len(augmenter.augmented_results):
+            return augmenter.augmented_results[evt.index]['image']
+        return None
+
+    with gr.Blocks(title="Dynamic Donut Polygon Augmentation") as demo:
+        gr.Markdown("# 🌀 Dynamic Donut Polygon Augmentation Tool")
+        gr.Markdown("Upload multiple images and JSON files to apply batch augmentation with configurable parameter ranges")
         
         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=[v["display_name"] for v in aug_options.values()],
-                    label="Augmentation Type",
-                    value=aug_options["rotate"]["display_name"]
+            with gr.Column(scale=1):
+                gr.Markdown("## 📁 Input Files")
+                
+                images_input = gr.File(
+                    file_count="multiple",
+                    file_types=["image"],
+                    label="Upload Images"
                 )
-                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"],
-                    step=0.01
+                
+                json_input = gr.File(
+                    file_count="multiple",
+                    file_types=[".json"],
+                    label="Upload LabelMe JSON Files"
                 )
                 
-                def update_slider(display_name):
-                    aug_key = next(k for k, v in aug_options.items() if v["display_name"] == display_name)
-                    return {
-                        aug_param: gr.update(
-                            minimum=aug_options[aug_key]["range"][0],
-                            maximum=aug_options[aug_key]["range"][1],
-                            value=aug_options[aug_key]["default"],
-                            label=aug_options[aug_key]["param_name"],
-                            step=0.01 if aug_key in ["pixel_dropout", "brightness_contrast", "scale"] else 1
-                        )
-                    }
+                num_augmentations = gr.Slider(
+                    minimum=1, maximum=5, value=2, step=1,
+                    label="Augmentations per configuration"
+                )
+                
+                gr.Markdown("## ⚙️ Augmentation Configuration")
+                
+                # Rotation parameters
+                with gr.Group():
+                    rotate_enabled = gr.Checkbox(label="Enable Rotation", value=True)
+                    with gr.Row():
+                        rotate_min = gr.Slider(-45, 45, -15, label="Min Rotation (degrees)")
+                        rotate_max = gr.Slider(-45, 45, 15, label="Max Rotation (degrees)")
+                
+                # Flip parameters
+                with gr.Group():
+                    hflip_enabled = gr.Checkbox(label="Enable Horizontal Flip", value=True)
+                    vflip_enabled = gr.Checkbox(label="Enable Vertical Flip", value=False)
                 
-                aug_type.change(fn=update_slider, inputs=aug_type, outputs=[aug_param])
+                # Scale parameters
+                with gr.Group():
+                    scale_enabled = gr.Checkbox(label="Enable Scale", value=True)
+                    with gr.Row():
+                        scale_min = gr.Slider(0.7, 1.3, 0.9, label="Min Scale")
+                        scale_max = gr.Slider(0.7, 1.3, 1.1, label="Max Scale")
                 
-                submit_btn = gr.Button("Apply Augmentation")
+                # Brightness parameters
+                with gr.Group():
+                    brightness_enabled = gr.Checkbox(label="Enable Brightness/Contrast", value=True)
+                    with gr.Row():
+                        brightness_min = gr.Slider(-0.3, 0.3, -0.1, label="Min Brightness")
+                        brightness_max = gr.Slider(-0.3, 0.3, 0.1, label="Max Brightness")
+                
+                # Dropout parameters
+                with gr.Group():
+                    dropout_enabled = gr.Checkbox(label="Enable Pixel Dropout", value=False)
+                    with gr.Row():
+                        dropout_min = gr.Slider(0.01, 0.1, 0.02, label="Min Dropout")
+                        dropout_max = gr.Slider(0.01, 0.1, 0.05, label="Max Dropout")
+                
+                generate_btn = gr.Button("🚀 Generate Augmentations", variant="primary")
+                status_text = gr.Textbox(label="Status", interactive=False)
             
-            with gr.Column():
-                output_image = gr.Image(type="pil", label="Augmented Image with Colored Polygons and Masks")
-                output_json = gr.Textbox(label="Augmented LabelMe JSON", lines=10, max_lines=20)
-        
-        def submit(image, json_file, display_name, aug_param):
-            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)
-        
-        submit_btn.click(
-            fn=submit,
-            inputs=[image_input, json_input, aug_type, aug_param],
-            outputs=[output_image, output_json]
+            with gr.Column(scale=2):
+                gr.Markdown("## 🖼️ Augmented Results")
+                gr.Markdown("*Click on any image to view with enhanced mask overlay*")
+                
+                augmented_gallery = gr.Gallery(
+                    label="Augmented Images with Polygon Masks",
+                    show_label=False,
+                    elem_id="gallery",
+                    columns=3,
+                    rows=3,
+                    height="auto"
+                )
+                
+                with gr.Row():
+                    download_btn = gr.Button("📥 Download All (ZIP)", variant="secondary")
+                    download_file = gr.File(label="Download Package", visible=False)
+                
+                gr.Markdown("## 📋 Augmentation Metadata")
+                json_output = gr.Code(
+                    label="Generated Metadata JSON",
+                    language="json",
+                    lines=15
+                )
+                
+                gr.Markdown("## 🎭 Enhanced Preview")
+                mask_preview = gr.Image(label="Selected Image with Mask Overlay")
+        
+        # Event handlers
+        generate_btn.click(
+            process_batch_augmentation,
+            inputs=[
+                images_input, json_input, num_augmentations,
+                rotate_enabled, rotate_min, rotate_max,
+                hflip_enabled, vflip_enabled,
+                scale_enabled, scale_min, scale_max,
+                brightness_enabled, brightness_min, brightness_max,
+                dropout_enabled, dropout_min, dropout_max
+            ],
+            outputs=[augmented_gallery, json_output, status_text]
+        )
+        
+        download_btn.click(
+            download_package,
+            outputs=download_file
+        )
+        
+        augmented_gallery.select(
+            show_mask_overlay,
+            outputs=mask_preview
         )
     
     return demo