File size: 4,865 Bytes
d3a15f1 |
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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
import os
import h5py
import numpy as np
import gradio as gr
import plotly.graph_objects as go
from railnet_model import RailNetSystem
from huggingface_hub import hf_hub_download
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
model = RailNetSystem.from_pretrained("Tournesol-Saturday/railNet-tooth-segmentation-in-CBCT-image").cuda()
model.load_weights(from_hub=True, repo_id="Tournesol-Saturday/railNet-tooth-segmentation-in-CBCT-image")
def render_plotly_volume(pred, x_eye=1.25, y_eye=1.25, z_eye=1.25):
downsample_factor = 2
pred_ds = pred[::downsample_factor, ::downsample_factor, ::downsample_factor]
fig = go.Figure(data=go.Volume(
x=np.repeat(np.arange(pred_ds.shape[0]), pred_ds.shape[1] * pred_ds.shape[2]),
y=np.tile(np.repeat(np.arange(pred_ds.shape[1]), pred_ds.shape[2]), pred_ds.shape[0]),
z=np.tile(np.arange(pred_ds.shape[2]), pred_ds.shape[0] * pred_ds.shape[1]),
value=pred_ds.flatten(),
isomin=0.5,
isomax=1.0,
opacity=0.1,
surface_count=1,
colorscale=[[0, 'rgb(255, 0, 0)'], [1, 'rgb(255, 0, 0)']],
showscale=False
))
fig.update_layout(
scene=dict(
xaxis=dict(visible=False),
yaxis=dict(visible=False),
zaxis=dict(visible=False),
camera=dict(eye=dict(x=x_eye, y=y_eye, z=z_eye))
),
margin=dict(l=0, r=0, b=0, t=0)
)
return fig
def handle_example(filename):
repo_id = "Tournesol-Saturday/railNet-tooth-segmentation-in-CBCT-image"
h5_path = hf_hub_download(repo_id=repo_id, filename=f"example_input_file/{filename}")
with h5py.File(h5_path, "r") as f:
image = f["image"][:]
label = f["label"][:]
name = filename.replace(".h5", "")
pred, dice, jc, hd, asd = model(image, label, "./output", name)
fig = render_plotly_volume(pred)
img_path = f"./output/{name}_img.nii.gz"
pred_path = f"./output/{name}_pred.nii.gz"
metrics = f"Dice: {dice:.4f}, Jaccard: {jc:.4f}, 95HD: {hd:.2f}, ASD: {asd:.2f}"
return metrics, pred, fig, img_path, pred_path
def clear_all():
return "", None, None, None, None
with gr.Blocks() as demo:
gr.HTML("<div style='text-align: center; font-size: 22px; font-weight: bold;'>🦷 Demo of RailNet: A CBCT Tooth Segmentation System</div>")
gr.HTML("<div style='text-align: center; font-size: 15px'>✅ Steps: Select a CBCT example file (.h5) → Automatic inference and metrics display → View 3D segmentation result (Mouse drag and scroll wheel zooming)</div>")
gr.HTML("""
<style>
.code-style {
font-family: monospace;
background-color: #2f363d;
color: #ffffff;
padding: 2px 6px;
border-radius: 4px;
font-size: 90%;
}
</style>
<div style='font-size: 15px; font-weight: bold;'>
📂 Step 1: Select a <span class='code-style'>.h5</span> example file from the <span class='code-style'>example_input_file</span> folder in our
<a href='https://huggingface.co/Tournesol-Saturday/railNet-tooth-segmentation-in-CBCT-image' target='_blank' style='text-decoration: none; color: #1f6feb; font-weight: bold;'>
Hugging Face model
</a> repository.
</div>
""")
example_files = ["CBCT_01.h5", "CBCT_02.h5", "CBCT_03.h5", "CBCT_04.h5"]
dropdown = gr.Dropdown(choices=example_files, label="Example File", value=example_files[0])
with gr.Row():
clear_btn = gr.Button("清除", variant="secondary")
submit_btn = gr.Button("提交", variant="primary")
gr.HTML("<div style='font-size: 15px; font-weight: bold;'>📊 Step 2: Metrics (Dice, Jaccard, 95HD, ASD)</div>")
result_text = gr.Textbox()
hidden_pred = gr.State(value=None)
gr.HTML("<div style='font-size: 15px; font-weight: bold;'>👁️ Step 3: 3D Visualisation</div>")
plot_output = gr.Plot()
gr.HTML("<div style='font-size: 15px; font-weight: bold;'>⬇️ Step 4: Download <span class='code-style'>NIfTI</span> files for accurate 1:1 visualization using <span class='code-style'>ITK-SNAP</span> software</div>")
with gr.Row():
hidden_img_file = gr.File(label="Download Original Image", interactive=False)
hidden_pred_file = gr.File(label="Download Segmentation Result", interactive=False)
submit_btn.click(
fn=handle_example,
inputs=[dropdown],
outputs=[result_text, hidden_pred, plot_output, hidden_img_file, hidden_pred_file]
)
clear_btn.click(
fn=clear_all,
inputs=[],
outputs=[result_text, hidden_pred, plot_output, hidden_img_file, hidden_pred_file]
)
demo.launch()
|