sliders

app.py

@@ -4,9 +4,41 @@ import numpy as np
 import requests
 from PIL import Image
 import torch 
-from inference import predict, random_sample, overlay_images
+from inference import predict, random_sample, overlay_images, post_process
 
 def main(): 
+
+    pp_options = [
+        "None",
+        "Thresholding",
+        "Closing",
+        "Opening",
+        "Canny Edge",
+        "Gaussian Smoothing",
+        "Hysteresis"
+    ]
+
+    def update_slider(post_process):
+        visibility = [0 for option in pp_options]
+        if post_process ==pp_options[0]:
+            # None
+            pass
+        else: 
+            # Retrieve index of post_process
+            assert post_process in pp_options
+            index = pp_options.index(post_process)
+            visibility[index] = 1
+        
+        ret_updates = []
+        for vis in visibility:
+            if vis == 1: 
+                ret_updates.append(gr.update(visible=True))
+            else:
+                ret_updates.append(gr.update(visible=False))
+        
+        return ret_updates
+    
+
     with gr.Blocks() as demo:
         # Button to select task
         
@@ -40,11 +72,23 @@ def main():
                 overlay_button = gr.Button("Overlay Result", elem_id="overlay-button")
                 overlay_button.click(fn=overlay_images, inputs=[seismic_image, prediction_image], outputs=[overlay_image])
                 gr.Markdown("### Post Processing")
-                post_process = gr.Radio(choices=['Remove', 'Thresholding', 'Closing', 'Opening', 'Canny Edge', 'Gaussian Smoothing', 'Hysteresis'], 
-                                        value='None', elem_id="post-processing")
-
+                with gr.Row():
+                    post_process = gr.Radio(choices=pp_options, 
+                                            value='None', elem_id="post-processing", label="Post Processing Method")
+                slider_none = gr.Slider(minimum=0, maximum=255, value=128, label="None Value", visible=False)
+                slider_thresh = gr.Slider(minimum=0, maximum=255, value=128, label="Threshold Value", visible=False)
+                slider_close = gr.Slider(minimum=0, maximum=64, value=32, label="Closing Value", visible=False)
+                slider_open = gr.Slider(minimum=0, maximum=64, value=32, label="Opening Value", visible=False)
+                slider_canny = gr.Slider(minimum=0, maximum=255, value=128, label="Canny Edge Value", visible=False)
+                slider_gauss = gr.Slider(minimum=0, maximum=255, value=128, label="Sigma", visible=False)
+                slider_hyst = gr.Slider(minimum=0, maximum=255, value=128, label="Hysteresis Min Value", visible=False)
+                post_process.change(
+                        fn=update_slider,
+                        inputs=[post_process],
+                        outputs=[slider_none, slider_thresh, slider_close, slider_open, slider_canny, slider_gauss, slider_hyst]
+                    )
                 gr.Button("Download Processed Image", elem_id="download-processed-button")
-
+                
     demo.launch()
 
 

inference.py

@@ -9,7 +9,7 @@ from huggingface_hub import hf_hub_download
 from PIL import Image
 import numpy as np
 from matplotlib import cm 
-
+from PIL import ImageFilter
 
 HFACE_FAULTS = "checkpoint-24.pth"
 HFACE_FACIES = "checkpoint-49.pth"
@@ -105,9 +105,36 @@ def overlay_images(seismic_image: Image, prediction_image: Image, alpha = 0.5) -
     return overlay_image
 
 
-def post_process(processed_prediction_image: Image, prediction_image: Image, method: str = 'None') -> Image:
-    pass
-
-
-def apply_thresholding(prediction_image: Image) -> Image:
-    return prediction_image.point(lambda p: p > 128 and 255)
+def post_process(processed_prediction_image: Image, prediction_image: Image, method: str = 'None', value = None) -> Image:
+    if method == 'None':
+        return processed_prediction_image
+    elif method == 'Thresholding':
+        return apply_thresholding(processed_prediction_image)
+    elif method == 'Closing':
+        return apply_closing(processed_prediction_image, value)
+    elif method == 'Opening':
+        return apply_opening(processed_prediction_image, value)
+    elif method == 'Canny Edge':
+        return apply_canny_edge(processed_prediction_image, value)
+    elif method == 'Gaussian Smoothing':
+        return apply_gaussian_smoothing(processed_prediction_image, value)
+    elif method == 'Hysteresis':
+        return apply_hysteresis(processed_prediction_image, value)
+    else:
+        raise ValueError(f"Unknown post-processing method: {method}")
+
+
+def apply_thresholding(image: Image, value: int) -> Image:
+    return image.point(lambda p: p > value and 255)
+def apply_closing(image: Image, value: int) -> Image:
+    # Apply closing (dilation followed by erosion)
+    return image.filter(ImageFilter.MaxFilter(size=value)).filter(ImageFilter.MinFilter(size=value))
+def apply_opening(image: Image, value: int) -> Image:
+    # Apply opening (erosion followed by dilation)
+    return image.filter(ImageFilter.MinFilter(size=value)).filter(ImageFilter.MaxFilter(size=value))
+def apply_canny_edge(image: Image, value: int) -> Image:
+    return image.filter(ImageFilter.FIND_EDGES)
+def apply_gaussian_smoothing(image: Image, value: int) -> Image:
+    return image.filter(ImageFilter.GaussianBlur(radius=value))
+def apply_hysteresis(image: Image, value: int) -> Image:
+    return image.point(lambda p: p > value and 255)  # Simple thresholding for hysteresis