| import argparse |
| import gradio as gr |
| import numpy as np |
| import cv2 |
| import torch |
|
|
| from model import SimpleHRNet, ViTHeatmap |
| from heatmap_utils import heatmaps_to_coords_dark |
| from secure_torch_load import secure_torch_load |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser(description="Cephalogram landmark inference app") |
| parser.add_argument("--checkpoint", type=str, default="best.pt.enc", help="Path to model checkpoint") |
| parser.add_argument("--device", type=str, default=("cuda" if torch.cuda.is_available() else "cpu"), help="Torch device, e.g. cuda or cpu") |
| parser.add_argument("--server-port", type=int, default=44065, help="Port for Gradio app") |
| parser.add_argument("--server-name", type=str, default="127.0.0.1", help="Host for Gradio app") |
| parser.add_argument("--share", action="store_true", help="Enable public Gradio share link") |
| parser.add_argument("--inbrowser", action="store_true", help="Open app in browser on launch") |
| return parser.parse_args() |
|
|
|
|
| def load_model(checkpoint_path, device): |
| ckpt = secure_torch_load(checkpoint_path, map_location="cpu") |
| |
| args = ckpt["args"] |
| landmark_symbols = ckpt.get("landmark_symbols", None) |
|
|
| if args["model"] == "hrnet": |
| model = SimpleHRNet(num_landmarks=args["num_landmarks"]) |
| else: |
| model = ViTHeatmap( |
| num_landmarks=args["num_landmarks"], |
| model_name=args["vit_name"], |
| pretrained=False, |
| img_size=(args["input_height"], args["input_width"]), |
| ) |
|
|
| model.load_state_dict(ckpt["model_state_dict"]) |
| model.to(device) |
| model.eval() |
|
|
| return model, args, landmark_symbols |
|
|
|
|
| def get_symbols(n, checkpoint_symbols): |
| if checkpoint_symbols is not None and len(checkpoint_symbols) == n: |
| return checkpoint_symbols |
| return [f"LM_{i}" for i in range(n)] |
|
|
|
|
| def preprocess(image, model_args, device): |
| h_orig, w_orig = image.shape[:2] |
| h_in = model_args["input_height"] |
| w_in = model_args["input_width"] |
|
|
| resized = cv2.resize(image, (w_in, h_in)) |
| tensor = torch.from_numpy(resized).permute(2, 0, 1).float() / 255.0 |
| tensor = tensor.unsqueeze(0).to(device) |
| return tensor, (h_orig, w_orig), (h_in, w_in) |
|
|
|
|
| def decode(pred_heatmaps, orig_size, input_size): |
| h_orig, w_orig = orig_size |
| h_in, w_in = input_size |
| h_hm, w_hm = pred_heatmaps.shape[2], pred_heatmaps.shape[3] |
|
|
| coords_hm = heatmaps_to_coords_dark(pred_heatmaps)[0] |
|
|
| coords_in = coords_hm.clone() |
| coords_in[:, 0] *= (w_in / w_hm) |
| coords_in[:, 1] *= (h_in / h_hm) |
|
|
| coords_orig = coords_in.clone() |
| coords_orig[:, 0] *= (w_orig / w_in) |
| coords_orig[:, 1] *= (h_orig / h_in) |
|
|
| return coords_orig.cpu().numpy() |
|
|
|
|
| def compute_confidence(heatmaps): |
| hm = heatmaps[0].detach().cpu().numpy() |
| return hm.reshape(hm.shape[0], -1).max(axis=1) |
|
|
|
|
| def draw_points(image, coords, symbols, color=(255, 0, 0)): |
| out = image.copy() |
| h, w = out.shape[:2] |
| for i, (x, y) in enumerate(coords): |
| x, y = int(round(float(x))), int(round(float(y))) |
| if 0 <= x < w and 0 <= y < h: |
| cv2.circle(out, (x, y), 4, color, -1, lineType=cv2.LINE_AA) |
| cv2.putText( |
| out, |
| symbols[i], |
| (x + 5, y - 5), |
| cv2.FONT_HERSHEY_SIMPLEX, |
| 0.4, |
| color, |
| 1, |
| cv2.LINE_AA, |
| ) |
| return out |
|
|
|
|
| def heatmap_overlay(image, heatmap): |
| h, w = image.shape[:2] |
| hm = cv2.resize(heatmap, (w, h), interpolation=cv2.INTER_LINEAR) |
| hm = (hm - hm.min()) / (hm.max() - hm.min() + 1e-6) |
| hm_color = cv2.applyColorMap((hm * 255).astype(np.uint8), cv2.COLORMAP_JET) |
| hm_color = cv2.cvtColor(hm_color, cv2.COLOR_BGR2RGB) |
| return cv2.addWeighted(image, 0.6, hm_color, 0.4, 0) |
|
|
|
|
| def make_single_landmark_view(orig, coords, symbols, hm_np, idx): |
| out = heatmap_overlay(orig, hm_np[idx]) |
| out = draw_points( |
| out, |
| np.array([coords[idx]], dtype=np.float32), |
| [symbols[idx]], |
| color=(255, 255, 255), |
| ) |
| return out |
|
|
|
|
| def build_demo(model, model_args, checkpoint_symbols, device): |
| default_symbols = get_symbols(model_args["num_landmarks"], checkpoint_symbols) |
|
|
| def run_inference(image): |
| if image is None: |
| return None, None, None, None, None, None, gr.Dropdown() |
|
|
| orig = image.copy() |
| tensor, orig_size, input_size = preprocess(orig, model_args, device) |
|
|
| with torch.no_grad(): |
| heatmaps = model(tensor) |
|
|
| coords = decode(heatmaps, orig_size, input_size) |
| hm_np = heatmaps[0].detach().cpu().numpy() |
| conf = compute_confidence(heatmaps) |
| symbols = get_symbols(len(coords), checkpoint_symbols) |
|
|
| pred_overlay = draw_points(orig, coords, symbols) |
| summed_overlay = heatmap_overlay(orig, hm_np.sum(axis=0)) |
| single_overlay = make_single_landmark_view(orig, coords, symbols, hm_np, 0) |
|
|
| table = [ |
| [symbols[i], float(coords[i, 0]), float(coords[i, 1]), float(conf[i])] |
| for i in range(len(symbols)) |
| ] |
|
|
| cache = { |
| "orig": orig, |
| "coords": coords, |
| "symbols": symbols, |
| "heatmaps": hm_np, |
| "pred_overlay": pred_overlay, |
| "summed_overlay": summed_overlay, |
| "table": table, |
| } |
|
|
| dropdown_update = gr.Dropdown(choices=symbols, value=symbols[0]) |
|
|
| return orig, pred_overlay, summed_overlay, single_overlay, table, cache, dropdown_update |
|
|
| def update_selected_landmark(selected_landmark, cache): |
| if cache is None: |
| return None |
|
|
| symbols = cache["symbols"] |
| idx = symbols.index(selected_landmark) if selected_landmark in symbols else 0 |
|
|
| return make_single_landmark_view( |
| cache["orig"], |
| cache["coords"], |
| cache["symbols"], |
| cache["heatmaps"], |
| idx, |
| ) |
|
|
| with gr.Blocks() as demo: |
| gr.Markdown("## Cephalogram Landmark Inference") |
|
|
| cache_state = gr.State() |
|
|
| with gr.Row(): |
| with gr.Column(scale=1, min_width=320): |
| input_image = gr.Image(type="numpy", label="Input Image", height=420) |
| run_button = gr.Button("Run Inference", variant="primary") |
| selected_landmark = gr.Dropdown( |
| choices=default_symbols, |
| value=default_symbols[0], |
| label="Landmark Heatmap Selector", |
| ) |
|
|
| with gr.Column(scale=2): |
| with gr.Row(): |
| out_orig = gr.Image(label="Original", height=284) |
| out_pred = gr.Image(label="Predictions", height=284) |
| with gr.Row(): |
| out_sum = gr.Image(label="All-Landmark Heatmap Overlay", height=284) |
| out_single = gr.Image(label="Selected Landmark Heatmap Overlay", height=284) |
|
|
| out_table = gr.Dataframe( |
| headers=["Landmark", "X", "Y", "Confidence"], |
| label="Predictions", |
| interactive=False, |
| wrap=True, |
| ) |
|
|
| run_button.click( |
| fn=run_inference, |
| inputs=[input_image], |
| outputs=[ |
| out_orig, |
| out_pred, |
| out_sum, |
| out_single, |
| out_table, |
| cache_state, |
| selected_landmark, |
| ], |
| ) |
|
|
| selected_landmark.change( |
| fn=update_selected_landmark, |
| inputs=[selected_landmark, cache_state], |
| outputs=[out_single], |
| ) |
|
|
| return demo |
|
|
|
|
| if __name__ == "__main__": |
| cli_args = parse_args() |
| model, model_args, checkpoint_symbols = load_model(cli_args.checkpoint, cli_args.device) |
| checkpoint_symbols = [ |
| "A", "ANS", "B", "Me", "N", "Or", "Pog", "PNS", "Pn", "R", |
| "S", "Ar", "Co", "Gn", "Go", "Po", "LPM", "LIT", "LMT", "UPM", |
| "UIA", "UIT", "UMT", "LIA", "Li", "Ls", "N`", "Pog`", "Sn" |
| ] |
|
|
| demo = build_demo(model, model_args, checkpoint_symbols, cli_args.device) |
| demo.launch( |
| |
| |
| |
| |
| ) |
|
|
