Spaces:
Sleeping
Sleeping
File size: 3,170 Bytes
3ecb175 | 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 | """Run YOLOv11-seg inference on tuber images and emit per-image .txt files.
Default predict parameters: conf=0.25, imgsz=1280. Output .txt files
follow the YOLO-seg format consumed by the morphometrics stage.
Single-trial usage:
uv run python scripts/run_yolo_inference.py \\
--weights path/to/bestModel.pt \\
--images path/to/2025NCPT \\
--out-dir outputs/preds/NCPT
All three trials usage:
uv run python scripts/run_yolo_inference.py \\
--weights path/to/bestModel.pt \\
--images-root path/to/rawData/images \\
--out-root outputs/preds
"""
from __future__ import annotations
import argparse
from pathlib import Path
from potatopheno.segmentation import load_model, predict_directory
TRIALS = ("NCPT", "ADC", "NCR")
def _parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--weights", type=Path, required=True, help="Path to bestModel.pt.")
single = p.add_argument_group("single-trial mode")
single.add_argument(
"--images",
type=Path,
help="Directory of input images (single trial).",
)
single.add_argument(
"--out-dir",
type=Path,
help="Directory to write .txt files into.",
)
multi = p.add_argument_group("all-trials mode")
multi.add_argument(
"--images-root",
type=Path,
help="Parent directory with 2025<TRIAL>/ subdirs.",
)
multi.add_argument(
"--out-root",
type=Path,
help="Parent directory; <TRIAL>/preds/ will be created under it.",
)
multi.add_argument(
"--trials",
type=lambda s: tuple(t.strip() for t in s.split(",") if t.strip()),
default=TRIALS,
help=f"Comma-separated trial names (default: {','.join(TRIALS)}).",
)
p.add_argument("--conf", type=float, default=0.25, help="Confidence threshold.")
p.add_argument("--imgsz", type=int, default=1280, help="Inference image size.")
p.add_argument("--chunk-size", type=int, default=15, help="Images per streaming chunk.")
return p.parse_args()
def main() -> None:
args = _parse_args()
single_mode = args.images is not None and args.out_dir is not None
multi_mode = args.images_root is not None and args.out_root is not None
if single_mode == multi_mode:
raise SystemExit(
"Pass EITHER (--images and --out-dir) OR (--images-root and --out-root), not both."
)
model = load_model(args.weights)
predict_kwargs = {"conf": args.conf, "imgsz": args.imgsz, "verbose": False}
if single_mode:
predict_directory(
model=model,
input_dir=args.images,
output_dir=args.out_dir,
chunk_size=args.chunk_size,
**predict_kwargs,
)
return
for trial in args.trials:
predict_directory(
model=model,
input_dir=args.images_root / f"2025{trial}",
output_dir=args.out_root / trial / "preds",
chunk_size=args.chunk_size,
**predict_kwargs,
)
if __name__ == "__main__":
main()
|