specs/2_code-level-spec.md

Code-Level Specification

1. Implementation Directive for Cursor

Implement the repository exactly according to this document and design-level-spec.md.

When the two documents conflict, this code-level specification takes precedence.

Do not add a backend, Gradio UI, database, training pipeline, MedSAM, or report generation in this phase.

Keep the implementation modular, explicit, and easy to extend.

Target Python version: Python 3.11+.

2. Required Repository Structure

cxr-medgemma-finding-localization-demo/
├── app.py
├── requirements.txt
├── README.md
├── design-level-spec.md
├── code-level-spec.md
├── data/
│   └── demo_cases.jsonl
├── outputs/
│   └── .gitkeep
├── tests/
│   └── test_multiview_pipeline.py
└── src/
    ├── __init__.py
    ├── config.py
    ├── dataset.py
    ├── inference.py
    ├── model.py
    ├── parsing.py
    ├── preprocessing.py
    ├── prompts.py
    ├── schemas.py
    └── visualization.py

Do not collapse the entire implementation into app.py.

app.py is orchestration only.

3. Dependencies

requirements.txt:

accelerate
numpy
Pillow
pydantic>=2
python-dotenv
torch
transformers

Do not hardcode a Hugging Face access token.

Authentication is expected to come from the user's Hugging Face environment or HF_TOKEN.

4. Configuration

File: src/config.py

Required constants:

MODEL_ID = os.getenv(
    "MEDGEMMA_MODEL_ID",
    "google/medgemma-1.5-4b-it",
)

BBOX_NORMALIZATION_SCALE = 1000
DEFAULT_MAX_FINDINGS = 5
DISCOVERY_MAX_NEW_TOKENS = 512
LOCALIZATION_MAX_NEW_TOKENS = 768

No secrets belong in this file.

5. Pydantic Schemas

File: src/schemas.py

Finding

class Finding(BaseModel):
    finding: str
    anatomical_location: str
    certainty: Literal[
        "positive",
        "probable",
        "questionable",
    ]

FindingDiscovery

class FindingDiscovery(BaseModel):
    findings: list[Finding]

Default findings to an empty list.

LocalizationBox

class LocalizationBox(BaseModel):
    box_2d: list[int]
    label: str

Validation requirements:

• exactly four coordinates;
• every coordinate between 0 and 1000;
• interpret order as [y0, x0, y1, x1];
• require x0 < x1;
• require y0 < y1.

LocalizationStatus

LocalizationStatus = Literal[
    "localized",
    "abstained",
    "rejected_by_quality_gate",
    "parser_error",
]

ViewLocalization

class ViewLocalization(BaseModel):
    image_index: int
    image_path: str
    status: LocalizationStatus
    boxes: list[LocalizationBox] = []
    candidate_boxes: list[LocalizationBox] = []
    rejection_reasons: list[str] = []

• boxes — accepted boxes only (MedSAM input);
• candidate_boxes — raw MedGemma proposals before gating;
• rejection_reasons — populated when status is rejected_by_quality_gate or parser_error.

LocalizedFinding

class LocalizedFinding(BaseModel):
    finding: str
    anatomical_location: str
    certainty: Certainty
    localizations: list[ViewLocalization]

Default localizations to an empty list.

Every finding must contain one ViewLocalization per study view.

PipelineResult

class PipelineResult(BaseModel):
    case_id: str | None
    input_images: list[str]
    findings: list[LocalizedFinding]

6. Study-Level Finding Discovery Prompt

File: src/prompts.py

Use a study-level multi-view prompt that:

• treats all provided images as complementary views of one examination;
• returns one consolidated study-level finding list as strict raw JSON only;
• does not duplicate the same finding across views;
• does not infer temporal change without prior-study images;
• remains valid when only one image is provided;
• forbids prose, headings, markdown fences, and top-level JSON lists in the model response.

Do not include the IU X-Ray reference report in this prompt.

7. Localization Prompt Builder

File: src/prompts.py

