Update app.py

app.py

@@ -8,9 +8,8 @@ import numpy as np
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 st.title("Image blur using segmentation and depth estimation.")
-st.markdown("""Upload your image and choose a blur style.
-
-            **Note**: The lens blur option takes a long time to process (>5min) since this space isn't linked to a GPU.""")
+st.info("Upload your image and choose a blur style")
+st.warning("**Note**: The lens blur option takes a long time to process (>5min) since this space isn't linked to a GPU.")
 
 @st.cache_resource(show_spinner="Pushing pixels...")
 def load_gblur_model():
@@ -28,7 +27,7 @@ def load_lblur_model():
 gblur_model = load_gblur_model()
 lblur_model, lblur_img_proc = load_lblur_model()
 
-def gaussian_blur(image):   
+def gaussian_blur(image, blur_str):   
     # Image transform
     image_size = (512, 512)
     transform_image = transforms.Compose([
@@ -46,13 +45,13 @@ def gaussian_blur(image):
     mask = np.array(pred_pil.resize(image.size))
 
     # Blurring
-    blur = np.array(image.filter(ImageFilter.GaussianBlur(radius=15)))
+    blur = np.array(image.filter(ImageFilter.GaussianBlur(radius=blur_str)))
     mask = np.expand_dims(mask, axis=2)
     output_image = np.where(mask, np.array(image), blur)
     
     return Image.fromarray(output_image)
 
-def lens_blur(image):
+def lens_blur(image, blur_str):
     # Process image
     inputs = lblur_img_proc(images=image, return_tensors="pt").to(device)
 
@@ -75,7 +74,7 @@ def lens_blur(image):
     
     # No of discrete blurs and max blur intensity
     num_levels = 15
-    max_radius = 30.0
+    max_radius = blur_str
     
     # Pre-compute all blur images
     blurred_images = []
@@ -127,18 +126,21 @@ if up_img:
     image = Image.open(up_img).convert("RGB")
     
     gblur, lblur = st.columns(2)
+    blur_str = st.slider(
+        "Blur Strength", value=(5, 50)
+    )
 
     disp_left, disp_right = st.columns(2)
     with disp_left:
         og_img = image.copy()
         st.image(og_img, caption="Original Image",  use_container_width=True)
     with disp_right:
-        if gblur.button("Gaussian Blur", use_container_width="True"):
+        if gblur.button("Gaussian Blur", use_container_width=True):
             with st.spinner(f"Spinning violently around the y-axis..."):
-                result = gaussian_blur(image)
-        if lblur.button("Lens Blur", use_container_width="True"):
+                result = gaussian_blur(image, blur_str)
+        if lblur.button("Lens Blur", use_container_width=True):
             with st.spinner(f"One mississippi, two mississippi..."):
-                result = lens_blur(image)
+                result = lens_blur(image, blur_str)
         st.image(result, "Blurred Image", use_container_width=True)
     # apply_func(image, options)