Spaces:
Running on Zero
Running on Zero
| """FleXray: Universal Clinical X-ray Segmentation — interactive Gradio demo. | |
| Loads the flagship FleXray UNet (plus its 4 ensemble siblings for the higher | |
| quality modes) from `VictorButoi/flexray` through the official `flexray` PyPI | |
| package and segments 60 anatomical structures from any clinical radiograph. | |
| """ | |
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # noqa: E402 (must precede torch) | |
| import colorsys # noqa: E402 | |
| import time # noqa: E402 | |
| import gradio as gr # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| from PIL import Image # noqa: E402 | |
| from fxr.inference import FleXraySegmenter # noqa: E402 | |
| MODEL_ID = "VictorButoi/flexray" | |
| # --------------------------------------------------------------------------- | |
| # Model loading (module scope, eager .to("cuda") — ZeroGPU intercepts this) | |
| # --------------------------------------------------------------------------- | |
| _t0 = time.perf_counter() | |
| _flagship = FleXraySegmenter.from_pretrained(MODEL_ID, device="cuda") | |
| _ensemble = FleXraySegmenter.from_pretrained(MODEL_ID, ensemble=True, device="cuda") | |
| print(f"[flexray] loaded flagship + {len(_ensemble.models)} ensemble members " | |
| f"in {time.perf_counter() - _t0:.1f}s", flush=True) | |
| LABEL_NAMES = [n for n in _flagship.label_names if n != "background"] | |
| MODES = ["Low", "Normal", "High", "X-High"] | |
| # --------------------------------------------------------------------------- | |
| # Label colors — the official FleXray protocol colormap, gaps filled deterministically | |
| # --------------------------------------------------------------------------- | |
| _PROTOCOL_COLORS = { | |
| "skull": "#8B4EC6", "scapulae": "#00FF00", "clavicles": "#0000FF", | |
| "humeri": "#FFFF00", "phalanges": "#00FFFF", "metacarpals": "#FF00FF", | |
| "carpals": "#FF6B6B", "radii": "#0000CD", "ulnae": "#CD853F", | |
| "tibiae": "#D2B48C", "femurs": "#CC66A0", "patellae": "#FF0000", | |
| "fibulae": "#008B8B", "metatarsals": "#2E8B57", "tarsals": "#4D96FF", | |
| "toes": "#6A5ACD", | |
| "rib_1": "#FFE0B2", "rib_2": "#FFCC80", "rib_3": "#FFB74D", | |
| "rib_4": "#FFA726", "rib_5": "#FF9800", "rib_6": "#FB8C00", | |
| "rib_7": "#F57C00", "rib_8": "#EF6C00", "rib_9": "#E65100", | |
| "rib_10": "#BF360C", "rib_11": "#991D0A", "rib_12": "#7D0C08", | |
| "sternum": "#B5179E", | |
| "vertebra_c1": "#BBDEFB", "vertebra_c2": "#A9D2F6", "vertebra_c3": "#97C6F1", | |
| "vertebra_c4": "#85BAEC", "vertebra_c5": "#73AEE7", "vertebra_c6": "#61A2E2", | |
| "vertebra_c7": "#4F96DD", "vertebra_t1": "#3D8AD8", "vertebra_t2": "#2B7ED3", | |
| "vertebra_t3": "#1976D2", "vertebra_t4": "#176DC4", "vertebra_t5": "#1565C0", | |
| "vertebra_t6": "#145DB2", "vertebra_t7": "#1255A4", "vertebra_t8": "#104D96", | |
| "vertebra_t9": "#0E4588", "vertebra_t10": "#0C3D7A", "vertebra_t11": "#0A356C", | |
| "vertebra_t12": "#082D5E", | |
| "vertebra_l1": "#6D28D9", "vertebra_l2": "#8B5CF6", "vertebra_l3": "#A855F7", | |
| "vertebra_l4": "#D946EF", "vertebra_l5": "#EC4899", | |
| "hips": "#90EE90", "sacrum": "#ADFF2F", | |
| "kidneys": "#2A9D8F", "liver": "#800000", "spleen": "#556B2F", | |
| "lungs": "#32CD32", "heart": "#F4A460", | |
| } | |
| def _fallback_color(index: int) -> str: | |
| """Deterministic, well-spread color for labels outside the official map.""" | |
| hue = (index * 0.61803398875) % 1.0 # golden-ratio spacing | |
| r, g, b = colorsys.hls_to_rgb(hue, 0.55, 0.75) | |
| return "#%02X%02X%02X" % (int(r * 255), int(g * 255), int(b * 255)) | |
| LABEL_COLORS = { | |
| name: _PROTOCOL_COLORS.get(name, _fallback_color(i)) | |
| for i, name in enumerate(LABEL_NAMES) | |
| } | |
| LABEL_HEX = [LABEL_COLORS[n][1:] for n in LABEL_NAMES] | |
| _CHANNELS = {name: i + 1 for i, name in enumerate(LABEL_NAMES)} # model channel idx (background=0) | |
| def _hex_to_rgb(h: str) -> tuple[int, int, int]: | |
| return (int(h[1:3], 16), int(h[3:5], 16), int(h[5:7], 16)) | |
| LABEL_RGB = np.array([_hex_to_rgb(LABEL_COLORS[n]) for n in LABEL_NAMES], dtype=np.float32) | |
| CHANNEL_TO_SLOT = np.array([_CHANNELS[name] for name in LABEL_NAMES], dtype=np.int64) | |
| # --------------------------------------------------------------------------- | |
| # Rendering | |
| # --------------------------------------------------------------------------- | |
| def _render_overlay(pil_image: Image.Image, masks: np.ndarray, selected: list[str]) -> Image.Image: | |
| """Blend a color-coded mask over the original image at original resolution. | |
| masks: uint8 array shaped (61, H, W) — channel 0 is background. | |
| selected: list of label names to draw; empty list draws all. | |
| """ | |
| width, height = pil_image.size | |
| base_rgb = np.repeat( | |
| np.asarray(pil_image.convert("L").resize((width, height)), dtype=np.float32)[:, :, None], | |
| 3, | |
| axis=2, | |
| ) | |
| present = masks[CHANNEL_TO_SLOT].astype(bool) # (60, h, w) ordered like LABEL_NAMES | |
| keep = np.ones(len(LABEL_NAMES), dtype=bool) | |
| if selected: | |
| keep = np.array([name in selected for name in LABEL_NAMES], dtype=bool) | |
| present = present & keep[:, None, None] | |
| # Class-index map: later labels overwrite earlier ones on overlap. | |
| h, w = masks.shape[-2:] | |
| idx = np.zeros((h, w), dtype=np.int64) | |
| for i in range(present.shape[0]): | |
| if keep[i]: | |
| idx[present[i]] = i + 1 | |
| # One RGB color per class index (0 = background, never drawn). | |
| lookup = np.zeros((len(LABEL_NAMES) + 1, 3), dtype=np.float32) | |
| lookup[1:] = LABEL_RGB | |
| rgb = np.zeros((h, w, 3), dtype=np.float32) | |
| nonzero = idx > 0 | |
| rgb[nonzero] = lookup[idx[nonzero]] | |
| # Nearest-neighbor upscale keeps color regions crisp; alpha comes from class presence. | |
| overlay_np = np.asarray( | |
| Image.fromarray(rgb.astype(np.uint8), mode="RGB").resize( | |
| (width, height), resample=Image.Resampling.NEAREST | |
| ), | |
| dtype=np.float32, | |
| ) | |
| alpha_img = np.asarray( | |
| Image.fromarray((nonzero * 255).astype(np.uint8)).resize( | |
| (width, height), resample=Image.Resampling.NEAREST | |
| ), | |
| dtype=np.float32, | |
| ) / 255.0 | |
| alpha3 = alpha_img[:, :, None] | |
| out = overlay_np * (0.55 * alpha3) + base_rgb * (1.0 - 0.55 * alpha3) | |
| return Image.fromarray(np.clip(out, 0, 255).astype(np.uint8), mode="RGB") | |
| def _build_selected_list(show_all: bool, label_selection: list[str]) -> list[str]: | |
| """Resolve the checkbox-group selection into a concrete label list.""" | |
| if show_all or not label_selection: | |
| return list(LABEL_NAMES) | |
| return [n for n in label_selection if n in LABEL_NAMES] | |
| # --------------------------------------------------------------------------- | |
| # Inference | |
| # --------------------------------------------------------------------------- | |
| def segment( | |
| image: Image.Image, | |
| mode: str = "Low", | |
| show_all_labels: bool = True, | |
| label_selection: list[str] | None = None, | |
| threshold: float = 0.5, | |
| ) -> tuple[Image.Image, str, str]: | |
| """Segment anatomical structures in a clinical X-ray. | |
| Args: | |
| image: Grayscale or color X-ray image (PNG/JPG; DICOM works via the CLI). | |
| mode: Quality mode — Low (flagship, 1 pass), Normal (flagship, 8-pass TTA), | |
| High (5-model ensemble), X-High (ensemble + 8-pass TTA). | |
| show_all_labels: When True, draw every predicted structure. | |
| label_selection: Specific structures to draw when show_all_labels is False. | |
| threshold: Probability threshold for a structure to count as present. | |
| Returns: | |
| (color overlay, detected-structure summary, timing note). | |
| """ | |
| if image is None: | |
| raise gr.Error("Please upload an X-ray image first.") | |
| start = time.perf_counter() | |
| selected = _build_selected_list(show_all_labels, label_selection or []) | |
| segmenter = _ensemble if mode in ("High", "X-High") else _flagship | |
| tta_samples = 8 if mode in ("Normal", "X-High") else 1 | |
| prediction = segmenter.predict(image, threshold=float(threshold), tta_samples=tta_samples) | |
| probs = prediction.probabilities[0].detach().cpu().numpy() # (61, 256, 256) | |
| masks = (probs >= float(threshold)).astype(np.uint8) | |
| elapsed = time.perf_counter() - start | |
| overlay = _render_overlay(image, masks, selected) | |
| found = [] | |
| for i, name in enumerate(LABEL_NAMES): | |
| frac = float(masks[_CHANNELS[name]].mean()) | |
| if frac > 0.001: | |
| found.append((frac, name, i)) | |
| found.sort(reverse=True) | |
| if found: | |
| summary = "<br>".join( | |
| f'<span style="display:inline-block;width:12px;height:12px;' | |
| f'background:#{LABEL_HEX[i]};margin-right:6px;border-radius:3px;"></span>' | |
| f"{name} — {frac * 100:.1f}% of image" | |
| for frac, name, i in found[:12] | |
| ) | |
| if len(found) > 12: | |
| summary += f"<br>… and {len(found) - 12} more" | |
| else: | |
| summary = "No structures detected above threshold — try another image or lower the threshold." | |
| note = f"mode: {mode} · {len(segmenter.models)} model(s) · TTA×{tta_samples} · {elapsed:.1f}s on GPU" | |
| return overlay, summary, note | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| .legend { display: flex; flex-wrap: wrap; gap: 6px 14px; font-size: 12px; } | |
| .legend span { display: inline-flex; align-items: center; gap: 5px; } | |
| .legend i { display: inline-block; width: 10px; height: 10px; border-radius: 2px; } | |
| """ | |
| legend_html = '<div class="legend">' + "".join( | |
| f'<span><i style="background:#{LABEL_HEX[i]}"></i>{name}</span>' | |
| for i, name in enumerate(LABEL_NAMES) | |
| ) + "</div>" | |
| def _change_mode(mode: str): | |
| """Placeholder kept for API stability; currently unused.""" | |
| return gr.update() | |
| with gr.Blocks() as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| """ | |
| # 🦴 FleXray: Universal Clinical X-ray Segmentation | |
| One 2D UNet that segments **60 anatomical structures** from any clinical | |
| radiograph — chest, spine, pelvis, hands, knees, and more. | |
| [Paper](https://arxiv.org/abs/2609.26756) · | |
| [Code](https://github.com/VictorButoi/FleXray) · | |
| [Weights](https://huggingface.co/VictorButoi/flexray) · | |
| [Website](https://flexray.csail.mit.edu/) | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| image_in = gr.Image( | |
| label="X-ray image", type="pil", image_mode="L", height=420 | |
| ) | |
| mode = gr.Radio( | |
| MODES, value="Low", label="Quality mode", | |
| info="Low: flagship · Normal: +8-pass TTA · High: 5-model ensemble · X-High: both", | |
| ) | |
| threshold = gr.Slider( | |
| 0.1, 0.9, value=0.5, step=0.05, label="Probability threshold" | |
| ) | |
| run_btn = gr.Button("Segment", variant="primary") | |
| with gr.Accordion("Focus on specific structures", open=False): | |
| show_all = gr.Checkbox(value=True, label="Show all structures") | |
| label_sel = gr.CheckboxGroup( | |
| LABEL_NAMES, label="Structures to display (when 'show all' is off)" | |
| ) | |
| with gr.Column(scale=1): | |
| overlay_out = gr.Image( | |
| label="Segmentation overlay", type="pil", format="png", height=420 | |
| ) | |
| stats_out = gr.HTML(label="Detected structures") | |
| note_out = gr.Textbox(label="Run info", interactive=False) | |
| with gr.Accordion("Label legend (60 structures)", open=False): | |
| gr.HTML(legend_html) | |
| gr.Examples( | |
| examples=[["examples/chest_lateral.png"], ["examples/pelvis_ap.png"]], | |
| inputs=[image_in], | |
| outputs=[overlay_out, stats_out, note_out], | |
| fn=segment, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Try a sample X-ray (radiographs courtesy of Radiopaedia.org, used with permission)", | |
| ) | |
| gr.Markdown( | |
| """ | |
| ⚠️ **Research demo** — not a medical device. Weights are CC-BY-NC-4.0. | |
| """ | |
| ) | |
| run_btn.click( | |
| segment, | |
| inputs=[image_in, mode, show_all, label_sel, threshold], | |
| outputs=[overlay_out, stats_out, note_out], | |
| api_name="segment", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS) |