working v1

.gradio/flagged/dataset1.csv

@@ -0,0 +1,2 @@
+seismic,Select Task,output,timestamp
+,Fault,,2025-08-20 13:37:37.675220

app.py

@@ -6,8 +6,26 @@ from PIL import Image
 import torch 
 from inference import predict, random_sample
 
-
 def main(): 
-    gr.Interface(fn=predict,
-             inputs=gr.Image(type="pil"),
-             outputs=gr.Image(type="pil"),).launch()
+    with gr.Blocks() as demo:
+        # Button to select task
+        
+        seismic_data = gr.State()
+        gr.Markdown("## SFM Inference Demo")
+        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")
+            prediction_image = gr.Image(label="Prediction Result")
+
+        with gr.Row():
+            random_sample_button = gr.Button("Upload Random Sample", elem_id="random-sample-button")
+            random_sample_button.click(fn=random_sample, inputs=[task], outputs=[seismic_image, seismic_data])
+
+        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])
+
+    demo.launch()

inference.py

@@ -4,9 +4,20 @@ import timm
 from util.datasets import ThebeSet, P3DFaciesSet
 from util.pos_embed import interpolate_pos_embed
 import random
+import huggingface_hub
+from huggingface_hub import hf_hub_download
+from PIL import Image
+import numpy as np
+from matplotlib import cm 
 
 
-def predict(seis: torch.Tensor, finetune: str, task='Fault', model_type='vit_large_patch16', device = 'cpu', thresh = 0.5): 
+HFACE_FAULTS = "checkpoint-24.pth"
+HFACE_FACIES = "checkpoint-49.pth"
+
+FAULT_DATA_PATH = "C:\\Users\\abhalekar\\Desktop\\DATASETS\\Thebe_DATASET\\crossline_combined_data"
+FACIES_DATA_PATH = "C:\\Users\\abhalekar\\Desktop\\DATASETS\\P3D_Vol_DATASET"
+
+def predict(seismic: torch.Tensor, task='Fault', model_type='vit_large_patch16', device = 'cpu', hface = True, thresh = 0.5): 
     if task == 'Fault':
         model = models_Fault.__dict__[model_type](
             img_size=768,
@@ -14,19 +25,21 @@ def predict(seis: torch.Tensor, finetune: str, task='Fault', model_type='vit_lar
             drop_path_rate=0.1,
             in_chans=1,
         )
+        checkpoint_path = hf_hub_download(repo_id="Ani24/SFM_Finetuned", filename=HFACE_FAULTS, subfolder="ckpts-Tversky-Neut")
     elif task == 'Facies':
         model = models_Facies.__dict__[model_type](
-            img_size=768,
+            img_size=128,
             num_classes=6,
             drop_path_rate=0.1,
             in_chans=1,
         )
+        checkpoint_path = hf_hub_download(repo_id="Ani24/SFM_Finetuned", filename=HFACE_FACIES, subfolder="ckpts-RSVSFacies-P3D")
     else:
         raise ValueError(f"Task not configured yet: {task}")
 
     model.to(device)
-    checkpoint = torch.load(finetune, map_location=device, weights_only=True)
-    print("Load pre-trained checkpoint from: %s" % finetune)
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+
     checkpoint_model = checkpoint['model']
     state_dict = model.state_dict()
 
@@ -40,10 +53,10 @@ def predict(seis: torch.Tensor, finetune: str, task='Fault', model_type='vit_lar
     print(msg)
 
 
-    print("Seismic data shape:", seis.shape)
+    print("Seismic data shape:", seismic.shape)
 
     with torch.no_grad():
-        output = model(seis.unsqueeze(0))
+        output = model(seismic.unsqueeze(0))
     output = output.squeeze(0)
     
     if task in ['Fault']: 
@@ -55,19 +68,27 @@ def predict(seis: torch.Tensor, finetune: str, task='Fault', model_type='vit_lar
         output = output.detach().cpu().numpy()
 
     print("Model output shape:", output.shape)
-
+    output = 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
 
 
-def random_sample(data_path: str, task = 'Fault', batch_size=1, num_workers=0):
+def random_sample(task = 'Fault', data_path = None, batch_size=1, num_workers=0):
     if task == 'Fault':
+        data_path = FAULT_DATA_PATH
         dataset = ThebeSet(data_path, [768, 768], 'test')
     elif task == 'Facies':
+        data_path = FACIES_DATA_PATH
         dataset = P3DFaciesSet(data_path, mode = 'train')
     else:
         raise ValueError(f"Task not configured yet: {task}")
     
     index = random.randint(0, len(dataset) - 1)
     seis, label = dataset[index]
+    
+    seis_image = seis.detach().cpu().numpy().squeeze(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, label
+    return seis_image, seis

test.ipynb

@@ -2,13 +2,52 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "aaadf81b",
    "metadata": {},
-   "outputs": [],
+   "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"
+    }
+   ],
    "source": [
-    "import app\n"
+    "import app\n",
+    "\n",
+    "app.main()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5f4991c2",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -18,7 +57,15 @@
    "name": "python3"
   },
   "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
    "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
    "version": "3.11.9"
   }
  },