demo_log.py

# #!/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. <interpret> 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:
# <tool type="anatomical_roi" label=["{region}"]><image>
# 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:
# <tool type="pathological_roi" label=["{sub-part}"]><image>
# 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. <finding> section

# Summarize all inspected observations grouped by anatomical region.

# Each line must follow:
# - **{Region Name}**: {finding}

# 3. <impression> section

# Provide clinically meaningful diagnostic conclusions.

# Each line must follow:
# - {impression}

# 4. <report> section

# Compose a complete, fluent narrative report summarizing all findings and suggest impressions.


# Required output structure:

# <interpret>
# Reviewing {region}...
# <tool type="anatomical_roi" label=["{region}"]><image>
# Inspecting {sub-part}...
# <tool type="pathological_roi" label=["{sub-part}"]><image>
# {observation}
# [repeat for all regions and sub-parts]
# </interpret>

# <finding>
# - {region}: {summary of observations}
# [repeat for all regions]
# </finding>

# <impression>
# {concise diagnostic conclusions}
# </impression>


# 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'<tool\s+type="([^"]+)"\s+label=\[([^\]]*)\]>', re.I)
# REPORT_RE       = re.compile(r'^(.*?\.)\n', re.S)
# INTERPRET_END   = re.compile(r'/interpret>', re.I)

# FINDING_RE      = re.compile(r'<finding>(.*?)</finding>', re.I | re.S)
# IMPRESSION_RE   = re.compile(r'<impression>(.*?)</impression>', re.I | re.S)
# FINAL_REPORT_RE = re.compile(r'<report>(.*?)</report>', re.I | re.S)

# UNFINISHED      = re.compile(r"(<tool[^>]*\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 ="""
#     <label class="sub-card" for="sub-card{i}-{j}">
#         <input class="sub-card-toggle" id="sub-card{i}-{j}" i="{i}" j="{j}" type="checkbox" {check_state}>
#         <div class="sub-heading">{heading}</div>
#         <div class="sub-content">{content}</div>
#     </label>
#     """
#     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 = """
#     <label class="card" for="card{i}">
#     <input class="card-toggle" id="card{i}" i="{i}" type="checkbox" {check_state}>
#     <div class="heading">
#         Reviewing <span class="highlight">{heading}</span> ...
#     </div>
#     <div class="content">
#         {sub_cards}
#     </div>
#     </label>
#     """
#     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 = """
#     <svg class="html-icon" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 25.5 18.14" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round">
#         <polyline points="7.39 15.16 .75 9.07 .75 9.06 7.39 2.97"/>
#         <polyline points="18.11 15.16 24.75 9.07 24.75 9.06 18.11 2.97"/>
#         <line x1="10.79" y1="17.39" x2="14.71" y2=".75"/>
#     </svg>
#     """
#     template = """
#     <div id="col3">
#         <div class="block-label">
#             {icon}
#             Reasoning Process
#         </div>
#         <div class="reasoning-wrapper">
#             {cards}
#         </div>
#     </div>
#     """
#     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 = """
#     <svg class="img-icon" xmlns='http://www.w3.org/2000/svg' viewBox='0 0 24 24' fill='none' stroke='currentColor' stroke-width='1.5' stroke-linecap='round' stroke-linejoin='round'>
#         <rect x='3' y='3' width='18' height='18' rx='2' ry='2'/>
#         <circle cx='8.5' cy='8.5' r='1.5'/>
#         <polyline points='21 15 16 10 5 21'/>
#     </svg>
#     """

#     template = """
#     <div id="col2">
#         <div class="block-label">
#             {label_icon}
#             Demo Image
#         </div>
#         <div class="img-wrapper">
#             {icon}
#             {img_tag}
#             {masks}
#         </div>
#     </div>
#     """
#     img_tag_template = '<img src="{img_base64}" alt="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 = """
#     <svg class="img-icon" xmlns='http://www.w3.org/2000/svg' viewBox='0 0 24 24' fill='none' stroke='currentColor' stroke-width='1.5' stroke-linecap='round' stroke-linejoin='round'>
#         <rect x='3' y='3' width='18' height='18' rx='2' ry='2'/>
#         <circle cx='8.5' cy='8.5' r='1.5'/>
#         <polyline points='21 15 16 10 5 21'/>
#     </svg>
#     """
#     iframe = """
#     <iframe id="diagram-html"
#         onload='this.style.height=this.contentWindow.document.body.scrollHeight + 20 + "px";'
#         src="{html_path}?t={time}"
#     ></iframe>
#     """
#     template = """
#     <div id="row2" style="{style}">
#         <div class="block-label">
#             {label_icon}
#             Interactive Demo
#         </div>
#         <div class="img-wrapper">
#             {icon}
#             {iframe}
#         </div>
#     </div>
#     """

