Spaces:
Runtime error
Runtime error
File size: 3,506 Bytes
e0e2b27 | 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 103 104 105 106 107 108 109 110 111 | """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()
|