Implement:

def build_localization_prompt(finding: Finding) -> str:
    ...

The prompt must inject:

• finding.finding
• finding.anatomical_location

The prompt must retain the localization contract adapted from the uploaded MedGemma notebook:

• bounding-box order [y0, x0, y1, x1];
• normalized coordinates [0, 1000];
• top-left and bottom-right semantics;
• patient anatomical left/right convention.

Additional required rules:

1. localize the abnormal finding itself;
2. do not merely box the entire normal anatomical region;
3. do not add a different finding;
4. allow multiple boxes for spatially separate regions;
5. return [] when the target cannot be localized in the current view;
6. include study-level context explaining that the finding was identified from a multi-view review but must only be localized if visible in the current image;
7. prefer abstention over speculative or anatomy-only boxes;
8. return strict raw JSON only — no prose, headings, or markdown fences.

Required output shape:

[
  {
    "box_2d": [120, 610, 290, 790],
    "label": "pulmonary nodule"
  }
]

7.1 Localization Quality Gate

File: src/localization_quality.py

Implement geometry helpers in normalized 1000×1000 space:

def calculate_bbox_area_ratio(box: LocalizationBox) -> float: ...
def calculate_bbox_width_ratio(box: LocalizationBox) -> float: ...
def calculate_bbox_height_ratio(box: LocalizationBox) -> float: ...

Finding profiles (demo heuristics):

FOCAL_PROFILE = LocalizationProfile(
    name="focal",
    max_area_ratio=0.12,
    max_width_ratio=0.60,
    max_height_ratio=0.60,
)

DIFFUSE_PROFILE = LocalizationProfile(
    name="diffuse",
    max_area_ratio=None,
    max_width_ratio=None,
    max_height_ratio=None,
)

Finding groups via exact finding.strip().lower() match:

• Focal: nodule, pulmonary nodule, nodular opacity, focal opacity, focal airspace opacity, fracture, rib fracture
• Diffuse: opacity, pleural effusion, atelectasis, atelectatic opacity, consolidation
• Unknown → DIFFUSE_PROFILE

Gate entry point:

def evaluate_localization_boxes(
    finding_name: str,
    boxes: list[LocalizationBox],
) -> LocalizationGateResult:
    ...

Reject a box when any non-None profile threshold is exceeded. Example reason:

bbox_area_ratio=0.2400 exceeds focal max_area_ratio=0.1200

Do not use reference_report or reference_findings in the gate.

8. Image Preprocessing

File: src/preprocessing.py

Implement:

def pad_image_to_square(image: Image.Image) -> Image.Image:
    ...

Rules:

1. convert input to RGB;
2. convert to numpy.uint8;
3. calculate max(height, width);
4. apply symmetric black padding;
5. return a PIL image.

Do not resize the image during this stage.

Do not multiply a uint8 array by 255 after padding.

The returned square image is the single coordinate space for:

• finding discovery;
• finding localization;
• visualization.

9. MedGemma Client

File: src/model.py

Implement:

class MedGemmaClient:
    def __init__(self, model_id: str = MODEL_ID) -> None:
        ...

    def generate(
        self,
        *,
        images: Sequence[Image.Image],
        prompt: str,
        max_new_tokens: int,
    ) -> str:
        ...

Model Loading

Use the notebook's Transformers pipeline pattern:

pipeline(
    "image-text-to-text",
    model=model_id,
    model_kwargs={
        "dtype": selected_dtype,
        "device_map": "auto",
    },
)

Load the pipeline exactly once in MedGemmaClient.__init__.

Dtype Selection

Implement:

CUDA + BF16 support -> torch.bfloat16
CUDA without BF16   -> torch.float16
CPU                 -> torch.float32

Message Format

Build content as:

[
    {"type": "image", "image": image},
    ...,
    {"type": "text", "text": prompt},
]

Then:

messages = [
    {
        "role": "user",
        "content": content,
    }
]

Inference settings:

do_sample=False

Use the stage-specific max_new_tokens.

