scripts/predict_segment_areas.py

"""Predict food segment area fractions for one image using a trained YOLO model.

Example:
    python scripts/predict_segment_areas.py --image path/to/meal.jpg --output data/nutrition/segments.json
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

import cv2
import numpy as np


PROJECT_ROOT = Path(__file__).resolve().parents[1]
SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))
os.environ.setdefault("YOLO_CONFIG_DIR", str(PROJECT_ROOT / ".ultralytics"))

from ultralytics import YOLO
from segmentation_utils import load_data_yaml_class_names


DEFAULT_WEIGHTS = (
    PROJECT_ROOT
    / "runs"
    / "foodseg103_target"
    / "yolov8s_target_e10_w02"
    / "weights"
    / "best.pt"
)
DEFAULT_DATA = PROJECT_ROOT / "data" / "processed" / "foodseg103_target_yolo" / "data.yaml"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image", type=Path, required=True)
    parser.add_argument("--weights", type=Path, default=DEFAULT_WEIGHTS)
    parser.add_argument("--data", type=Path, default=DEFAULT_DATA)
    parser.add_argument("--imgsz", type=int, default=512)
    parser.add_argument("--conf", type=float, default=0.05)
    parser.add_argument("--output", type=Path)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    class_names = load_data_yaml_class_names(args.data)
    image_bgr = cv2.imread(str(args.image))
    if image_bgr is None:
        raise RuntimeError(f"Could not read image: {args.image}")

    image = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    height, width = image.shape[:2]
    model = YOLO(str(args.weights))
    result = model.predict(image, imgsz=args.imgsz, conf=args.conf, verbose=False)[0]

    area_by_class = {name: 0 for name in class_names.values()}
    occupied = np.zeros((height, width), dtype=bool)

    if result.masks is not None:
        masks = result.masks.data.cpu().numpy()
        classes = result.boxes.cls.cpu().numpy().astype(int)
        confs = result.boxes.conf.cpu().numpy()

        for idx in np.argsort(confs):
            item = masks[idx]
            if item.shape != (height, width):
                item = cv2.resize(item, (width, height), interpolation=cv2.INTER_NEAREST)
            binary = item > 0.5
            class_name = class_names[int(classes[idx])]
            area_by_class[class_name] += int(np.logical_and(binary, ~occupied).sum())
            occupied |= binary

    total_food_area = sum(area_by_class.values())
    segments = []
    for class_name, pixel_area in area_by_class.items():
        if pixel_area <= 0 or total_food_area <= 0:
            continue
        segments.append(
            {
                "class_name": class_name,
                "pixel_area": pixel_area,
                "area_fraction": round(pixel_area / total_food_area, 6),
            }
        )

    payload = {
        "image": str(args.image),
        "weights": str(args.weights),
        "conf": args.conf,
        "total_food_area": total_food_area,
        "segments": segments,
    }

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(payload, indent=2), encoding="utf-8")
        print(f"Wrote segments: {args.output}")
    else:
        print(json.dumps(payload, indent=2))


if __name__ == "__main__":
    main()