Update app.py

app.py

@@ -422,7 +422,8 @@ class PolygonAugmentation:
         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)
+            # Use angle directly (not abs) and set limit as tuple for specific angle
+            aug_transform = A.Rotate(limit=(rotation_angle, 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)
@@ -746,94 +747,124 @@ class PolygonAugmentation:
     def create_download_package(self):
         """Create a zip file with all augmented images and proper LabelMe JSON files"""
         if not self.augmented_results:
+            logger.warning("No augmented results available for download")
             return None
             
+        logger.info(f"Creating download package with {len(self.augmented_results)} results")
         zip_buffer = io.BytesIO()
         
-        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
-            # Add all augmented images and their corresponding LabelMe JSON files
-            for idx, result in enumerate(self.augmented_results):
-                # Save augmented image
-                img_buffer = io.BytesIO()
-                
-                # 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']
-                }
+        try:
+            with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
+                # Add all augmented images and their corresponding LabelMe JSON files
+                for idx, result in enumerate(self.augmented_results):
+                    filename = result['metadata']['filename']
+                    
+                    # Save augmented image
+                    try:
+                        # Convert PIL image to RGB if needed
+                        if result['image'].mode != 'RGB':
+                            img_rgb = result['image'].convert('RGB')
+                        else:
+                            img_rgb = result['image']
+                        
+                        # Save as PNG bytes
+                        img_buffer = io.BytesIO()
+                        img_rgb.save(img_buffer, format='PNG', optimize=True)
+                        zip_file.writestr(filename, img_buffer.getvalue())
+                        logger.info(f"Added image: {filename}")
+                        
+                    except Exception as e:
+                        logger.error(f"Error saving image {filename}: {str(e)}")
+                        continue
+                    
+                    # Save corresponding LabelMe JSON file
+                    json_filename = filename.replace('.png', '.json')
+                    
+                    try:
+                        # Create a clean LabelMe JSON structure
+                        clean_json_data = {
+                            "version": "5.0.1",
+                            "flags": {},
+                            "shapes": [],
+                            "imagePath": filename,
+                            "imageData": None,  # No embedded image data as requested
+                            "imageHeight": result['json_data']['imageHeight'],
+                            "imageWidth": result['json_data']['imageWidth'],
+                            "imageDepth": 3
+                        }
+                        
+                        # Copy shapes with proper LabelMe format
+                        for shape in result['json_data']['shapes']:
+                            clean_shape = {
+                                "label": shape['label'],
+                                "points": shape['points'],
+                                "group_id": shape.get('group_id'),
+                                "shape_type": "polygon",
+                                "flags": shape.get('flags', {}),
+                                "description": shape.get('description', ''),
+                                "iscrowd": shape.get('iscrowd', 0),
+                                "attributes": shape.get('attributes', {})
+                            }
+                            clean_json_data['shapes'].append(clean_shape)
+                        
+                        # Write JSON file
+                        json_str = json.dumps(clean_json_data, indent=2, ensure_ascii=False)
+                        zip_file.writestr(json_filename, json_str)
+                        logger.info(f"Added JSON: {json_filename} with {len(clean_json_data['shapes'])} shapes")
+                        
+                    except Exception as e:
+                        logger.error(f"Error saving JSON {json_filename}: {str(e)}")
+                        continue
                 
-                json_str = json.dumps(clean_json_data, indent=2, ensure_ascii=False)
-                zip_file.writestr(json_filename, json_str)
-            
-            # Add comprehensive summary metadata
-            summary = {
-                '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']]))
+                # Add comprehensive summary metadata
+                summary = {
+                    '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) if self.augmented_results else 0
                     }
-                    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, ensure_ascii=False))
-            
-            # Add README for the package
-            readme_content = f"""# Augmented Dataset Package
+                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
+- **Annotations**: LabelMe JSON format annotation files (standard format)
 - **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
+- Compatible with LabelMe, CVAT, and other annotation tools
 
 ## Statistics
 - Total augmented images: {len(self.augmented_results)}
@@ -849,11 +880,19 @@ This package contains {len(self.augmented_results)} augmented images with their
 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()
+                zip_file.writestr('README.md', readme_content)
+                
+                logger.info("Successfully created ZIP package with all files")
+            
+            zip_buffer.seek(0)
+            logger.info(f"Created download package with {len(self.augmented_results)} image-annotation pairs")
+            return zip_buffer.getvalue()
+            
+        except Exception as e:
+            logger.error(f"Error creating ZIP package: {str(e)}")
+            import traceback
+            logger.error(traceback.format_exc())
+            return None
 
 def create_interface():
     augmenter = PolygonAugmentation(tolerance=2.0, area_threshold=0.01, debug=True)
@@ -957,7 +996,34 @@ def create_interface():
             return [], error_msg, None
     
     def download_package():
-        return augmenter.create_download_package()
+        """Handle download package creation and return proper file data"""
+        try:
+            package_data = augmenter.create_download_package()
+            if package_data is None:
+                return None
+                
+            # Save the package to a temporary file for download
+            import tempfile
+            import os
+            
+            # Create temporary file with proper name
+            temp_file = tempfile.NamedTemporaryFile(
+                delete=False, 
+                suffix='.zip', 
+                prefix='augmented_dataset_'
+            )
+            
+            with open(temp_file.name, 'wb') as f:
+                f.write(package_data)
+            
+            logger.info(f"Created download package: {temp_file.name}")
+            return temp_file.name
+            
+        except Exception as e:
+            logger.error(f"Error creating download package: {str(e)}")
+            import traceback
+            logger.error(traceback.format_exc())
+            return None
     
     def show_mask_overlay(evt: gr.SelectData):
         if evt.index < len(augmenter.augmented_results):
@@ -1042,7 +1108,7 @@ def create_interface():
                 
                 with gr.Row():
                     download_btn = gr.Button("📥 Download All (ZIP)", variant="secondary")
-                    download_file = gr.File(label="Download Package", visible=False)
+                    download_file = gr.File(label="Download Package", visible=True)
                 
                 gr.Markdown("## 📋 Augmentation Metadata")
                 json_output = gr.Code(