Update app.py

app.py

@@ -176,30 +176,63 @@ class PolygonAugmentation:
         for poly, label in zip(aug_polygons, aug_labels):
             if not poly or len(poly) < 3:
                 continue
-            new_shapes.append({
+            
+            # Create LabelMe format shape
+            shape_data = {
                 "label": label,
                 "points": poly,
                 "group_id": None,
                 "shape_type": "polygon",
                 "flags": {},
-                "confidence": 1.0
-            })
+                "description": "",
+                "attributes": {},
+                "iscrowd": 0,
+                "difficult": 0
+            }
+            
+            # Add additional metadata for special polygon types
+            if label.lower() in ['ring', 'donut', 'annulus', 'circle', 'round']:
+                shape_data["attributes"]["polygon_type"] = "ring"
+            elif label.lower() in ['background', 'bg', 'back']:
+                shape_data["attributes"]["polygon_type"] = "background"
+            else:
+                shape_data["attributes"]["polygon_type"] = "object"
+            
+            new_shapes.append(shape_data)
         
         # Get actual dimensions from augmented image
         aug_height, aug_width = aug_image.shape[:2]
         
+        # Create LabelMe compatible JSON structure
         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_height,  # Use actual augmented image height
-            "imageWidth": aug_width     # Use actual augmented image width
+            "imageData": None,  # Explicitly set to None as requested
+            "imageHeight": aug_height,
+            "imageWidth": aug_width,
+            "imageDepth": 3 if len(aug_image.shape) == 3 else 1,
+            
+            # Additional LabelMe metadata
+            "lineColor": [0, 255, 0, 128],
+            "fillColor": [255, 0, 0, 128],
+            "textSize": 10,
+            "textColor": [0, 0, 0, 255],
+            
+            # Augmentation metadata
+            "augmentation": {
+                "augmented": True,
+                "augmentation_id": aug_id,
+                "original_file": original_data.get("imagePath", "unknown"),
+                "augmentation_timestamp": datetime.now().isoformat(),
+                "augmentation_tool": "PolygonAugmentation v1.0"
+            }
         }
         
         if self.debug:
-            logger.info(f"[DEBUG] Created augmented data: {len(new_shapes)} shapes, size: {aug_width}x{aug_height}")
+            logger.info(f"[DEBUG] Created LabelMe JSON: {len(new_shapes)} shapes, size: {aug_width}x{aug_height}")
+            logger.info(f"[DEBUG] Shape types: {[s['attributes'].get('polygon_type', 'unknown') for s in new_shapes]}")
         
         return aug_data
 
@@ -574,26 +607,90 @@ class PolygonAugmentation:
         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
+        """Create visualization with colored polygon masks and outlines for each class"""
+        # Create a dynamic color map for unique labels with better color distribution
         unique_labels = list(set(shape['label'] for shape in aug_data['shapes']))
         if not unique_labels:
             label_color_map = {"unknown": (0, 255, 0)}
         else:
             num_labels = len(unique_labels)
-            hues = [i / num_labels for i in range(num_labels)]
+            # Create more distinct colors using different hue ranges
             label_color_map = {}
-            for label, hue in zip(unique_labels, hues):
-                rgb = colorsys.hsv_to_rgb(hue, 1.0, 1.0)
-                rgb = tuple(int(c * 255) for c in rgb)
+            for i, label in enumerate(unique_labels):
+                if label.lower() in ['background', 'bg', 'back']:
+                    # Background gets a neutral gray-blue color
+                    rgb = (100, 149, 237)  # Cornflower blue with low opacity
+                elif 'ring' in label.lower() or 'donut' in label.lower():
+                    # Ring/donut shapes get purple-pink colors
+                    hue = 0.8 + (i * 0.1) % 0.2  # Purple range
+                    rgb = colorsys.hsv_to_rgb(hue, 0.8, 0.9)
+                    rgb = tuple(int(c * 255) for c in rgb)
+                else:
+                    # Regular objects get distributed colors across the spectrum
+                    hue = (i * 0.618033988749895) % 1.0  # Golden ratio for better distribution
+                    saturation = 0.7 + (i % 3) * 0.1  # Vary saturation
+                    value = 0.8 + (i % 2) * 0.15  # Vary brightness
+                    rgb = colorsys.hsv_to_rgb(hue, saturation, value)
+                    rgb = tuple(int(c * 255) for c in rgb)
+                
                 label_color_map[label] = rgb
         
         # Convert augmented image to RGB for visualization
         aug_image_rgb = cv2.cvtColor(aug_image, cv2.COLOR_BGR2RGB)
         overlay = aug_image_rgb.copy()
