Delete app.py
Browse files
app.py
DELETED
@@ -1,135 +0,0 @@
|
|
1 |
-
import os
|
2 |
-
import h5py
|
3 |
-
import numpy as np
|
4 |
-
import gradio as gr
|
5 |
-
import plotly.graph_objects as go
|
6 |
-
from railnet_model import RailNetSystem
|
7 |
-
|
8 |
-
from huggingface_hub import hf_hub_download
|
9 |
-
|
10 |
-
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
|
11 |
-
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
12 |
-
|
13 |
-
model = RailNetSystem.from_pretrained("Tournesol-Saturday/railNet-tooth-segmentation-in-CBCT-image").cuda()
|
14 |
-
|
15 |
-
model.load_weights(from_hub=True, repo_id="Tournesol-Saturday/railNet-tooth-segmentation-in-CBCT-image")
|
16 |
-
|
17 |
-
def render_plotly_volume(pred, x_eye=1.25, y_eye=1.25, z_eye=1.25):
|
18 |
-
downsample_factor = 2
|
19 |
-
pred_ds = pred[::downsample_factor, ::downsample_factor, ::downsample_factor]
|
20 |
-
|
21 |
-
fig = go.Figure(data=go.Volume(
|
22 |
-
x=np.repeat(np.arange(pred_ds.shape[0]), pred_ds.shape[1] * pred_ds.shape[2]),
|
23 |
-
y=np.tile(np.repeat(np.arange(pred_ds.shape[1]), pred_ds.shape[2]), pred_ds.shape[0]),
|
24 |
-
z=np.tile(np.arange(pred_ds.shape[2]), pred_ds.shape[0] * pred_ds.shape[1]),
|
25 |
-
value=pred_ds.flatten(),
|
26 |
-
isomin=0.5,
|
27 |
-
isomax=1.0,
|
28 |
-
opacity=0.1,
|
29 |
-
surface_count=1,
|
30 |
-
colorscale=[[0, 'rgb(255, 0, 0)'], [1, 'rgb(255, 0, 0)']],
|
31 |
-
showscale=False
|
32 |
-
))
|
33 |
-
|
34 |
-
fig.update_layout(
|
35 |
-
scene=dict(
|
36 |
-
xaxis=dict(visible=False),
|
37 |
-
yaxis=dict(visible=False),
|
38 |
-
zaxis=dict(visible=False),
|
39 |
-
camera=dict(eye=dict(x=x_eye, y=y_eye, z=z_eye))
|
40 |
-
),
|
41 |
-
margin=dict(l=0, r=0, b=0, t=0)
|
42 |
-
)
|
43 |
-
return fig
|
44 |
-
|
45 |
-
def handle_example(filename):
|
46 |
-
repo_id = "Tournesol-Saturday/railNet-tooth-segmentation-in-CBCT-image"
|
47 |
-
h5_path = hf_hub_download(repo_id=repo_id, filename=f"example_input_file/{filename}")
|
48 |
-
|
49 |
-
with h5py.File(h5_path, "r") as f:
|
50 |
-
image = f["image"][:]
|
51 |
-
label = f["label"][:]
|
52 |
-
|
53 |
-
name = filename.replace(".h5", "")
|
54 |
-
pred, dice, jc, hd, asd = model(image, label, "./output", name)
|
55 |
-
|
56 |
-
fig = render_plotly_volume(pred)
|
57 |
-
|
58 |
-
img_path = f"./output/{name}_img.nii.gz"
|
59 |
-
pred_path = f"./output/{name}_pred.nii.gz"
|
60 |
-
|
61 |
-
metrics = f"Dice: {dice:.4f}, Jaccard: {jc:.4f}, 95HD: {hd:.2f}, ASD: {asd:.2f}"
|
62 |
-
|
63 |
-
return metrics, pred, fig, img_path, pred_path
|
64 |
-
|
65 |
-
def clear_all():
|
66 |
-
return "", None, None, None, None
|
67 |
-
|
68 |
-
with gr.Blocks() as demo:
|
69 |
-
gr.HTML("<div style='text-align: center; font-size: 22px; font-weight: bold;'>🦷 Demo of RailNet: A CBCT Tooth Segmentation System</div>")
|
70 |
-
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>")
|
71 |
-
|
72 |
-
gr.HTML("""
|
73 |
-
<style>
|
74 |
-
.code-style {
|
75 |
-
font-family: monospace;
|
76 |
-
background-color: #2f363d;
|
77 |
-
color: #ffffff;
|
78 |
-
padding: 2px 6px;
|
79 |
-
border-radius: 4px;
|
80 |
-
font-size: 90%;
|
81 |
-
}
|
82 |
-
</style>
|
83 |
-
|
84 |
-
<div style='font-size: 15px; font-weight: bold;'>
|
85 |
-
📂 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
|
86 |
-
<a href='https://huggingface.co/Tournesol-Saturday/railNet-tooth-segmentation-in-CBCT-image' target='_blank' style='text-decoration: none; color: #1f6feb; font-weight: bold;'>
|
87 |
-
Hugging Face model
|
88 |
-
</a> repository.
|
89 |
-
</div>
|
90 |
-
""")
|
91 |
-
|
92 |
-
example_files = ["CBCT_01.h5", "CBCT_02.h5", "CBCT_03.h5", "CBCT_04.h5"]
|
93 |
-
dropdown = gr.Dropdown(choices=example_files, label="Example File", value=example_files[0])
|
94 |
-
|
95 |
-
|
96 |
-
with gr.Row():
|
97 |
-
clear_btn = gr.Button("清除", variant="secondary")
|
98 |
-
submit_btn = gr.Button("提交", variant="primary")
|
99 |
-
|
100 |
-
gr.HTML("<div style='font-size: 15px; font-weight: bold;'>📊 Step 2: Metrics (Dice, Jaccard, 95HD, ASD)</div>")
|
101 |
-
result_text = gr.Textbox()
|
102 |
-
hidden_pred = gr.State(value=None)
|
103 |
-
|
104 |
-
gr.HTML("<div style='font-size: 15px; font-weight: bold;'>👁️ Step 3: 3D Visualisation</div>")
|
105 |
-
plot_output = gr.Plot()
|
106 |
-
|
107 |
-
# hidden_img_file = gr.File(visible=False)
|
108 |
-
# hidden_pred_file = gr.File(visible=False)
|
109 |
-
|
110 |
-
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>")
|
111 |
-
with gr.Row():
|
112 |
-
hidden_img_file = gr.File(label="Download Original Image", interactive=False)
|
113 |
-
hidden_pred_file = gr.File(label="Download Segmentation Result", interactive=False)
|
114 |
-
|
115 |
-
submit_btn.click(
|
116 |
-
fn=handle_example,
|
117 |
-
inputs=[dropdown],
|
118 |
-
outputs=[result_text, hidden_pred, plot_output, hidden_img_file, hidden_pred_file]
|
119 |
-
)
|
120 |
-
|
121 |
-
# def update_view(pred, x_eye, y_eye, z_eye):
|
122 |
-
# if pred is None:
|
123 |
-
# return gr.update()
|
124 |
-
# return render_plotly_volume(pred, x_eye, y_eye, z_eye)
|
125 |
-
|
126 |
-
clear_btn.click(
|
127 |
-
fn=clear_all,
|
128 |
-
inputs=[],
|
129 |
-
outputs=[result_text, hidden_pred, plot_output, hidden_img_file, hidden_pred_file]
|
130 |
-
)
|
131 |
-
|
132 |
-
# download_img_btn.click(fn=lambda f: f, inputs=[hidden_img_file], outputs=[hidden_img_file])
|
133 |
-
# download_pred_btn.click(fn=lambda f: f, inputs=[hidden_pred_file], outputs=[hidden_pred_file])
|
134 |
-
|
135 |
-
demo.launch()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|