Thinking Trace Removal

Preserve the uploaded notebook's behavior:

if "<unused95>" in response:
    response = response.split("<unused95>", 1)[1].lstrip()

Do not print or save the removed thinking trace.

10. JSON Parsing

File: src/parsing.py

Implement:

def strip_medgemma_thinking_trace(response: str) -> str:
    ...

def extract_finding_discovery_payload(response: str) -> dict[str, Any]:
    ...

def extract_json_payload(
    response: str,
    expected_type: type,
) -> Any:
    ...

Finding discovery parsing

extract_finding_discovery_payload is used only for study-level finding discovery. It normalizes known MedGemma structural deviations into:

{"findings": [...]}

Accepted inputs include raw JSON objects, top-level finding lists, fenced JSON, and prose followed by Consolidated list of findings: plus a JSON array.

Parser failure must raise ValueError with a response preview. It must not silently return {"findings": []}.

Localization parsing

extract_json_payload remains the generic parser for localization responses expecting a JSON list of bounding boxes.

Parsing order for extract_json_payload:

1. strip MedGemma thinking trace;
2. search for a fenced ```json ... ``` payload;
3. if no code fence exists:
   • for dict, extract from first { to last };
   • for list, extract from first [ to last ];
4. parse the first JSON value with JSONDecoder.raw_decode;
5. validate the top-level Python type;
6. raise ValueError when parsing fails.

Do not silently fabricate a fallback finding.

11. Demo Dataset Loader

File: src/dataset.py

Implement:

def load_demo_cases(path: Path) -> list[dict[str, Any]]:
    ...

def get_demo_case(
    cases: list[dict[str, Any]],
    case_id: str,
) -> dict[str, Any]:
    ...

def resolve_case_images(
    case: dict[str, Any],
    dataset_root: Path,
) -> list[Path]:
    ...

Path resolution:

dataset_root / raw_path.lstrip("/")

Example:

raw JSONL path:
  /iu_xray/image/CXR3281_IM-1562/0.png

dataset_root:
  /data/iu-xray-dataset

resolved:
  /data/iu-xray-dataset/iu_xray/image/CXR3281_IM-1562/0.png

The dataset loader may read:

• case_id
• images
• reference_report
• reference_findings

However, reference_report and reference_findings may only be used by app.py after the MedGemma pipeline finishes.

They must never be passed to MedGemmaClient.generate.

12. Inference Service

File: src/inference.py

Implement:

class CXRFindingsLocalizer:
    def __init__(
        self,
        medgemma: MedGemmaClient,
        max_findings: int = DEFAULT_MAX_FINDINGS,
    ) -> None:
        ...

    def discover_findings(
        self,
        images: Sequence[Image.Image],
    ) -> FindingDiscovery:
        ...

    def localize_finding(
        self,
        image: Image.Image,
        finding: Finding,
        image_index: int | None = None,
    ) -> list[LocalizationBox]:
        ...

    def run(
        self,
        *,
        images: Sequence[Image.Image],
        input_images: list[str],
        case_id: str | None = None,
    ) -> tuple[PipelineResult, list[Image.Image]]:
        ...

discover_findings

Flow:

all square study images
    ↓
FINDING_DISCOVERY_PROMPT
    ↓
MedGemmaClient.generate (one call)
    ↓
extract_finding_discovery_payload
    ↓
FindingDiscovery.model_validate

localize_finding

Flow:

square image + Finding
    ↓
build_localization_prompt(finding)
    ↓
MedGemmaClient.generate
    ↓
extract_json_payload(..., list)
    ↓
LocalizationBox.model_validate for each item

Print debug output including finding, location, image index, and raw response.

run

Required flow:

processed_images = [pad_image_to_square(image) for image in images]

discovery = discover_findings(processed_images)