-        
-        # Create masks and outlines for visualization
         height, width = aug_image.shape[:2]
+        
+        # Create a composite mask to handle overlapping polygons
+        composite_mask = np.zeros((height, width, 3), dtype=np.uint8)
+        
+        # Group shapes by label for better visualization
+        shapes_by_label = {}
+        for shape in aug_data['shapes']:
+            label = shape['label']
+            if label not in shapes_by_label:
+                shapes_by_label[label] = []
+            shapes_by_label[label].append(shape)
+        
+        # Process each label group
+        for label, shapes in shapes_by_label.items():
+            color = label_color_map.get(label, (0, 255, 0))
+            
+            # Create mask for all polygons of this label
+            label_mask = np.zeros((height, width), dtype=np.uint8)
+            
+            for shape in shapes:
+                points = np.array(shape['points'], dtype=np.int32)
+                if len(points) < 3:
+                    continue
+                
+                # Fill the polygon area
+                cv2.fillPoly(label_mask, [points], 255)
+            
+            # Apply color to the mask areas
+            if label_mask.sum() > 0:  # Only if mask has content
+                # Determine alpha based on label type
+                if label.lower() in ['background', 'bg', 'back']:
+                    alpha = 0.15  # Lower opacity for background
+                elif 'ring' in label.lower() or 'donut' in label.lower():
+                    alpha = 0.4   # Medium opacity for rings
+                else:
+                    alpha = 0.35  # Standard opacity for objects
+                
+                # Create colored mask
+                colored_mask = np.zeros_like(aug_image_rgb)
+                colored_mask[label_mask == 255] = color
+                
+                # Blend with overlay
+                mask_area = label_mask == 255
+                overlay[mask_area] = cv2.addWeighted(
+                    overlay[mask_area], 
+                    1.0 - alpha, 
+                    colored_mask[mask_area], 
+                    alpha, 
+                    0
+                )
+        
+        # Draw polygon outlines with thicker lines for better visibility
         for shape in aug_data['shapes']:
             label = shape['label']
             color = label_color_map.get(label, (0, 255, 0))
@@ -601,49 +698,161 @@ class PolygonAugmentation:
             if len(points) < 3:
                 continue
             
-            # 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
-            overlay = cv2.addWeighted(overlay, 1.0, colored_mask, alpha, 0.0)
+            # Determine line thickness based on polygon type
+            if label.lower() in ['background', 'bg', 'back']:
+                thickness = 1  # Thinner lines for background
+            elif 'ring' in label.lower() or 'donut' in label.lower():
+                thickness = 3  # Thicker lines for rings to show structure
+            else:
+                thickness = 2  # Standard thickness
             
             # Draw polygon outline
-            cv2.polylines(overlay, [points], isClosed=True, color=color, thickness=2)
+            cv2.polylines(overlay, [points], isClosed=True, color=color, thickness=thickness)
+            
+            # Add label text near the polygon
+            if len(points) > 0:
+                # Find a good position for the label
+                moments = cv2.moments(points)
+                if moments['m00'] != 0:
+                    cx = int(moments['m10'] / moments['m00'])
+                    cy = int(moments['m01'] / moments['m00'])
+                else:
+                    cx, cy = points[0][0], points[0][1]
+                
+                # Ensure text position is within image bounds
+                cx = max(10, min(cx, width - 50))
+                cy = max(20, min(cy, height - 10))
+                
+                # Add text background for better readability
+                font = cv2.FONT_HERSHEY_SIMPLEX
+                font_scale = 0.4
+                text_thickness = 1
+                text_size = cv2.getTextSize(label, font, font_scale, text_thickness)[0]
+                
+                # Draw background rectangle
+                cv2.rectangle(overlay, 
+                            (cx - 2, cy - text_size[1] - 4), 
+                            (cx + text_size[0] + 2, cy + 2), 
+                            (0, 0, 0), -1)
+                
+                # Draw text
+                cv2.putText(overlay, label, (cx, cy - 2), font, font_scale, color, text_thickness)
+        
+        if self.debug:
+            logger.info(f"[DEBUG] Created visualization with {len(unique_labels)} unique labels: {list(unique_labels)}")
         
         return Image.fromarray(overlay)
 
     def create_download_package(self):
