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

```text
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`:

```text
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:

```python
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`

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

### `FindingDiscovery`

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

Default `findings` to an empty list.

### `LocalizationBox`

```python
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`

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

### `ViewLocalization`

```python
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`

```python
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`

```python
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:

```python
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:

```json
[
  {
    "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:

```python
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):

```python
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:

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

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

```text
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:

```python
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:

```python
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:

```python
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:

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

### Message Format

Build content as:

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

Then:

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

Inference settings:

```python
do_sample=False
```

Use the stage-specific `max_new_tokens`.

### Thinking Trace Removal

Preserve the uploaded notebook's behavior:

```python
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:

```python
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:

```json
{"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:

```python
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:

```python
dataset_root / raw_path.lstrip("/")
```

Example:

```text
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:

```python
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:

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

### `localize_finding`

Flow:

```text
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:

```python
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:

```python
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`.

```python
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

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

#### Demo Case Mode

```bash
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

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

Defaults:

```text
--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:

```text
=== 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:

```text
=== Structured JSON ===
```

and pretty-print `PipelineResult.model_dump()`.

In demo-case mode, print only after inference:

```text
=== 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:

```text
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:

```json
{
  "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:

```text
[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`:

```bash
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.