Upload 4 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+best.pt.enc filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,67 +1,250 @@
-import os
-import traceback
+import argparse
 import gradio as gr
+import numpy as np
+import cv2
 import torch
-from huggingface_hub import hf_hub_download
 
-# Change this
-MODEL_REPO_ID = "farrell236/CephVIT"
-MODEL_FILENAME = "best.pt"
+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")
+    # ckpt = 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()
 
-def inspect_checkpoint():
-    try:
-        hf_token = os.getenv("HF_TOKEN")
-        if not hf_token:
-            return "ERROR: HF_TOKEN is missing. Add it in Space Settings -> Secrets."
+    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()
 
-        local_path = hf_hub_download(
-            repo_id=MODEL_REPO_ID,
-            filename=MODEL_FILENAME,
-            token=hf_token,
-        )
 
-        lines = []
-        lines.append("Download successful.")
-        lines.append(f"Local path: {local_path}")
+def compute_confidence(heatmaps):
+    hm = heatmaps[0].detach().cpu().numpy()
+    return hm.reshape(hm.shape[0], -1).max(axis=1)
 
-        ckpt = torch.load(local_path, map_location="cpu")
 
-        lines.append("")
-        lines.append(f"Top-level object type: {type(ckpt).__name__}")
+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
 
-        if isinstance(ckpt, dict):
-            top_keys = list(ckpt.keys())
-            lines.append(f"Top-level key count: {len(top_keys)}")
-            lines.append("Top-level keys:")
-            for k in top_keys[:50]:
-                lines.append(f"  - {k}")
 
-            if "state_dict" in ckpt and isinstance(ckpt["state_dict"], dict):
-                sd_keys = list(ckpt["state_dict"].keys())
-                lines.append("")
-                lines.append(f"state_dict key count: {len(sd_keys)}")
-                lines.append("First 20 state_dict keys:")
-                for k in sd_keys[:20]:
-                    lines.append(f"  - {k}")
+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)
 
-        else:
-            lines.append("Checkpoint is not a dict, so no keys to print.")
 
-        return "\n".join(lines)
+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
 
-    except Exception as e:
-        return f"ERROR:\n{type(e).__name__}: {e}\n\n{traceback.format_exc()}"
 
+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],
+        )
 
-demo = gr.Interface(
-    fn=inspect_checkpoint,
-    inputs=None,
-    outputs=gr.Textbox(label="Checkpoint inspection", lines=30),
-    title="Private checkpoint test",
-    description="Checks whether best.pt can be downloaded from a private Hugging Face repo and inspected.",
-)
+    return demo
 
 
 if __name__ == "__main__":
-    demo.launch()
+    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"
+    ]  # TEMPORARY HARD CODE
+
+    demo = build_demo(model, model_args, checkpoint_symbols, cli_args.device)
+    demo.launch(
+        # server_name=cli_args.server_name,
+        # server_port=cli_args.server_port,
+        # share=cli_args.share,
+        # inbrowser=cli_args.inbrowser,
+    )

best.pt.enc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:018b89108a18be63689aea6ce7d2cedbc22c09e1fe28558b5bddd901efb3f558
+size 976710027

requirements.txt

@@ -1,3 +1,16 @@
-gradio
-torch
-huggingface_hub
+--extra-index-url https://download.pytorch.org/whl/cu121
+
+albumentations==1.3.1
+cryptography==46.0.6
+gradio==4.44.1
+huggingface_hub==0.31.2
+numpy==1.26.3
+opencv-python==4.11.0.86
+pandas==2.3.3
+pillow==10.4.0
+pydantic==2.10.6
+timm==1.0.9
+torch==2.5.1
+torchvision==0.20.1
+torchaudio==2.5.1
+tqdm==4.66.5

secure_torch_load.py

@@ -0,0 +1,56 @@
+import gzip
+import io
+import os
+import torch
+from typing import Optional
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+
+def _parse_key(key_str: str) -> bytes:
+    key_str = key_str.strip()
+
+    try:
+        key = bytes.fromhex(key_str)
+        if len(key) == 32:
+            return key
+    except ValueError:
+        pass
+
+    key = key_str.encode("utf-8")
+    if len(key) == 32:
+        return key
+
+    raise ValueError("Key must be either a 64-character hex string or a 32-character raw string.")
+
+
+def _get_key(key: Optional[str] = None, env_var: str = "MODEL_KEY") -> bytes:
+    if key is not None:
+        return _parse_key(key)
+
+    env_value = os.environ.get(env_var)
+    if not env_value:
+        raise RuntimeError("Missing key. Provide key=... or set environment variable {}.".format(env_var))
+    return _parse_key(env_value)
+
+
+def decrypt_and_decompress_to_bytes(path: str, key: Optional[str] = None, env_var: str = "MODEL_KEY") -> bytes:
+    key_bytes = _get_key(key=key, env_var=env_var)
+    aesgcm = AESGCM(key_bytes)
+
+    with open(path, "rb") as f:
+        data = f.read()
+
+    if len(data) < 13:
+        raise ValueError("Encrypted file is too short or invalid.")
+
+    nonce = data[:12]
+    ciphertext = data[12:]
+    compressed = aesgcm.decrypt(nonce, ciphertext, None)
+    plaintext = gzip.decompress(compressed)
+    return plaintext
+
+
+def secure_torch_load(path: str, *args, key: Optional[str] = None, env_var: str = "MODEL_KEY", **kwargs):
+    plaintext = decrypt_and_decompress_to_bytes(path, key=key, env_var=env_var)
+    buffer = io.BytesIO(plaintext)
+    return torch.load(buffer, *args, **kwargs)