Update app.py

app.py

@@ -5,7 +5,8 @@ import albumentations as A
 import random
 
 def apply_augmentations(image, flip_h, flip_v, rotate, crop, gray, scale,
-                       prob_flip_h, prob_flip_v, prob_rotate, prob_crop, prob_gray, prob_scale):
+                       prob_flip_h, prob_flip_v, prob_rotate, prob_crop, prob_gray, prob_scale,
+                       rotation_limit):
     augmentations = []
     
     if flip_h:
@@ -13,7 +14,7 @@ def apply_augmentations(image, flip_h, flip_v, rotate, crop, gray, scale,
     if flip_v:
         augmentations.append(A.VerticalFlip(p=float(prob_flip_v)))
     if rotate:
-        augmentations.append(A.Rotate(limit=90, p=float(prob_rotate)))
+        augmentations.append(A.Rotate(limit=(-rotation_limit, rotation_limit), p=float(prob_rotate)))
     if crop:
         augmentations.append(A.RandomResizedCrop(
             size=(image.shape[0], image.shape[1]),
@@ -89,6 +90,7 @@ def main():
             with gr.Column():
                 rotate = gr.Checkbox(label="Rotate")
                 prob_rotate = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")
+                rotation_limit = gr.Slider(minimum=0, maximum=90, value=10, label="Rotation Limit (±degrees)")
             
             with gr.Column():
                 crop = gr.Checkbox(label="Random Crop")
@@ -110,7 +112,8 @@ def main():
             inputs=[
                 input_image,
                 flip_h, flip_v, rotate, crop, gray, scale,
-                prob_flip_h, prob_flip_v, prob_rotate, prob_crop, prob_gray, prob_scale
+                prob_flip_h, prob_flip_v, prob_rotate, prob_crop, prob_gray, prob_scale,
+                rotation_limit
             ],
             outputs=output_image
         )