Spaces:
Sleeping
Sleeping
| """Read YOLO-seg .txt predictions and build a per-tuber morphometrics dataframe. | |
| YOLO-seg .txt format (one detection per line): | |
| class_id x1 y1 x2 y2 ... xN yN # coords normalized to [0, 1] | |
| """ | |
| from __future__ import annotations | |
| from pathlib import Path | |
| import numpy as np | |
| import pandas as pd | |
| from PIL import Image | |
| from .morphometrics import polygon_shape_metrics | |
| IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"} | |
| def read_yolo_seg_txt(txt_path: str | Path) -> list[dict]: | |
| """Parse a YOLO-seg .txt file into a list of {class_id, xy_norm}. | |
| Lines with fewer than 3 vertices (6 coordinates) or an odd coordinate | |
| count are skipped. | |
| """ | |
| polys: list[dict] = [] | |
| with open(txt_path) as f: | |
| for line in f: | |
| line = line.strip() | |
| if not line: | |
| continue | |
| parts = line.split() | |
| class_id = int(float(parts[0])) | |
| coords = np.array(list(map(float, parts[1:])), dtype=float) | |
| if coords.size < 6 or coords.size % 2 != 0: | |
| continue | |
| polys.append({"class_id": class_id, "xy_norm": coords.reshape(-1, 2)}) | |
| return polys | |
| def _list_images(folder: Path) -> list[Path]: | |
| return sorted(p for p in folder.rglob("*") if p.suffix.lower() in IMG_EXTS) | |
| def build_morphometrics_dataframe( | |
| trial_name: str, | |
| images_dir: str | Path, | |
| preds_dir: str | Path, | |
| ) -> pd.DataFrame: | |
| """Produce one row per detected tuber with all pixel-scale shape metrics. | |
| Pairs each .txt file in `preds_dir` with the image of the same stem in | |
| `images_dir` so we can read the image dimensions (needed both for | |
| denormalizing polygon coordinates and for the downstream calibration). | |
| """ | |
| images_dir = Path(images_dir) | |
| preds_dir = Path(preds_dir) | |
| if not preds_dir.exists(): | |
| raise FileNotFoundError(f"[{trial_name}] preds_dir not found: {preds_dir}") | |
| if not images_dir.exists(): | |
| raise FileNotFoundError(f"[{trial_name}] images_dir not found: {images_dir}") | |
| image_map = {p.stem: p for p in _list_images(images_dir)} | |
| txt_files = sorted(preds_dir.glob("*.txt")) | |
| if not txt_files: | |
| raise ValueError(f"[{trial_name}] No .txt files found in {preds_dir}") | |
| rows: list[dict] = [] | |
| skipped_txt_no_image = 0 | |
| for txt in txt_files: | |
| img_stem = txt.stem | |
| img_path = image_map.get(img_stem) | |
| if img_path is None: | |
| skipped_txt_no_image += 1 | |
| continue | |
| with Image.open(img_path) as im: | |
| img_width, img_height = im.size | |
| for line_idx, poly in enumerate(read_yolo_seg_txt(txt), start=1): | |
| xy = poly["xy_norm"] | |
| xs = xy[:, 0] * img_width | |
| ys = xy[:, 1] * img_height | |
| metrics = polygon_shape_metrics(xs, ys) | |
| rows.append( | |
| { | |
| "trial": trial_name, | |
| "image_id": img_stem, | |
| "detection_id": f"{img_stem}_{line_idx}", | |
| "class_id": poly["class_id"], | |
| "img_width_px": img_width, | |
| "img_height_px": img_height, | |
| "n_points": len(xs), | |
| **metrics, | |
| } | |
| ) | |
| df = pd.DataFrame(rows) | |
| if skipped_txt_no_image: | |
| print( | |
| f"[{trial_name}] images={len(image_map)} txt={len(txt_files)} " | |
| f"instances={len(df)} skipped_txt_no_image={skipped_txt_no_image}" | |
| ) | |
| return df | |
| def plot_id_from_image_id(image_id: pd.Series) -> pd.Series: | |
| """Derive a raw plot identifier from the last 8 characters of the image stem. | |
| The result (e.g. "NCPT-129", "5_ADC101") is further normalized by | |
| standardize_plot_id in the yield-prediction stage. | |
| """ | |
| return image_id.str[-8:] | |