for finding in discovery.findings[:max_findings]:
    for image_index, processed_image in enumerate(processed_images):
        try:
            candidate_boxes = localize_finding(
                processed_image, finding, image_index=image_index
            )
        except (ValueError, ValidationError) as error:
            ViewLocalization(status="parser_error", ...)
        elif not candidate_boxes:
            ViewLocalization(status="abstained", ...)
        else:
            gate = evaluate_localization_boxes(finding.finding, candidate_boxes)
            if gate.accepted_boxes:
                ViewLocalization(status="localized", boxes=gate.accepted_boxes, ...)
            else:
                ViewLocalization(status="rejected_by_quality_gate", boxes=[], ...)

If one localization cannot be parsed or validated for one finding/view pair:

• print a warning;
• preserve the finding;
• set status="parser_error" with empty boxes and candidate_boxes;
• continue to the next finding/view pair.

Do not catch finding-discovery JSON errors as an empty normal result. A malformed discovery response should fail visibly because otherwise the demo could present a false "no abnormality" state.

13. Visualization

File: src/visualization.py

Implement:

def draw_view_localizations(
    image: Image.Image,
    image_index: int,
    findings: list[LocalizedFinding],
) -> Image.Image:
    ...

For each finding, locate the ViewLocalization whose image_index matches the current view and draw only its accepted boxes. Never draw candidate_boxes.

y0, x0, y1, x1 = box_2d

left = x0 / 1000 * width
top = y0 / 1000 * height
right = x1 / 1000 * width
bottom = y1 / 1000 * height

Draw:

• red rectangle;
• readable text label;
• black label background;
• white label text.

Draw on the square-padded image returned by the inference service.

Do not transform boxes back to the original non-square coordinate space in this phase.

14. CLI Application

File: app.py

app.py must remain a thin orchestration layer.

Required Input Modes

Use an argparse mutually exclusive group.

Direct Image Mode

python app.py   --image /path/to/chest_xray.png   --show

Demo Case Mode

python app.py   --case-id CXR3281_IM-1562   --dataset-root /path/to/dataset/root   --show

Demo-case mode loads all images listed in the case.

Required Arguments

--image PATH
--case-id STRING
--demo-cases PATH
--dataset-root PATH
--output-dir PATH
--max-findings INT
--model-id STRING
--show

Defaults:

--demo-cases data/demo_cases.jsonl
--output-dir outputs
--max-findings 5
--model-id google/medgemma-1.5-4b-it

--dataset-root is required only with --case-id.

15. Console Output

Implement a readable text summary:

=== Study Input ===
Case ID       : CXR3281_IM-1562
Study views   : 2
  [0] /.../0.png
  [1] /.../1.png

=== MedGemma Findings and Localization ===

[1] pulmonary nodule
    Location  : right upper lung
    Certainty : positive

    View [0]
      Image  : /.../0.png
      Status : localized
      BBox 1 : [120, 610, 290, 790] (label=pulmonary nodule)
      Area   : 0.0306

    View [1]
      Image  : /.../1.png
      Status : abstained
      BBox   : not localized

[2] fracture
    Location  : right humerus
    Certainty : positive

    View [0]
      Image     : /.../0.png
      Status    : rejected_by_quality_gate
      Candidate : [100, 100, 700, 500]
      Reason    : bbox_area_ratio=0.2400 exceeds focal max_area_ratio=0.1200
      BBox      : not accepted

Always print Status for each view. For localized, also print accepted bbox coordinates and area ratio. For abstained, print BBox : not localized. For rejected_by_quality_gate, print candidate boxes, reasons, and BBox : not accepted. For parser_error, print the reason.

Then print:

=== Structured JSON ===

and pretty-print PipelineResult.model_dump().

In demo-case mode, print only after inference:

=== Hidden IU X-Ray Reference ===
The following reference data was NOT sent to MedGemma.

Then display every reference_findings item in a concise text form.

Do not calculate an automatic match score in this phase.

16. Saved Output Contract

Use:

outputs/<run_name>/
├── result.json
├── annotated_0.png
├── annotated_1.png
└── ...

run_name:

• demo mode: case_id
• direct-image mode: input image stem

