Added few more features

app.py

@@ -9,6 +9,7 @@ import supervision as sv
 import uuid
 import random
 from pathlib import Path
+import colorsys
 
 class PolygonAugmentation:
     def __init__(self, tolerance=0.2, area_threshold=0.01, debug=False):
@@ -34,12 +35,10 @@ class PolygonAugmentation:
         else:
             data = json.load(json_file)
         
-        # 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"),
@@ -52,7 +51,7 @@ class PolygonAugmentation:
                 for seg in data['segments']
             ]
         else:
-            raise ValueError("Invalid JSON: Neither 'shapes' nor ' segments' key found or not a list")
+            raise ValueError("Invalid JSON: Neither 'shapes' nor 'segments' key found or not a list")
         
         polygons = []
         labels = []
@@ -156,7 +155,7 @@ class PolygonAugmentation:
                 "group_id": None,
                 "shape_type": "polygon",
                 "flags": {},
-                "confidence": 1.0  # Default confidence for augmented shapes
+                "confidence": 1.0
             })
         
         aug_data = {
@@ -343,7 +342,7 @@ class PolygonAugmentation:
                 contrast_limit=aug_param,
                 p=1.0
             ),
-            "pixel_dropout": A.PixelDropout(dropout_prob=aug_param, p=1.0)
+            "pixel_dropout": A.PixelDropout(dropout_prob=min(max(aug_param, 0.0), 1.0), p=1.0)
         }
         
         if aug_type not in aug_dict:
@@ -372,6 +371,21 @@ class PolygonAugmentation:
 
 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
         img_np = np.array(image)
         img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
@@ -387,13 +401,18 @@ def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param:
             img_np, polygons, labels, original_areas, original_data, aug_type, aug_param
         )
         
-        # Create a color map for unique labels
+        # Create a dynamic color map for unique labels
         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
-        ]
-        label_color_map = {label: colors[i % len(colors)] for i, label in enumerate(unique_labels)}
+        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)]
+            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)
+                label_color_map[label] = rgb
         
         # Convert augmented image to RGB for visualization
         aug_image_rgb = cv2.cvtColor(aug_image, cv2.COLOR_BGR2RGB)
@@ -403,7 +422,7 @@ def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param:
         height, width = aug_image.shape[:2]
         for shape in aug_data['shapes']:
             label = shape['label']
-            color = label_color_map[label]
+            color = label_color_map.get(label, (0, 255, 0))
             points = np.array(shape['points'], dtype=np.int32)
             
             # Draw filled mask with transparency
@@ -411,7 +430,7 @@ 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
             cv2.addWeighted(colored_mask, alpha, overlay, 1 - alpha, 0, overlay)
             
             # Draw polygon outline
@@ -455,7 +474,8 @@ def create_interface():
                     minimum=aug_options["Rotate"]["range"][0],
                     maximum=aug_options["Rotate"]["range"][1],
                     value=aug_options["Rotate"]["default"],
-                    label=aug_options["Rotate"]["param_name"]
+                    label=aug_options["Rotate"]["param_name"],
+                    step=0.01
                 )
                 
                 def update_slider(aug_type):
@@ -464,7 +484,8 @@ def create_interface():
                             minimum=aug_options[aug_type]["range"][0],
                             maximum=aug_options[aug_type]["range"][1],
                             value=aug_options[aug_type]["default"],
-                            label=aug_options[aug_type]["param_name"]
+                            label=aug_options[aug_type]["param_name"],
+                            step=0.01 if aug_type in ["Pixel Dropout", "Brightness/Contrast", "Scale"] else 1
                         )
                     }