Added few more features

app.py

@@ -373,12 +373,12 @@ def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param:
     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)
+            "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}")
@@ -448,12 +448,12 @@ def augment_image(image: Image.Image, json_file: Any, aug_type: str, aug_param:
 
 # Define augmentation types and parameter ranges
 aug_options = {
-    "Rotate": {"param_name": "Angle (degrees)", "range": (-30, 30), "default": 0},
-    "Horizontal Flip": {"param_name": "Apply Flip (0 or 1)", "range": (0, 1), "default": 0},
-    "Vertical Flip": {"param_name": "Apply Flip (0 or 1)", "range": (0, 1), "default": 0},
-    "Scale": {"param_name": "Scale Factor", "range": (0.5, 1.5), "default": 1.0},
-    "Brightness/Contrast": {"param_name": "Brightness/Contrast Limit", "range": (-0.3, 0.3), "default": 0},
-    "Pixel Dropout": {"param_name": "Dropout Probability", "range": (0.01, 0.1), "default": 0.05}
+    "rotate": {"display_name": "Rotate", "param_name": "Angle (degrees)", "range": (-30, 30), "default": 0},
+    "horizontal_flip": {"display_name": "Horizontal Flip", "param_name": "Apply Flip (0 or 1)", "range": (0, 1), "default": 0},
+    "vertical_flip": {"display_name": "Vertical Flip", "param_name": "Apply Flip (0 or 1)", "range": (0, 1), "default": 0},
+    "scale": {"display_name": "Scale", "param_name": "Scale Factor", "range": (0.5, 1.5), "default": 1.0},
+    "brightness_contrast": {"display_name": "Brightness/Contrast", "param_name": "Brightness/Contrast Limit", "range": (-0.3, 0.3), "default": 0},
+    "pixel_dropout": {"display_name": "Pixel Dropout", "param_name": "Dropout Probability", "range": (0.01, 0.1), "default": 0.05}
 }
 
 def create_interface():
@@ -466,26 +466,28 @@ def create_interface():
                 image_input = gr.Image(type="pil", label="Input Image")
                 json_input = gr.File(label="LabelMe JSON File", file_types=[".json"])
                 aug_type = gr.Dropdown(
-                    choices=list(aug_options.keys()),
+                    choices=[v["display_name"] for v in aug_options.values()],
                     label="Augmentation Type",
-                    value="Rotate"
+                    value=aug_options["rotate"]["display_name"]
                 )
                 aug_param = gr.Slider(
-                    minimum=aug_options["Rotate"]["range"][0],
-                    maximum=aug_options["Rotate"]["range"][1],
-                    value=aug_options["Rotate"]["default"],
-                    label=aug_options["Rotate"]["param_name"],
+                    minimum=aug_options["rotate"]["range"][0],
+                    maximum=aug_options["rotate"]["range"][1],
+                    value=aug_options["rotate"]["default"],
+                    label=aug_options["rotate"]["param_name"],
                     step=0.01
                 )
                 
-                def update_slider(aug_type):
+                def update_slider(display_name):
+                    # Map display_name back to internal key
+                    aug_key = next(k for k, v in aug_options.items() if v["display_name"] == display_name)
                     return {
                         aug_param: gr.update(
-                            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"],
-                            step=0.01 if aug_type in ["Pixel Dropout", "Brightness/Contrast", "Scale"] else 1
+                            minimum=aug_options[aug_key]["range"][0],
+                            maximum=aug_options[aug_key]["range"][1],
+                            value=aug_options[aug_key]["default"],
+                            label=aug_options[aug_key]["param_name"],
+                            step=0.01 if aug_key in ["pixel_dropout", "brightness_contrast", "scale"] else 1
                         )
                     }
                 
@@ -497,8 +499,13 @@ def create_interface():
                 output_image = gr.Image(type="pil", label="Augmented Image with Colored Polygons and Masks")
                 output_json = gr.Textbox(label="Augmented LabelMe JSON", lines=10, max_lines=20)
         
+        def submit(image, json_file, display_name, aug_param):
+            # Map display_name to internal key
+            aug_key = next(k for k, v in aug_options.items() if v["display_name"] == display_name)
+            return augment_image(image, json_file, aug_key, aug_param)
+        
         submit_btn.click(
-            fn=augment_image,
+            fn=submit,
             inputs=[image_input, json_input, aug_type, aug_param],
             outputs=[output_image, output_json]
         )