Tournesol-Saturday commited on
Commit
56f16b8
·
verified ·
1 Parent(s): 704e0dd

Delete app.py

Browse files
Files changed (1) hide show
  1. app.py +0 -135
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()