#     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 = """
#     <svg class="html-icon" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 25.5 18.14" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round">
#         <polyline points="7.39 15.16 .75 9.07 .75 9.06 7.39 2.97"/>
#         <polyline points="18.11 15.16 24.75 9.07 24.75 9.06 18.11 2.97"/>
#         <line x1="10.79" y1="17.39" x2="14.71" y2=".75"/>
#     </svg>
#     """
#     list_template = """
#     <div class="text-box" style="{style}">
#         <div class="block-label">
#             {icon}
#             {label}
#         </div>
#         <ul class="text-wrapper">
#             {value}
#         </ul>
#     </div>
#     """
#     paragraph_template = """
#     <div class="text-box" style="{style}">
#         <div class="block-label">
#             {icon}
#             {label}
#         </div>
#         <div class="text-wrapper">
#             {value}
#         </div>
#     </div>
#     """
#     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'</?(finding|impression|report)>', '', value, flags=re.I).strip()

#         # Add bold text and <li> 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"<b>{key}:</b> {value}"
#             if self.type == "list":
#                 line = f"<li>{line}</li>"
#             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(["<script>", f.read(), "</script>"])

#     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. <interpret> 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:
# <tool type="anatomical_roi" label=["{region}"]><image>
# 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:
# <tool type="pathological_roi" label=["{sub-part}"]><image>
# 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. <finding> section

# Summarize all inspected observations grouped by anatomical region.

# Each line must follow:
# - **{Region Name}**: {finding}

# 3. <impression> section

# Provide clinically meaningful diagnostic conclusions.

# Each line must follow:
# - {impression}

# 4. <report> section

# Compose a complete, fluent narrative report summarizing all findings and suggest impressions.


# Required output structure:

# <interpret>
# Reviewing {region}...
# <tool type="anatomical_roi" label=["{region}"]><image>
# Inspecting {sub-part}...
# <tool type="pathological_roi" label=["{sub-part}"]><image>
# {observation}
# [repeat for all regions and sub-parts]
# </interpret>

# <finding>
# - {region}: {summary of observations}
# [repeat for all regions]
# </finding>

# <impression>
# {concise diagnostic conclusions}
# </impression>


# 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. <interpret> 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:
<tool type="anatomical_roi" label=["{region}"]><image>
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:
<tool type="pathological_roi" label=["{sub-part}"]><image>
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. <finding> section

Summarize all inspected observations grouped by anatomical region.

Each line must follow:
- **{Region Name}**: {finding}

3. <impression> section

Provide clinically meaningful diagnostic conclusions.

Each line must follow:
- {impression}

4. <report> section

Compose a complete, fluent narrative report summarizing all findings and suggest impressions.


Required output structure:

<interpret>
Reviewing {region}...
<tool type="anatomical_roi" label=["{region}"]><image>
Inspecting {sub-part}...
<tool type="pathological_roi" label=["{sub-part}"]><image>
{observation}
[repeat for all regions and sub-parts]
</interpret>

<finding>
- {region}: {summary of observations}
[repeat for all regions]
</finding>

<impression>
{concise diagnostic conclusions}
</impression>


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 <interpret> section.  
# 2. For EACH anatomical region:  
#    a. Output: Reviewing {region}...  
#    b. Immediately output a tool call in this form:  
#       <tool type="anatomical_roi" label=["{region}"]><image>  
#    c. After the anatomical tool, output: inspecting {sub-part}...  
#    d. Immediately follow with a pathological_roi tool call in this form:  
#       <tool type="pathological_roi" label=["{sub-part}"]><image>  
#    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 <finding>, <impression>, or <report> until the entire <interpret> section is complete.  

# The required final structure is:

# <interpret>
# Reviewing {region}...
# <tool type="anatomical_roi" label=["{region}"]><image>
# inspecting {sub-part}...
# <tool type="pathological_roi" label=["{sub-part}"]><image>
# {observation}
# [repeat for all regions and sub-parts]
# </interpret>

# <finding>
# - {region}: {summary of observations}
# [repeat for all regions]
# </finding>

# <impression>
# {concise diagnostic conclusions}
# </impression>

# <report>
# {full narrative radiology report}
# </report>

# Rules:  
# - Each "Reviewing ..." MUST be followed by an anatomical_roi <tool>.  
# - Each "inspecting ..." MUST be followed by a pathological_roi <tool>.  
# - Skipping tool calls is invalid.  
# - Jumping to <finding>, <impression>, or <report> 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'<tool\s+type="([^"]+)"\s+label=\[([^\]]*)\]>', re.I)
REPORT_RE       = re.compile(r'^(.*?\.)\n', re.S)
INTERPRET_END   = re.compile(r'/interpret>', re.I)

