File size: 3,435 Bytes
e93bfbd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """
Batch prediction script for non-interactive use.
Processes many X-ray images (or a single one) and outputs results as CSV.
"""
import argparse
import csv
import os
import sys
from pathlib import Path
from optimize import (
clear_memory,
infer_blip,
infer_fusion,
infer_fusion_onnx,
set_cpu_threads,
)
CHECKPOINT_PATH = "./checkpoints/fusion_model.pth"
ONNX_FULL_DIR = "./checkpoints/onnx_full"
HAS_FUSION = os.path.exists(CHECKPOINT_PATH)
HAS_ONNX = os.path.exists(os.path.join(ONNX_FULL_DIR, "fusion_full.onnx"))
def find_images(path):
path = Path(path)
if path.is_file():
return [path]
exts = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".dcm"}
return sorted([p for p in path.rglob("*") if p.suffix.lower() in exts])
def process_image(image_path, symptoms, use_vision, use_onnx):
result = {"image": str(image_path), "symptoms": symptoms}
if use_vision:
try:
caption = infer_blip(str(image_path), use_onnx=HAS_ONNX)
result["caption"] = caption
except Exception as e:
result["caption"] = f"ERROR: {e}"
if HAS_FUSION and symptoms:
try:
if use_onnx and HAS_ONNX:
diagnosis, confidence = infer_fusion_onnx(str(image_path), symptoms)
else:
diagnosis, confidence = infer_fusion(str(image_path), symptoms)
result["diagnosis"] = diagnosis if diagnosis else "INCONCLUSIVE"
result["confidence"] = f"{confidence:.4f}" if isinstance(confidence, float) else confidence
except Exception as e:
result["diagnosis"] = f"ERROR: {e}"
result["confidence"] = ""
clear_memory()
return result
def main():
set_cpu_threads()
parser = argparse.ArgumentParser(description="Batch X-ray prediction")
parser.add_argument("input", help="Path to image file or directory")
parser.add_argument("--symptoms", default="What abnormality is present in this chest X-ray?",
help="Symptoms text (same for all images)")
parser.add_argument("--output", "-o", default="predictions.csv",
help="Output CSV path")
parser.add_argument("--vision", action="store_true",
help="Run vision captioning (BLIP)")
parser.add_argument("--onnx", action="store_true",
help="Use ONNX runtime (if exported)")
args = parser.parse_args()
images = find_images(args.input)
if not images:
print(f"No images found at: {args.input}")
sys.exit(1)
print(f"Found {len(images)} image(s)")
print(f"Symptoms: {args.symptoms}")
print(f"Vision: {'ON' if args.vision else 'OFF'}")
print(f"ONNX: {'ON' if args.onnx and HAS_ONNX else 'OFF'}")
print()
results = []
for i, img_path in enumerate(images):
print(f"[{i+1}/{len(images)}] Processing {img_path.name}...")
result = process_image(str(img_path), args.symptoms, args.vision, args.onnx)
results.append(result)
fieldnames = ["image", "symptoms", "caption", "diagnosis", "confidence"]
with open(args.output, "w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
writer.writeheader()
writer.writerows(results)
print(f"\nDone. Results saved to {args.output}")
if __name__ == "__main__":
main()
|