Spaces:
Sleeping
Sleeping
| """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() | |