Spaces:
Sleeping
Sleeping
File size: 4,770 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """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}
|