FINDING_RE      = re.compile(r'<finding>(.*?)</finding>', re.I | re.S)
IMPRESSION_RE   = re.compile(r'<impression>(.*?)</impression>', re.I | re.S)
FINAL_REPORT_RE = re.compile(r'<report>(.*?)</report>', re.I | re.S)

UNFINISHED      = re.compile(r"(<tool[^>]*\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 ="""
    <label class="sub-card" for="sub-card{i}-{j}">
        <input class="sub-card-toggle" id="sub-card{i}-{j}" i="{i}" j="{j}" type="checkbox" {check_state}>
        <div class="sub-heading">{heading}</div>
        <div class="sub-content">{content}</div>
    </label>
    """
    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 = """
    <label class="card" for="card{i}">
    <input class="card-toggle" id="card{i}" i="{i}" type="checkbox" {check_state}>
    <div class="heading">
        Reviewing <span class="highlight">{heading}</span> ...
    </div>
    <div class="content">
        {sub_cards}
    </div>
    </label>
    """
    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 = """
    <svg class="html-icon" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 25.5 18.14" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round">
        <polyline points="7.39 15.16 .75 9.07 .75 9.06 7.39 2.97"/>
        <polyline points="18.11 15.16 24.75 9.07 24.75 9.06 18.11 2.97"/>
        <line x1="10.79" y1="17.39" x2="14.71" y2=".75"/>
    </svg>
    """
    template = """
    <div id="col3">
        <div class="block-label">
            {icon}
            Reasoning Process
        </div>
        <div class="reasoning-wrapper">
            {cards}
        </div>
    </div>
    """
    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 = """
    <svg class="img-icon" xmlns='http://www.w3.org/2000/svg' viewBox='0 0 24 24' fill='none' stroke='currentColor' stroke-width='1.5' stroke-linecap='round' stroke-linejoin='round'>
        <rect x='3' y='3' width='18' height='18' rx='2' ry='2'/>
        <circle cx='8.5' cy='8.5' r='1.5'/>
        <polyline points='21 15 16 10 5 21'/>
    </svg>
    """

    template = """
    <div id="col2">
        <div class="block-label">
            {label_icon}
            Demo Image
        </div>
        <div class="img-wrapper">
            {icon}
            {img_tag}
            {masks}
        </div>
    </div>
    """
    img_tag_template = '<img src="{img_base64}" alt="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 = """
    <svg class="img-icon" xmlns='http://www.w3.org/2000/svg' viewBox='0 0 24 24' fill='none' stroke='currentColor' stroke-width='1.5' stroke-linecap='round' stroke-linejoin='round'>
        <rect x='3' y='3' width='18' height='18' rx='2' ry='2'/>
        <circle cx='8.5' cy='8.5' r='1.5'/>
        <polyline points='21 15 16 10 5 21'/>
    </svg>
    """
    iframe = """
    <iframe id="diagram-html"
        onload='this.style.height=this.contentWindow.document.body.scrollHeight + 20 + "px";'
        srcdoc="{html_content}"
    ></iframe>
    """
    template = """
    <div id="row2" style="{style}">
        <div class="block-label">
            {label_icon}
            Interactive Demo
        </div>
        <div class="img-wrapper">
            {icon}
            {iframe}
        </div>
    </div>
    """

    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 = """
    <svg class="html-icon" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 25.5 18.14" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round">
        <polyline points="7.39 15.16 .75 9.07 .75 9.06 7.39 2.97"/>
        <polyline points="18.11 15.16 24.75 9.07 24.75 9.06 18.11 2.97"/>
        <line x1="10.79" y1="17.39" x2="14.71" y2=".75"/>
    </svg>
    """
    list_template = """
    <div class="text-box" style="{style}">
        <div class="block-label">
            {icon}
            {label}
        </div>
        <ul class="text-wrapper">
            {value}
        </ul>
    </div>
    """
    paragraph_template = """
    <div class="text-box" style="{style}">
        <div class="block-label">
            {icon}
            {label}
        </div>
        <div class="text-wrapper">
            {value}
        </div>
    </div>
    """
    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'</?(finding|impression|report)>', '', 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"<b>{key}:</b> {value}"
            if self.type == "list":
                line = f"<li>{line}</li>"
            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'</?report>', '', 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 = '<div class="text-wrapper">Generation stopped due to per case time limit.</div>'
            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(["<script>", js_source, "</script>"])

    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)