precomputed all rotations

.gitignore

@@ -1,2 +1 @@
-flagged/
-*.ckpt
+app_copy.py

app.py

@@ -3,95 +3,12 @@ import os
 import gradio as gr
 import matplotlib.pyplot as plt
 import numpy as np
-import pandas as pd 
 import skimage
-from mediffusion import DiffusionModule
 import monai as mn
 import torch
 
 from io_utils import LoadImageD
 
-# Loading the model for inference
-
-model = DiffusionModule("./diffusion_configs.yaml")
-model.load_ckpt("./data/model.ckpt")
-model.eval();
-
-# Loading a baseline noise for making predictions
-
-seed = 3407
-np.random.seed(seed)
-torch.random.manual_seed(seed)
-torch.backends.cudnn.deterministic = True
-BASELINE_NOISE = torch.randn(1, 1, 256, 256).half()
-
-# Model helper functions
-
-def create_ds(img_paths):
-    if type(img_paths) == str:
-        img_paths = [img_paths]
-    data_list = [{"img": img_path} for img_path in img_paths]
-        
-    # Get the transforms            
-    Ts_list = [
-            LoadImageD(keys=["img"], transpose=True, normalize=True),
-            mn.transforms.EnsureChannelFirstD(
-                keys=["img"], channel_dim="no_channel"
-            ),
-            mn.transforms.ResizeD(
-                keys=["img"],
-                spatial_size=(256, 256),
-                mode=["bicubic"],
-            ),
-            mn.transforms.ScaleIntensityD(keys=["img"], minv=0, maxv=1),
-            mn.transforms.ToTensorD(keys=["img"], track_meta=None),
-            mn.transforms.SelectItemsD(keys=["img"]),
-        ]
-    return mn.data.Dataset(data_list, transform=mn.transforms.Compose(Ts_list))
-
-def make_predictions(img_path, angles=None, cls_batch=None, rotate_to_standard=False, sampler="DDIM5"):
-    
-    global model
-    global BASELINE_NOISE
-    
-    # Create the image dataset
-    if cls_batch is not None:
-        ds = create_ds([img_path]*len(cls_batch))
-    else:
-        ds = create_ds(img_path)
-    dl = mn.data.DataLoader(ds, batch_size=len(ds), num_workers=0 if len(ds)==1 else 4, shuffle=False)
-    input_batch = next(iter(dl))
-    original_imgs = input_batch["img"].detach().cpu().numpy()
-    
-    # Create the classifier condition if not provided
-    if cls_batch is None:
-        fp = torch.zeros(768)
-        if rotate_to_standard or angles is None:
-            angles = [1000, 1000, 1000]
-            cls_value = torch.tensor([2, *angles, *fp])
-        else:
-            cls_value = torch.tensor([1, *angles, *fp])
-        cls_batch = cls_value.unsqueeze(0).repeat(input_batch["img"].shape[0], 1)
-
-    # Generate noise
-    noise = BASELINE_NOISE.repeat(input_batch["img"].shape[0], 1, 1, 1)
-    model_kwargs = {
-        "cls": cls_batch,
-        "concat": input_batch["img"]
-    }
-    
-    # Make predictions
-    preds = model.predict(
-        noise, model_kwargs=model_kwargs, classifier_cond_scale=4, inference_protocol=sampler 
-    )
-    adjusted_preds = list()
-    for pred, original_img in zip(preds, original_imgs):
-        adjusted_pred = pred.detach().cpu().numpy().squeeze()
-        original_img = original_img.squeeze()
-        adjusted_pred = skimage.exposure.match_histograms(adjusted_pred, original_img)
-        adjusted_preds.append(adjusted_pred)
-    return adjusted_preds
-
 # Gradio helper functions
 
 current_img = None
@@ -101,65 +18,52 @@ def rotate_btn_fn(img_path, xt, yt, zt, add_bone_cmap=False):
     
     global current_img
     
-    angles = [float(xt), float(yt), float(zt)]
-    out_img = make_predictions(img_path, angles)[0]
+    angles = (xt, yt, zt)
+    out_img_path = f'data/cached_outputs/{os.path.basename(img_path)[:-4]}_{angles}.png'
+    out_img = skimage.io.imread(out_img_path)
     if not add_bone_cmap:
