diff --git a/.gitignore b/.gitignore index 5ada641..4fa1fd6 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,17 @@ *$py.class OAI_CONFIG_LIST +# webapp demo: per-job uploads/outputs and downloaded smoke-test fixtures +webapp/data/* +!webapp/data/.gitkeep +!webapp/data/samples/ +webapp/data/samples/* +!webapp/data/samples/.gitkeep + +# unrelated local scratch content, not part of this project +/IVUS/ +/.claude/ + # C extensions *.so diff --git a/README.md b/README.md index bd5f3d0..7a7b705 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,15 @@ The skill is intended for environments that already have nnUNet v2 available. To run inference, the target environment must provide a configured `nnUNet_results` path and usable nnUNet CLI or Python API support. +## Demo Webapp +[`webapp/`](webapp) contains a standalone, ChatGPT-style chat demo: type a free-text segmentation request and upload a CT/MRI `.nii.gz` volume, and it runs real [TotalSegmentator](https://github.com/wasserth/TotalSegmentator) inference (modality auto-detected from the volume) and shows the result overlaid in-browser via [Niivue](https://niivue.com/). See [`webapp/README.md`](webapp/README.md) for setup and details. Once dependencies are installed (`pip install -r webapp/requirements.txt`), launch it with: + +```bash +uvicorn webapp.backend.main:app --host 0.0.0.0 --port 8000 --reload +``` + +Then open `http://localhost:8000/`. + ## Supported Datasets The repository currently includes metadata for the following datasets and targets. The summary below is adapted from the dataset table in the paper. The links point to official dataset access pages rather than direct-download mirrors hosted by this repository; some datasets still require registration, challenge participation, or a signed data-use agreement before download. diff --git a/webapp/README.md b/webapp/README.md new file mode 100644 index 0000000..d37c266 --- /dev/null +++ b/webapp/README.md @@ -0,0 +1,137 @@ +# MedSegAgent Demo Webapp + +A ChatGPT-style chat UI where you type a natural-language segmentation +request (e.g. "肝臓をセグメンテーションして" / "segment the liver") and +upload a CT or MRI volume (`.nii.gz`). The app parses the query with an LLM, +resolves the target structure against this repo's +[`dataset/TotalSegmentator_v2.json`](../dataset/TotalSegmentator_v2.json) / +[`TotalSegmentator_MRI.json`](../dataset/TotalSegmentator_MRI.json) label +maps, runs real [TotalSegmentator](https://github.com/wasserth/TotalSegmentator) +inference (`total` / `total_mr` tasks, Apache-2.0, no license key needed), +and shows the result overlaid on the volume with +[Niivue](https://niivue.com/). + +This is purely additive: it does not modify anything in `agents/`, `utils/`, +`dataset/`, or the root `requirements.txt`. + +Scope: single-target queries only (e.g. "segment the liver"), CT and MRI +only. TotalSegmentator's `fast` (3mm) mode is used for every request; there +is no quality toggle. + +Result display follows the familiar image-chat pattern (ChatGPT/Gemini +style): each completed turn shows a small static thumbnail inline in its +chat bubble; tapping it opens a modal with the full interactive Niivue +viewer (multiplanar view, scrollable slices). There is one shared Niivue +canvas for the whole page, reused for both off-screen thumbnail rendering +and the on-screen modal. + +## Features + +- Natural-language query in English or Japanese ("segment the liver" / + "肝臓をセグメンテーションして") -- an LLM extracts the target structure. +- Modality (CT vs. MRI) is auto-detected from the uploaded volume's own + voxel intensities, not from the query text (a query often won't mention + modality at all even though the file is unambiguously one or the other). +- Real segmentation via TotalSegmentator, not a mock -- the whole task is + run and the requested structure is isolated into its own mask file. +- Chat-style result: a tappable thumbnail per turn, expanding into an + interactive multiplanar volume viewer with an overlay on/off toggle. + +## Project layout + +``` +webapp/ +├── backend/ +│ ├── main.py # FastAPI app: /api/chat, /api/jobs/{id}, static file serving +│ ├── config.py # paths, model name, fast-mode default +│ ├── query_parser.py # LLM call: free text -> {modality, target} +│ ├── modality_detector.py # CT vs MR from voxel intensities +│ ├── label_matcher.py # target string -> TotalSegmentator label name/id +│ ├── segmentation_runner.py # runs TotalSegmentator (GPU-serialized) +│ ├── mask_extractor.py # isolates just the target label into its own file +│ └── jobs.py # in-memory job status store +├── frontend/ +│ ├── index.html / app.js / style.css # vanilla JS chat UI + Niivue viewer +└── data/ # per-job uploads/outputs (gitignored) + smoke-test samples +``` + +## Setup + +```bash +cd /home/nlab/MedSegAgent +source .venv/bin/activate + +# Check for dependency conflicts with the pinned research deps +# (pyautogen==0.2.35, spacy==3.8.7, ...) before installing for real. +# If this wants to change already-installed torch/nnunetv2/spacy/pyautogen +# versions, stop and reconsider (e.g. a separate venv) rather than proceeding. +pip install --dry-run -r webapp/requirements.txt + +pip install -r webapp/requirements.txt +``` + +`OAI_CONFIG_LIST` at the repo root must already contain a working entry for +the model named in `webapp/backend/config.py` (`MODEL_NAME`, default +`"deepseek-v4-flash"`). + +## Run + +```bash +cd /home/nlab/MedSegAgent +source .venv/bin/activate +uvicorn webapp.backend.main:app --host 0.0.0.0 --port 8000 --reload +``` + +Then open `http://localhost:8000/` (WSL2 forwards this to the Windows host +automatically). + +## Smoke test + +Grab TotalSegmentator's own bundled test fixtures (small, real NIfTI +volumes, safe to use for a demo): + +```bash +mkdir -p webapp/data/samples +curl -L -o webapp/data/samples/example_ct_sm.nii.gz \ + https://raw.githubusercontent.com/wasserth/TotalSegmentator/master/tests/reference_files/example_ct_sm.nii.gz +curl -L -o webapp/data/samples/example_mr_sm.nii.gz \ + https://raw.githubusercontent.com/wasserth/TotalSegmentator/master/tests/reference_files/example_mr_sm.nii.gz +``` + +Backend only (no browser needed): + +```bash +curl -F "query=segment the liver" \ + -F "file=@webapp/data/samples/example_ct_sm.nii.gz" \ + http://localhost:8000/api/chat +# -> {"job_id": "...", ...} +curl http://localhost:8000/api/jobs/ +# poll until "status": "done" +``` + +Then repeat with a Japanese query (`肝臓をセグメンテーションして`) and with +the MRI fixture (`query=segment the spleen`, +`file=@webapp/data/samples/example_mr_sm.nii.gz`) to exercise both language +handling and the MRI path. + +Finally, drive the same three queries from the browser UI and confirm the +Niivue canvas shows a plausible overlay, and the "オーバーレイを表示" +checkbox toggles it on/off. + +## Known limitations + +- Every request segments the *entire* task (all ~117 CT / ~56 MRI + structures), not just the requested one. An earlier version used + TotalSegmentator's `roi_subset` to restrict inference to just the target + structure for speed, but testing found that `roi_subset`'s preliminary + low-resolution "rough crop" detection pass can silently miss small + structures (reproduced: `roi_subset=["spleen"]` on `total_mr` came back + completely empty, while `roi_subset=["liver"]` on the same task worked). + Running the full task every time is slower but has been reliably correct; + see `segmentation_runner.py`'s docstring. +- Single-user/single-GPU demo: concurrent requests are serialized on a GPU + lock (`segmentation_runner.GPU_LOCK`), not parallelized. +- Job data (`webapp/data//`) is cleaned up automatically after ~1 + hour (`config.JOB_MAX_AGE_SECONDS`), but is otherwise unbounded during + that window -- fine for a demo, not for unattended long-term operation. +- No auth, no persistence across server restarts (jobs are in-memory). diff --git a/webapp/__init__.py b/webapp/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/webapp/__init__.py diff --git a/webapp/backend/__init__.py b/webapp/backend/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/webapp/backend/__init__.py diff --git a/webapp/backend/config.py b/webapp/backend/config.py new file mode 100644 index 0000000..d6810e2 --- /dev/null +++ b/webapp/backend/config.py @@ -0,0 +1,38 @@ +"""Central configuration for the webapp demo backend. + +Deliberately does not import anything from the repo's agents/ or utils/ +packages (see webapp/backend/query_parser.py for why) -- this module only +knows about plain paths and constants. +""" + +from pathlib import Path + +# webapp/backend/config.py -> backend -> webapp -> repo root +BASE_DIR = Path(__file__).resolve().parents[2] +WEBAPP_DIR = BASE_DIR / "webapp" + +DATA_DIR = WEBAPP_DIR / "data" +FRONTEND_DIR = WEBAPP_DIR / "frontend" + +DATASET_DIR = BASE_DIR / "dataset" +CT_LABELS_PATH = DATASET_DIR / "TotalSegmentator_v2.json" +MRI_LABELS_PATH = DATASET_DIR / "TotalSegmentator_MRI.json" + +OAI_CONFIG_PATH = BASE_DIR / "OAI_CONFIG_LIST" + +# Model used for query parsing and label disambiguation. Confirmed present +# and working in OAI_CONFIG_LIST during development of this app -- change +# here if the entry in OAI_CONFIG_LIST changes. +MODEL_NAME = "deepseek-v4-flash" + +# TotalSegmentator quality mode. Fixed to fast (3mm) for demo responsiveness; +# no UI toggle by design (confirmed with the user). +FAST_MODE = True + +# Reject uploads larger than this (bytes). +MAX_UPLOAD_BYTES = 1024 * 1024 * 1024 # 1 GiB + +# Delete job directories older than this many seconds on each new request. +JOB_MAX_AGE_SECONDS = 60 * 60 # 1 hour + +ALLOWED_UPLOAD_EXTENSIONS = (".nii.gz", ".nii") diff --git a/webapp/backend/jobs.py b/webapp/backend/jobs.py new file mode 100644 index 0000000..579ae73 --- /dev/null +++ b/webapp/backend/jobs.py @@ -0,0 +1,32 @@ +"""In-memory job store. Single-process, single-user demo -- no DB/Redis +needed. Guarded by a plain lock since FastAPI's BackgroundTasks run in a +threadpool, not the event loop thread. +""" + +import threading +import time + +_jobs = {} +_jobs_lock = threading.Lock() + + +def create_job(job_id, **fields): + with _jobs_lock: + _jobs[job_id] = { + "status": "queued", + "message": "", + "created_at": time.time(), + **fields, + } + + +def update_job(job_id, **fields): + with _jobs_lock: + if job_id in _jobs: + _jobs[job_id].update(fields) + + +def get_job(job_id): + with _jobs_lock: + job = _jobs.get(job_id) + return dict(job) if job is not None else None diff --git a/webapp/backend/label_matcher.py b/webapp/backend/label_matcher.py new file mode 100644 index 0000000..712a8c8 --- /dev/null +++ b/webapp/backend/label_matcher.py @@ -0,0 +1,70 @@ +"""Resolve a free-text segmentation target (e.g. "liver") to a canonical +TotalSegmentator label name + integer id, using this repo's own +dataset/TotalSegmentator_v2.json (CT) / dataset/TotalSegmentator_MRI.json +(MRI) as the source of truth for the name<->id mapping. + +Resolution order: exact match -> normalized match -> unique substring match +-> LLM disambiguation against the full label list. +""" + +import json +import re + +from . import config +from .query_parser import extract_json_from_output, get_oai_client_for_model + + +def _labels_path_for_modality(modality): + return config.CT_LABELS_PATH if modality.upper().startswith("CT") else config.MRI_LABELS_PATH + + +def _load_labels(modality): + path = _labels_path_for_modality(modality) + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + return data["labels"] # {name: int_id}, includes "background": 0 + + +def _normalize(name): + return re.sub(r"[\s\-]+", "_", name.strip().lower()) + + +def _llm_disambiguate(target, labels, model): + names = [n for n in labels.keys() if n != "background"] + prompt = ( + "You are a medical imaging expert. Given the segmentation target " + f'"{target}", pick the single closest matching name from this list ' + f"of available structures: {json.dumps(names, ensure_ascii=False)}\n" + "Respond ONLY in this exact format:\n" + 'Final Answer: {"label": "chosen_name"}' + ) + client = get_oai_client_for_model(model) + response = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + ) + text = response.choices[0].message.content or "" + parsed = extract_json_from_output(text, expected_type="dict") + label_name = parsed.get("label") if parsed else None + if not label_name or label_name not in labels: + raise ValueError( + f"「{target}」に該当する構造が見つかりませんでした（LLM応答: {text!r}）。" + ) + return label_name + + +def resolve_target(target, modality, model): + """Returns {"name": str, "id": int}. Raises ValueError if unresolved.""" + labels = _load_labels(modality) + norm_target = _normalize(target) + + if norm_target in labels: + return {"name": norm_target, "id": labels[norm_target]} + + candidates = [name for name in labels if norm_target in name or name in norm_target] + candidates = [name for name in candidates if name != "background"] + if len(candidates) == 1: + return {"name": candidates[0], "id": labels[candidates[0]]} + + label_name = _llm_disambiguate(target, labels, model) + return {"name": label_name, "id": labels[label_name]} diff --git a/webapp/backend/main.py b/webapp/backend/main.py new file mode 100644 index 0000000..ebf4642 --- /dev/null +++ b/webapp/backend/main.py @@ -0,0 +1,222 @@ +"""FastAPI app for the MedSegAgent chat-based segmentation demo. + +Run from the repo root (so OAI_CONFIG_LIST/dataset/ paths in config.py and +plain package imports resolve correctly): + + cd /home/nlab/MedSegAgent + source .venv/bin/activate + uvicorn webapp.backend.main:app --host 0.0.0.0 --port 8000 --reload +""" + +import logging +import shutil +import time +import uuid + +from fastapi import BackgroundTasks, FastAPI, File, Form, HTTPException, UploadFile +from fastapi.responses import FileResponse +from fastapi.staticfiles import StaticFiles + +from . import config, jobs +from .label_matcher import resolve_target +from .mask_extractor import extract_single_label +from .modality_detector import detect_modality +from .query_parser import parse_query +from .segmentation_runner import GPU_LOCK, run_totalsegmentator + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger("medsegagent_webapp") + +app = FastAPI(title="MedSegAgent Demo") + +config.DATA_DIR.mkdir(parents=True, exist_ok=True) + + +class NoCacheStaticFiles(StaticFiles): + """StaticFiles with caching disabled. + + Without this, browsers can skip revalidating index.html/app.js/style.css + on a normal reload (heuristic caching kicks in when no Cache-Control + header is set) and silently keep serving a stale build after an edit -- + a real footgun during active development of this demo. /static (the + frontend source) always uses this; /data (per-job NIfTI results) does + not need it since a given job_id's files never change after creation. + """ + + def file_response(self, *args, **kwargs): + response = super().file_response(*args, **kwargs) + response.headers["Cache-Control"] = "no-store" + return response + + +app.mount("/data", StaticFiles(directory=str(config.DATA_DIR)), name="data") +app.mount("/static", NoCacheStaticFiles(directory=str(config.FRONTEND_DIR)), name="static") + + +@app.get("/") +def index(): + return FileResponse(str(config.FRONTEND_DIR / "index.html"), headers={"Cache-Control": "no-store"}) + + +def _cleanup_old_job_dirs(): + """Delete job directories older than config.JOB_MAX_AGE_SECONDS. + Called on every new /api/chat request; cheap for a handful of jobs.""" + now = time.time() + if not config.DATA_DIR.exists(): + return + for entry in config.DATA_DIR.iterdir(): + if entry.name == "samples" or not entry.is_dir(): + continue + try: + if now - entry.stat().st_mtime > config.JOB_MAX_AGE_SECONDS: + shutil.rmtree(entry, ignore_errors=True) + except FileNotFoundError: + pass + + +def _run_segmentation_job(job_id, input_path, seg_path, mask_path, modality, label_id): + """Background task: waits for the GPU lock, runs inference, extracts just + the target label into its own file, and updates job status. + + Status transitions: queued -> running -> done | error. The "queued" + state persists (accurately) while waiting on GPU_LOCK for a second + concurrent request, not just while waiting for FastAPI to schedule the + background task. + + seg_path ends up holding TotalSegmentator's full multi-label output + (every structure it found); mask_path holds a derived single-label mask + containing only `label_id` -- see mask_extractor.py for why the frontend + is served mask_path, not seg_path, as the overlay. + """ + try: + jobs.update_job(job_id, status="queued", message="GPUの空き待ちです...") + with GPU_LOCK: + jobs.update_job(job_id, status="running", message="セグメンテーション実行中...") + run_totalsegmentator(input_path, seg_path, modality) + extract_single_label(seg_path, label_id, mask_path) + jobs.update_job(job_id, status="done", message="完了しました。") + except Exception as e: # noqa: BLE001 - surface any failure to the UI + logger.exception("segmentation job %s failed", job_id) + jobs.update_job(job_id, status="error", message=f"推論に失敗しました: {e}") + + +@app.post("/api/chat") +async def chat( + background_tasks: BackgroundTasks, + query: str = Form(...), + file: UploadFile = File(...), +): + _cleanup_old_job_dirs() + + filename = file.filename or "" + if not filename.lower().endswith(config.ALLOWED_UPLOAD_EXTENSIONS): + raise HTTPException( + status_code=400, + detail="対応していないファイル形式です（.nii.gz または .nii のみアップロードできます）。", + ) + + job_id = uuid.uuid4().hex + job_dir = config.DATA_DIR / job_id + job_dir.mkdir(parents=True, exist_ok=True) + input_path = job_dir / "input.nii.gz" + + size = 0 + try: + with open(input_path, "wb") as f: + while True: + chunk = await file.read(1024 * 1024) + if not chunk: + break + size += len(chunk) + if size > config.MAX_UPLOAD_BYTES: + raise HTTPException( + status_code=413, + detail="ファイルサイズが大きすぎます（上限1GB）。", + ) + f.write(chunk) + except HTTPException: + shutil.rmtree(job_dir, ignore_errors=True) + raise + + try: + parsed = parse_query(query) + except Exception as e: + shutil.rmtree(job_dir, ignore_errors=True) + raise HTTPException(status_code=400, detail=f"クエリを解析できませんでした: {e}") + + target = parsed["target"] + + # Modality routing is decided from the uploaded volume's own voxel + # intensities, NOT from the LLM's parsed query text -- a query like + # "segment the liver" very often won't mention the modality at all even + # though the file itself is unambiguously a CT or MRI. See + # modality_detector.py for why. parsed["modality"] (e.g. "not specified", + # or whatever the user's wording implied) is kept only for display. + try: + modality = detect_modality(input_path) + except Exception as e: + shutil.rmtree(job_dir, ignore_errors=True) + raise HTTPException(status_code=400, detail=f"画像のモダリティを判定できませんでした: {e}") + + try: + resolved_label = resolve_target(target, modality, model=config.MODEL_NAME) + except Exception as e: + shutil.rmtree(job_dir, ignore_errors=True) + raise HTTPException(status_code=400, detail=str(e)) + + seg_path = job_dir / "segmentation.nii.gz" + mask_path = job_dir / "target_mask.nii.gz" + jobs.create_job( + job_id, + status="queued", + message="GPUの空き待ちです...", + modality=modality, + modality_from_query=parsed.get("modality"), + target=target, + resolved_label=resolved_label, + ) + background_tasks.add_task( + _run_segmentation_job, + job_id, + input_path, + seg_path, + mask_path, + modality, + resolved_label["id"], + ) + + return { + "job_id": job_id, + "modality": modality, + "modality_from_query": parsed.get("modality"), + "target": target, + "resolved_label": resolved_label, + "status": "queued", + } + + +@app.get("/api/jobs/{job_id}") +def get_job_status(job_id: str): + job = jobs.get_job(job_id) + if job is None: + raise HTTPException(status_code=404, detail="ジョブが見つかりません。") + + result = None + if job["status"] == "done": + result = { + "input_url": f"/data/{job_id}/input.nii.gz", + # target_mask.nii.gz (not the raw multi-label segmentation.nii.gz) + # -- contains only 0 and label_id, see mask_extractor.py. + "seg_url": f"/data/{job_id}/target_mask.nii.gz", + "label_id": job["resolved_label"]["id"], + "label_name": job["resolved_label"]["name"], + } + + return { + "status": job["status"], + "message": job.get("message", ""), + "modality": job.get("modality"), + "target": job.get("target"), + "resolved_label": job.get("resolved_label"), + "result": result, + } diff --git a/webapp/backend/mask_extractor.py b/webapp/backend/mask_extractor.py new file mode 100644 index 0000000..d9be511 --- /dev/null +++ b/webapp/backend/mask_extractor.py @@ -0,0 +1,31 @@ +"""Post-processes a TotalSegmentator multi-label output into a single-label +mask containing only the requested structure. + +Why this exists: since segmentation_runner.py always runs the full task +(all ~117 CT / ~56 MRI structures, see its docstring for why roi_subset was +dropped), the raw output NIfTI contains every structure TotalSegmentator +found, not just the one the user asked about. Isolating "just this one +label" was first attempted purely client-side in app.js via Niivue's +setColormapLabel, but that did not behave like a sparse index->color lookup +in testing -- every structure rendered in the overlay color, not just the +target. Zeroing out everything but the target label here, server-side, +sidesteps that entirely: the file handed to the browser only ever contains +0 and the target's label id, so however the client colors it, only the +target structure can show up. +""" + +import nibabel as nib +import numpy as np + + +def extract_single_label(segmentation_path, label_id, output_path): + """Reads segmentation_path, zeroes every voxel not equal to label_id, + and writes the result to output_path. Returns output_path.""" + img = nib.load(str(segmentation_path)) + data = np.asanyarray(img.dataobj) + mask = (np.rint(data).astype(np.int32) == int(label_id)).astype(np.uint8) * int(label_id) + + out_img = nib.Nifti1Image(mask, img.affine, header=img.header) + out_img.set_data_dtype(np.uint8) + nib.save(out_img, str(output_path)) + return output_path diff --git a/webapp/backend/modality_detector.py b/webapp/backend/modality_detector.py new file mode 100644 index 0000000..4b4d6c6 --- /dev/null +++ b/webapp/backend/modality_detector.py @@ -0,0 +1,26 @@ +"""Detect CT vs MRI directly from the uploaded volume's voxel intensities, +rather than relying on the user's query text mentioning the modality. + +Why this exists: a user asking "segment the liver" while uploading a CT +volume very often won't say the word "CT" at all -- the modality is implicit +in the file, not necessarily in the sentence. Falling back to the LLM's +parsed `modality` field (which can legitimately be "not specified") for +task routing would silently misroute CT volumes to the MRI model. Image +intensity is a much more reliable signal: CT values are calibrated +Hounsfield units where air is always ~-1000 HU, while MRI intensities are +arbitrary per-scan units that are essentially never that negative. +""" + +import nibabel as nib +import numpy as np + +# Air in CT is ~ -1000 HU; real MRI images essentially never go this negative. +CT_MIN_HU_THRESHOLD = -500 + + +def detect_modality(path): + """Returns "CT" or "MR" based on the 1st percentile voxel intensity.""" + img = nib.load(str(path)) + data = img.get_fdata() + p1 = float(np.percentile(data, 1)) + return "CT" if p1 < CT_MIN_HU_THRESHOLD else "MR" diff --git a/webapp/backend/query_parser.py b/webapp/backend/query_parser.py new file mode 100644 index 0000000..5c864e1 --- /dev/null +++ b/webapp/backend/query_parser.py @@ -0,0 +1,117 @@ +"""Natural-language query parsing: free text -> {"modality": ..., "target": ...}. + +NOTE on why this doesn't import from agents/ or utils/: +`utils/chat_utils.py` does `from agents import *`, which (via agents/__init__.py) +eagerly imports all ~10 agent pattern modules and the full pyautogen stack just +to reuse a couple of small, self-contained pieces. To keep this webapp +decoupled from that research-code import chain (today it's harmless, but it's +incidental, not structural -- see the implementation plan), the prompt text +and JSON-extraction helper are copied here verbatim from +agents/c2f_tree_search.py (lines ~35-91 at time of writing), and the OpenAI +client lookup is reimplemented directly instead of importing +utils.chat_utils.get_OAI_client_for_model. +""" + +import json +import logging +import re + +from autogen import config_list_from_json +from openai import OpenAI + +from . import config + +# Copied verbatim from agents/c2f_tree_search.py (SYSTEM_PROMPT_QUERY_PARSER). +SYSTEM_PROMPT_QUERY_PARSER = """ +# Role: +You are an AI assistant specialized in parsing queries related to medical image segmentation. Your task is to split the query into its imaging modality and segmentation target with as much specificity as possible. + +## Skills: +1. Query Parsing: +- Identify and extract the imaging modality, including specific subtypes (e.g., CT contrast, CT non-contrast, MR ADC, MR FLAIR, etc.) from the query. +- If specific subtypes are not mentioned, extract only the main modality (e.g., CT, MR). +- Identify and extract the segmentation target (e.g., brain, liver, tumor, etc.). + +## Constraints: +1. Response Format +- MUST use the following JSON format for the Final Answer: +Final Answer: {"modality": "modality_name", "target": "segmentation_target"} +Here are some examples: +{"modality": "CT", "target": "liver"} +{"modality": "MR", "target": "spleen"} +{"modality": "MR ADC", "target": "brain tumor"} +""" + + +def extract_json_from_output(llm_output, expected_type="dict"): + """Adapted from agents/c2f_tree_search.py's extract_json_from_output. + + One deliberate change from the original: the default expected_type is + "dict" here (the original repo function defaults to "list", which is a + footgun for this exact prompt -- SYSTEM_PROMPT_QUERY_PARSER above always + answers with a JSON *object*, so relying on the original default would + silently fail to match every single response). + """ + if expected_type == "list": + pattern = r"Final Answer:\s*(\[[\s\S]*?\])" + elif expected_type == "dict": + pattern = r"Final Answer:\s*(\{[\s\S]*?\})" + else: + raise ValueError("expected_type must be 'list' or 'dict'") + + match = re.search(pattern, llm_output) + if not match: + logging.error("No 'Final Answer' found in LLM output: %r", llm_output) + return None + + json_str = match.group(1) + json_str = re.sub(r"//.*", "", json_str) + json_str = re.sub(r"/\*[\s\S]*?\*/", "", json_str) + try: + return json.loads(json_str) + except json.JSONDecodeError: + logging.info("Initial JSON parse failed, retrying with quote fix...") + json_str_fixed = json_str.replace("'", '"') + try: + return json.loads(json_str_fixed) + except json.JSONDecodeError: + logging.error("Failed to parse JSON from LLM output: %r", json_str) + return None + + +def get_oai_client_for_model(model, config_path=None): + """Equivalent to utils/chat_utils.get_OAI_client_for_model, reimplemented + locally to avoid importing utils.chat_utils (see module docstring).""" + path = str(config_path or config.OAI_CONFIG_PATH) + cfg = config_list_from_json(env_or_file=path, filter_dict={"model": [model, "default"]})[0] + return OpenAI(api_key=cfg["api_key"], base_url=cfg["base_url"]) + + +def parse_query(query, model=None): + """Parse a free-text segmentation request into {"modality": str, "target": str}. + + Raises ValueError if the LLM response couldn't be parsed or is missing + required fields. + """ + model = model or config.MODEL_NAME + client = get_oai_client_for_model(model) + + response = client.chat.completions.create( + model=model, + messages=[ + {"role": "system", "content": SYSTEM_PROMPT_QUERY_PARSER}, + {"role": "user", "content": query}, + ], + ) + text = response.choices[0].message.content or "" + parsed = extract_json_from_output(text, expected_type="dict") + + if not parsed: + raise ValueError(f"クエリを解析できませんでした（LLM応答: {text!r}）") + + modality = parsed.get("modality") + target = parsed.get("target") or parsed.get("targets") + if not modality or not target: + raise ValueError(f"クエリからmodality/targetを抽出できませんでした: {parsed!r}") + + return {"modality": str(modality), "target": str(target)} diff --git a/webapp/backend/segmentation_runner.py b/webapp/backend/segmentation_runner.py new file mode 100644 index 0000000..5f0b849 --- /dev/null +++ b/webapp/backend/segmentation_runner.py @@ -0,0 +1,53 @@ +"""Runs TotalSegmentator inference on an uploaded volume. + +GPU_LOCK is exported (not used internally) so callers can distinguish +"queued, waiting for the GPU" from "actually running" in job status -- +see main.py's background task wrapper. +""" + +import threading + +from totalsegmentator.python_api import totalsegmentator + +from . import config + +# Single GPU in this deployment: serialize inference across concurrent +# requests to avoid two totalsegmentator() calls fighting over the same +# CUDA context / VRAM. +GPU_LOCK = threading.Lock() + + +def run_totalsegmentator(input_path, output_path, modality, fast=None): + """Runs TotalSegmentator and returns output_path. + + modality: a free-text modality string parsed from the user's query + (e.g. "CT", "CT contrast", "MR", "MR ADC"). Anything starting with + "CT" routes to task="total"; anything else routes to task="total_mr". + fast: overrides config.FAST_MODE if given. + + Deliberately does NOT use `roi_subset` to restrict inference to just the + requested structure, even though that would be faster: roi_subset adds a + preliminary low-resolution "rough crop" detection pass, and in testing + that pass silently missed a small structure (spleen) on the total_mr + task -- roi_subset=["liver"] worked, roi_subset=["spleen"] came back + completely empty on the exact same volume, while running the full task + (no roi_subset) found spleen correctly every time. Always segmenting + everything is slower but has been reliably correct in testing across + CT/MRI and multiple target organs; the target label is still extracted + afterward the same way. + """ + if fast is None: + fast = config.FAST_MODE + + task = "total" if modality.upper().startswith("CT") else "total_mr" + + totalsegmentator( + input=str(input_path), + output=str(output_path), + task=task, + fast=fast, + ml=True, # single multi-label NIfTI output, not one file per structure + device="gpu", + quiet=True, + ) + return output_path diff --git a/webapp/data/.gitkeep b/webapp/data/.gitkeep new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/webapp/data/.gitkeep diff --git a/webapp/data/samples/.gitkeep b/webapp/data/samples/.gitkeep new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/webapp/data/samples/.gitkeep diff --git a/webapp/frontend/app.js b/webapp/frontend/app.js new file mode 100644 index 0000000..54728a1 --- /dev/null +++ b/webapp/frontend/app.js @@ -0,0 +1,273 @@ +// MedSegAgent demo: ChatGPT-style chat UI driving a real TotalSegmentator run, +// visualized with Niivue. No build step, no framework -- see the implementation +// plan (webapp is a from-scratch additive feature) for context. +// +// Result display pattern: each completed turn shows a small static thumbnail +// inline in its chat bubble (like an image-generation chat) rather than a +// permanently visible viewer panel. Tapping a thumbnail opens a modal with +// the full interactive Niivue viewer (multiplanar view, scrollable slices). +// There is exactly one Niivue instance/canvas for the whole page, reused for +// both off-screen thumbnail rendering and the on-screen modal -- see +// style.css's .viewer-modal (visibility-toggled, never display:none, so the +// canvas keeps a stable non-zero size and never loses its WebGL context). + +const chatLog = document.getElementById("chat-log"); +const composer = document.getElementById("composer"); +const queryInput = document.getElementById("query-input"); +const fileInput = document.getElementById("file-input"); +const fileChip = document.getElementById("file-chip"); +const sendButton = document.getElementById("send-button"); +const overlayToggle = document.getElementById("overlay-toggle"); +const viewerModal = document.getElementById("viewer-modal"); +const viewerModalClose = document.getElementById("viewer-modal-close"); + +const POLL_INTERVAL_MS = 2000; + +// --- Niivue setup ----------------------------------------------------- + +const nv = new niivue.Niivue({ backColor: [0, 0, 0, 1] }); +nv.attachTo("gl"); + +const SEG_VOLUME_INDEX = 1; + +// Identifies whichever turn's volumes are currently loaded into the shared +// canvas, so opening a thumbnail can skip re-fetching the NIfTI files when +// it's already showing that same turn's result. +let loadedJobId = null; + +overlayToggle.addEventListener("change", () => { + if (!nv.volumes[SEG_VOLUME_INDEX]) return; + nv.volumes[SEG_VOLUME_INDEX].opacity = overlayToggle.checked ? 0.7 : 0; + nv.updateGLVolume(); +}); + +function openViewerModal() { + viewerModal.classList.remove("closed"); +} + +function closeViewerModal() { + viewerModal.classList.add("closed"); +} + +viewerModalClose.addEventListener("click", closeViewerModal); +viewerModal.addEventListener("click", (e) => { + if (e.target === viewerModal) closeViewerModal(); // click on the backdrop +}); + +// Loads a turn's result into the shared canvas and isolates the requested +// structure (background transparent, target label colored). No-op if this +// job's volumes are already the ones loaded. +async function loadResultIntoViewer(jobId, result) { + if (loadedJobId === jobId) return; + + await nv.loadVolumes([ + { url: result.input_url, colormap: "gray", opacity: 1 }, + { url: result.seg_url, colormap: "gray", opacity: 0.7 }, + ]); + + // TotalSegmentator was run with roi_subset=[target], so the segmentation + // volume should only ever contain 0 and this one label id anyway -- this + // isolation is just to guarantee background is fully transparent. + const cmap = { + R: [0, 255], + G: [0, 70], + B: [0, 70], + A: [0, 200], + I: [0, result.label_id], + labels: ["background", result.label_name], + }; + // setColormapLabel is a method on the NVImage instance (nv.volumes[i]), + // not on the Niivue instance itself -- confirmed against niivue source. + nv.volumes[SEG_VOLUME_INDEX].setColormapLabel(cmap); + nv.updateGLVolume(); + + overlayToggle.checked = true; + loadedJobId = jobId; +} + +// Renders a turn's result and captures a static PNG snapshot for the chat +// bubble thumbnail. Waits a couple of animation frames after loading so the +// GPU has actually painted before we read the canvas back. +async function renderThumbnail(jobId, result) { + await loadResultIntoViewer(jobId, result); + await new Promise((resolve) => requestAnimationFrame(() => requestAnimationFrame(resolve))); + return nv.canvas.toDataURL("image/png"); +} + +// --- Chat state --------------------------------------------------------- + +let selectedFile = null; +let busy = false; + +fileInput.addEventListener("change", () => { + selectedFile = fileInput.files[0] || null; + if (selectedFile) { + fileChip.textContent = selectedFile.name; + fileChip.hidden = false; + } else { + fileChip.hidden = true; + } +}); + +queryInput.addEventListener("input", () => { + queryInput.style.height = "auto"; + queryInput.style.height = Math.min(queryInput.scrollHeight, 128) + "px"; +}); + +queryInput.addEventListener("keydown", (e) => { + if (e.key === "Enter" && !e.shiftKey) { + e.preventDefault(); + composer.requestSubmit(); + } +}); + +composer.addEventListener("submit", async (e) => { + e.preventDefault(); + if (busy) return; + + const query = queryInput.value.trim(); + if (!query) return; + if (!selectedFile) { + addBubble("assistant", "先に .nii.gz / .nii ファイルをアップロードしてください。", { error: true }); + return; + } + + addBubble("user", query, { fileName: selectedFile.name }); + + const fileToSend = selectedFile; + queryInput.value = ""; + queryInput.style.height = "auto"; + fileInput.value = ""; + selectedFile = null; + fileChip.hidden = true; + + // A new turn is starting: don't leave the modal open showing a stale + // result while this turn's completion is about to reuse the same canvas. + closeViewerModal(); + + await runChatTurn(query, fileToSend); +}); + +function addBubble(role, text, opts = {}) { + const row = document.createElement("div"); + row.className = `bubble-row ${role}`; + + const bubble = document.createElement("div"); + bubble.className = "bubble" + (opts.error ? " error" : ""); + bubble.textContent = text; + + if (opts.fileName) { + const fileNameEl = document.createElement("span"); + fileNameEl.className = "file-name"; + fileNameEl.textContent = `📎 ${opts.fileName}`; + bubble.appendChild(fileNameEl); + } + + row.appendChild(bubble); + chatLog.appendChild(row); + chatLog.scrollTop = chatLog.scrollHeight; + return bubble; +} + +function addThumbnailToBubble(bubble, jobId, result, dataUrl) { + const wrap = document.createElement("div"); + wrap.className = "msg-thumb-wrap"; + + const img = document.createElement("img"); + img.className = "msg-thumb"; + img.src = dataUrl; + img.alt = `${result.label_name} segmentation preview`; + img.addEventListener("click", async () => { + await loadResultIntoViewer(jobId, result); + openViewerModal(); + }); + wrap.appendChild(img); + + const caption = document.createElement("div"); + caption.className = "msg-thumb-caption"; + caption.textContent = "タップして拡大 (断面をスクロール表示)"; + wrap.appendChild(caption); + + bubble.appendChild(wrap); + chatLog.scrollTop = chatLog.scrollHeight; +} + +function setBusy(isBusy) { + busy = isBusy; + sendButton.disabled = isBusy; +} + +async function runChatTurn(query, file) { + setBusy(true); + const statusBubble = addBubble("assistant", "クエリを解析中..."); + + const form = new FormData(); + form.append("query", query); + form.append("file", file, file.name); + + let submitResult; + try { + const res = await fetch("/api/chat", { method: "POST", body: form }); + submitResult = await res.json(); + if (!res.ok) { + throw new Error(submitResult.detail || "リクエストに失敗しました。"); + } + } catch (err) { + statusBubble.classList.add("error"); + statusBubble.textContent = `エラー: ${err.message}`; + setBusy(false); + return; + } + + const { job_id, modality, resolved_label } = submitResult; + statusBubble.textContent = + `対象を検出しました: ${resolved_label.name} (label ${resolved_label.id}, ${modality})\n` + + `セグメンテーション実行を待機中...`; + + const startedAt = Date.now(); + await pollJob(job_id, statusBubble, resolved_label, startedAt); + setBusy(false); +} + +function pollJob(jobId, statusBubble, resolvedLabel, startedAt) { + return new Promise((resolve) => { + const timer = setInterval(async () => { + let job; + try { + const res = await fetch(`/api/jobs/${jobId}`); + job = await res.json(); + if (!res.ok) throw new Error(job.detail || "ジョブ情報の取得に失敗しました。"); + } catch (err) { + clearInterval(timer); + statusBubble.classList.add("error"); + statusBubble.textContent = `エラー: ${err.message}`; + resolve(); + return; + } + + const elapsedSec = Math.round((Date.now() - startedAt) / 1000); + + if (job.status === "queued") { + statusBubble.textContent = `GPUの空き待ちです... (${elapsedSec}秒)`; + } else if (job.status === "running") { + statusBubble.textContent = `セグメンテーション実行中... (${elapsedSec}秒)`; + } else if (job.status === "done") { + clearInterval(timer); + statusBubble.textContent = + `完了しました（${elapsedSec}秒）: ${resolvedLabel.name} (label ${resolvedLabel.id}) を検出しました。`; + try { + const dataUrl = await renderThumbnail(jobId, job.result); + addThumbnailToBubble(statusBubble, jobId, job.result, dataUrl); + } catch (err) { + statusBubble.textContent += `\n(プレビュー画像の生成に失敗しました: ${err.message})`; + } + resolve(); + } else if (job.status === "error") { + clearInterval(timer); + statusBubble.classList.add("error"); + statusBubble.textContent = `エラー: ${job.message}`; + resolve(); + } + }, POLL_INTERVAL_MS); + }); +} diff --git a/webapp/frontend/index.html b/webapp/frontend/index.html new file mode 100644 index 0000000..0025868 --- /dev/null +++ b/webapp/frontend/index.html @@ -0,0 +1,56 @@ + + + + + + MedSegAgent Demo + + + + +
+
+ MedSegAgent + CT / MRI segmentation demo (TotalSegmentator) +
+ +
+
+ +
+ + + + +
+
+
+ + +
+
+ + +
+ +
+
+
+ + + + diff --git a/webapp/frontend/style.css b/webapp/frontend/style.css new file mode 100644 index 0000000..043bf0c --- /dev/null +++ b/webapp/frontend/style.css @@ -0,0 +1,290 @@ +:root { + --bg: #212121; + --panel-bg: #2a2a2a; + --bubble-user: #3a3a3a; + --bubble-assistant: #2f2f2f; + --bubble-error: #4a2626; + --text: #ececec; + --text-dim: #9b9b9b; + --accent: #10a37f; + --accent-dim: #0d8a6c; + --border: #3d3d3d; +} + +* { + box-sizing: border-box; +} + +html, body { + height: 100%; + margin: 0; + background: var(--bg); + color: var(--text); + font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", "Hiragino Sans", "Yu Gothic", sans-serif; +} + +.app { + display: flex; + flex-direction: column; + height: 100vh; +} + +.app-header { + display: flex; + align-items: baseline; + gap: 0.6em; + padding: 0.75em 1.2em; + border-bottom: 1px solid var(--border); + flex-shrink: 0; +} + +.app-title { + font-weight: 700; + font-size: 1.1em; +} + +.app-subtitle { + color: var(--text-dim); + font-size: 0.85em; +} + +.chat-panel { + flex: 1; + display: flex; + flex-direction: column; + min-height: 0; +} + +/* Result viewer modal. Uses visibility (not display:none) so the Niivue + canvas inside keeps a stable, non-zero layout size at all times -- it + needs to be paintable off-screen to produce thumbnails even while the + modal itself isn't shown to the user. */ +.viewer-modal { + position: fixed; + inset: 0; + display: flex; + align-items: center; + justify-content: center; + background: rgba(0, 0, 0, 0.75); + z-index: 1000; +} + +.viewer-modal.closed { + visibility: hidden; + opacity: 0; + pointer-events: none; +} + +.viewer-modal-inner { + position: relative; + width: min(1100px, 92vw); + height: min(720px, 86vh); + background: #000; + border-radius: 12px; + overflow: hidden; +} + +.viewer-modal-inner canvas { + width: 100%; + height: 100%; + display: block; +} + +.viewer-modal-close { + position: absolute; + top: 0.6em; + right: 0.7em; + width: 2em; + height: 2em; + border-radius: 50%; + border: none; + background: rgba(0, 0, 0, 0.55); + color: var(--text); + font-size: 1em; + cursor: pointer; + z-index: 1; +} + +.viewer-modal-close:hover { + background: rgba(0, 0, 0, 0.8); +} + +.viewer-controls { + position: absolute; + left: 0.8em; + bottom: 0.6em; + display: flex; + align-items: center; + gap: 1em; + background: rgba(0, 0, 0, 0.55); + padding: 0.4em 0.8em; + border-radius: 8px; + font-size: 0.85em; +} + +.overlay-toggle { + display: flex; + align-items: center; + gap: 0.4em; + cursor: pointer; + user-select: none; +} + +/* Inline result thumbnail shown inside an assistant chat bubble. */ +.msg-thumb-wrap { + margin-top: 0.5em; +} + +.msg-thumb { + display: block; + width: 220px; + max-width: 100%; + border-radius: 10px; + border: 1px solid var(--border); + background: #000; + cursor: pointer; +} + +.msg-thumb:hover { + filter: brightness(1.15); +} + +.msg-thumb-caption { + margin-top: 0.3em; + font-size: 0.78em; + color: var(--text-dim); +} + +.chat-log { + flex: 1; + overflow-y: auto; + padding: 1em 1.2em; + display: flex; + flex-direction: column; + gap: 0.9em; +} + +.bubble-row { + display: flex; +} + +.bubble-row.user { + justify-content: flex-end; +} + +.bubble-row.assistant { + justify-content: flex-start; +} + +.bubble { + max-width: 70%; + padding: 0.6em 0.9em; + border-radius: 14px; + line-height: 1.5; + white-space: pre-wrap; + font-size: 0.95em; +} + +.bubble-row.user .bubble { + background: var(--bubble-user); + border-bottom-right-radius: 4px; +} + +.bubble-row.assistant .bubble { + background: var(--bubble-assistant); + border-bottom-left-radius: 4px; +} + +.bubble.error { + background: var(--bubble-error); +} + +.bubble .file-name { + display: block; + margin-top: 0.35em; + font-size: 0.8em; + color: var(--text-dim); +} + +.bubble .status-line { + color: var(--text-dim); + font-size: 0.85em; +} + +.composer { + flex-shrink: 0; + display: flex; + align-items: flex-end; + gap: 0.5em; + padding: 0.8em 1.2em 1.1em; + border-top: 1px solid var(--border); +} + +.attach-button { + flex-shrink: 0; + display: flex; + align-items: center; + justify-content: center; + width: 2.4em; + height: 2.4em; + border-radius: 50%; + background: var(--panel-bg); + cursor: pointer; + font-size: 1.1em; +} + +.attach-button:hover { + background: var(--border); +} + +.file-chip { + flex-shrink: 0; + max-width: 10em; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; + background: var(--panel-bg); + border-radius: 999px; + padding: 0.35em 0.8em; + font-size: 0.8em; + color: var(--text-dim); +} + +.query-input { + flex: 1; + resize: none; + max-height: 8em; + border-radius: 18px; + border: 1px solid var(--border); + background: var(--panel-bg); + color: var(--text); + padding: 0.6em 1em; + font-size: 0.95em; + font-family: inherit; +} + +.query-input:focus { + outline: none; + border-color: var(--accent); +} + +.send-button { + flex-shrink: 0; + border: none; + border-radius: 18px; + background: var(--accent); + color: #fff; + padding: 0.6em 1.2em; + font-size: 0.95em; + cursor: pointer; +} + +.send-button:hover { + background: var(--accent-dim); +} + +.send-button:disabled { + background: var(--border); + color: var(--text-dim); + cursor: not-allowed; +} diff --git a/webapp/requirements.txt b/webapp/requirements.txt new file mode 100644 index 0000000..669af37 --- /dev/null +++ b/webapp/requirements.txt @@ -0,0 +1,12 @@ +# Dependencies for the webapp demo only. Installed into the same .venv as the +# repo root requirements.txt, but kept in a separate file so the research +# code's pinned deps (pyautogen==0.2.35, spacy==3.8.7, ...) stay easy to +# reason about on their own. Before installing, run: +# pip install --dry-run -r webapp/requirements.txt +# and check whether it wants to change torch/nnunetv2/spacy/pyautogen's +# already-installed versions -- if so, stop and decide (e.g. a separate venv) +# instead of installing blindly. +fastapi +uvicorn[standard] +python-multipart +TotalSegmentator