overlays

.env

@@ -1 +0,0 @@
-PYTHONPATH="C:\Users\abhalekar\AppData\Local\Microsoft\WindowsApps\python3.11.exe"

app.py

@@ -4,20 +4,24 @@ import numpy as np
 import requests
 from PIL import Image
 import torch 
-from inference import predict, random_sample
+from inference import predict, random_sample, overlay_images
 
 def main(): 
     with gr.Blocks() as demo:
         # Button to select task
         
         seismic_data = gr.State()
+        prediction_data = gr.State()
+        processed_prediction_data = gr.State()
+
         gr.Markdown("## SFM Inference Demo")
+        gr.Markdown("### Select a task and run inference on seismic data")
         with gr.Row():
             task = gr.Radio(choices=['Fault', 'Facies'], label="Select Task", value='Fault')
         gr.Markdown("### Upload your seismic data or sample from dataset")
 
         with gr.Row():
-            seismic_image = gr.Image(label="Upload Seismic Data")
+            seismic_image = gr.Image(label="Seismic Data")
             prediction_image = gr.Image(label="Prediction Result")
 
         with gr.Row():
@@ -26,6 +30,26 @@ def main():
 
         with gr.Row():
             predict_button = gr.Button("Run Inference", elem_id="predict-button")
-            predict_button.click(fn=predict, inputs=[seismic_data, task], outputs=[prediction_image])
+            predict_button.click(fn=predict, inputs=[seismic_data, task], outputs=[prediction_image, prediction_data])
+            processed_prediction_data = prediction_data
+
+        with gr.Row(): 
+            overlay_image = gr.Image(label="Overlay Result")
+            with gr.Column():
+                gr.Markdown("### Overlay Seismic Data with Prediction Result")
+                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")
+
+                post_process.change(fn=show_slider, inputs=[post_process], outputs=[])
+
+                
+                gr.Button("Download Processed Image", elem_id="download-processed-button")
+
+    demo.launch()
+
 
-    demo.launch()
+if __name__ == "__main__":
+    main()

inference.py

@@ -67,11 +67,11 @@ def predict(seismic: torch.Tensor, task='Fault', model_type='vit_large_patch16',
         output = output.argmax(dim=0)
         output = output.detach().cpu().numpy()
 
-    print("Model output shape:", output.shape)
-    output = output/ output.max()  # Normalize output to [0, 1] range
+    output_image = output/ output.max()  # Normalize output to [0, 1] range
     # output is numpy 2d array - convert to pil RGB image
-    output = Image.fromarray((output * 255).astype(np.uint8)).convert("RGB")
-    return output
+    output_image = Image.fromarray((output_image * 255).astype(np.uint8)).convert("RGB")
+    
+    return output_image, output
 
 
 def random_sample(task = 'Fault', data_path = None, batch_size=1, num_workers=0):
@@ -91,4 +91,44 @@ def random_sample(task = 'Fault', data_path = None, batch_size=1, num_workers=0)
     seis_image = (seis_image - seis_image.min()) / (seis_image.max() - seis_image.min())  # Normalize to [0, 1] range
     seis_image = Image.fromarray(np.uint8(cm.seismic(seis_image) * 255))  # Convert to PIL Image
 
-    return seis_image, seis
+    return seis_image, seis
+
+
+def overlay_images(seismic_image: Image, prediction_image: Image, alpha = 0.5) -> Image:
+    # Create an overlay of the predicted facies/faults on the original seismic image
+    prediction_image = Image.fromarray(np.array(prediction_image).astype(np.uint8)).convert("RGBA")
+    seismic_image = Image.fromarray(np.array(seismic_image).astype(np.uint8)).convert("RGBA")
+    
+    prediction_image.putalpha(int(255 * alpha))  # Set alpha for overlay
+
+    overlay_image = Image.alpha_composite(seismic_image, prediction_image)
+    return overlay_image
+
+
+def post_process(processed_prediction_image: Image, prediction_image: Image, method: str = 'None') -> Image:
+    if method == 'Remove':
+        # Remove the predicted regions from the original image
+        output = prediction_image.copy()
+    elif method == 'Thresholding':
+        # Apply thresholding to the predicted image
+        output = prediction_image.point(lambda p: p > 128 and 255)
+    elif method == 'Closing':
+        # Apply closing operation to the predicted image
+        output = prediction_image.filter(ImageFilter.MaxFilter(size=5))
+    elif method == 'Opening':
+        # Apply opening operation to the predicted image
+        output = prediction_image.filter(ImageFilter.MinFilter(size=5))
+    elif method == 'Canny Edge':
+        # Apply Canny edge detection to the predicted image
+        output = prediction_image.filter(ImageFilter.FIND_EDGES)
+    elif method == 'Gaussian Smoothing':
+        # Apply Gaussian smoothing to the predicted image
+        output = prediction_image.filter(ImageFilter.GaussianBlur(radius=2))
+    elif method == 'Hysteresis':
+        # Apply hysteresis thresholding to the predicted image
+        output = prediction_image.point(lambda p: p > 128 and 255)
+    return output
+
+
+def apply_thresholding(prediction_image: Image) -> Image:
+    return prediction_image.point(lambda p: p > 128 and 255)

test.ipynb

@@ -2,39 +2,10 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "aaadf81b",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "C:\\Users\\abhalekar\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
-      "  from .autonotebook import tqdm as notebook_tqdm\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "* Running on local URL:  http://127.0.0.1:7860\n",
-      "* To create a public link, set `share=True` in `launch()`.\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "<div><iframe src=\"http://127.0.0.1:7860/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
+   "outputs": [],
    "source": [
     "import app\n",
     "\n",