Added few more features

app.py

@@ -33,26 +33,44 @@ class PolygonAugmentation:
                 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")
+        
+        # Check for 'shapes' (LabelMe) or 'segments' (custom format)
+        shapes = []
+        if 'shapes' in data and isinstance(data['shapes'], list):
+            shapes = data['shapes']
+        elif 'segments' in data and isinstance(data['segments'], list):
+            # Convert custom 'segments' to LabelMe 'shapes' format
+            shapes = [
+                {
+                    "label": seg.get("class", "unknown"),
+                    "points": seg.get("polygon", []),
+                    "shape_type": "polygon",
+                    "group_id": None,
+                    "flags": {},
+                    "confidence": seg.get("confidence", 1.0)
+                }
+                for seg in data['segments']
+            ]
+        else:
+            raise ValueError("Invalid JSON: Neither 'shapes' nor ' segments' key found or not a list")
         
         polygons = []
         labels = []
         original_areas = []
         
-        for shape in data['shapes']:
+        for shape in 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")
+                    print(f"[DEBUG] Skipping invalid shape: {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):
+            except (ValueError, TypeError) as e:
                 if self.debug:
-                    print(f"[DEBUG] Error processing points: {shape['points']}")
+                    print(f"[DEBUG] Error processing points: {shape['points']}, error: {str(e)}")
                 continue   
         
         if not polygons and self.debug:
@@ -137,7 +155,8 @@ class PolygonAugmentation:
                 "points": poly,
                 "group_id": None,
                 "shape_type": "polygon",
-                "flags": {}
+                "flags": {},
+                "confidence": 1.0  # Default confidence for augmented shapes
             })
         
         aug_data = {
@@ -369,7 +388,7 @@ def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param:
         )
         
         # Create a color map for unique labels
-        unique_labels = list(set(aug_data['shapes'][i]['label'] for i in range(len(aug_data['shapes']))))
+        unique_labels = list(set(shape['label'] for shape in aug_data['shapes']))
         colors = [
             (255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255),
             (128, 0, 0), (0, 128, 0), (0, 0, 128), (128, 128, 0)  # Add more if needed
@@ -382,10 +401,10 @@ def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param:
         
         # Create masks for visualization
         height, width = aug_image.shape[:2]
-        for poly, shape in zip(aug_data['shapes'], aug_data['shapes']):
+        for shape in aug_data['shapes']:
             label = shape['label']
             color = label_color_map[label]
-            points = np.array(poly['points'], dtype=np.int32)
+            points = np.array(shape['points'], dtype=np.int32)
             
             # Draw filled mask with transparency
             mask = np.zeros((height, width), dtype=np.uint8)