Update app.py

app.py

@@ -1,45 +1,30 @@
-"""
-OrthoTimes QuickCephTool — HRNet Landmark Detection API
-Hugging Face Space backend.
-
-Loads cwlachap/hrnet-cephalometric-landmark-detection, runs inference,
-and returns 19 landmark coordinates as normalised (0-1) JSON.
-
-The /detect endpoint is called by the HTML tool with a base64-encoded image.
-"""
-
 import io
 import json
 import base64
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageDraw
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from huggingface_hub import hf_hub_download
 import gradio as gr
 
-# ─── Model architecture (HRNet-W32 heatmap output) ───────────────────────────
-# Minimal re-implementation matching the checkpoint structure.
-# Adapted from the official HRNet codebase.
+# ── HRNet-W32 ─────────────────────────────────────────────────────────────────
 
 class BasicBlock(nn.Module):
     expansion = 1
     def __init__(self, inplanes, planes, stride=1, downsample=None):
         super().__init__()
         self.conv1 = nn.Conv2d(inplanes, planes, 3, stride=stride, padding=1, bias=False)
-        self.bn1   = nn.BatchNorm2d(planes, momentum=0.1)
-        self.relu  = nn.ReLU(inplace=True)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
         self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
-        self.bn2   = nn.BatchNorm2d(planes, momentum=0.1)
+        self.bn2 = nn.BatchNorm2d(planes)
         self.downsample = downsample
 
     def forward(self, x):
-        residual = x
         out = self.relu(self.bn1(self.conv1(x)))
         out = self.bn2(self.conv2(out))
-        if self.downsample: residual = self.downsample(x)
-        return self.relu(out + residual)
+        return self.relu(out + (self.downsample(x) if self.downsample else x))
 
 
 class Bottleneck(nn.Module):
@@ -47,346 +32,227 @@ class Bottleneck(nn.Module):
     def __init__(self, inplanes, planes, stride=1, downsample=None):
         super().__init__()
         self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
-        self.bn1   = nn.BatchNorm2d(planes, momentum=0.1)
+        self.bn1 = nn.BatchNorm2d(planes)
         self.conv2 = nn.Conv2d(planes, planes, 3, stride=stride, padding=1, bias=False)
-        self.bn2   = nn.BatchNorm2d(planes, momentum=0.1)
+        self.bn2 = nn.BatchNorm2d(planes)
         self.conv3 = nn.Conv2d(planes, planes * 4, 1, bias=False)
-        self.bn3   = nn.BatchNorm2d(planes * 4, momentum=0.1)
-        self.relu  = nn.ReLU(inplace=True)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
         self.downsample = downsample
 
     def forward(self, x):
-        residual = x
         out = self.relu(self.bn1(self.conv1(x)))
         out = self.relu(self.bn2(self.conv2(out)))
         out = self.bn3(self.conv3(out))
-        if self.downsample: residual = self.downsample(x)
-        return self.relu(out + residual)
+        return self.relu(out + (self.downsample(x) if self.downsample else x))
 
 
-class HRModule(nn.Module):
-    def __init__(self, num_branches, block, num_blocks, num_inchannels, num_channels, fuse_method='SUM'):
-        super().__init__()
-        self.num_branches   = num_branches
-        self.fuse_method    = fuse_method
-        self.branches       = self._make_branches(num_branches, block, num_blocks, num_channels)
-        self.fuse_layers    = self._make_fuse_layers(num_inchannels, num_channels)
-        self.relu           = nn.ReLU(True)
-
-    def _make_one_branch(self, branch_index, block, num_blocks, num_channels, stride=1):
-        layers = []
-        for i in range(num_blocks[branch_index]):
-            layers.append(block(num_channels[branch_index], num_channels[branch_index]))
-        return nn.Sequential(*layers)
-
-    def _make_branches(self, num_branches, block, num_blocks, num_channels):
-        return nn.ModuleList([
-            self._make_one_branch(i, block, num_blocks, num_channels)
-            for i in range(num_branches)
-        ])
+def make_layer(block, inplanes, planes, blocks, stride=1):
+    downsample = None
+    if stride != 1 or inplanes != planes * block.expansion:
+        downsample = nn.Sequential(
+            nn.Conv2d(inplanes, planes * block.expansion, 1, stride=stride, bias=False),
+            nn.BatchNorm2d(planes * block.expansion)
+        )
+    layers = [block(inplanes, planes, stride, downsample)]
+    for _ in range(1, blocks):
+        layers.append(block(planes * block.expansion, planes))
+    return nn.Sequential(*layers)
 
