added LoadImageD from osail-utils

app.py

@@ -9,6 +9,8 @@ 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")
@@ -25,22 +27,6 @@ BASELINE_NOISE = torch.randn(1, 1, 256, 256).half()
 
 # Model helper functions
 
-class LoadImageD(mn.transforms.Transform):
-    def __init__(self, keys, transpose=False, normalize=False):
-        self.keys = keys
-        self.transpose = transpose
-        self.normalize = normalize
-    def __call__(self, data):
-        for key in self.keys:
-            img = skimage.io.imread(data[key])
-            if self.transpose:
-                img = img.transpose(0, 1)
-            if self.normalize:
-                img -= img.min()
-                img /= (img.max()+1e-6)
-            data[key] = img
-        return data
-
 def create_ds(img_paths):
     if type(img_paths) == str:
         img_paths = [img_paths]
@@ -125,58 +111,15 @@ def rotate_btn_fn(img_path, xt, yt, zt, add_bone_cmap=False):
     out_img = (out_img[..., :3] * 255).astype(np.uint8)
     current_img = out_img
     return out_img
-    
-def rotate_to_standard_btn_fn(img_path, add_bone_cmap=False):
-    
-    global current_img
-    
-    out_img = make_predictions(img_path, rotate_to_standard=True)[0]
-    if not add_bone_cmap:
-        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 make_live_btn_fn(img_path, axis, add_bone_cmap=False):
-    
-    global live_preds
-    
-    base_angles = list(range(-20, 21, 1))
-    base_angles = [float(i) for i in base_angles]
-    if axis.lower() == "axis x":
-        all_angles = [[i, 0, 0] for i in base_angles]
-    elif axis.lower() == "axis y":
-        all_angles = [[0, i, 0] for i in base_angles]
-    elif axis.lower() == "axis z":
-        all_angles = [[0, 0, i] for i in base_angles]
-    fp = torch.zeros(768)
-    cls_batch = torch.tensor([[1, *angles, *fp] for angles in all_angles])
-    
-    live_preds = make_predictions(img_path, cls_batch=cls_batch)
-    live_preds = {angle: live_preds[i] for i, angle in enumerate(base_angles)}
-    return img_path    
-    
-def rotate_live_img_fn(angle, add_bone_cmap=False):
-    
-    global live_img
-    global live_preds
-    
-    if live_img is not None:
-        if angle == 0:
-            return live_img
-        return live_preds[float(angle)]
     
 css_style = "./style.css"
 callback = gr.CSVLogger()
 with gr.Blocks(css=css_style) as app:
-    gr.HTML("VCNet: A Deep Learning Solution for Roating RadioGraphs in 3D Space", elem_classes="title")
-    gr.HTML("Developed by the Orthopedics Surgery Artificial Intelligence Lab (OSAIL)", elem_classes="note")
+    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")
     
     with gr.TabItem("Single Rotation"):
@@ -207,41 +150,10 @@ with gr.Blocks(css=css_style) as app:
                 zt = gr.Slider(label='Rotation angle in z axis:', elem_classes='angle', value=0, minimum=-20, maximum=20, step=1)
         with gr.Row():
             rotate_btn = gr.Button("Rotate!", elem_classes='rotate_button')
