run.py

import argparse
import json
from pathlib import Path

from pipeline_aaai import MedDetectPipeline


def parse_args():
    parser = argparse.ArgumentParser(description="Run the minimal MedDetect-Graph pipeline.")
    parser.add_argument("--config", required=True)
    parser.add_argument("--input", required=True, help="JSON list, JSON object, or JSONL file")
    parser.add_argument("--output", required=True, help="Single final JSON result file")
    parser.add_argument("--image-root", default="")
    return parser.parse_args()


def load_records(path):
    path = Path(path)
    if path.suffix.lower() == ".jsonl":
        return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
    data = json.loads(path.read_text(encoding="utf-8"))
    return data if isinstance(data, list) else [data]


def resolve_image(value, image_root):
    if isinstance(value, dict):
        value = value.get("file") or value.get("path") or value.get("image_id")
    if isinstance(value, list):
        value = value[0] if value else None
    if not value:
        return None
    path = Path(str(value))
    if path.is_file():
        return str(path.resolve())
    rooted = Path(image_root) / path
    return str(rooted.resolve()) if rooted.is_file() else None


def normalize(record, image_root, default_task):
    response = record.get("response")
    if response is None:
        response = record.get("hallucinated_report", record.get("report", ""))
    raw_segments = record.get("segments") or record.get("hallucinated_report_sentence_level") or []
    segments = []
    for item in raw_segments:
        if isinstance(item, dict):
            section = item.get("section", "")
            text = item.get("text", "")
            segments.append(f"{section}: {text}" if section else text)
        else:
            segments.append(str(item))
    target_image = record.get("target_image")
    image_value = target_image or record.get("image") or record.get("image_path")
    image_modality = (
        target_image.get("image_type") if isinstance(target_image, dict) else None
    ) or record.get("image_type") or record.get("image_modality")
    return {
        "uid": record.get("uid") or record.get("case_id") or record.get("id"),
        "task": record.get("task") or default_task,
        "query": record.get("query"),
        "response": response,
        "segments": segments or None,
        "image_path": resolve_image(image_value, image_root),
        "image_modality": image_modality,
    }


def main():
    args = parse_args()
    pipeline = MedDetectPipeline(args.config)
    task = pipeline.config["general"]["task"]
    results = [
        pipeline.run(normalize(record, args.image_root, task))
        for record in load_records(args.input)
    ]
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")


if __name__ == "__main__":
    main()