Spaces:
Sleeping
Sleeping
File size: 4,221 Bytes
f161c8a 5a72667 eb51f25 5a72667 eb51f25 0d88515 5a72667 16cf656 0d88515 eb51f25 0d88515 16cf656 eb51f25 16cf656 5a72667 16cf656 273ebc0 5a72667 aee3bc1 1d9018e 5a72667 f80665b 5a72667 0d88515 5a72667 01e170d 5a72667 1d9018e 5a72667 1d9018e 5a72667 1d9018e 5a72667 16cf656 5a72667 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
# Archive paper: https://arxiv.org/abs/2404.13000
import os
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import skimage
from skimage import io, transform
from io_utils import LoadImageD
# Gradio helper functions
current_img = None
live_preds = None
def rotate_img(image, angles):
"""Rotate the input image based on the given angles."""
xt, yt, zt = angles
# Apply rotations (this is just a placeholder, replace with your actual logic)
rotated_img = skimage.transform.rotate(image, angle=xt, mode='edge')
rotated_img = skimage.transform.rotate(rotated_img, angle=yt, mode='edge')
rotated_img = skimage.transform.rotate(rotated_img, angle=zt, mode='edge')
return rotated_img
def rotate_btn_fn(img, xt, yt, zt, add_bone_cmap=False):
global current_img
try:
angles = (xt, yt, zt)
print(f"Rotating with angles: {angles}")
if isinstance(img, np.ndarray):
input_img_path = "uploaded_image.png"
skimage.io.imsave(input_img_path, img)
elif isinstance(img, str) and os.path.exists(img):
input_img_path = img
img = skimage.io.imread(input_img_path)
else:
raise ValueError("Invalid input image")
# Apply the rotation
out_img = rotate_img(img, angles)
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
except Exception as e:
print(f"Error in rotate_btn_fn: {e}")
return None
css_style = "./style.css"
callback = gr.CSVLogger()
with gr.Blocks(css=css_style, title="RadRotator") as app:
gr.HTML("RadRotator: 3D Rotation of Radiographs with Diffusion Models", elem_classes="title")
gr.HTML("Developed by:<br>Pouria Rouzrokh, Bardia Khosravi, Shahriar Faghani, Kellen Mulford, Michael J. Taunton, Bradley J. Erickson, Cody C. Wyles<br><a href='https://pouriarouzrokh.github.io/RadRotator'>[Our website]</a>, <a href='https://arxiv.org/abs/2404.13000'>[arXiv Paper]</a>", elem_classes="note")
gr.HTML("Note: The demo operates on a CPU, and since diffusion models require more computational capacity to function, all predictions are precomputed.", elem_classes="note")
with gr.TabItem("Demo"):
with gr.Row():
input_img = gr.Image(type='numpy', label='Input image', interactive=True, elem_classes='imgs')
output_img = gr.Image(type='numpy', label='Output image', interactive=False, elem_classes='imgs')
with gr.Row():
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='x axis (medial/lateral rotation):', elem_classes='angle', value=0, minimum=-15, maximum=15, step=5)
with gr.Column(scale=1):
yt = gr.Slider(label='y axis (inlet/outlet rotation):', elem_classes='angle', value=0, minimum=-15, maximum=15, step=5)
with gr.Column(scale=1):
zt = gr.Slider(label='z axis (plane rotation):', elem_classes='angle', value=0, minimum=-15, maximum=15, step=5)
with gr.Row():
rotate_btn = gr.Button("Rotate!", elem_classes='rotate_button')
rotate_btn.click(fn=rotate_btn_fn, inputs=[input_img, xt, yt, zt], outputs=output_img)
try:
app.close()
gr.close_all()
except Exception as e:
print(f"Error closing app: {e}")
demo = app.launch(
max_threads=4,
share=True,
inline=False,
show_api=False,
show_error=False,
) |