Update app.py

app.py

@@ -41,7 +41,11 @@ class PolygonAugmentation:
         if isinstance(json_file, str):
             with open(json_file, 'r', encoding='utf-8') as f:
                 data = json.load(f)
+        elif isinstance(json_file, dict):
+            # Handle dictionary data directly
+            data = json_file
         else:
+            # Handle file object
             data = json.load(json_file)
         
         shapes = []
@@ -344,8 +348,8 @@ class PolygonAugmentation:
         
         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),
+            "horizontal_flip": A.HorizontalFlip(p=1.0 if aug_param == 1 else 0.0),
+            "vertical_flip": A.VerticalFlip(p=1.0 if aug_param == 1 else 0.0),
             "scale": A.Affine(scale=aug_param, p=1.0),
             "brightness_contrast": A.RandomBrightnessContrast(
                 brightness_limit=aug_param,
@@ -361,16 +365,39 @@ 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)
+        # Convert masks array to list for albumentations
+        masks_list = [masks[i] for i in range(masks.shape[0])]
+        
+        # Create additional targets for each mask
+        additional_targets = {f'mask{i}': 'mask' for i in range(len(masks_list))}
+        
+        # Update transform with additional targets
+        transform = A.Compose([
+            aug_dict[aug_type],
+            A.RandomCrop(width=crop_width, height=crop_height, p=0.8)
+        ], additional_targets=additional_targets)
+        
+        # Prepare input dictionary
+        input_dict = {'image': image}
+        for i, mask in enumerate(masks_list):
+            input_dict[f'mask{i}'] = mask
+        
+        # Apply augmentation
+        aug_result = transform(**input_dict)
         aug_image = aug_result['image']
-        aug_masks = aug_result['masks']
+        
+        # Collect augmented masks
+        aug_masks_list = []
+        for i in range(len(masks_list)):
+            aug_masks_list.append(aug_result[f'mask{i}'])
+        
+        aug_masks = np.array(aug_masks_list, dtype=np.uint8)
         
         # Validate augmented image
         if aug_image is None or aug_image.size == 0:
@@ -387,23 +414,28 @@ class PolygonAugmentation:
 
     def batch_augment_images(self, image_json_pairs, aug_configs, num_augmentations):
         """Batch process multiple images with multiple augmentation configurations"""
+        logger.info(f"Starting batch augmentation with {len(image_json_pairs)} pairs, {len(aug_configs)} configs, {num_augmentations} augmentations each")
         self.augmented_results = []
         results = []
         
-        for pair_idx, (image, json_file) in enumerate(image_json_pairs):
-            if image is None or json_file is None:
+        for pair_idx, (image, json_data) in enumerate(image_json_pairs):
+            if image is None or json_data is None:
+                logger.warning(f"Skipping pair {pair_idx}: missing image or JSON data")
                 continue
                 
             try:
+                logger.info(f"Processing image pair {pair_idx}")
                 # 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)
+                # Load data - pass the JSON data directly
+                img_np, polygons, labels, original_areas, original_data, _ = self.load_labelme_data(json_data, img_np)
+                logger.info(f"Loaded {len(polygons)} polygons for image {pair_idx}")
                 
                 # Apply each augmentation configuration
-                for config in aug_configs:
+                for config_idx, config in enumerate(aug_configs):
+                    logger.info(f"Applying config {config_idx}: {config['aug_type']}")
                     for aug_idx in range(num_augmentations):
                         # Generate random parameter within range
                         min_val, max_val = config['param_range']
@@ -413,6 +445,7 @@ class PolygonAugmentation:
                             aug_param = random.uniform(min_val, max_val)
                         
                         try:
+                            logger.info(f"Generating augmentation {aug_idx} with {config['aug_type']}, param: {aug_param}")
                             aug_image, aug_data = self.augment_single_image(
                                 img_np, polygons, labels, original_areas, 
                                 original_data, config['aug_type'], aug_param
@@ -438,15 +471,21 @@ class PolygonAugmentation:
                             
                             self.augmented_results.append(result_data)
                             results.append(aug_image_vis)
+                            logger.info(f"Successfully generated augmentation {aug_idx} for image {pair_idx}")
                             
                         except Exception as e:
                             logger.error(f"Error augmenting image {pair_idx} with {config['aug_type']}: {str(e)}")
+                            import traceback
+                            logger.error(traceback.format_exc())
                             continue
                             
             except Exception as e:
                 logger.error(f"Error processing image pair {pair_idx}: {str(e)}")
+                import traceback
+                logger.error(traceback.format_exc())
                 continue
         
+        logger.info(f"Batch augmentation completed. Generated {len(results)} total results.")
         return results
 
     def create_visualization(self, aug_image, aug_data):
@@ -544,9 +583,17 @@ def create_interface():
             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))
+                    # Load JSON file content properly
+                    json_path = json_files[i].name
+                    logger.info(f"Loading JSON from: {json_path}")
+                    with open(json_path, 'r', encoding='utf-8') as f:
+                        json_data = json.load(f)
+                    logger.info(f"Successfully loaded JSON with keys: {list(json_data.keys())}")
+                    image_json_pairs.append((image, json_data))
                 except Exception as e:
-                    logger.error(f"Error loading image {i}: {str(e)}")
+                    logger.error(f"Error loading image/JSON pair {i}: {str(e)}")
+                    import traceback
+                    logger.error(traceback.format_exc())
                     continue
         
         if not image_json_pairs:
@@ -596,6 +643,7 @@ def create_interface():
         
         # Process augmentations
         try:
+            logger.info(f"Starting batch augmentation with {len(image_json_pairs)} image pairs and {len(aug_configs)} configurations")
             augmented_images = augmenter.batch_augment_images(
                 image_json_pairs, aug_configs, num_augmentations
             )
@@ -604,11 +652,15 @@ def create_interface():
             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"
+            logger.info(status)
             return augmented_images, json_summary, status
             
         except Exception as e:
-            logger.error(f"Batch augmentation error: {str(e)}")
-            return [], f"Error: {str(e)}", None
+            error_msg = f"Batch augmentation error: {str(e)}"
+            logger.error(error_msg)
+            import traceback
+            logger.error(traceback.format_exc())
+            return [], error_msg, None
     
     def download_package():
         return augmenter.create_download_package()