Spaces:
Sleeping
Sleeping
| """YOLOv11-seg instance segmentation for tuber detection. | |
| Wraps the Ultralytics YOLO API with conf=0.25, imgsz=1280, default NMS | |
| IoU, no augmentation. Polygon coordinates are normalized to [0, 1] and | |
| written in the YOLO-seg .txt format consumed by io.read_yolo_seg_txt. | |
| """ | |
| from __future__ import annotations | |
| import gc | |
| from collections.abc import Iterable, Iterator | |
| from pathlib import Path | |
| from typing import Any | |
| import numpy as np | |
| from ultralytics import YOLO | |
| from .io import IMG_EXTS | |
| YOLO_PREDICT_DEFAULTS: dict[str, Any] = { | |
| "conf": 0.25, | |
| "imgsz": 1280, | |
| "verbose": False, | |
| } | |
| def load_model(weights_path: str | Path) -> YOLO: | |
| """Load YOLOv11-seg weights via the Ultralytics API.""" | |
| return YOLO(str(weights_path)) | |
| def predict_image( | |
| model: YOLO, | |
| image: str | Path | np.ndarray, | |
| **predict_kwargs: Any, | |
| ) -> list[tuple[int, np.ndarray]]: | |
| """Run inference on a single image; return [(class_id, polygon_normalized)]. | |
| `polygon_normalized` is an (N, 2) float array of (x, y) coordinates in | |
| [0, 1]. Returns an empty list if no instances were detected. | |
| """ | |
| params = {**YOLO_PREDICT_DEFAULTS, **predict_kwargs} | |
| results = model.predict(source=image, stream=False, **params) | |
| r = results[0] | |
| if r.masks is None: | |
| return [] | |
| h, w = r.orig_shape | |
| detections: list[tuple[int, np.ndarray]] = [] | |
| for poly_xy, box in zip(r.masks.xy, r.boxes): | |
| class_id = int(box.cls[0]) | |
| normalized = poly_xy / np.array([w, h], dtype=poly_xy.dtype) | |
| detections.append((class_id, normalized)) | |
| return detections | |
| def write_yolo_seg_txt(detections: Iterable[tuple[int, np.ndarray]], txt_path: str | Path) -> None: | |
| """Write detections in the legacy YOLO-seg format. | |
| One line per detection: `class_id x1 y1 x2 y2 ...` with coords in [0, 1]. | |
| Creates an empty file if `detections` is empty (matches legacy behavior). | |
| """ | |
| with open(txt_path, "w") as f: | |
| for class_id, poly in detections: | |
| coords = poly.reshape(-1) | |
| f.write(f"{class_id} " + " ".join(map(str, coords)) + "\n") | |
| def _list_images(folder: Path) -> list[Path]: | |
| return sorted(p for p in folder.rglob("*") if p.suffix.lower() in IMG_EXTS) | |
| def _chunked(lst: list[Path], n: int) -> Iterator[list[Path]]: | |
| for i in range(0, len(lst), n): | |
| yield lst[i : i + n] | |
| def predict_directory( | |
| model: YOLO, | |
| input_dir: str | Path, | |
| output_dir: str | Path, | |
| chunk_size: int = 15, | |
| **predict_kwargs: Any, | |
| ) -> dict[str, int]: | |
| """Run inference over every image under input_dir, writing one .txt per image. | |
| Streams predictions in chunks of chunk_size and triggers GC between | |
| chunks. Empty .txt files are still created for images with no detections. | |
| Returns a dict with n_images, n_with_masks, n_seen for QC. | |
| """ | |
| input_dir = Path(input_dir) | |
| output_dir = Path(output_dir) | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| image_paths = _list_images(input_dir) | |
| print(f"[INFO] {input_dir.name}: {len(image_paths)} images") | |
| if not image_paths: | |
| return {"n_images": 0, "n_seen": 0, "n_with_masks": 0} | |
| params = {**YOLO_PREDICT_DEFAULTS, **predict_kwargs} | |
| seen = 0 | |
| with_masks = 0 | |
| for ci, part in enumerate(_chunked(image_paths, chunk_size), start=1): | |
| gen = model.predict(source=[p.as_posix() for p in part], stream=True, **params) | |
| # zip() rather than trusting r.path: in some Ultralytics versions | |
| # r.path falls back to "imageN" inside a batched/streamed call, | |
| # which causes chunks to overwrite each other's outputs. stream=True | |
| # preserves input order, so part[i] matches the i-th yielded result. | |
| for img_path, r in zip(part, gen): | |
| seen += 1 | |
| txt_path = output_dir / f"{img_path.stem}.txt" | |
| with open(txt_path, "w") as f: | |
| if r.masks is not None: | |
| with_masks += 1 | |
| h, w = r.orig_shape | |
| for poly_xy, box in zip(r.masks.xy, r.boxes): | |
| class_id = int(box.cls[0]) | |
| coords: list[float] = [] | |
| for x, y in poly_xy: | |
| coords.append(x / w) | |
| coords.append(y / h) | |
| f.write(f"{class_id} " + " ".join(map(str, coords)) + "\n") | |
| del r | |
| gc.collect() | |
| if ci == 1 or ci % 10 == 0: | |
| print( | |
| f"[INFO] chunk {ci}: {min(ci * chunk_size, len(image_paths))}/" | |
| f"{len(image_paths)} images" | |
| ) | |
| print(f"[DONE] {output_dir}: seen={seen}, with_masks={with_masks}") | |
| return {"n_images": len(image_paths), "n_seen": seen, "n_with_masks": with_masks} | |