-        with gr.Row():
-            rotate_to_standard_btn = gr.Button("Rotate to standard view!", elem_classes='rotate_to_standard_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)
-        rotate_to_standard_btn.click(fn=rotate_to_standard_btn_fn, inputs=[input_img], outputs=output_img)
         use_current_btn.click(fn=use_current_btn_fn, inputs=[output_img], outputs=input_img)
-    
-    with gr.TabItem("Live Rotation"):
-        with gr.Row():
-            live_img = gr.Image(type='filepath', label='Live Image', sources='upload', 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 = [live_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 = [live_img],
-                label = "DRR Examples",
-                elem_id='examples'
-            )
-        with gr.Row():
-            gr.Markdown('Please select an example image, an axis, and then press Make Live!', elem_classes='text')
-        with gr.Row():
-            axis = gr.Dropdown(choices=['Axis X', 'Axis Y', 'Axis Z'], show_label=False, elem_classes='angle', value='Axis X')
-            live_btn = gr.Button("Make Live!", elem_classes='make_live_button')
-        with gr.Row():
-             gr.Markdown('You can now rotate the radiograph in your selected axis using the scaler.', elem_classes='text')
-        with gr.Row():
-             slider = gr.Slider(show_label=False, minimum=-20, maximum=20, step=1, value=0, elem_classes='slider', interactive=True)
-        live_btn.click(fn=make_live_btn_fn, inputs=[live_img, axis], outputs=live_img)
-        slider.change(fn=rotate_live_img_fn, inputs=[slider], outputs=live_img)
         
 try:
     app.close()

io_utils.py

@@ -0,0 +1,121 @@
+################################################################################
+# This files contains OSAIL utils to read and write files.
+################################################################################
+
+from .data import pad_to_square
+import copy
+import monai as mn
+import numpy as np
+import os
+import skimage
+
+################################################################################
+# -F: load_image
+
+def load_image(input_object, pad=False, normalize=True, standardize=False, 
+               dtype=np.float32, percentile_clip=None, target_shape=None, 
+               transpose=False, ensure_grayscale=True, LoadImage_args=[], LoadImage_kwargs={}):
+    """A helper function to load different input types.
+
+    Args:
+        input_object (Union[np.ndarray, str]): 
+            a 2D NumPy array of X-ray an image, a DICOM file of an X-ray image, 
+            or a string path to a .npy, any regular image file format 
+            saved on disk that skimage.io can load.
+        pad (bool, optional): whether to pad the image to square shape. 
+            Defaults to True.
+        normalize (bool, optional): whether to normalize the image. 
+            Defaults to True.
+        standardize (bool, optional): whether to standardize the image.
+            Defaults to False.
+        dtype (np.dtype, optional): the data type of the output image. 
+            Defaults to np.float32.
+        percentile_clip (float, optional): the percentile to clip the image. 
+            Defaults to 2.5.
+        target_shape (tuple, optional): the target shape of the output image. 
+            Defaults to None, which means no resizing.
+        transpose (bool, optional): whether to transpose the image.
+            Defaults to False.
+        ensure_grayscale (bool, optional): whether to make the image grayscale.
+            Defaults to True.
+        LoadImg_args: a list of keyword arguments to pass to  mn.transforms.LoadImage.
+        LoadImg_kwargs: a dictionary of keyword arguments to pass to  mn.transforms.LoadImage.
+            
+    Returns:
+        the loaded image array.
+    """
+    # Load the image.
+    if isinstance(input_object, np.ndarray):
+        image = input_object
+    elif isinstance(input_object, str):
+        assert os.path.exists(input_object), f"File not found: {input_object}"
+        reader = mn.transforms.LoadImage(image_only=True, *LoadImage_args, **LoadImage_kwargs)
+        image = reader(input_object)
+
+    # Make the image 2D.
+    if ensure_grayscale:
+        if image.shape[-1] == 3:
+            image = np.mean(image, axis=-1)  
+        elif image.shape[0] == 3:
+            image = np.mean(image, axis=0)
+        elif image.shape[-1] == 4:
+            image = np.mean(image[...,:3], axis=-1)  
+        elif image.shape[0] == 4:
+            image = np.mean(image[:3,...], axis=0)  
+        assert len(image.shape) == 2, f"Image must be 2D: {image.shape}"
+    
+    # Transpose the image.
+    if transpose:
+        image = np.transpose(image, axes=(1,0))
+    
+    # Clip the image.
+    if percentile_clip is not None:
+        percentile_low = np.percentile(image, percentile_clip)
+        percentile_high = np.percentile(image, 100-percentile_clip)
+        image = np.clip(image, percentile_low, percentile_high)
+        
+    # Standardize the image.
+    if standardize:
+        image = image.astype(np.float32)
+        image -= image.mean()
+        image /= (image.std() + 1e-8)
+        
+    # Normalize the image.
+    if normalize:
+        image = image.astype(np.float32)
+        image -= image.min()
+        image /= (image.max() + 1e-8)
+    
+    # Pad the image to square shape.
+    if pad:
+        image = pad_to_square(image)   
+    
+    # Resize the image.
+    if target_shape is not None:
+        image = skimage.transform.resize(image, target_shape, preserve_range=True)
+        
+    # Cast the image to the target data type.
+    if dtype is np.uint8:
+        image = (image * 255).astype(np.uint8)
+    else:
+        image = image.astype(dtype)  
+    
+    return image
+
+################################################################################
+# -C: LoadImageD
+
+class LoadImageD(mn.transforms.Transform):
+    """A MONAI transform to load input image using load_image function.
+    """
+    def __init__(self, keys, *to_pass_keys, **to_pass_kwargs) -> None:
+        super().__init__()
+        self.keys = keys
+        self.to_pass_keys = to_pass_keys
+        self.to_pass_kwargs = to_pass_kwargs
+        
+    def __call__(self, data):
+        data_copy = copy.deepcopy(data)
+        for key in self.keys:
+            data_copy[key] = load_image(data[key], *self.to_pass_keys, **self.to_pass_kwargs)
+        return data_copy