-    def _make_fuse_layers(self, num_inchannels, num_channels):
-        fuse_layers = []
-        for i in range(self.num_branches):
-            fuse_layer = []
-            for j in range(self.num_branches):
+
+class FuseLayer(nn.Module):
+    def __init__(self, num_branches, num_channels):
+        super().__init__()
+        self.num_branches = num_branches
+        fuse = []
+        for i in range(num_branches):
+            row = []
+            for j in range(num_branches):
                 if j > i:
-                    fuse_layer.append(nn.Sequential(
-                        nn.Conv2d(num_inchannels[j], num_inchannels[i], 1, bias=False),
-                        nn.BatchNorm2d(num_inchannels[i], momentum=0.1),
-                        nn.Upsample(scale_factor=2**(j-i), mode='nearest')
+                    row.append(nn.Sequential(
+                        nn.Conv2d(num_channels[j], num_channels[i], 1, bias=False),
+                        nn.BatchNorm2d(num_channels[i]),
+                        nn.Upsample(scale_factor=2 ** (j - i), mode='nearest')
                     ))
                 elif j == i:
-                    fuse_layer.append(None)
+                    row.append(nn.Identity())
                 else:
-                    conv3x3s = []
+                    convs = []
                     for k in range(i - j):
-                        if k == i - j - 1:
-                            conv3x3s.append(nn.Sequential(
-                                nn.Conv2d(num_inchannels[j], num_inchannels[i], 3, stride=2, padding=1, bias=False),
-                                nn.BatchNorm2d(num_inchannels[i], momentum=0.1)
-                            ))
-                        else:
-                            conv3x3s.append(nn.Sequential(
-                                nn.Conv2d(num_inchannels[j], num_inchannels[j], 3, stride=2, padding=1, bias=False),
-                                nn.BatchNorm2d(num_inchannels[j], momentum=0.1),
-                                nn.ReLU(True)
-                            ))
-                    fuse_layer.append(nn.Sequential(*conv3x3s))
-            fuse_layers.append(nn.ModuleList(fuse_layer))
-        return nn.ModuleList(fuse_layers)
+                        inc = num_channels[j] if k == 0 else num_channels[i]
+                        convs += [nn.Conv2d(inc, num_channels[i], 3, stride=2, padding=1, bias=False),
+                                  nn.BatchNorm2d(num_channels[i])]
+                        if k < i - j - 1:
+                            convs.append(nn.ReLU(True))
+                    row.append(nn.Sequential(*convs))
+            fuse.append(nn.ModuleList(row))
+        self.fuse = nn.ModuleList(fuse)
+        self.relu = nn.ReLU(True)
 
     def forward(self, x):
-        for i, branch in enumerate(self.branches):
-            x[i] = branch(x[i])
-        x_fuse = []
-        for i in range(len(self.fuse_layers)):
-            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
+        out = []
+        for i in range(self.num_branches):
+            y = x[0] if i == 0 else self.fuse[i][0](x[0])
             for j in range(1, self.num_branches):
-                if i == j:
-                    y = y + x[j]
-                elif j > i:
-                    y = y + self.fuse_layers[i][j](x[j])
-                else:
-                    y = y + self.fuse_layers[i][j](x[j])
-            x_fuse.append(self.relu(y))
-        return x_fuse
+                y = y + (x[j] if i == j else self.fuse[i][j](x[j]))
+            out.append(self.relu(y))
+        return out
+
+
+class HRStage(nn.Module):
+    def __init__(self, num_branches, block, num_blocks, num_channels):
+        super().__init__()
+        self.branches = nn.ModuleList([
+            nn.Sequential(*[block(num_channels[i], num_channels[i]) for _ in range(num_blocks)])
+            for i in range(num_branches)
+        ])
+        self.fuse = FuseLayer(num_branches, num_channels)
+
+    def forward(self, x):
+        x = [b(xi) for b, xi in zip(self.branches, x)]
+        return self.fuse(x)
 
 
 class HRNet(nn.Module):
-    """HRNet-W32 for heatmap-based landmark detection."""
     def __init__(self, num_joints=19):
         super().__init__()
-        self.num_joints = num_joints
-        # Stem
-        self.conv1  = nn.Conv2d(3, 64, 3, stride=2, padding=1, bias=False)
-        self.bn1    = nn.BatchNorm2d(64, momentum=0.1)
-        self.conv2  = nn.Conv2d(64, 64, 3, stride=2, padding=1, bias=False)
-        self.bn2    = nn.BatchNorm2d(64, momentum=0.1)
-        self.relu   = nn.ReLU(inplace=True)
-        # Layer1
-        self.layer1 = self._make_layer(Bottleneck, 64, 64, 4)
-        # Transition1 (256 → [32, 64])
-        self.transition1 = nn.ModuleList([
-            nn.Sequential(nn.Conv2d(256,32,3,padding=1,bias=False), nn.BatchNorm2d(32,momentum=0.1), nn.ReLU(True)),
-            nn.Sequential(nn.Sequential(nn.Conv2d(256,64,3,stride=2,padding=1,bias=False), nn.BatchNorm2d(64,momentum=0.1), nn.ReLU(True)))
-        ])
-        # Stage2
-        self.stage2 = nn.Sequential(HRModule(2, BasicBlock, [4,4], [32,64], [32,64]))
-        # Transition2
-        self.transition2 = nn.ModuleList([
-            None, None,
-            nn.Sequential(nn.Conv2d(64,128,3,stride=2,padding=1,bias=False), nn.BatchNorm2d(128,momentum=0.1), nn.ReLU(True))
-        ])
-        # Stage3
-        self.stage3 = nn.Sequential(*[HRModule(3, BasicBlock, [4,4,4], [32,64,128], [32,64,128]) for _ in range(4)])
-        # Transition3
-        self.transition3 = nn.ModuleList([
-            None, None, None,
-            nn.Sequential(nn.Conv2d(128,256,3,stride=2,padding=1,bias=False), nn.BatchNorm2d(256,momentum=0.1), nn.ReLU(True))
+        self.stem = nn.Sequential(
+            nn.Conv2d(3, 64, 3, stride=2, padding=1, bias=False), nn.BatchNorm2d(64), nn.ReLU(True),
+            nn.Conv2d(64, 64, 3, stride=2, padding=1, bias=False), nn.BatchNorm2d(64), nn.ReLU(True),
+        )
+        self.layer1 = make_layer(Bottleneck, 64, 64, 4)
+        self.trans1 = nn.ModuleList([
+            nn.Sequential(nn.Conv2d(256, 32, 3, padding=1, bias=False), nn.BatchNorm2d(32), nn.ReLU(True)),
+            nn.Sequential(nn.Conv2d(256, 64, 3, stride=2, padding=1, bias=False), nn.BatchNorm2d(64), nn.ReLU(True)),
         ])
-        # Stage4
-        self.stage4 = nn.Sequential(*[HRModule(4, BasicBlock, [4,4,4,4], [32,64,128,256], [32,64,128,256]) for _ in range(3)])
-        # Head
-        self.final_layer = nn.Conv2d(32, num_joints, 1)
-
-    def _make_layer(self, block, inplanes, planes, blocks, stride=1):
-        downsample = None
-        if stride != 1 or inplanes != planes * block.expansion:
-            downsample = nn.Sequential(
-                nn.Conv2d(inplanes, planes*block.expansion, 1, stride=stride, bias=False),
-                nn.BatchNorm2d(planes*block.expansion, momentum=0.1)
-            )
-        layers = [block(inplanes, planes, stride, downsample)]
-        for _ in range(1, blocks):
-            layers.append(block(planes*block.expansion, planes))
-        return nn.Sequential(*layers)
+        self.stage2 = HRStage(2, BasicBlock, 4, [32, 64])
+        self.trans2 = nn.Sequential(nn.Conv2d(64, 128, 3, stride=2, padding=1, bias=False), nn.BatchNorm2d(128), nn.ReLU(True))
+        self.stage3 = nn.Sequential(*[HRStage(3, BasicBlock, 4, [32, 64, 128]) for _ in range(4)])
+        self.trans3 = nn.Sequential(nn.Conv2d(128, 256, 3, stride=2, padding=1, bias=False), nn.BatchNorm2d(256), nn.ReLU(True))
+        self.stage4 = nn.Sequential(*[HRStage(4, BasicBlock, 4, [32, 64, 128, 256]) for _ in range(3)])
+        self.head = nn.Conv2d(32, num_joints, 1)
 
     def forward(self, x):
-        x = self.relu(self.bn1(self.conv1(x)))
-        x = self.relu(self.bn2(self.conv2(x)))
+        x = self.stem(x)
         x = self.layer1(x)
-        xl = [t(x) if t else x for t in self.transition1]
-        xl = list(self.stage2[0](xl))
-        xl2 = []
-        for i, t in enumerate(self.transition2):
-            if t is None:
-                xl2.append(xl[i] if i < len(xl) else xl[-1])
-            else:
-                xl2.append(t(xl[-1]))
-        xl = xl2
+        x = [t(x) for t in self.trans1]
+        x = self.stage2(x)
+        x = [x[0], x[1], self.trans2(x[1])]
         for m in self.stage3:
-            xl = m(xl)
-        xl3 = []
-        for i, t in enumerate(self.transition3):
-            if t is None:
-                xl3.append(xl[i] if i < len(xl) else xl[-1])
-            else:
-                xl3.append(t(xl[-1]))
-        xl = xl3
+            x = m(x)
+        x = [x[0], x[1], x[2], self.trans3(x[2])]
         for m in self.stage4:
-            xl = m(xl)
-        return self.final_layer(xl[0])
+            x = m(x)
+        return self.head(x[0])
 
 
-# ─── Load model ──────────────────────────────────────────────────────────────
+# ── Load weights ──────────────────────────────────────────────────────────────
 print("Downloading model weights...")
 model_path = hf_hub_download(
     repo_id="cwlachap/hrnet-cephalometric-landmark-detection",
     filename="best_model.pth"
 )
+checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
+state_dict = checkpoint.get("model_state_dict", checkpoint.get("state_dict", checkpoint))
 model = HRNet(num_joints=19)
-checkpoint = torch.load(model_path, map_location='cpu')
-state = checkpoint.get('model_state_dict', checkpoint)
-model.load_state_dict(state, strict=False)
+missing, unexpected = model.load_state_dict(state_dict, strict=False)
+print(f"Loaded. Missing: {len(missing)} | Unexpected: {len(unexpected)}")
 model.eval()
-print("Model loaded.")
-
-# ─── Landmark mapping ────────────────────────────────────────────────────────
-# HRNet model outputs 19 landmarks in this order (ISBI 2015 dataset standard)
-HRNET_LANDMARKS = [
-    'S',     # 0  Sella
-    'N',     # 1  Nasion
-    'Or',    # 2  Orbitale
-    'Po',    # 3  Porion
-    'ANS',   # 4  ANS
-    'PNS',   # 5  PNS
-    'A',     # 6  Point A
-    'U1tip', # 7  Upper incisor tip
-    'L1tip', # 8  Lower incisor tip
-    'B',     # 9  Point B
-    'Pog',   # 10 Pogonion
-    'Me',    # 11 Menton
-    'Gn',    # 12 Gnathion
-    'Go',    # 13 Gonion
-    'Co',    # 14 Condylion  (some datasets use Ar here)
-    'L1ap',  # 15 Lower incisor apex
-    'U1ap',  # 16 Upper incisor apex
-    'U6',    # 17 Upper molar
-    'L6',    # 18 Lower molar
-]
-
-# ─── Preprocessing ───────────────────────────────────────────────────────────
-INPUT_W, INPUT_H = 256, 320   # model input size
-
-def preprocess(img_pil):
-    """Convert PIL image → normalised tensor [1,3,H,W]."""
-    img = img_pil.convert('RGB').resize((INPUT_W, INPUT_H), Image.BILINEAR)
-    arr = np.array(img, dtype=np.float32) / 255.0
-    mean = np.array([0.485, 0.456, 0.406])
-    std  = np.array([0.229, 0.224, 0.225])
-    arr  = (arr - mean) / std
-    tensor = torch.from_numpy(arr).permute(2,0,1).unsqueeze(0).float()
-    return tensor
-
-def heatmap_to_coords(heatmaps, orig_w, orig_h):
-    """
-    heatmaps: [1, num_joints, H, W]
-    Returns list of (x_norm, y_norm) tuples in original image space.
-    """
-    hm = heatmaps[0]          # [num_joints, H, W]
-    num_joints, hm_h, hm_w = hm.shape
-    coords = []
-    for j in range(num_joints):
-        flat = hm[j].reshape(-1)
-        idx  = int(flat.argmax())
-        py   = idx // hm_w
-        px   = idx  % hm_w
-        # Sub-pixel refinement: nudge toward neighbouring maxima
-        if 1 <= px < hm_w-1 and 1 <= py < hm_h-1:
-            dx = float(hm[j, py, px+1] - hm[j, py, px-1])
-            dy = float(hm[j, py+1, px] - hm[j, py-1, px])
-            px += 0.25 * np.sign(dx)
-            py += 0.25 * np.sign(dy)
-        # Normalise to 0-1 in original image space
-        x_norm = (px / hm_w)  * (INPUT_W  / orig_w)
-        y_norm = (py / hm_h)  * (INPUT_H  / orig_h)
-        # Clamp
-        x_norm = float(np.clip(x_norm, 0.0, 1.0))
-        y_norm = float(np.clip(y_norm, 0.0, 1.0))
-        coords.append((x_norm, y_norm))
-    return coords
+print("Model ready.")
 
-# ─── Inference function ───────────────────────────────────────────────────────
-def detect_landmarks(image_b64: str, mime_type: str = "image/jpeg") -> str:
-    """
-    Accepts base64-encoded image string.
-    Returns JSON: {"landmarks": {"S": {"x":..,"y":..}, ...}}
-    """
-    try:
-        img_bytes = base64.b64decode(image_b64)
-        img_pil   = Image.open(io.BytesIO(img_bytes)).convert('RGB')
-        orig_w, orig_h = img_pil.size
+# ── Constants ─────────────────────────────────────────────────────────────────
+LM_IDS = ['S', 'N', 'Or', 'Po', 'ANS', 'PNS', 'A', 'U1tip', 'L1tip', 'B',
+           'Pog', 'Me', 'Gn', 'Go', 'Co', 'L1ap', 'U1ap', 'U6', 'L6']
 
-        tensor = preprocess(img_pil)
-        with torch.no_grad():
-            heatmaps = model(tensor)
+LM_COLORS = {
+    'S': '#58a6ff', 'N': '#58a6ff', 'Or': '#58a6ff', 'Po': '#58a6ff',
+    'ANS': '#3fb950', 'PNS': '#3fb950', 'A': '#3fb950',
+    'U1tip': '#3fb950', 'U1ap': '#3fb950', 'U6': '#3fb950',
+    'B': '#f0883e', 'L1tip': '#f0883e', 'L1ap': '#f0883e',
+    'Pog': '#f0883e', 'Me': '#f0883e', 'Gn': '#f0883e',
+    'Go': '#f0883e', 'Co': '#f0883e', 'L6': '#f0883e'
+}
 
-        coords = heatmap_to_coords(heatmaps.numpy(), orig_w, orig_h)
+INPUT_W, INPUT_H = 256, 320
 
-        result = {}
-        for i, lm_id in enumerate(HRNET_LANDMARKS):
-            if i < len(coords):
-                result[lm_id] = {"x": round(coords[i][0], 4),
-                                  "y": round(coords[i][1], 4),
-                                  "confidence": 0.85}
 
-        return json.dumps({"landmarks": result, "notes": "HRNet-W32 detection"})
+# ── Helpers ───────────────────────────────────────────────────────────────────
+def preprocess(pil_img):
+    img = pil_img.convert('RGB').resize((INPUT_W, INPUT_H), Image.BILINEAR)
+    arr = np.array(img, dtype=np.float32) / 255.0
+    arr = (arr - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
+    return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).float()
+
+
+def heatmap_to_coords(hm_np, orig_w, orig_h):
+    coords = {}
+    nj, hh, hw = hm_np.shape
+    for j in range(min(nj, len(LM_IDS))):
+        hm = hm_np[j]
+        idx = int(hm.argmax())
+        py, px = divmod(idx, hw)
+        if 1 <= px < hw - 1 and 1 <= py < hh - 1:
+            px += 0.25 * np.sign(float(hm[py, px + 1] - hm[py, px - 1]))
+            py += 0.25 * np.sign(float(hm[py + 1, px] - hm[py - 1, px]))
+        x_norm = float(np.clip((px / hw) * (INPUT_W / orig_w), 0, 1))
+        y_norm = float(np.clip((py / hh) * (INPUT_H / orig_h), 0, 1))
+        coords[LM_IDS[j]] = {"x": round(x_norm, 4), "y": round(y_norm, 4), "confidence": 0.85}
+    return coords
 
-    except Exception as e:
-        return json.dumps({"error": str(e)})
 
+def run_detection(pil_img):
+    orig_w, orig_h = pil_img.size
+    tensor = preprocess(pil_img)
+    with torch.no_grad():
+        hm = model(tensor)[0].numpy()
+    return heatmap_to_coords(hm, orig_w, orig_h)
 
-# ─── Gradio interface ─────────────────────────────────────────────────────────
-# We expose both a visual demo AND a pure API endpoint
 
-def demo_fn(image):
-    """Visual demo: accepts PIL image, returns annotated image + JSON."""
-    if image is None:
+# ── Gradio functions ──────────────────────────────────────────────────────────
+def detect_visual(pil_img):
+    if pil_img is None:
         return None, "{}"
-    buf = io.BytesIO()
-    image.save(buf, format='JPEG')
-    b64 = base64.b64encode(buf.getvalue()).decode()
-    result_json = detect_landmarks(b64)
-    result = json.loads(result_json)
-
-    # Draw landmarks on image for visual output
-    from PIL import ImageDraw, ImageFont
-    draw = ImageDraw.Draw(image)
-    w, h = image.size
-    colors = {'S':'#58a6ff','N':'#58a6ff','Or':'#58a6ff','Po':'#58a6ff',
-              'ANS':'#3fb950','PNS':'#3fb950','A':'#3fb950',
-              'U1tip':'#3fb950','U1ap':'#3fb950','U6':'#3fb950',
-              'B':'#f0883e','L1tip':'#f0883e','L1ap':'#f0883e',
-              'Pog':'#f0883e','Me':'#f0883e','Gn':'#f0883e',
-              'Go':'#f0883e','Co':'#f0883e','L6':'#f0883e'}
-
-    for lm_id, pt in result.get('landmarks', {}).items():
-        cx = int(pt['x'] * w)
-        cy = int(pt['y'] * h)
-        col = colors.get(lm_id, '#ffffff')
-        r = 5
-        draw.ellipse([cx-r, cy-r, cx+r, cy+r], fill=col, outline='black')
-        draw.text((cx+7, cy-8), lm_id, fill=col)
-
-    return image, json.dumps(result, indent=2)
-
-def api_fn(image_b64, mime_type="image/jpeg"):
-    """Pure JSON API — called by the HTML tool."""
-    return detect_landmarks(image_b64, mime_type)
-
-with gr.Blocks(title="OrthoTimes — Ceph Landmark Detection") as demo:
-    gr.Markdown("""
-    # OrthoTimes QuickCephTool — Landmark Detection API
-    **HRNet-W32** pretrained on cephalometric radiographs.
-    Detects 19 landmarks with MRE ~1.2–1.6 mm.
-
-    ### API Usage (from JavaScript):
-    ```js
-    const result = await fetch(
-      "https://YOUR-SPACE.hf.space/run/api",
-      { method:"POST", headers:{"Content-Type":"application/json"},
-        body: JSON.stringify({ data: [base64String, "image/jpeg"] }) }
-    );
-    const json = await result.json();
-    const landmarks = JSON.parse(json.data[0]).landmarks;
-    ```
-    """)
+    coords = run_detection(pil_img)
+    out = pil_img.copy().convert("RGB")
+    draw = ImageDraw.Draw(out)
+    w, h = out.size
+    r = max(4, w // 120)
+    for lm_id, pt in coords.items():
+        cx, cy = int(pt['x'] * w), int(pt['y'] * h)
+        col = LM_COLORS.get(lm_id, '#ffffff')
+        draw.ellipse([cx - r, cy - r, cx + r, cy + r], fill=col, outline='black', width=1)
+        draw.text((cx + r + 2, cy - r), lm_id, fill=col)
+    return out, json.dumps({"landmarks": coords}, indent=2)
+
+
+def detect_api(image_b64: str) -> str:
+    try:
+        img_bytes = base64.b64decode(image_b64)
+        pil_img = Image.open(io.BytesIO(img_bytes)).convert('RGB')
+        coords = run_detection(pil_img)
+        return json.dumps({"landmarks": coords})
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+
+# ── UI ────────────────────────────────────────────────────────────────────────
+with gr.Blocks(title="OrthoTimes Landmark Detection") as demo:
+    gr.Markdown("## OrthoTimes QuickCephTool — HRNet Landmark Detection\nDetects 19 cephalometric landmarks automatically.")
 
     with gr.Row():
-        inp = gr.Image(type='pil', label='Upload lateral cephalogram')
+        img_in = gr.Image(type="pil", label="Upload lateral cephalogram")
         with gr.Column():
-            out_img  = gr.Image(type='pil', label='Detected landmarks')
-            out_json = gr.Textbox(label='JSON output', lines=20)
+            img_out = gr.Image(type="pil", label="Detected landmarks")
+            json_out = gr.Textbox(label="JSON output", lines=12)
 
-    gr.Button("Detect Landmarks").click(demo_fn, inputs=inp, outputs=[out_img, out_json])
+    gr.Button("▶ Detect Landmarks").click(fn=detect_visual, inputs=img_in, outputs=[img_out, json_out])
 
-    # Headless API endpoint
+    # Headless API endpoint — called by the HTML tool
     gr.Interface(
-        fn=api_fn,
-        inputs=[gr.Textbox(label="base64 image"), gr.Textbox(label="mime type", value="image/jpeg")],
+        fn=detect_api,
+        inputs=gr.Textbox(label="Base64 image"),
         outputs=gr.Textbox(label="JSON result"),
-        api_name="api"
+        api_name="detect"
     )
 
-demo.launch()
+demo.launch()