# #!/usr/bin/env python3 # import os, re, cv2, time, base64, uuid, threading, argparse # import json, shutil # from datetime import datetime # from io import BytesIO # from pathlib import Path # from functools import partial # from typing import List, Optional, Iterable, Tuple # import numpy as np # import gradio as gr # from PIL import Image, ImageFilter # import torch # from transformers import ( # AutoProcessor, AutoTokenizer, AutoModelForVision2Seq, # TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList # ) # import cxas # from cxas.label_mapper import name2id as cxas_name2id # # --------------------------------------------- # # argparse # # --------------------------------------------- # def parse_args(): # p = argparse.ArgumentParser() # p.add_argument("--model", default="vreason") # p.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"]) # p.add_argument("--viz_mode", choices=["blur", "crop", "blurcrop"], default="blurcrop") # p.add_argument("--context_ring", type=int, default=8) # p.add_argument("--blur_radius", type=int, default=31) # p.add_argument("--feather", type=int, default=6) # p.add_argument("--resize_base_to", nargs=2, type=int, default=[512, 512], metavar=("W", "H")) # p.add_argument("--resize_roi_to", nargs=2, type=int, default=[256, 256], metavar=("W", "H")) # p.add_argument("--cxas_gpus", default="0") # p.add_argument("--title", default="Interactive Chest X ray Demo") # p.add_argument("--share", action="store_true") # p.add_argument("--server_name", default=None) # p.add_argument("--server_port", type=int, default=None) # p.add_argument("--static_dir", default="./static") # p.add_argument("--tmp_dir", default="./tmp") # p.add_argument("--css_path", default="layout.css") # p.add_argument("--js_path", default="script.js") # p.add_argument("--display_shortest", type=int, default=512, help="Shortest side for display images in the web interface, preserving aspect ratio") # p.add_argument("--log_dir", default="./logs") # return p.parse_args() # class SessionLogger: # def __init__(self, base_dir: Path): # self.base_dir = Path(base_dir) # self.base_dir.mkdir(parents=True, exist_ok=True) # stamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") # self.dir = self.base_dir / stamp # self.dir.mkdir(parents=True, exist_ok=False) # self.paths = {"frontal": None, "lateral": None} # self.anat_to_path = {} # region name -> mask jpg filename # self.patho_to_path = {} # (region, sub) -> mask jpg filename # self._counters = {"anatomy": 0, "pathology": 0} # def save_input_images(self, pa_path: str, lat_path: Optional[str]): # frontal_dst = self.dir / "frontal.jpg" # Image.open(pa_path).convert("RGB").save(frontal_dst, "JPEG", quality=95, optimize=True, subsampling=1) # self.paths["frontal"] = frontal_dst.name # if lat_path: # lateral_dst = self.dir / "lateral.jpg" # Image.open(lat_path).convert("RGB").save(lateral_dst, "JPEG", quality=95, optimize=True, subsampling=1) # self.paths["lateral"] = lateral_dst.name # def _save_mask_jpg(self, mask_bool: np.ndarray, kind: str) -> str: # # mask_bool shape H x W at original size # name = f"{kind}_mask_{self._counters[kind]:03d}.jpg" # self._counters[kind] += 1 # out = self.dir / name # img = Image.fromarray(mask_bool.astype(np.uint8) * 255, mode="L") # img.save(out, "JPEG", quality=95, optimize=True, subsampling=1) # return name # def add_anatomy_mask(self, anatomy_heading: str, mask_bool: np.ndarray): # fname = self._save_mask_jpg(mask_bool, "anatomy") # self.anat_to_path[anatomy_heading] = fname # def add_pathology_mask(self, anatomy_heading: str, sub_heading: str, mask_bool: np.ndarray): # fname = self._save_mask_jpg(mask_bool, "pathology") # self.patho_to_path[(anatomy_heading, sub_heading)] = fname # def finalize_reasoning(self, reasoning_wrapper: "ReasoningWrapper"): # reasoning_list = [] # for card in reasoning_wrapper.cards: # region = card.heading # region_path = self.anat_to_path.get(region) # patho_list = [] # for sub in card.sub_cards: # patho_list.append({ # "anatomies": sub.heading, # "reason": sub.content, # "path": self.patho_to_path.get((region, sub.heading)) # }) # reasoning_list.append({ # "region": region, # "path": region_path, # "pathological": patho_list # }) # payload = { # "input_image": self.paths["frontal"], # "lateral_image": self.paths["lateral"], # "reasoning": reasoning_list # } # out_json = self.dir / "reasoning.json" # with open(out_json, "w", encoding="utf8") as f: # json.dump(payload, f, ensure_ascii=False, indent=2) # return out_json # # --------------------------------------------- # # system message and regex # # --------------------------------------------- # SYSTEM_MESSAGE = """You are a radiologist assistant. You must strictly follow the inspection and reasoning protocol described below. # You perform a step-by-step anatomical inspection, followed by a summary of findings, diagnostic impression, and final report. # Inspection procedure: # 1. section # For each anatomical region that is commonly examined: # a. You will describe your focus by saying: Reviewing {region}... # b. Then, you bring attention to this area by calling the anatomical tool in the form: # # c. After examining the anatomical region, move on to its relevant pathological sub-parts. # d. Indicate this step by saying: Inspecting {sub-part}... # e. Immediately follow with a pathological tool call in the form: # # f. Then describe what you observe for that sub-part in clear clinical terms. # Continue this process until all regions and their sub-parts have been reviewed. # 2. section # Summarize all inspected observations grouped by anatomical region. # Each line must follow: # - **{Region Name}**: {finding} # 3. section # Provide clinically meaningful diagnostic conclusions. # Each line must follow: # - {impression} # 4. section # Compose a complete, fluent narrative report summarizing all findings and suggest impressions. # Required output structure: # # Reviewing {region}... # # Inspecting {sub-part}... # # {observation} # [repeat for all regions and sub-parts] # # # - {region}: {summary of observations} # [repeat for all regions] # # # {concise diagnostic conclusions} # # Guidence: # 1. Use the most specific anatomical label possible from the internal knowledge graph. # 2. The pathological region for inspection should be within the anatomical region for review. Each anatomical region must has its own pathological region. # 3. Avoid inspecting the same region twice. # 4. Observations should be relevant to the region of image, and clinically accurate. # 5. You must base all judgments strictly on the visual evidence in the provided image. Don't rely on general statistical expectations unless the evidence clearly supports them. # """ # # SYSTEM_MESSAGE = "You are a radiologist assistant." # INTERPRET_START = re.compile(r'interpret>', re.I) # ANATOMY_RE = re.compile(r"Reviewing\s+(.+?)\.\.\.", re.I) # SUB_ANATOMY_RE = re.compile(r"Inspecting\s+(.+?)\.\.\.", re.I) # TOOL_RE = re.compile(r'', re.I) # REPORT_RE = re.compile(r'^(.*?\.)\n', re.S) # INTERPRET_END = re.compile(r'/interpret>', re.I) # FINDING_RE = re.compile(r'(.*?)', re.I | re.S) # IMPRESSION_RE = re.compile(r'(.*?)', re.I | re.S) # FINAL_REPORT_RE = re.compile(r'(.*?)', re.I | re.S) # UNFINISHED = re.compile(r"(]*\blabel=\[[^\]]+\][^>]*>)", re.I) # LABEL_RE = re.compile(r'label=\s*(\[[^\]]+\])', re.I) # # --------------------------------------------- # # ROI pipeline identical to dataset builder # # --------------------------------------------- # def ids_for_many(names: Iterable[str]) -> List[int]: # out, seen = [], set() # for n in names or []: # n = str(n).strip().lower() # if n in cxas_name2id: # for i in cxas_name2id[n]: # i = int(i) # if i not in seen: # seen.add(i) # out.append(i) # return out # def mask_union(arr: Optional[np.ndarray], idx: List[int]) -> Optional[np.ndarray]: # if arr is None or not idx: # return None # safe = [i for i in idx if 0 <= i < arr.shape[0]] # if not safe: # return None # return arr[safe].any(axis=0) # def bbox_from_mask(mask: Optional[np.ndarray], w: int, h: int, pad: int = 0) -> Tuple[int, int, int, int]: # if mask is None or not mask.any(): # return 0, w, 0, h # ys, xs = np.where(mask) # x0 = max(0, int(xs.min()) - pad) # x1 = min(w, int(xs.max()) + pad) # y0 = max(0, int(ys.min()) - pad) # y1 = min(h, int(ys.max()) + pad) # return x0, x1, y0, y1 # def to_alpha(mask_bool: Optional[np.ndarray], invert: bool, feather: int, size_wh: Tuple[int, int]) -> Image.Image: # if mask_bool is None: # return Image.new("L", size_wh, color=255 if invert else 0) # sel = (~mask_bool if invert else mask_bool).astype(np.uint8) * 255 # if size_wh != (mask_bool.shape[1], mask_bool.shape[0]): # sel = cv2.resize(sel, size_wh, interpolation=cv2.INTER_NEAREST) # if feather > 0: # sel = cv2.GaussianBlur(sel, ksize=(0, 0), sigmaX=feather, sigmaY=feather) # return Image.fromarray(sel, mode="L") # def optionally_resize(img: Image.Image, target_wh: Optional[Tuple[int, int]]) -> Image.Image: # if not target_wh: # return img # w, h = int(target_wh[0]), int(target_wh[1]) # return img.resize((w, h), Image.BICUBIC) # def save_jpg(img: Image.Image, path: Path, quality: int = 95): # path.parent.mkdir(parents=True, exist_ok=True) # if img.mode not in ("L", "RGB"): # img = img.convert("RGB") # img.save(path, format="JPEG", quality=quality, subsampling=1, optimize=True) # def save_blur(img_base: Image.Image, mask_bool: Optional[np.ndarray], out_path: Path, # blur_radius: int, feather: int, roi_wh: Optional[Tuple[int, int]]): # base = img_base.convert("RGB") # blurred = base.filter(ImageFilter.GaussianBlur(radius=blur_radius)) # alpha = to_alpha(mask_bool, invert=True, feather=feather, size_wh=base.size) # comp = Image.composite(blurred, base, alpha) # comp = optionally_resize(comp, roi_wh) # save_jpg(comp, out_path) # def save_crop(img_base: Image.Image, mask_bool: Optional[np.ndarray], out_path: Path, ring: int, # roi_wh: Optional[Tuple[int, int]]): # base = img_base # w, h = base.size # x0, x1, y0, y1 = bbox_from_mask(mask_bool, w, h, pad=ring) # crop = base.crop((x0, y0, x1, y1)) # crop = optionally_resize(crop, roi_wh) # save_jpg(crop, out_path) # def save_blurcrop(img_base: Image.Image, mask_bool: Optional[np.ndarray], out_path: Path, # blur_radius: int, feather: int, ring: int, roi_wh: Optional[Tuple[int, int]]): # base = img_base.convert("RGB") # w, h = base.size # if mask_bool is not None and (mask_bool.shape[1], mask_bool.shape[0]) != (w, h): # mask_resized = cv2.resize(mask_bool.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST).astype(bool) # else: # mask_resized = mask_bool # if mask_resized is None: # crop = base # else: # x0, x1, y0, y1 = bbox_from_mask(mask_resized, w, h, pad=ring) # crop = base.crop((x0, y0, x1, y1)) # mask_resized = mask_resized[y0:y1, x0:x1] # blurred = crop.filter(ImageFilter.GaussianBlur(radius=blur_radius)) # alpha = to_alpha(mask_resized, invert=True, feather=feather, size_wh=crop.size) # comp = Image.composite(blurred, crop, alpha) # comp = optionally_resize(comp, roi_wh) # save_jpg(comp, out_path) # def save_viz(img_base: Image.Image, arr0: Optional[np.ndarray], idx: List[int], out_path: Path, # mode: str, blur_radius: int, feather: int, ring: int, roi_wh: Optional[Tuple[int, int]]): # mask = mask_union(arr0, idx) if idx else None # if mode == "blur": # save_blur(img_base, mask, out_path, blur_radius=blur_radius, feather=feather, roi_wh=roi_wh) # elif mode == "crop": # save_crop(img_base, mask, out_path, ring=ring, roi_wh=roi_wh) # else: # save_blurcrop(img_base, mask, out_path, blur_radius=blur_radius, feather=feather, ring=ring, roi_wh=roi_wh) # # --------------------------------------------- # # your UI classes # # paste your Card, SubAnatomyCard, AnatomyCard, ReasoningWrapper, # # Mask, SubAnatomyMask, AnatomyMask, ImageWrapper, IFrame, TextBox # # exactly as you shared earlier # # --------------------------------------------- # from abc import abstractmethod # class Card: # @property # def check_state(self): # if self.checked: # return "checked" # return "" # @abstractmethod # def uncheck(self): # pass # @abstractmethod # def check(self): # pass # @staticmethod # def decode_match(match): # return [match.group(1)] # class SubAnatomyCard(Card): # template =""" # # """ # def __init__(self, i, j, heading): # self.i = i # self.j = j # self.checked = True # self.heading = heading # self.content = "" # def uncheck(self): # self.checked = False # def check(self): # self.checked = True # def update_content(self, content): # self.content = content # return self.content # def render(self): # return self.template.format(i=self.i, j=self.j, check_state=self.check_state, heading=self.heading, content=self.content) # class AnatomyCard(Card): # template = """ # # """ # def __init__(self, i, heading): # self.i = i # self.j = 0 # self.checked = True # self.heading = heading # self.sub_cards = [] # def uncheck(self): # self.checked = False # for sub_card in self.sub_cards: # sub_card.checked = False # def check(self): # self.checked = True # # for sub_card in self.sub_cards: # # sub_card.checked = True # def add_sub_card(self, sub_card_heading): # sub_card = SubAnatomyCard(self.i, self.j, sub_card_heading) # self.sub_cards.append(sub_card) # self.j += 1 # return sub_card # def render(self): # return self.template.format(i=self.i, check_state=self.check_state, heading=self.heading, sub_cards="\n".join([x.render() for x in self.sub_cards])) # class ReasoningWrapper: # html_icon = """ # # # # # # """ # template = """ #
#
# {icon} # Reasoning Process #
#
# {cards} #
#
# """ # def __init__(self): # self.i = 0 # self.cards = [] # def add_card(self, card_heading): # card = AnatomyCard(self.i, card_heading) # self.cards.append(card) # self.i += 1 # return card # def select(self, card, sub_card=None): # if sub_card is not None: # sub_card.check() # for x in card.sub_cards: # if x != sub_card: # x.uncheck() # card.check() # for x in self.cards: # if x != card: # x.uncheck() # def unselect(self, card): # card.uncheck() # def render(self): # return self.template.format(icon=self.html_icon, cards="\n".join([x.render() for x in self.cards])) # def output(self): # explain = [] # for card in self.cards: # roi = card.heading # reason = [] # for sub_card in card.sub_cards: # reason.append(sub_card.content) # reason = " ".join(reason) # explain.append({"roi": roi, "reason": reason}) # return {"explain": explain} # from utils import smooth, mask_to_svg # class Mask: # @property # def hidden_style(self): # if self.to_show: # return "" # return "display:none;" # @property # def svg(self): # return mask_to_svg(self.mask) # @abstractmethod # def hide(self): # pass # @abstractmethod # def show(self): # pass # @staticmethod # def preprocess(mask): # # m = np.asarray(mask) # # if m.ndim == 3: # # m = m[..., 0] # # if m.dtype == bool: # # m = m.astype(np.uint8) * 255 # # else: # # m = m.astype(np.uint8) # m = smooth(mask) # return m # # return mask # @staticmethod # def decode_match(match, anatomy_masks=None): # tool_type = match.group(1) # tool_labels = match.group(2) # tool_labels = [x.strip().strip('"').strip("'").lower() for x in tool_labels.split("\", \"")] # Split labels by comma, strip quotes and spaces # # # Get BBox through mask # # bbox = mask.convert("L").getbbox() # # Get AnatomyMask through mask # idxs = [i for label in tool_labels for i in cxas_name2id[label]] # mask = sum(anatomy_masks[i] for i in idxs) > 0 # return idxs, mask # class SubAnatomyMask(Mask): # def __init__(self, i, j, mask): # self.i = i # self.j = j # self.mask = mask # self.to_show = True # def hide(self): # self.to_show = False # def show(self): # self.to_show = True # def render(self): # return self.svg.format(sub_class="sub-svg", prefix="sub-", index=f"{self.i}-{self.j}", hidden_style=self.hidden_style, sub_svgs="", extra_data="") # class AnatomyMask(Mask): # def __init__(self, i, mask): # self.i = i # self.j = 0 # self.mask = mask # self.to_show = True # self.sub_masks = [] # @property # def inner_mask(self): # inner_mask = np.zeros_like(self.mask) # if self.sub_masks: # inner_mask = sum(sub_mask.mask for sub_mask in self.sub_masks) > 0 # inner_mask = inner_mask.astype(np.uint8) * 255 # return inner_mask # def hide(self): # self.to_show = False # for sub_mask in self.sub_masks: # sub_mask.to_show = False # def show(self): # self.to_show = True # # for sub_mask in self.sub_masks: # # sub_mask.to_show = True # def add_sub_mask(self, mask): # # mask = Mask.preprocess(mask) # sub_mask = SubAnatomyMask(self.i, self.j, mask) # self.sub_masks.append(sub_mask) # self.j += 1 # return sub_mask # def render(self): # return self.svg.format(sub_class="", prefix="", index=self.i, hidden_style=self.hidden_style, sub_svgs="\n".join([x.render() for x in self.sub_masks]), extra_data="") # def to_b64(path): # with open(path, "rb") as f: # return "data:image/png;base64," + base64.b64encode(f.read()).decode() # class ImageWrapper: # img_icon = """ # # # # # # """ # template = """ #
#
# {label_icon} # Demo Image #
#
# {icon} # {img_tag} # {masks} #
#
# """ # img_tag_template = 'Demo Image' # # TODO: image-size # def __init__(self, img_path=""): # self.i = 0 # self.j = 0 # self.set_img(img_path) # self.masks = [] # @property # def icon(self): # if self.img_tag: # return "" # return ImageWrapper.img_icon # def set_img(self, img_path): # if img_path is None or img_path == '': # self.img_path = "" # self.img_tag = "" # else: # self.img_path = img_path # img_base64 = to_b64(img_path) # self.img_tag = self.img_tag_template.format(img_base64=img_base64) # def add_mask(self, mask): # mask = AnatomyMask(self.i, mask) # self.masks.append(mask) # self.i += 1 # return mask # def select(self, mask, sub_mask=None): # if sub_mask is not None: # sub_mask.show() # for x in mask.sub_masks: # if x != sub_mask: # x.hide() # mask.show() # for x in self.masks: # if x != mask: # x.hide() # def unselect(self, mask): # mask.hide() # def render(self): # return self.template.format(label_icon=self.img_icon, icon=self.icon, img_tag=self.img_tag, masks="\n".join([x.render() for x in self.masks])) # def output(self): # return { # "image_path": self.img_path, # "explain": [{"mask": mask.inner_mask} for mask in self.masks] # } # class IFrame: # icon = """ # # # # # # """ # iframe = """ # # """ # template = """ #
#
# {label_icon} # Interactive Demo #
#
# {icon} # {iframe} #
#
# """ # def __init__(self, html_path=""): # self.html_path = html_path # self.icon = IFrame.icon # self.iframe = "" # if self.html_path: # self.iframe = IFrame.iframe.format(html_path=self.html_path, time=time.time()) # self.icon = "" # @property # def style(self): # if self.html_path: # return "" # return "min-height: 400px;" # def render(self): # return self.template.format(style=self.style, label_icon=IFrame.icon, icon=self.icon, iframe=self.iframe) # class TextBox: # icon = """ # # # # # # """ # list_template = """ #
#
# {icon} # {label} #
#
    # {value} #