-        """Create a zip file with all augmented images and JSON files"""
+        """Create a zip file with all augmented images and proper LabelMe 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
+            # Add all augmented images and their corresponding LabelMe JSON files
             for idx, result in enumerate(self.augmented_results):
-                # Save image
+                # Save augmented 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)
+                # Convert PIL image back to OpenCV format for saving
+                img_cv = np.array(result['image'])
+                img_bgr = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)
+                
+                # Encode image as PNG
+                success, encoded_img = cv2.imencode('.png', img_bgr)
+                if success:
+                    zip_file.writestr(result['metadata']['filename'], encoded_img.tobytes())
+                else:
+                    # Fallback to PIL saving
+                    result['image'].save(img_buffer, format='PNG')
+                    zip_file.writestr(result['metadata']['filename'], img_buffer.getvalue())
+                
+                # Save corresponding LabelMe JSON file
+                json_filename = result['metadata']['filename'].replace('.png', '.json')
+                
+                # Create a clean LabelMe JSON without internal metadata
+                clean_json_data = {
+                    "version": result['json_data'].get("version", "5.0.1"),
+                    "flags": result['json_data'].get("flags", {}),
+                    "shapes": result['json_data']['shapes'],
+                    "imagePath": result['metadata']['filename'],
+                    "imageData": None,
+                    "imageHeight": result['json_data']['imageHeight'],
+                    "imageWidth": result['json_data']['imageWidth']
+                }
+                
+                json_str = json.dumps(clean_json_data, indent=2, ensure_ascii=False)
                 zip_file.writestr(json_filename, json_str)
             
-            # Add summary metadata
+            # Add comprehensive summary metadata
             summary = {
-                'total_augmentations': len(self.augmented_results),
-                'generation_timestamp': datetime.now().isoformat(),
-                'augmentation_summary': [result['metadata'] for result in self.augmented_results]
+                'package_info': {
+                    'total_augmentations': len(self.augmented_results),
+                    'generation_timestamp': datetime.now().isoformat(),
+                    'generator': 'PolygonAugmentation v1.0',
+                    'format': 'LabelMe JSON + PNG images'
+                },
+                'augmentation_summary': [
+                    {
+                        'filename': result['metadata']['filename'],
+                        'json_file': result['metadata']['filename'].replace('.png', '.json'),
+                        'augmentation_type': result['metadata']['augmentation_type'],
+                        'parameter_value': result['metadata']['parameter_value'],
+                        'polygon_count': len(result['json_data']['shapes']),
+                        'image_size': f"{result['json_data']['imageWidth']}x{result['json_data']['imageHeight']}",
+                        'timestamp': result['metadata']['timestamp'],
+                        'labels': list(set([shape['label'] for shape in result['json_data']['shapes']]))
+                    }
+                    for result in self.augmented_results
+                ],
+                'statistics': {
+                    'unique_augmentation_types': list(set([r['metadata']['augmentation_type'] for r in self.augmented_results])),
+                    'total_polygons': sum([len(r['json_data']['shapes']) for r in self.augmented_results]),
+                    'unique_labels': list(set([
+                        shape['label'] 
+                        for result in self.augmented_results 
+                        for shape in result['json_data']['shapes']
+                    ])),
+                    'average_polygons_per_image': sum([len(r['json_data']['shapes']) for r in self.augmented_results]) / len(self.augmented_results)
+                }
             }
-            zip_file.writestr('augmentation_summary.json', json.dumps(summary, indent=2))
+            zip_file.writestr('augmentation_summary.json', json.dumps(summary, indent=2, ensure_ascii=False))
+            
+            # Add README for the package
+            readme_content = f"""# Augmented Dataset Package
+
+## Overview
+This package contains {len(self.augmented_results)} augmented images with their corresponding LabelMe annotation files.
+
+## Contents
+- **Images**: PNG format augmented images
+- **Annotations**: LabelMe JSON format annotation files
+- **Summary**: augmentation_summary.json with detailed metadata
+
+## File Structure
+- Each image file (*.png) has a corresponding annotation file (*.json) with the same base name
+- All annotations are in standard LabelMe format without embedded image data
+
+## Statistics
+- Total augmented images: {len(self.augmented_results)}
+- Total polygons: {sum([len(r['json_data']['shapes']) for r in self.augmented_results])}
+- Unique labels: {list(set([shape['label'] for result in self.augmented_results for shape in result['json_data']['shapes']]))}
+- Augmentation types used: {list(set([r['metadata']['augmentation_type'] for r in self.augmented_results]))}
+
+## Usage
+1. Extract the ZIP file
+2. Load images and annotations using any tool that supports LabelMe format
+3. Use the augmentation_summary.json for batch processing or analysis
+
+Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
+Tool: PolygonAugmentation v1.0
+"""
+            zip_file.writestr('README.md', readme_content)
         
         zip_buffer.seek(0)
+        logger.info(f"Created download package with {len(self.augmented_results)} image-annotation pairs")
         return zip_buffer.getvalue()
 
 def create_interface():