92 lines
3.0 KiB
Python
Executable File
92 lines
3.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Create a metadata-free, consent-traceable Timmy training record."""
|
|
import argparse
|
|
import hashlib
|
|
import hmac
|
|
import io
|
|
import json
|
|
import os
|
|
import sys
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
from PIL import Image, ImageOps
|
|
|
|
COLORS = ("brown", "green", "yellow", "pale", "red", "black")
|
|
QUALITIES = ("good", "fair", "poor")
|
|
REVIEWERS = ("user-confirmed", "clinician-reviewed")
|
|
|
|
|
|
def parser() -> argparse.ArgumentParser:
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--input", required=True)
|
|
p.add_argument("--output-dir", required=True)
|
|
p.add_argument("--subject-id", required=True)
|
|
p.add_argument("--consent-version", required=True)
|
|
p.add_argument("--bristol-type", required=True, type=int, choices=range(1, 8))
|
|
p.add_argument("--color", required=True, choices=COLORS)
|
|
p.add_argument("--quality", required=True, choices=QUALITIES)
|
|
p.add_argument("--reviewer", required=True, choices=REVIEWERS)
|
|
return p
|
|
|
|
|
|
def main() -> int:
|
|
args = parser().parse_args()
|
|
salt = os.environ.get("TIMMY_DATASET_SALT", "")
|
|
if len(salt) < 8:
|
|
raise SystemExit("TIMMY_DATASET_SALT must contain at least 8 characters")
|
|
if not args.consent_version.strip():
|
|
raise SystemExit("consent version is required")
|
|
|
|
try:
|
|
with Image.open(args.input) as source:
|
|
image = ImageOps.exif_transpose(source).convert("RGB")
|
|
image.thumbnail((1200, 1200))
|
|
buffer = io.BytesIO()
|
|
image.save(buffer, "JPEG", quality=85, optimize=True, exif=b"")
|
|
except Exception as exc:
|
|
raise SystemExit(f"input is not a decodable image: {exc}") from exc
|
|
|
|
data = buffer.getvalue()
|
|
image_id = hashlib.sha256(data).hexdigest()
|
|
subject_key = hmac.new(salt.encode(), args.subject_id.encode(), hashlib.sha256).hexdigest()[:16]
|
|
bucket = int(subject_key[:2], 16)
|
|
split = "train" if bucket < 205 else "validation" if bucket < 230 else "test"
|
|
|
|
output = Path(args.output_dir)
|
|
images = output / "images"
|
|
images.mkdir(parents=True, exist_ok=True)
|
|
relative_path = f"images/{image_id}.jpg"
|
|
target = output / relative_path
|
|
if not target.exists():
|
|
target.write_bytes(data)
|
|
|
|
record = {
|
|
"schemaVersion": 1,
|
|
"imageId": image_id,
|
|
"relativePath": relative_path,
|
|
"subjectKey": subject_key,
|
|
"split": split,
|
|
"consentVersion": args.consent_version.strip(),
|
|
"ingestedAt": datetime.now(timezone.utc).isoformat(),
|
|
"labels": {
|
|
"bristolType": args.bristol_type,
|
|
"color": args.color,
|
|
"quality": args.quality,
|
|
},
|
|
"review": {"status": args.reviewer},
|
|
"derivation": {
|
|
"format": "jpeg",
|
|
"maxDimension": 1200,
|
|
"metadataStripped": True,
|
|
},
|
|
}
|
|
with (output / "manifest.jsonl").open("a", encoding="utf-8") as manifest:
|
|
manifest.write(json.dumps(record, separators=(",", ":")) + "\n")
|
|
print(json.dumps(record))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|