#
# """ # paragraph_template = """ #
#
# {icon} # {label} #
#
# {value} #
#
# """ # def __init__(self, label, type="list", value=""): # self.type = type # self.template = getattr(TextBox, f"{type}_template", self.paragraph_template) # self.icon = TextBox.icon # self.label = label # self.value = self.preprocess(value) # @property # def style(self): # if self.value: # return "" # return "min-height: 300px;" # def preprocess(self, value): # # Remove tag if any # value = re.sub(r'', '', value, flags=re.I).strip() # # Add bold text and
  • tag accordingly # lines = [] # pattern = r'^\s*\*+\s*(.*?)\s*\*+\s*:\s*(.*)$' # for line in value.splitlines(): # line = line.strip().lstrip("-").strip() # m = re.match(pattern, line) # if m: # key, value = m.groups() # line = f"{key}: {value}" # if self.type == "list": # line = f"
  • {line}
    • " # lines.append(line) # return '\n'.join(lines) # def update_content(self, value): # self.value = self.preprocess(value) # return value # def render(self): # return self.template.format(style=self.style, icon=self.icon, label=self.label, value=self.value) # # --------------------------------------------- # # small helpers # # --------------------------------------------- # def to_b64(path: str) -> str: # with open(path, "rb") as f: # return "data:image/png;base64," + base64.b64encode(f.read()).decode() # # --------------------------------------------- # # build models # # --------------------------------------------- # def build_models(args): # device = "cuda" if (args.device == "auto" and torch.cuda.is_available()) else args.device # proc = AutoProcessor.from_pretrained(args.model, trust_remote_code=True) # tok = AutoTokenizer.from_pretrained(args.model) # vlm = AutoModelForVision2Seq.from_pretrained( # args.model, # torch_dtype=torch.float16 if device == "cuda" else torch.float32, # trust_remote_code=True, # ).to(device).eval() # cxas_model = cxas.CXAS(gpus=args.cxas_gpus if device == "cuda" else "").eval() # return device, proc, tok, vlm, cxas_model # # --------------------------------------------- # # app factory so args are captured # # --------------------------------------------- # def build_app(args, device, proc, tok, vlm, cxas_model): # RESIZE_BASE_TO = tuple(map(int, args.resize_base_to)) # RESIZE_ROI_TO = tuple(map(int, args.resize_roi_to)) # VIZ_MODE = args.viz_mode # CONTEXT_RING = int(args.context_ring) # BLUR_RADIUS = int(args.blur_radius) # FEATHER_SIGMA = int(args.feather) # TMP_DIR = Path(args.tmp_dir); TMP_DIR.mkdir(parents=True, exist_ok=True) # class _GenStop(StoppingCriteria): # def __init__(self, tokenizer, pattern): # super().__init__() # self.tok = tokenizer # self.pattern = pattern # self.decoded = "" # def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs): # self.decoded += self.tok.decode(input_ids[0][-1:], skip_special_tokens=False) # return bool(self.pattern.search(self.decoded)) # def generate(pa_path: str, lat_path: Optional[str]): # if not pa_path: # raise gr.Error("Upload at least a frontal image") # logger = SessionLogger(base_dir=Path(args.log_dir)) # logger.save_input_images(pa_path, lat_path) # img_wrapper = ImageWrapper(pa_path) # reasoning_wrapper = ReasoningWrapper() # iframe = IFrame() # finding_wrapper = TextBox(label="Findings", value="", type="list") # impression_wrapper = TextBox(label="Impressions", value="", type="list") # report_wrapper = TextBox(label="Report", value="", type="paragraph") # yield ( # gr.update(value=img_wrapper.render()), # gr.update(value=reasoning_wrapper.render()), # gr.update(value=iframe.render()), # gr.update(value=finding_wrapper.render()), # gr.update(value=impression_wrapper.render()), # gr.update(value=report_wrapper.render()), # ) # # CXAS masks # anatomy_masks = cxas_model.seg(pa_path) # anatomy_masks = np.asarray(anatomy_masks, dtype=bool) # # original and base images # pa_orig = Image.open(pa_path).convert("RGB") # # model base is 512 x 512 or whatever RESIZE_BASE_TO is # pa_base = pa_orig.resize(RESIZE_BASE_TO, Image.BICUBIC) # used for model tokens and ROI building # ##### TO REMOVE # ses_dir = TMP_DIR / f"s_{uuid.uuid4().hex}" # (ses_dir / "masks").mkdir(parents=True, exist_ok=True) # pa_tmp = ses_dir / "0.jpg" # pa_orig.save(pa_tmp) # cxas_model.process_file( # filename=str(pa_tmp), # do_store=True, # output_directory=str(ses_dir / "masks"), # storage_type="npy", # ) # npy_path = ses_dir / "masks" / "0.npy" # if not npy_path.exists(): # raise gr.Error("CXAS did not produce mask npy") # raw = np.load(npy_path).astype(bool) # anatomy_masks_orig = raw # C, H, W = raw.shape # bw, bh = RESIZE_BASE_TO # if (W, H) != (bw, bh): # resized = np.zeros((C, bh, bw), dtype=bool) # for i in range(C): # resized[i] = cv2.resize( # raw[i].astype(np.uint8), # (bw, bh), # interpolation=cv2.INTER_NEAREST, # ).astype(bool) # anatomy_masks = resized # else: # anatomy_masks = raw # #### TO REMOVE # # display copy keeps aspect ratio with shortest side args.display_shortest # orig_w, orig_h = pa_orig.size # disp_short = args.display_shortest # if orig_w < orig_h: # disp_w = disp_short # disp_h = int(orig_h * disp_short / orig_w) # else: # disp_h = disp_short # disp_w = int(orig_w * disp_short / orig_h) # pa_disp = pa_orig.resize((disp_w, disp_h), Image.BICUBIC) # disp_path = TMP_DIR / f"display_{uuid.uuid4().hex}.jpg" # pa_disp.save(disp_path) # img_wrapper = ImageWrapper(str(disp_path)) # resize_map = { # "orig": (orig_w, orig_h), # "display": (disp_w, disp_h), # "model": tuple(RESIZE_BASE_TO), # } # # ensure mask stack matches the model base resolution # C, H, W = anatomy_masks.shape # if (W, H) != pa_base.size: # resized = np.zeros((C, RESIZE_BASE_TO[1], RESIZE_BASE_TO[0]), dtype=bool) # for i in range(C): # resized[i] = cv2.resize( # anatomy_masks[i].astype(np.uint8), # RESIZE_BASE_TO, # interpolation=cv2.INTER_NEAREST, # ).astype(bool) # anatomy_masks = resized # # prepare an optional lateral copy for the model at model size # lat_base = None # if lat_path: # lat_base = Image.open(lat_path).convert("RGB").resize(RESIZE_BASE_TO, Image.BICUBIC) # convo = [ # {"role": "system", "content": [{"type": "text", "text": SYSTEM_MESSAGE}]}, # {"role": "user", "content": [ # {"type": "image", "image": pa_path}, # *([{"type": "image", "image": lat_path}] if lat_path else []), # {"type": "text", "text": "Based on the provided chest radiographs, explain your diagnosis procedure and write a report."} # ]} # ] # from utils import RegexStopper # UI state transitions # stoppers = { # "interpret_start": RegexStopper(INTERPRET_START), # "anatomy": RegexStopper(ANATOMY_RE), # "sub_anatomy": RegexStopper(SUB_ANATOMY_RE), # "anatomy_tool": RegexStopper(TOOL_RE), # "sub_anatomy_tool": RegexStopper(TOOL_RE), # "report" : RegexStopper(REPORT_RE), # "interpret_end": RegexStopper(INTERPRET_END), # "finding": RegexStopper(FINDING_RE), # "impression": RegexStopper(IMPRESSION_RE), # "final_report": RegexStopper(FINAL_REPORT_RE), # } # sequence = { # "interpret_start": ["anatomy"], # "anatomy": ["anatomy_tool"], # "anatomy_tool": ["sub_anatomy"], # "sub_anatomy": ["sub_anatomy_tool"], # "sub_anatomy_tool": ["report"], # "report": ["anatomy", "sub_anatomy", "interpret_end"], # "interpret_end": ["finding"], # "finding": ["impression"], # "impression": ["final_report"], # "final_report": [], # } # cur_items = { # "interpret_start": [], # "anatomy": None, # "anatomy_tool": None, # "sub_anatomy": None, # "sub_anatomy_tool": None, # "report": "", # "interpret_end": [], # "finding": "", # "impression": "", # "final_report": "", # } # match_decoder_fn = { # "interpret_start": lambda x: [None], # "anatomy": lambda m: [m.group(1)], # "anatomy_tool": partial(Mask.decode_match, anatomy_masks=anatomy_masks), # "sub_anatomy": lambda m: [m.group(1)], # "sub_anatomy_tool": partial(Mask.decode_match, anatomy_masks=anatomy_masks), # "report": lambda m: [m.group(1)], # "interpret_end": lambda x: [None], # "finding": lambda m: [m.group(1)], # "impression": lambda m: [m.group(1)], # "final_report": lambda m: [m.group(1)], # } # cur_sequences = ["interpret_start"] # reply_full = "" # while True: # # Build the image tensor list that matches the number and order of image tokens # images_for_proc = [] # for m in convo: # for p in m["content"]: # if p.get("type") == "image": # path = p["image"] # if path == pa_path: # images_for_proc.append(pa_base) # 512 x 512 for model # elif lat_path and path == lat_path: # images_for_proc.append(lat_base) # 512 x 512 for model # else: # # ROI images created during the loop are already sized by resize_roi_to # images_for_proc.append(Image.open(path).convert("RGB")) # prompt = proc.apply_chat_template(convo, tokenize=False, add_generation_prompt=True) # inputs = proc(text=[prompt], images=[images_for_proc], padding=True, return_tensors="pt").to(device) # streamer = TextIteratorStreamer( # tok, # skip_prompt=True, # skip_special_tokens=False, # decode_kwargs={"skip_special_tokens": False}, # ) # stopper = _GenStop(tok, UNFINISHED) # stopping = StoppingCriteriaList([stopper]) # gen_kwargs = dict( # **inputs, # max_new_tokens=256, # eos_token_id=tok.eos_token_id, # streamer=streamer, # stopping_criteria=stopping # ) # thread = threading.Thread(target=vlm.generate, kwargs=gen_kwargs) # thread.start() # response = "" # for chunk in streamer: # response += chunk # for ch in chunk: # if "report" in cur_sequences: # cur_items["report"] += ch # cur_items["sub_anatomy"].update_content(cur_items["report"]) # elif "finding" in cur_sequences: # cur_items["finding"] += ch # finding_wrapper.update_content(cur_items["finding"]) # elif "impression" in cur_sequences: # cur_items["impression"] += ch # impression_wrapper.update_content(cur_items["impression"]) # elif "final_report" in cur_sequences: # cur_items["final_report"] += ch # report_wrapper.update_content(cur_items["final_report"]) # for name in list(cur_sequences): # if not stoppers[name](ch): # continue # cur_stopper = stoppers[name] # decoded = match_decoder_fn[name](cur_stopper.match) # if name == "anatomy": # cur_items[name] = reasoning_wrapper.add_card(*decoded) # reasoning_wrapper.select(cur_items["anatomy"]) # if cur_items["anatomy_tool"]: # img_wrapper.unselect(cur_items["anatomy_tool"]) # elif name == "anatomy_tool": # idxs, mask_model = decoded # # For next round / inference: resize to 256×256 (roi) # mask_for_next = cv2.resize( # mask_model.astype(np.uint8), # tuple(RESIZE_ROI_TO), # interpolation=cv2.INTER_NEAREST # ).astype(bool) # # For display overlay: resize to display size keeping aspect ratio # mask_for_display = cv2.resize( # mask_model.astype(np.uint8), # resize_map["display"], # interpolation=cv2.INTER_NEAREST # ).astype(bool) # # Update UI # cur_items[name] = img_wrapper.add_mask(Mask.preprocess(mask_for_display)) # img_wrapper.select(cur_items["anatomy_tool"]) # # Save ROI image for the next round (use model-sized mask) # roi_path = TMP_DIR / f"roi_{uuid.uuid4().hex}.jpg" # save_viz( # img_base=pa_base, # arr0=anatomy_masks, # idx=idxs, # out_path=roi_path, # mode=VIZ_MODE, # blur_radius=BLUR_RADIUS, # feather=FEATHER_SIGMA, # ring=CONTEXT_RING, # roi_wh=tuple(RESIZE_ROI_TO) # ) # mask_orig = mask_union(anatomy_masks_orig, idxs) if idxs else np.zeros(anatomy_masks_orig.shape[1:], dtype=bool) # logger.add_anatomy_mask(cur_items["anatomy"].heading, mask_orig) # elif name == "sub_anatomy": # cur_items[name] = cur_items["anatomy"].add_sub_card(*decoded) # reasoning_wrapper.select(cur_items["anatomy"], cur_items["sub_anatomy"]) # elif name == "sub_anatomy_tool": # idxs, mask_model = decoded # # For next round / inference: resize to 256×256 (roi) # mask_for_next = cv2.resize( # mask_model.astype(np.uint8), # tuple(RESIZE_ROI_TO), # interpolation=cv2.INTER_NEAREST # ).astype(bool) # # For display overlay: resize to display size keeping aspect ratio # mask_for_display = cv2.resize( # mask_model.astype(np.uint8), # resize_map["display"], # interpolation=cv2.INTER_NEAREST # ).astype(bool) # # Update UI # cur_items[name] = cur_items["anatomy_tool"].add_sub_mask(Mask.preprocess(mask_for_display)) # img_wrapper.select(cur_items["anatomy_tool"], cur_items["sub_anatomy_tool"]) # # Save ROI image for the next round (use model-sized mask) # roi_path = TMP_DIR / f"roi_{uuid.uuid4().hex}.jpg" # save_viz( # img_base=pa_base, # arr0=anatomy_masks, # idx=idxs, # out_path=roi_path, # mode=VIZ_MODE, # blur_radius=BLUR_RADIUS, # feather=FEATHER_SIGMA, # ring=CONTEXT_RING, # roi_wh=tuple(RESIZE_ROI_TO) # ) # mask_orig = mask_union(anatomy_masks_orig, idxs) if idxs else np.zeros(anatomy_masks_orig.shape[1:], dtype=bool) # logger.add_pathology_mask(cur_items["anatomy"].heading, cur_items["sub_anatomy"].heading, mask_orig) # elif name == "report": # cur_items[name] = "" # elif name == "interpret_end": # from render import render_diagram_html # img_wrapper_json = img_wrapper.output() # reasoning_wrapper_json = reasoning_wrapper.output() # explain_json = {"explain": []} # for iw, rw in zip(img_wrapper_json['explain'], reasoning_wrapper_json["explain"]): # explain_json["explain"].append({ # "roi": rw["roi"], # "mask": iw["mask"], # "reason": rw["reason"] # }) # img_wrapper_json.update(reasoning_wrapper_json) # img_wrapper_json.update(explain_json) # json_output = img_wrapper_json # html_path = render_diagram_html(json_output) # iframe = IFrame(html_path=html_path) # logger.finalize_reasoning(reasoning_wrapper) # cur_sequences = sequence[name] # yield ( # gr.update(value=img_wrapper.render()), # gr.update(value=reasoning_wrapper.render()), # gr.update(value=iframe.render()), # gr.update(value=finding_wrapper.render()), # gr.update(value=impression_wrapper.render()), # gr.update(value=report_wrapper.render()), # ) # thread.join() # reply_full += response # # If we stopped on an unfinished tool tag, fulfill it now # m_unfinished = UNFINISHED.search(response) # if m_unfinished: # tool_call_end = m_unfinished.end() # text_to_send = response[:tool_call_end] # # # parse labels directly from the unfinished tag # # m_labels = LABEL_RE.search(m_unfinished.group(0)) # # print("*"*30) # # print(m_unfinished.group(0), m_unfinished.group(1)) # # labels = [] # # if m_labels: # # try: # # print("raw", raw) # # raw = ast.literal_eval(m_labels.group(1)) # # labels = [str(x).strip().strip('"').strip("'").lower() for x in raw] # # except Exception: # # labels = [] # # idxs = ids_for_many(labels) # mask_bool = mask_union(anatomy_masks, idxs) if idxs else np.zeros_like(anatomy_masks[0], bool) # # # update the on screen overlay immediately # # new_mask = img_wrapper.add_mask(mask_bool) # # img_wrapper.select(new_mask) # # cur_items["anatomy_tool"] = new_mask # # create ROI using the same pipeline as the dataset builder # roi_path = TMP_DIR / f"roi_{uuid.uuid4().hex}.jpg" # save_viz( # img_base=pa_base, arr0=anatomy_masks, idx=idxs, out_path=roi_path, # mode=VIZ_MODE, blur_radius=BLUR_RADIUS, feather=FEATHER_SIGMA, # ring=CONTEXT_RING, roi_wh=tuple(RESIZE_ROI_TO) # ) # # append the assistant message with all text up to the tool tag and the real image # convo.append({ # "role": "assistant", # "content": [ # {"type": "text", "text": text_to_send}, # {"type": "image", "image": str(roi_path)}, # ] # }) # # push a live UI update so the user sees the overlay right away # yield ( # gr.update(value=img_wrapper.render()), # gr.update(value=reasoning_wrapper.render()), # gr.update(value=iframe.render()), # gr.update(value=finding_wrapper.render()), # gr.update(value=impression_wrapper.render()), # gr.update(value=report_wrapper.render()), # ) # # ask the model for the next chunk # continue # # no unfinished tool tag means we are done # break # # static files and Blocks # with open(args.css_path, "r") as f: # css = f.read() # with open(args.js_path, "r") as f: # js = "\n".join([""]) # gr.set_static_paths(paths=[Path(args.static_dir).resolve()]) # gr.set_static_paths(paths=[Path(args.tmp_dir).resolve()]) # with gr.Blocks(title=args.title, css=css, head=js) as demo: # gr.Markdown(f"## {args.title}") # with gr.Row(elem_id="row1"): # with gr.Column(elem_id="col1"): # with gr.Column(elem_id="col1-1"): # pa_in = gr.Image(elem_classes=["image"], label="Upload Frontal", type="filepath") # lat_in = gr.Image(elem_classes=["image"], label="Upload Lateral optional", type="filepath") # btn = gr.Button("Generate", variant="primary") # img_out = gr.HTML(ImageWrapper().render()) # reasoning_out = gr.HTML(ReasoningWrapper().render()) # diagram_out = gr.HTML(IFrame().render()) # with gr.Row(elem_id="row3"): # with gr.Column(scale=2): # findings_out = gr.HTML(TextBox(label="Findings", value="", type="list").render()) # with gr.Column(scale=1): # impressions_out = gr.HTML(TextBox(label="Impressions", value="", type="list").render()) # with gr.Row(elem_id="row4"): # report_out = gr.HTML(TextBox(label="Report", value="", type="paragraph").render()) # btn.click( # generate, # inputs=[pa_in, lat_in], # outputs=[img_out, reasoning_out, diagram_out, findings_out, impressions_out, report_out] # ) # return demo # # --------------------------------------------- # # main # # --------------------------------------------- # if __name__ == "__main__": # args = parse_args() # device, proc, tok, vlm, cxas_model = build_models(args) # demo = build_app(args, device, proc, tok, vlm, cxas_model) # demo.launch(share=args.share, server_name=args.server_name, server_port=args.server_port) #!/usr/bin/env python3 import os, re, cv2, time, base64, uuid, threading, argparse import json, shutil from datetime import datetime from io import BytesIO from pathlib import Path from functools import partial from typing import List, Optional, Iterable, Tuple import numpy as np import gradio as gr from PIL import Image, ImageFilter import torch from transformers import ( AutoProcessor, AutoTokenizer, AutoModelForVision2Seq, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList ) try: import spaces except Exception: spaces = None import cxas from cxas.label_mapper import name2id as cxas_name2id # --------------------------------------------- # argparse # --------------------------------------------- def parse_args(): p = argparse.ArgumentParser() p.add_argument("--model", default="EvidenceAIResearch/VReason-QwenVL") p.add_argument("--hf_token", default=None, help="HF token; defaults to HF_TOKEN/HUGGINGFACEHUB_API_TOKEN env") p.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"]) p.add_argument("--viz_mode", choices=["blur", "crop", "blurcrop"], default="blurcrop") p.add_argument("--context_ring", type=int, default=8) p.add_argument("--blur_radius", type=int, default=31) p.add_argument("--feather", type=int, default=6) p.add_argument("--resize_base_to", nargs=2, type=int, default=[512, 512], metavar=("W", "H")) p.add_argument("--resize_roi_to", nargs=2, type=int, default=[256, 256], metavar=("W", "H")) p.add_argument("--cxas_gpus", default="0") p.add_argument("--title", default="Interactive Chest X ray Demo") p.add_argument("--share", action="store_true") p.add_argument("--server_name", default="0.0.0.0") p.add_argument("--server_port", type=int, default=None) p.add_argument("--static_dir", default="./static") p.add_argument("--tmp_dir", default="./tmp") p.add_argument("--css_path", default="layout.css") p.add_argument("--js_path", default="script.js") p.add_argument("--display_shortest", type=int, default=512, help="Shortest side for display images in the web interface, preserving aspect ratio") p.add_argument("--log_dir", default="./logs") p.add_argument("--case_timeout", type=float, default=600.0, help="max seconds per request from click to final yield") return p.parse_args() class SessionLogger: def __init__(self, base_dir: Path): self.base_dir = Path(base_dir) self.base_dir.mkdir(parents=True, exist_ok=True) stamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") self.dir = self.base_dir / stamp self.dir.mkdir(parents=True, exist_ok=False) self.paths = {"frontal": None, "lateral": None} self.anat_to_path = {} self.patho_to_path = {} self._counters = {"anatomy": 0, "pathology": 0} def save_input_images(self, pa_path: str, lat_path: Optional[str]): frontal_dst = self.dir / "frontal.jpg" Image.open(pa_path).convert("RGB").save(frontal_dst, "JPEG", quality=95, optimize=True, subsampling=1) self.paths["frontal"] = frontal_dst.name if lat_path: lateral_dst = self.dir / "lateral.jpg" Image.open(lat_path).convert("RGB").save(lateral_dst, "JPEG", quality=95, optimize=True, subsampling=1) self.paths["lateral"] = lateral_dst.name def _save_mask_jpg(self, mask_bool: np.ndarray, kind: str) -> str: name = f"{kind}_mask_{self._counters[kind]:03d}.jpg" self._counters[kind] += 1 out = self.dir / name img = Image.fromarray(mask_bool.astype(np.uint8) * 255, mode="L") img.save(out, "JPEG", quality=95, optimize=True, subsampling=1) return name def add_anatomy_mask(self, anatomy_heading: str, mask_bool: np.ndarray): fname = self._save_mask_jpg(mask_bool, "anatomy") self.anat_to_path[anatomy_heading] = fname def add_pathology_mask(self, anatomy_heading: str, sub_heading: str, mask_bool: np.ndarray): fname = self._save_mask_jpg(mask_bool, "pathology") self.patho_to_path[(anatomy_heading, sub_heading)] = fname def finalize_reasoning(self, reasoning_wrapper: "ReasoningWrapper", report_text: Optional[str] = None): reasoning_list = [] for card in reasoning_wrapper.cards: region = card.heading region_path = self.anat_to_path.get(region) patho_list = [] for sub in card.sub_cards: patho_list.append({ "anatomies": sub.heading, "reason": sub.content, "path": self.patho_to_path.get((region, sub.heading)) }) reasoning_list.append({ "region": region, "path": region_path, "pathological": patho_list }) payload = { "input_image": self.paths["frontal"], "lateral_image": self.paths["lateral"], "reasoning": reasoning_list, "report": report_text } out_json = self.dir / "reasoning.json" with open(out_json, "w", encoding="utf8") as f: json.dump(payload, f, ensure_ascii=False, indent=2) return out_json # --------------------------------------------- # system message and regex # --------------------------------------------- # SYSTEM_MESSAGE = """You are a radiologist assistant. You must strictly follow the inspection and reasoning protocol described below. # You perform a step-by-step inspection, followed by a summary of findings, diagnostic impression, and final report. # Inspection procedure: # 1. section # For each anatomical region that is commonly examined: # a. You will describe your focus by saying: Reviewing {region}... # b. Then, you bring attention to this area by calling the anatomical tool in the form: # # c. Immediately after the anatomical review, you must inspect at least one pathological sub-part belonging to that region. For each sub-part: # Write: Inspecting {sub-part}... # Then call the pathological tool in the exact form: # # Describe the observation for that sub-part in precise clinical terms. # d. Each anatomical region is considered incomplete unless it includes at least one “Inspecting {sub-part}...” step with its corresponding pathological tool call. # e. After completing all inspections for one region, proceed to the next anatomical region and repeat the same sequence. # Continue this process until all anatomical regions and their pathological sub-parts have been both reviewed and inspected. # 2. section # Summarize all inspected observations grouped by anatomical region. # Each line must follow: # - **{Region Name}**: {finding} # 3. section # Provide clinically meaningful diagnostic conclusions. # Each line must follow: # - {impression} # 4. section # Compose a complete, fluent narrative report summarizing all findings and suggest impressions. # Required output structure: # # Reviewing {region}... # # Inspecting {sub-part}... # # {observation} # [repeat for all regions and sub-parts] # # # - {region}: {summary of observations} # [repeat for all regions] # # # {concise diagnostic conclusions} # # Guidence: # 1. Use the most specific anatomical label possible from the internal knowledge graph. # 2. The pathological region for inspection should be within the anatomical region for review. Each anatomical region must has its own pathological region. # 3. Avoid inspecting the same region twice. # 4. Observations should be relevant to the region of image, and clinically accurate. # 5. You must base all judgments strictly on the visual evidence in the provided image. Don't rely on general statistical expectations unless the evidence clearly supports them. # """ SYSTEM_MESSAGE = """You are a radiologist assistant. You must strictly follow the inspection and reasoning protocol described below. You perform a step-by-step inspection, followed by a summary of findings, diagnostic impression, and final report. Inspection procedure: 1. section For each anatomical region that is commonly examined: a. You will describe your focus by saying: Reviewing {region}... b. Then, you bring attention to this area by calling the anatomical tool in the form: c. After examining the anatomical region, move on to its relevant pathological sub-parts. d. Indicate this step by saying: Inspecting {sub-part}... e. Immediately follow with a pathological tool call in the form: f. Then describe what you observe for that sub-part in clear clinical terms. Continue this process until all regions and their sub-parts have been reviewed. 2. section Summarize all inspected observations grouped by anatomical region. Each line must follow: - **{Region Name}**: {finding} 3. section Provide clinically meaningful diagnostic conclusions. Each line must follow: - {impression} 4. section Compose a complete, fluent narrative report summarizing all findings and suggest impressions. Required output structure: Reviewing {region}... Inspecting {sub-part}... {observation} [repeat for all regions and sub-parts] - {region}: {summary of observations} [repeat for all regions] {concise diagnostic conclusions} Guidence: 1. Use the most specific anatomical label possible from the internal knowledge graph. 2. The pathological region for inspection should be within the anatomical region for review. Each anatomical region must has its own pathological region. 3. Avoid inspecting the same region twice. 4. Observations should be relevant to the region of image, and clinically accurate. 5. You must base all judgments strictly on the visual evidence in the provided image. Don't rely on general statistical expectations unless the evidence clearly supports them. 6. Do not use comparative, temporal, or change-related language. """ # SYSTEM_MESSAGE = """You are a radiologist assistant. You must follow a strict inspection protocol. # 1. Always begin with an section. # 2. For EACH anatomical region: # a. Output: Reviewing {region}... # b. Immediately output a tool call in this form: # # c. After the anatomical tool, output: inspecting {sub-part}... # d. Immediately follow with a pathological_roi tool call in this form: # # e. Then describe the observation for that sub-part. # 3. Repeat the cycle until ALL regions and their sub-parts have been inspected. # 4. Do not output , , or until the entire section is complete. # The required final structure is: # # Reviewing {region}... # # inspecting {sub-part}... # # {observation} # [repeat for all regions and sub-parts] # # # - {region}: {summary of observations} # [repeat for all regions] # # # {concise diagnostic conclusions} # # # {full narrative radiology report} # # Rules: # - Each "Reviewing ..." MUST be followed by an anatomical_roi . # - Each "inspecting ..." MUST be followed by a pathological_roi . # - Skipping tool calls is invalid. # - Jumping to , , or before all inspections is invalid. # """ INTERPRET_START = re.compile(r'interpret>', re.I) ANATOMY_RE = re.compile(r"Reviewing\s+(.+?)\.\.\.", re.I) SUB_ANATOMY_RE = re.compile(r"Inspecting\s+(.+?)\.\.\.", re.I) TOOL_RE = re.compile(r'', re.I) REPORT_RE = re.compile(r'^(.*?\.)\n', re.S) INTERPRET_END = re.compile(r'/interpret>', re.I) FINDING_RE = re.compile(r'(.*?)', re.I | re.S) IMPRESSION_RE = re.compile(r'(.*?)', re.I | re.S) FINAL_REPORT_RE = re.compile(r'(.*?)', re.I | re.S) UNFINISHED = re.compile(r"(]*\blabel=\[[^\]]+\][^>]*>)", re.I) LABEL_RE = re.compile(r'label=\s*(\[[^\]]+\])', re.I) # --------------------------------------------- # ROI pipeline identical to dataset builder # --------------------------------------------- def ids_for_many(names: Iterable[str]) -> List[int]: out, seen = [], set() for n in names or []: n = str(n).strip().lower() if n in cxas_name2id: for i in cxas_name2id[n]: i = int(i) if i not in seen: seen.add(i) out.append(i) return out def mask_union(arr: Optional[np.ndarray], idx: List[int]) -> Optional[np.ndarray]: if arr is None or not idx: return None safe = [i for i in idx if 0 <= i < arr.shape[0]] if not safe: return None return arr[safe].any(axis=0) def bbox_from_mask(mask: Optional[np.ndarray], w: int, h: int, pad: int = 0) -> Tuple[int, int, int, int]: if mask is None or not mask.any(): return 0, w, 0, h ys, xs = np.where(mask) x0 = max(0, int(xs.min()) - pad) x1 = min(w, int(xs.max()) + pad) y0 = max(0, int(ys.min()) - pad) y1 = min(h, int(ys.max()) + pad) return x0, x1, y0, y1 def to_alpha(mask_bool: Optional[np.ndarray], invert: bool, feather: int, size_wh: Tuple[int, int]) -> Image.Image: if mask_bool is None: return Image.new("L", size_wh, color=255 if invert else 0) sel = (~mask_bool if invert else mask_bool).astype(np.uint8) * 255 if size_wh != (mask_bool.shape[1], mask_bool.shape[0]): sel = cv2.resize(sel, size_wh, interpolation=cv2.INTER_NEAREST) if feather > 0: sel = cv2.GaussianBlur(sel, ksize=(0, 0), sigmaX=feather, sigmaY=feather) return Image.fromarray(sel, mode="L") def optionally_resize(img: Image.Image, target_wh: Optional[Tuple[int, int]]) -> Image.Image: if not target_wh: return img w, h = int(target_wh[0]), int(target_wh[1]) return img.resize((w, h), Image.BICUBIC) def save_jpg(img: Image.Image, path: Path, quality: int = 95): path.parent.mkdir(parents=True, exist_ok=True) if img.mode not in ("L", "RGB"): img = img.convert("RGB") img.save(path, format="JPEG", quality=quality, subsampling=1, optimize=True) def save_blur(img_base: Image.Image, mask_bool: Optional[np.ndarray], out_path: Path, blur_radius: int, feather: int, roi_wh: Optional[Tuple[int, int]]): base = img_base.convert("RGB") blurred = base.filter(ImageFilter.GaussianBlur(radius=blur_radius)) alpha = to_alpha(mask_bool, invert=True, feather=feather, size_wh=base.size) comp = Image.composite(blurred, base, alpha) comp = optionally_resize(comp, roi_wh) save_jpg(comp, out_path) def save_crop(img_base: Image.Image, mask_bool: Optional[np.ndarray], out_path: Path, ring: int, roi_wh: Optional[Tuple[int, int]]): base = img_base w, h = base.size x0, x1, y0, y1 = bbox_from_mask(mask_bool, w, h, pad=ring) crop = base.crop((x0, y0, x1, y1)) crop = optionally_resize(crop, roi_wh) save_jpg(crop, out_path) def save_blurcrop(img_base: Image.Image, mask_bool: Optional[np.ndarray], out_path: Path, blur_radius: int, feather: int, ring: int, roi_wh: Optional[Tuple[int, int]]): base = img_base.convert("RGB") w, h = base.size if mask_bool is not None and (mask_bool.shape[1], mask_bool.shape[0]) != (w, h): mask_resized = cv2.resize(mask_bool.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST).astype(bool) else: mask_resized = mask_bool if mask_resized is None: crop = base else: x0, x1, y0, y1 = bbox_from_mask(mask_resized, w, h, pad=ring) crop = base.crop((x0, y0, x1, y1)) mask_resized = mask_resized[y0:y1, x0:x1] blurred = crop.filter(ImageFilter.GaussianBlur(radius=blur_radius)) alpha = to_alpha(mask_resized, invert=True, feather=feather, size_wh=crop.size) comp = Image.composite(blurred, crop, alpha) comp = optionally_resize(comp, roi_wh) save_jpg(comp, out_path) def save_viz(img_base: Image.Image, arr0: Optional[np.ndarray], idx: List[int], out_path: Path, mode: str, blur_radius: int, feather: int, ring: int, roi_wh: Optional[Tuple[int, int]]): mask = mask_union(arr0, idx) if idx else None if mode == "blur": save_blur(img_base, mask, out_path, blur_radius=blur_radius, feather=feather, roi_wh=roi_wh) elif mode == "crop": save_crop(img_base, mask, out_path, ring=ring, roi_wh=roi_wh) else: save_blurcrop(img_base, mask, out_path, blur_radius=blur_radius, feather=feather, ring=ring, roi_wh=roi_wh) # --------------------------------------------- # UI classes # --------------------------------------------- from abc import abstractmethod class Card: @property def check_state(self): if self.checked: return "checked" return "" @abstractmethod def uncheck(self): pass @abstractmethod def check(self): pass @staticmethod def decode_match(match): return [match.group(1)] class SubAnatomyCard(Card): template =""" """ def __init__(self, i, j, heading): self.i = i self.j = j self.checked = True self.heading = heading self.content = "" def uncheck(self): self.checked = False def check(self): self.checked = True def update_content(self, content): self.content = content return self.content def render(self): return self.template.format(i=self.i, j=self.j, check_state=self.check_state, heading=self.heading, content=self.content) class AnatomyCard(Card): template = """ """ def __init__(self, i, heading): self.i = i self.j = 0 self.checked = True self.heading = heading self.sub_cards = [] def uncheck(self): self.checked = False for sub_card in self.sub_cards: sub_card.checked = False def check(self): self.checked = True def add_sub_card(self, sub_card_heading): sub_card = SubAnatomyCard(self.i, self.j, sub_card_heading) self.sub_cards.append(sub_card) self.j += 1 return sub_card def render(self): return self.template.format(i=self.i, check_state=self.check_state, heading=self.heading, sub_cards="\n".join([x.render() for x in self.sub_cards])) class ReasoningWrapper: html_icon = """ """ template = """
    {icon} Reasoning Process
    {cards}
    """ def __init__(self): self.i = 0 self.cards = [] def add_card(self, card_heading): card = AnatomyCard(self.i, card_heading) self.cards.append(card) self.i += 1 return card def select(self, card, sub_card=None): if sub_card is not None: sub_card.check() for x in card.sub_cards: if x != sub_card: x.uncheck() card.check() for x in self.cards: if x != card: x.uncheck() def unselect(self, card): card.uncheck() def render(self): return self.template.format(icon=self.html_icon, cards="\n".join([x.render() for x in self.cards])) def output(self): explain = [] for card in self.cards: roi = card.heading reason = [] for sub_card in card.sub_cards: reason.append(sub_card.content) reason = " ".join(reason) explain.append({"roi": roi, "reason": reason}) return {"explain": explain} from utils import smooth, mask_to_svg class Mask: @property def hidden_style(self): if self.to_show: return "" return "display:none;" @property def svg(self): return mask_to_svg(self.mask) @abstractmethod def hide(self): pass @abstractmethod def show(self): pass @staticmethod def preprocess(mask): m = smooth(mask) return m @staticmethod def decode_match(match, anatomy_masks=None): tool_type = match.group(1) tool_labels = match.group(2) tool_labels = [x.strip().strip('"').strip("'").lower() for x in tool_labels.split("\", \"")] idxs = [i for label in tool_labels for i in cxas_name2id[label]] mask = sum(anatomy_masks[i] for i in idxs) > 0 return idxs, mask class SubAnatomyMask(Mask): def __init__(self, i, j, mask): self.i = i self.j = j self.mask = mask self.to_show = True def hide(self): self.to_show = False def show(self): self.to_show = True def render(self): return self.svg.format(sub_class="sub-svg", prefix="sub-", index=f"{self.i}-{self.j}", hidden_style=self.hidden_style, sub_svgs="", extra_data="") class AnatomyMask(Mask): def __init__(self, i, mask): self.i = i self.j = 0 self.mask = mask self.to_show = True self.sub_masks = [] @property def inner_mask(self): inner_mask = np.zeros_like(self.mask) if self.sub_masks: inner_mask = sum(sub_mask.mask for sub_mask in self.sub_masks) > 0 inner_mask = inner_mask.astype(np.uint8) * 255 return inner_mask def hide(self): self.to_show = False for sub_mask in self.sub_masks: sub_mask.to_show = False def show(self): self.to_show = True def add_sub_mask(self, mask): sub_mask = SubAnatomyMask(self.i, self.j, mask) self.sub_masks.append(sub_mask) self.j += 1 return sub_mask def render(self): return self.svg.format(sub_class="", prefix="", index=self.i, hidden_style=self.hidden_style, sub_svgs="\n".join([x.render() for x in self.sub_masks]), extra_data="") def to_b64(path): with open(path, "rb") as f: return "data:image/png;base64," + base64.b64encode(f.read()).decode() class ImageWrapper: img_icon = """ """ template = """
    {label_icon} Demo Image
    {icon} {img_tag} {masks}
    """ img_tag_template = 'Demo Image' def __init__(self, img_path=""): self.i = 0 self.j = 0 self.set_img(img_path) self.masks = [] @property def icon(self): if self.img_tag: return "" return ImageWrapper.img_icon def set_img(self, img_path): if img_path is None or img_path == '': self.img_path = "" self.img_tag = "" else: self.img_path = img_path img_base64 = to_b64(img_path) self.img_tag = self.img_tag_template.format(img_base64=img_base64) def add_mask(self, mask): mask = AnatomyMask(self.i, mask) self.masks.append(mask) self.i += 1 return mask def select(self, mask, sub_mask=None): if sub_mask is not None: sub_mask.show() for x in mask.sub_masks: if x != sub_mask: x.hide() mask.show() for x in self.masks: if x != mask: x.hide() def unselect(self, mask): mask.hide() def render(self): return self.template.format(label_icon=self.img_icon, icon=self.icon, img_tag=self.img_tag, masks="\n".join([x.render() for x in self.masks])) def output(self): return { "image_path": self.img_path, "explain": [{"mask": mask.inner_mask} for mask in self.masks] } class IFrame: icon = """ """ iframe = """ """ template = """
    {label_icon} Interactive Demo
    {icon} {iframe}
    """ def __init__(self, html_content=""): self.html_content = html_content self.icon = IFrame.icon self.iframe = "" if self.html_content: import html as _html self.iframe = IFrame.iframe.format(html_content=_html.escape(html_content, quote=True)) self.icon = "" @property def style(self): if self.html_content: return "" return "min-height: 400px;" def render(self): return self.template.format(style=self.style, label_icon=IFrame.icon, icon=self.icon, iframe=self.iframe) class TextBox: icon = """ """ list_template = """
    {icon} {label}
      {value}
    """ paragraph_template = """
    {icon} {label}
    {value}
    """ def __init__(self, label, type="list", value=""): self.type = type self.template = getattr(TextBox, f"{type}_template", self.paragraph_template) self.icon = TextBox.icon self.label = label self.value = self.preprocess(value) @property def style(self): if self.value: return "" return "min-height: 300px;" def preprocess(self, value): value = re.sub(r'', '', value, flags=re.I).strip() lines = [] pattern = r'^\s*\*+\s*(.*?)\s*\*+\s*:\s*(.*)$' for line in value.splitlines(): line = line.strip().lstrip("-").strip() m = re.match(pattern, line) if m: key, value = m.groups() line = f"{key}: {value}" if self.type == "list": line = f"
  • {line}
    • " lines.append(line) return '\n'.join(lines) def update_content(self, value): self.value = self.preprocess(value) return value def render(self): return self.template.format(style=self.style, icon=self.icon, label=self.label, value=self.value) # --------------------------------------------- # small helpers # --------------------------------------------- def to_b64(path: str) -> str: with open(path, "rb") as f: return "data:image/png;base64," + base64.b64encode(f.read()).decode() # --------------------------------------------- # build models # --------------------------------------------- def build_models(args): if args.device == "auto": device = "cuda" if torch.cuda.is_available() else "cpu" else: device = args.device hf_token = ( args.hf_token or os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN") ) is_local_model = Path(args.model).exists() if not is_local_model and not hf_token: raise RuntimeError( "No Hugging Face token found for remote/private model loading. " "Set HF_TOKEN (or HUGGING_FACE_HUB_TOKEN) in Space Secrets." ) common_kwargs = {"trust_remote_code": True} if hf_token: # Transformers now requires using only `token`. common_kwargs["token"] = hf_token try: proc = AutoProcessor.from_pretrained(args.model, **common_kwargs) tok = AutoTokenizer.from_pretrained(args.model, token=hf_token) vlm = AutoModelForVision2Seq.from_pretrained( args.model, torch_dtype=torch.float16 if device == "cuda" else torch.float32, **common_kwargs, ).to(device).eval() except OSError as exc: raise RuntimeError( f"Failed to load model '{args.model}'. " "If it is private, ensure your Space secret HF_TOKEN has read access to that repo." ) from exc cxas_model = cxas.CXAS(gpus=args.cxas_gpus if device == "cuda" else "cpu").eval() return device, proc, tok, vlm, cxas_model # --------------------------------------------- # time budget stoppers # --------------------------------------------- class _GenStop(StoppingCriteria): def __init__(self, tokenizer, pattern): super().__init__() self.tok = tokenizer self.pattern = pattern self.decoded = "" def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs): self.decoded += self.tok.decode(input_ids[0][-1:], skip_special_tokens=False) return bool(self.pattern.search(self.decoded)) class _TimeBudgetStop(torch.nn.Module): def __init__(self, start_time: float, budget_sec: float): super().__init__() self.start_time = start_time self.budget = budget_sec self.triggered = False def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs): if time.time() - self.start_time >= self.budget: self.triggered = True return True return False # --------------------------------------------- # app factory so args are captured # --------------------------------------------- def build_app(args, device=None, proc=None, tok=None, vlm=None, cxas_model=None): RESIZE_BASE_TO = tuple(map(int, args.resize_base_to)) RESIZE_ROI_TO = tuple(map(int, args.resize_roi_to)) VIZ_MODE = args.viz_mode CONTEXT_RING = int(args.context_ring) BLUR_RADIUS = int(args.blur_radius) FEATHER_SIGMA = int(args.feather) TMP_DIR = Path(args.tmp_dir); TMP_DIR.mkdir(parents=True, exist_ok=True) LOG_DIR = Path(args.log_dir); LOG_DIR.mkdir(parents=True, exist_ok=True) def cleanup_dir(path: Path): try: if path.exists(): shutil.rmtree(path, ignore_errors=True) except Exception: pass def generate(pa_path: str, lat_path: Optional[str]): nonlocal device, proc, tok, vlm, cxas_model start_time = time.time() budget = float(args.case_timeout) request_id = uuid.uuid4().hex request_tmp_dir = TMP_DIR / request_id request_log_dir = LOG_DIR / request_id request_tmp_dir.mkdir(parents=True, exist_ok=True) request_log_dir.mkdir(parents=True, exist_ok=True) generated_report = False def ensure_time(need_sec: float = 0.0): if time.time() - start_time >= budget - max(0.0, need_sec): raise TimeoutError try: if any(x is None for x in (device, proc, tok, vlm, cxas_model)): device, proc, tok, vlm, cxas_model = build_models(args) if not pa_path: raise gr.Error("Upload at least a frontal image") logger = SessionLogger(base_dir=request_log_dir) logger.save_input_images(pa_path, lat_path) img_wrapper = ImageWrapper(pa_path) reasoning_wrapper = ReasoningWrapper() iframe = IFrame() finding_wrapper = TextBox(label="Findings", value="", type="list") impression_wrapper = TextBox(label="Impressions", value="", type="list") report_wrapper = TextBox(label="Report", value="", type="paragraph") yield ( gr.update(value=img_wrapper.render()), gr.update(value=reasoning_wrapper.render()), gr.update(value=iframe.render()), gr.update(value=finding_wrapper.render()), gr.update(value=impression_wrapper.render()), gr.update(value=report_wrapper.render()), ) ensure_time() anatomy_masks = cxas_model.seg(pa_path) ensure_time() anatomy_masks = np.asarray(anatomy_masks, dtype=bool) pa_orig = Image.open(pa_path).convert("RGB") ensure_time() pa_base = pa_orig.resize(RESIZE_BASE_TO, Image.BICUBIC) ses_dir = request_tmp_dir / "cxas" (ses_dir / "masks").mkdir(parents=True, exist_ok=True) pa_tmp = ses_dir / "0.jpg" pa_orig.save(pa_tmp) ensure_time() cxas_model.process_file( filename=str(pa_tmp), do_store=True, output_directory=str(ses_dir / "masks"), storage_type="npy", ) ensure_time() npy_path = ses_dir / "masks" / "0.npy" if not npy_path.exists(): raise gr.Error("CXAS did not produce mask npy") raw = np.load(npy_path).astype(bool) ensure_time() anatomy_masks_orig = raw C, H, W = raw.shape bw, bh = RESIZE_BASE_TO if (W, H) != (bw, bh): resized = np.zeros((C, bh, bw), dtype=bool) for i in range(C): resized[i] = cv2.resize( raw[i].astype(np.uint8), (bw, bh), interpolation=cv2.INTER_NEAREST, ).astype(bool) anatomy_masks = resized else: anatomy_masks = raw ensure_time() orig_w, orig_h = pa_orig.size disp_short = args.display_shortest if orig_w < orig_h: disp_w = disp_short disp_h = int(orig_h * disp_short / orig_w) else: disp_h = disp_short disp_w = int(orig_w * disp_short / orig_h) pa_disp = pa_orig.resize((disp_w, disp_h), Image.BICUBIC) ensure_time() disp_path = request_tmp_dir / "display.jpg" pa_disp.save(disp_path) img_wrapper = ImageWrapper(str(disp_path)) resize_map = { "orig": (orig_w, orig_h), "display": (disp_w, disp_h), "model": tuple(RESIZE_BASE_TO), } C, H, W = anatomy_masks.shape if (W, H) != pa_base.size: resized = np.zeros((C, RESIZE_BASE_TO[1], RESIZE_BASE_TO[0]), dtype=bool) for i in range(C): resized[i] = cv2.resize( anatomy_masks[i].astype(np.uint8), RESIZE_BASE_TO, interpolation=cv2.INTER_NEAREST, ).astype(bool) anatomy_masks = resized ensure_time() lat_base = None if lat_path: lat_base = Image.open(lat_path).convert("RGB").resize(RESIZE_BASE_TO, Image.BICUBIC) ensure_time() convo = [ {"role": "system", "content": [{"type": "text", "text": SYSTEM_MESSAGE}]}, {"role": "user", "content": [ {"type": "image", "image": pa_path}, *([{"type": "image", "image": lat_path}] if lat_path else []), {"type": "text", "text": "Based on the provided chest radiographs, explain your diagnosis procedure and write a report."} ]} ] from utils import RegexStopper stoppers = { "interpret_start": RegexStopper(INTERPRET_START), "anatomy": RegexStopper(ANATOMY_RE), "sub_anatomy": RegexStopper(SUB_ANATOMY_RE), "anatomy_tool": RegexStopper(TOOL_RE), "sub_anatomy_tool": RegexStopper(TOOL_RE), "report" : RegexStopper(REPORT_RE), "interpret_end": RegexStopper(INTERPRET_END), "finding": RegexStopper(FINDING_RE), "impression": RegexStopper(IMPRESSION_RE), "final_report": RegexStopper(FINAL_REPORT_RE), } sequence = { "interpret_start": ["anatomy"], "anatomy": ["anatomy_tool"], "anatomy_tool": ["sub_anatomy"], "sub_anatomy": ["sub_anatomy_tool"], "sub_anatomy_tool": ["report"], "report": ["anatomy", "sub_anatomy", "interpret_end"], "interpret_end": ["finding"], "finding": ["impression"], "impression": ["final_report"], "final_report": [], } cur_items = { "interpret_start": [], "anatomy": None, "anatomy_tool": None, "sub_anatomy": None, "sub_anatomy_tool": None, "report": "", "interpret_end": [], "finding": "", "impression": "", "final_report": "", } match_decoder_fn = { "interpret_start": lambda x: [None], "anatomy": lambda m: [m.group(1)], "anatomy_tool": partial(Mask.decode_match, anatomy_masks=anatomy_masks), "sub_anatomy": lambda m: [m.group(1)], "sub_anatomy_tool": partial(Mask.decode_match, anatomy_masks=anatomy_masks), "report": lambda m: [m.group(1)], "interpret_end": lambda x: [None], "finding": lambda m: [m.group(1)], "impression": lambda m: [m.group(1)], "final_report": lambda m: [m.group(1)], } cur_sequences = ["interpret_start"] reply_full = "" while True: images_for_proc = [] for m in convo: for p in m["content"]: if p.get("type") == "image": path = p["image"] if path == pa_path: images_for_proc.append(pa_base) elif lat_path and path == lat_path: images_for_proc.append(lat_base) else: images_for_proc.append(Image.open(path).convert("RGB")) ensure_time() prompt = proc.apply_chat_template(convo, tokenize=False, add_generation_prompt=True) inputs = proc(text=[prompt], images=[images_for_proc], padding=True, return_tensors="pt").to(device) ensure_time() streamer = TextIteratorStreamer( tok, skip_prompt=True, skip_special_tokens=False, decode_kwargs={"skip_special_tokens": False}, ) stopper = _GenStop(tok, UNFINISHED) time_stop = _TimeBudgetStop(start_time, budget) stopping = StoppingCriteriaList([stopper, time_stop]) gen_kwargs = dict( **inputs, max_new_tokens=1024, eos_token_id=tok.eos_token_id, streamer=streamer, stopping_criteria=stopping ) thread = threading.Thread(target=vlm.generate, kwargs=gen_kwargs) thread.start() response = "" for chunk in streamer: response += chunk ensure_time() for ch in chunk: if "report" in cur_sequences: cur_items["report"] += ch if cur_items["sub_anatomy"] is not None: cur_items["sub_anatomy"].update_content(cur_items["report"]) elif "finding" in cur_sequences: cur_items["finding"] += ch finding_wrapper.update_content(cur_items["finding"]) elif "impression" in cur_sequences: cur_items["impression"] += ch impression_wrapper.update_content(cur_items["impression"]) elif "final_report" in cur_sequences: cur_items["final_report"] += ch report_wrapper.update_content(cur_items["final_report"]) for name in list(cur_sequences): if not stoppers[name](ch): continue cur_stopper = stoppers[name] decoded = match_decoder_fn[name](cur_stopper.match) if name == "anatomy": cur_items[name] = reasoning_wrapper.add_card(*decoded) reasoning_wrapper.select(cur_items["anatomy"]) if cur_items["anatomy_tool"]: img_wrapper.unselect(cur_items["anatomy_tool"]) elif name == "anatomy_tool": idxs, mask_model = decoded mask_for_display = cv2.resize( mask_model.astype(np.uint8), resize_map["display"], interpolation=cv2.INTER_NEAREST ).astype(bool) cur_items[name] = img_wrapper.add_mask(Mask.preprocess(mask_for_display)) img_wrapper.select(cur_items["anatomy_tool"]) roi_path = request_tmp_dir / f"roi_{uuid.uuid4().hex}.jpg" save_viz( img_base=pa_base, arr0=anatomy_masks, idx=idxs, out_path=roi_path, mode=VIZ_MODE, blur_radius=BLUR_RADIUS, feather=FEATHER_SIGMA, ring=CONTEXT_RING, roi_wh=tuple(RESIZE_ROI_TO) ) mask_orig = mask_union(anatomy_masks_orig, idxs) if idxs else np.zeros(anatomy_masks_orig.shape[1:], dtype=bool) logger.add_anatomy_mask(cur_items["anatomy"].heading, mask_orig) elif name == "sub_anatomy": cur_items[name] = cur_items["anatomy"].add_sub_card(*decoded) reasoning_wrapper.select(cur_items["anatomy"], cur_items["sub_anatomy"]) elif name == "sub_anatomy_tool": idxs, mask_model = decoded mask_for_display = cv2.resize( mask_model.astype(np.uint8), resize_map["display"], interpolation=cv2.INTER_NEAREST ).astype(bool) cur_items[name] = cur_items["anatomy_tool"].add_sub_mask(Mask.preprocess(mask_for_display)) img_wrapper.select(cur_items["anatomy_tool"], cur_items["sub_anatomy_tool"]) roi_path = request_tmp_dir / f"roi_{uuid.uuid4().hex}.jpg" save_viz( img_base=pa_base, arr0=anatomy_masks, idx=idxs, out_path=roi_path, mode=VIZ_MODE, blur_radius=BLUR_RADIUS, feather=FEATHER_SIGMA, ring=CONTEXT_RING, roi_wh=tuple(RESIZE_ROI_TO) ) mask_orig = mask_union(anatomy_masks_orig, idxs) if idxs else np.zeros(anatomy_masks_orig.shape[1:], dtype=bool) logger.add_pathology_mask(cur_items["anatomy"].heading, cur_items["sub_anatomy"].heading, mask_orig) elif name == "report": cur_items[name] = "" elif name == "final_report": generated_report = True elif name == "interpret_end": from render import render_diagram_html img_wrapper_json = img_wrapper.output() reasoning_wrapper_json = reasoning_wrapper.output() explain_json = {"explain": []} for iw, rw in zip(img_wrapper_json['explain'], reasoning_wrapper_json["explain"]): explain_json["explain"].append({ "roi": rw["roi"], "mask": iw["mask"], "reason": rw["reason"] }) img_wrapper_json.update(reasoning_wrapper_json) img_wrapper_json.update(explain_json) json_output = img_wrapper_json html_content = render_diagram_html(json_output) iframe = IFrame(html_content=html_content) logger.finalize_reasoning(reasoning_wrapper) cur_sequences = sequence[name] logger.finalize_reasoning( reasoning_wrapper, report_text=re.sub(r'', '', cur_items["final_report"].strip(), flags=re.I).strip() ) yield ( gr.update(value=img_wrapper.render()), gr.update(value=reasoning_wrapper.render()), gr.update(value=iframe.render()), gr.update(value=finding_wrapper.render()), gr.update(value=impression_wrapper.render()), gr.update(value=report_wrapper.render()), ) thread.join(timeout=0.1) ensure_time() reply_full += response m_unfinished = UNFINISHED.search(response) if m_unfinished: tool_call_end = m_unfinished.end() text_to_send = response[:tool_call_end] idxs = [] # safe default mask_bool = mask_union(anatomy_masks, idxs) if idxs else np.zeros_like(anatomy_masks[0], bool) roi_path = request_tmp_dir / f"roi_{uuid.uuid4().hex}.jpg" save_viz( img_base=pa_base, arr0=anatomy_masks, idx=idxs, out_path=roi_path, mode=VIZ_MODE, blur_radius=BLUR_RADIUS, feather=FEATHER_SIGMA, ring=CONTEXT_RING, roi_wh=tuple(RESIZE_ROI_TO) ) convo.append({ "role": "assistant", "content": [ {"type": "text", "text": text_to_send}, {"type": "image", "image": str(roi_path)}, ] }) yield ( gr.update(value=img_wrapper.render()), gr.update(value=reasoning_wrapper.render()), gr.update(value=iframe.render()), gr.update(value=finding_wrapper.render()), gr.update(value=impression_wrapper.render()), gr.update(value=report_wrapper.render()), ) continue break except TimeoutError: msg = '
    Generation stopped due to per case time limit.
    ' yield ( gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(value=msg), ) return finally: if generated_report: cleanup_dir(request_tmp_dir) cleanup_dir(request_log_dir) if spaces is not None: generate = spaces.GPU(generate, duration=120) with open(args.css_path, "r") as f: css = f.read() js_path_candidates = [ Path(args.js_path), Path("script.js"), Path("templates/script.js"), Path("static/script.js"), ] js_source = "" for cand in js_path_candidates: if cand.exists(): js_source = cand.read_text() break js = "\n".join([""]) static_dir = Path(args.static_dir).resolve() tmp_dir_abs = Path(args.tmp_dir).resolve() static_dir.mkdir(parents=True, exist_ok=True) tmp_dir_abs.mkdir(parents=True, exist_ok=True) with gr.Blocks(title=args.title, css=css, head=js) as demo: gr.Markdown(f"## {args.title}") with gr.Row(elem_id="row1"): with gr.Column(elem_id="col1"): with gr.Column(elem_id="col1-1"): pa_in = gr.Image(elem_classes=["image"], label="Upload Frontal", type="filepath") lat_in = gr.Image(elem_classes=["image"], label="Upload Lateral optional", type="filepath") btn = gr.Button("Generate", variant="primary") img_out = gr.HTML(ImageWrapper().render()) reasoning_out = gr.HTML(ReasoningWrapper().render()) diagram_out = gr.HTML(IFrame().render()) with gr.Row(elem_id="row3"): with gr.Column(scale=2): findings_out = gr.HTML(TextBox(label="Findings", value="", type="list").render()) with gr.Column(scale=1): impressions_out = gr.HTML(TextBox(label="Impressions", value="", type="list").render()) with gr.Row(elem_id="row4"): report_out = gr.HTML(TextBox(label="Report", value="", type="paragraph").render()) btn.click( generate, inputs=[pa_in, lat_in], outputs=[img_out, reasoning_out, diagram_out, findings_out, impressions_out, report_out] ) return demo # --------------------------------------------- # main # --------------------------------------------- if __name__ == "__main__": args = parse_args() device, proc, tok, vlm, cxas_model = build_models(args) demo = build_app(args, device, proc, tok, vlm, cxas_model) demo.launch(share=args.share, server_name=args.server_name, server_port=args.server_port)