result.json example:

{
  "case_id": "CXR3281_IM-1562",
  "input_images": [
    "/data/iu-xray-dataset/iu_xray/image/CXR3281_IM-1562/0.png",
    "/data/iu-xray-dataset/iu_xray/image/CXR3281_IM-1562/1.png"
  ],
  "findings": [
    {
      "finding": "pulmonary nodule",
      "anatomical_location": "right upper lung",
      "certainty": "positive",
      "localizations": [
        {
          "image_index": 0,
          "image_path": "/data/iu-xray-dataset/iu_xray/image/CXR3281_IM-1562/0.png",
          "boxes": [
            {
              "box_2d": [120, 610, 290, 790],
              "label": "pulmonary nodule"
            }
          ]
        },
        {
          "image_index": 1,
          "image_path": "/data/iu-xray-dataset/iu_xray/image/CXR3281_IM-1562/1.png",
          "boxes": []
        }
      ]
    }
  ]
}

annotated_N.png must contain only boxes whose image_index == N.

17. Error Handling

Fail with a clear exception for:

• missing direct image;
• missing demo JSONL;
• unknown case ID;
• missing --dataset-root in demo mode;
• case with no images;
• resolved image path not found;
• discovery JSON parsing failure;
• discovery Pydantic validation failure.

For localization-only parsing/validation failure:

[warning] Could not parse localization for '<finding>': <error>

Continue processing remaining findings.

18. Acceptance Tests

Run mocked tests in tests/test_multiview_pipeline.py, tests/test_parsing.py, and tests/test_localization_quality.py:

python -m unittest discover -s tests -v

At minimum, verify:

Finding discovery parser tests (tests/test_parsing.py)

• strict {"findings":[...]} dict;
• top-level finding list normalization;
• real MedGemma prose plus Consolidated list of findings: array;
• fenced JSON;
• explicit empty findings;
• malformed prose raises ValueError;
• multiple standalone finding dicts raise ValueError.

Inference normalization test

discover_findings with the real MedGemma prose response preserves two findings after normalization.

Test A: one multi-view discovery call

For two study images, MedGemmaClient.generate(images=[img0, img1], ...) is called exactly once for discovery.

Test B: localization call count

For 2 findings and 2 images, localization is called 4 times.

Test C: bbox/image association

Every localization result contains image_index, image_path, status, and boxes.

Test D: finding visible in one view only

Mock one bbox on view 0 and [] on view 1. Final JSON must contain both ViewLocalization entries with statuses localized and abstained.

Localization quality gate tests (tests/test_localization_quality.py)

• Test A: small focal nodule box accepted;
• Test B: broad fracture box rejected with area-ratio reason;
• Test C: large diffuse opacity box not geometry-rejected.

Localization gate state tests (tests/test_multiview_pipeline.py)

• explicit MedGemma abstention ([]) → abstained;
• broad focal candidate fully rejected → rejected_by_quality_gate, candidate_boxes preserved;
• malformed localization response → parser_error;
• visualizer draws accepted boxes only, not rejected candidate_boxes.

Test E: annotated image consistency

annotated_0.png contains only boxes from image_index == 0.

Test F: direct image mode

One image produces one discovery call, one localization call per finding, and annotated_0.png.

Test G: hidden reference leakage

Verify reference_report and reference_findings are not passed to MedGemmaClient.generate.

19. Implementation Order for Cursor

Implement in this exact order:

1. create repository structure;
2. copy data/demo_cases.jsonl;
3. implement Pydantic schemas;
4. implement config;
5. implement prompts;
6. implement thinking-trace removal and JSON parser;
7. implement square padding;
8. implement MedGemma client;
9. implement dataset loader;
10. implement inference service;
11. implement visualization;
12. implement CLI orchestration;
13. add README commands;
14. run syntax checks;
15. run mocked parser/schema tests;
16. run one real MedGemma demo case.

Do not start MedSAM integration until this phase passes the acceptance tests.