-        print(out_img.shape)
         return out_img
     cmap = plt.get_cmap('bone')
     out_img = cmap(out_img)
     out_img = (out_img[..., :3] * 255).astype(np.uint8)
     current_img = out_img
     return out_img
-
-def use_current_btn_fn(input_img):
-    return input_img
-
-def retrieve_examples(examples, inputs):
-    global current_img
-    if current_img is not None:
-        return current_img
-    return examples[0]
     
 css_style = "./style.css"
 callback = gr.CSVLogger()
 with gr.Blocks(css=css_style) as app:
     gr.HTML("VCNet: A tool for 3D Rotation of Radiographs with Diffusion Models", elem_classes="title")
-    gr.HTML("Developed by: Pouria Rouzrokh, Bardia Khosravi, Shahriar Faghani, Kellen Mulford, Michael J. Taunton, Bradley J. Erickson, Cody C. Wyles", elem_classes="note")
-    gr.HTML("Note: This is a proof-of-concept demo of an AI tool that is not yet finalized. Please interpret with care!", elem_classes="note")
+    gr.HTML("Developed by: Pouria Rouzrokh, Bardia Khosravi, Shahriar Faghani, Kellen Mulford, Michael J. Taunton, Bradley J. Erickson, Cody C. Wyles", elem_classes="subtitle")
+    gr.HTML("Note: This is a proof-of-concept demo running on CPU. All predictions are pre-computed.", elem_classes="note")
     
-    with gr.TabItem("Single Rotation"):
+    with gr.TabItem("Demo"):
         with gr.Row():
             input_img = gr.Image(type='filepath', label='Input image', sources='upload', interactive=False, elem_classes='imgs')
             output_img = gr.Image(type='pil', label='Output image', interactive=False, elem_classes='imgs')
         with gr.Row():
-            gr.Examples(
-                examples = [os.path.join("./data/examples", f) for f in os.listdir("./data/examples") if "xr" in f], 
-                inputs = [input_img],
-                label = "Xray Examples",
-                elem_id='examples',
-            )
-            gr.Examples(
-                examples = [os.path.join("./data/examples", f) for f in os.listdir("./data/examples") if "drr" in f], 
-                inputs = [input_img],
-                label = "DRR Examples",
-                elem_id='examples',
-            )
+            with gr.Column(scale=0.25):
+                pass
+            with gr.Column(scale=1):
+                gr.Examples(
+                    examples = [os.path.join("./data/examples", f) for f in os.listdir("./data/examples") if "xr" in f], 
+                    inputs = [input_img],
+                    label = "Xray Examples",
+                    elem_id='examples',
+                )
+            with gr.Column(scale=0.25):
+                pass
         with gr.Row():
             gr.Markdown('Please select an example image, choose your rotation angles, and press Rotate!', elem_classes='text')
         with gr.Row():
             with gr.Column(scale=1):
-                xt = gr.Slider(label='Rotation angle in x axis:', elem_classes='angle', value=0, minimum=-20, maximum=20, step=1)
+                xt = gr.Slider(label='Rotation angle in x axis:', elem_classes='angle', value=0, minimum=-15, maximum=15, step=5)
             with gr.Column(scale=1):
-                yt = gr.Slider(label='Rotation angle in y axis:', elem_classes='angle', value=0, minimum=-20, maximum=20, step=1)
+                yt = gr.Slider(label='Rotation angle in y axis:', elem_classes='angle', value=0, minimum=-15, maximum=15, step=5)
             with gr.Column(scale=1):
-                zt = gr.Slider(label='Rotation angle in z axis:', elem_classes='angle', value=0, minimum=-20, maximum=20, step=1)
+                zt = gr.Slider(label='Rotation angle in z axis:', elem_classes='angle', value=0, minimum=-15, maximum=15, step=5)
         with gr.Row():
             rotate_btn = gr.Button("Rotate!", elem_classes='rotate_button')
-        with gr.Row():
-            use_current_btn = gr.Button("Use the current output as the new input!", elem_classes='use_current_button')
         rotate_btn.click(fn=rotate_btn_fn, inputs=[input_img, xt, yt, zt], outputs=output_img)
-        use_current_btn.click(fn=use_current_btn_fn, inputs=[output_img], outputs=input_img)
         
 try:
     app.close()