#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ prepare_taco.py Functionality: - Mode A (default): Local TACO images + a single COCO annotations file -> Map classes to Alami super-classes -> Stratified split (train/val/test) -> Write YOLOv8-seg labels + dataset.yaml + stats - Mode B (pre-split JSONs already exist): --annotations_train/val/test -> No stratification; use provided splits as-is -> Write YOLOv8-seg labels + dataset.yaml + stats - Optional in both modes: 1) If images are missing locally, download them directly from COCO image fields (flickr_640_url, flickr_url, coco_url, url) and save under image['file_name'] in --images_dir. 2) If still missing or explicitly desired: CSV fallback (--images_urls_csv), matching by basename (URL pathname) to image['file_name']. Notes: - Dedupe via MD5 available (default on). - RLE masks are skipped (polygons only). NEW (backwards compatible, default keeps behavior): - --min_poly_area_px: discard polygons smaller than this area (pixel^2). default 0 = off. """ import argparse import json import os import sys import shutil import random from collections import defaultdict, Counter from pathlib import Path from typing import Dict, List, Tuple, Any, Optional import re # ---- Required deps ---- try: from PIL import Image, ImageOps except ImportError: print("Please `pip install pillow`", file=sys.stderr); raise try: import numpy as np except ImportError: print("Please `pip install numpy`", file=sys.stderr); raise # --------------------------- # Helpers: IO / JSON / FS # --------------------------- def read_json(path: Path) -> Any: with path.open("r", encoding="utf-8") as f: return json.load(f) def write_json(path: Path, obj: Any) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as f: json.dump(obj, f, ensure_ascii=False, indent=2) def write_text(path: Path, s: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as f: f.write(s) def safe_symlink_or_copy(src: Path, dst: Path) -> None: dst.parent.mkdir(parents=True, exist_ok=True) if os.name == "nt": # Windows: Symlink pain -> copy if not dst.exists(): shutil.copy2(src, dst) else: shutil.copy2(src, dst) return try: if dst.exists() or dst.is_symlink(): dst.unlink() os.symlink(src, dst) except Exception: if not dst.exists(): shutil.copy2(src, dst) else: shutil.copy2(src, dst) # --------------------------- # Label map (rule-based) # --------------------------- def load_label_map(label_map_path: Path) -> Dict[str, Any]: lm = read_json(label_map_path) required = ["target_classes", "default", "rules"] for k in required: if k not in lm: raise ValueError(f"label_map.json missing key: {k}") # normalize for case-insensitive matching rules = [] for r in lm["rules"]: nr = { "to": r["to"], "match_any_substring": [s.lower() for s in r.get("match_any_substring", [])], "match_any_regex": r.get("match_any_regex", []), } rules.append(nr) lm["rules"] = rules return lm def map_class(name: str, label_map: Dict[str, Any]) -> str: n = name.lower().strip() for r in label_map["rules"]: for sub in r["match_any_substring"]: if sub in n: return r["to"] for pattern in r["match_any_regex"]: if re.search(pattern, n): return r["to"] return label_map["default"] # --------------------------- # COCO / TACO reading # --------------------------- def index_coco(coco: Dict[str, Any]) -> Tuple[Dict[int, dict], Dict[int, dict], Dict[int, List[dict]]]: images_by_id = {im["id"]: im for im in coco.get("images", [])} cats_by_id = {c["id"]: c for c in coco.get("categories", [])} anns_by_image = defaultdict(list) for ann in coco.get("annotations", []): anns_by_image[ann["image_id"]].append(ann) return images_by_id, cats_by_id, anns_by_image # --------------------------- # Geometry helpers (YOLOv8 seg expects normalized polygon points) # --------------------------- def coco_segmentation_to_yolo_polys(seg, img_w: int, img_h: int) -> List[List[float]]: """ COCO 'segmentation' can be list of polygons (each as flat list of x,y) or RLE. We only support polygons here; RLE is skipped. Returns list of [x1_norm,y1_norm, x2_norm,y2_norm, ...] per polygon. """ polys = [] if isinstance(seg, list): for poly in seg: if not isinstance(poly, list) or len(poly) < 6: continue norm = [] for i, v in enumerate(poly): if i % 2 == 0: x = max(0.0, min(float(v) / img_w, 1.0)) norm.append(x) else: y = max(0.0, min(float(v) / img_h, 1.0)) norm.append(y) polys.append(norm) # else: RLE not supported here return polys def _polygon_area_px(norm_poly: List[float], img_w: int, img_h: int) -> float: """ Shoelace area in pixels from normalized coordinates. norm_poly: [x1n,y1n,x2n,y2n,...] """ if len(norm_poly) < 6: return 0.0 xs = [norm_poly[i] * img_w for i in range(0, len(norm_poly), 2)] ys = [norm_poly[i] * img_h for i in range(1, len(norm_poly), 2)] area = 0.0 n = len(xs) for i in range(n): j = (i + 1) % n area += xs[i] * ys[j] - xs[j] * ys[i] return abs(area) * 0.5 # --------------------------- # Hashing / Dedupe (optional) # --------------------------- def image_md5(path: Path) -> Optional[str]: try: import hashlib h = hashlib.md5() with path.open("rb") as f: for chunk in iter(lambda: f.read(8192), b""): h.update(chunk) return h.hexdigest() except Exception: return None def image_phash(path: Path) -> Optional[str]: try: import scipy.fftpack except Exception: return None try: img = Image.open(path).convert("L") img = ImageOps.fit(img, (32, 32)) arr = np.asarray(img, dtype=np.float32) dct = scipy.fftpack.dct(scipy.fftpack.dct(arr.T, norm="ortho").T, norm="ortho") dct_low = dct[:8, :8] med = np.median(dct_low) bits = (dct_low > med).flatten() return "".join("1" if b else "0" for b in bits) except Exception: return None # --------------------------- # Split (stratify by primary class) # --------------------------- def stratified_split(items: List[dict], y: List[str], train_ratio=0.8, val_ratio=0.1, seed=42): rnd = random.Random(seed) by_class = defaultdict(list) for i, c in enumerate(y): by_class[c].append(i) train, val, test = [], [], [] for c, idx_list in by_class.items(): rnd.shuffle(idx_list) n = len(idx_list) n_train = int(round(n * train_ratio)) n_val = int(round(n * val_ratio)) n_test = n - n_train - n_val train += idx_list[:n_train] val += idx_list[n_train:n_train+n_val] test += idx_list[n_train+n_val:] for arr in (train, val, test): rnd.shuffle(arr) return train, val, test # --------------------------- # URL download support (CSV) # --------------------------- def read_urls_csv(csv_path: Path) -> List[str]: urls = [] with csv_path.open("r", encoding="utf-8") as f: for line in f: u = line.strip() if not u or u.lower().startswith("#"): continue urls.append(u) return urls def basename_from_url(u: str) -> str: try: from urllib.parse import urlparse, unquote p = urlparse(u) b = Path(unquote(p.path)).name return b except Exception: return Path(u).name def ensure_images_from_urls( urls_csv: Path, target_images_dir: Path, coco_images_by_id: Dict[int, dict] ) -> None: """ Download images from URLs into target_images_dir if they match COCO file_name by basename. - Heuristic: basename(URL) == basename(image['file_name']) - Non-matching URLs are ignored. - Existing files are not re-downloaded. """ urls = read_urls_csv(urls_csv) if not urls: print(f"[WARN] No URLs found in {urls_csv}", file=sys.stderr) return from collections import defaultdict by_base = defaultdict(list) for u in urls: by_base[basename_from_url(u).lower()].append(u) target_images_dir.mkdir(parents=True, exist_ok=True) q: List[Tuple[str, Path]] = [] for im in coco_images_by_id.values(): base = Path(im["file_name"]).name.lower() cand_urls = by_base.get(base) out = target_images_dir / Path(im["file_name"]).name if out.exists(): continue if cand_urls: q.append((cand_urls[0], out)) if not q: print("[INFO] No downloads required (all images present or no matches).") return print(f"[INFO] Downloading {len(q)} images from CSV into {target_images_dir} ...") import requests from concurrent.futures import ThreadPoolExecutor, as_completed def _fetch(u_dst: Tuple[str, Path]) -> Tuple[Path, bool, str]: u, dst = u_dst try: last = "" for attempt in range(3): try: r = requests.get(u, timeout=20) r.raise_for_status() dst.parent.mkdir(parents=True, exist_ok=True) with dst.open("wb") as f: f.write(r.content) return dst, True, "" except Exception as e: last = str(e) return dst, False, last except Exception as e: return dst, False, str(e) ok, fail = 0, 0 with ThreadPoolExecutor(max_workers=min(16, os.cpu_count() or 4)) as ex: futs = [ex.submit(_fetch, job) for job in q] for fu in as_completed(futs): dst, success, msg = fu.result() if success: ok += 1 else: fail += 1 print(f"[WARN] Download failed for {dst.name}: {msg}", file=sys.stderr) print(f"[INFO] CSV downloads done. ok={ok}, failed={fail}") # --------------------------- # NEW: URL download from COCO annotations (flickr_* / coco_url / url) # --------------------------- def _candidate_url_from_image(im: dict) -> Optional[str]: # Priority order for k in ("flickr_640_url", "flickr_url", "coco_url", "url"): u = im.get(k) if isinstance(u, str) and u.strip(): return u.strip() return None def ensure_images_from_ann_urls( coco_images_by_id: Dict[int, dict], target_images_dir: Path ) -> None: """ Download images based on URLs directly from COCO image objects. Target filename is ALWAYS image['file_name']. """ target_images_dir.mkdir(parents=True, exist_ok=True) jobs: List[Tuple[str, Path]] = [] for im in coco_images_by_id.values(): url = _candidate_url_from_image(im) if not url: continue dst = target_images_dir / Path(im["file_name"]).name if dst.exists(): continue jobs.append((url, dst)) if not jobs: print("[INFO] No annotation-based downloads required (all images present or no URLs).") return print(f"[INFO] Downloading {len(jobs)} images from annotation URLs into {target_images_dir} ...") import time as _time import requests from concurrent.futures import ThreadPoolExecutor, as_completed def _fetch(u_dst: Tuple[str, Path]) -> Tuple[Path, bool, str]: u, dst = u_dst try: last = "" for attempt in range(3): try: r = requests.get(u, timeout=25) r.raise_for_status() dst.parent.mkdir(parents=True, exist_ok=True) with dst.open("wb") as f: f.write(r.content) return dst, True, "" except Exception as e: last = str(e) _time.sleep(0.5 * (attempt + 1)) return dst, False, last except Exception as e: return dst, False, str(e) ok, fail = 0, 0 with ThreadPoolExecutor(max_workers=min(16, os.cpu_count() or 4)) as ex: futs = [ex.submit(_fetch, job) for job in jobs] for fu in as_completed(futs): dst, success, msg = fu.result() if success: ok += 1 else: fail += 1 print(f"[WARN] Annotation URL download failed for {dst.name}: {msg}", file=sys.stderr) print(f"[INFO] Annotation-based downloads done. ok={ok}, failed={fail}") # --------------------------- # Core helpers # --------------------------- def build_primary_label(anns: List[dict], cats_by_id: Dict[int, dict], label_map: Dict[str, Any]) -> Optional[str]: if not anns: return None mapped = [] for a in anns: cat = cats_by_id.get(a["category_id"]) if not cat: continue mapped.append(map_class(cat["name"], label_map)) if not mapped: return None cnt = Counter(mapped) return cnt.most_common(1)[0][0] def write_yolo_seg_label_file( label_path: Path, ann_list: List[dict], cats_by_id: Dict[int, dict], label_map: Dict[str, Any], img_w: int, img_h: int, class_to_index: Dict[str, int], min_poly_area_px: float = 0.0 ) -> int: """ Write YOLOv8-seg label file. NEW: min_poly_area_px > 0 filters small polygons (shoelace area in pixels). """ lines = [] for ann in ann_list: cat = cats_by_id.get(ann["category_id"]) if not cat: continue mapped = map_class(cat["name"], label_map) cls_id = class_to_index[mapped] polys = coco_segmentation_to_yolo_polys(ann.get("segmentation"), img_w, img_h) for poly in polys: if len(poly) < 6: continue if min_poly_area_px > 0.0: area = _polygon_area_px(poly, img_w, img_h) if area < min_poly_area_px: continue lines.append(" ".join([str(cls_id)] + [f"{p:.6f}" for p in poly])) if lines: label_path.parent.mkdir(parents=True, exist_ok=True) with label_path.open("w", encoding="utf-8") as f: f.write("\n".join(lines)) return len(lines) def items_from_coco( coco: Dict[str, Any], images_dir: Path, cats_by_id: Dict[int, dict], anns_by_image: Dict[int, List[dict]], label_map: Dict[str, Any] ) -> Tuple[List[dict], Counter, Counter]: items = [] raw_class_stats = Counter() unmapped = Counter() for img_id, img in {im["id"]: im for im in coco.get("images", [])}.items(): file_name = img["file_name"] width = img.get("width") height = img.get("height") # get dims if missing if width is None or height is None: try: with Image.open(images_dir / file_name) as im: width, height = im.size except Exception: # try with basename try: with Image.open(images_dir / Path(file_name).name) as im: width, height = im.size file_name = Path(file_name).name except Exception: continue ann_list = anns_by_image.get(img_id, []) for a in ann_list: cat = cats_by_id.get(a["category_id"]) if cat: raw_class_stats[cat["name"]] += 1 mapped = map_class(cat["name"], label_map) if mapped not in label_map["target_classes"]: unmapped[mapped] += 1 primary = build_primary_label(ann_list, cats_by_id, label_map) items.append({ "id": img_id, "file_name": file_name, "width": width, "height": height, "primary_label": primary, "anns": ann_list }) return items, raw_class_stats, unmapped def write_split( split_name: str, split_items: List[dict], out_root: Path, images_dir: Path, cats_by_id: Dict[int, dict], label_map: Dict[str, Any], class_to_index: Dict[str, int], min_poly_area_px: float = 0.0 ) -> Tuple[int, Counter, List[dict]]: (out_root / "images" / split_name).mkdir(parents=True, exist_ok=True) (out_root / "labels" / split_name).mkdir(parents=True, exist_ok=True) count_used = 0 class_presence = Counter() used_items: List[dict] = [] for it in split_items: # copy/symlink image src_img = images_dir / it["file_name"] if not src_img.exists(): # try basename fallback src_img = images_dir / Path(it["file_name"]).name if not src_img.exists(): # skip if missing even after downloads continue dst_img = out_root / "images" / split_name / Path(it["file_name"]).name safe_symlink_or_copy(src_img, dst_img) # write label file lbl_path = out_root / "labels" / split_name / (Path(it["file_name"]).stem + ".txt") n_written = write_yolo_seg_label_file( lbl_path, it["anns"], cats_by_id, label_map, it["width"], it["height"], class_to_index, min_poly_area_px=min_poly_area_px ) if n_written == 0: if lbl_path.exists(): lbl_path.unlink() if dst_img.exists(): try: dst_img.unlink() except Exception: pass continue mapped_classes = [] for a in it["anns"]: cat = cats_by_id.get(a["category_id"]) if cat: mapped_classes.append(map_class(cat["name"], label_map)) for mc in set(mapped_classes): class_presence[mc] += 1 count_used += 1 used_items.append(it) return count_used, class_presence, used_items # --------------------------- # Manifests (per split) # --------------------------- def write_split_manifest( manifests_root: Path, split_name: str, used_items: List[dict], images_by_id: Dict[int, dict] ) -> int: """ Write a manifest file per split: \t Preferred URL: flickr_640_url, then flickr_url, coco_url, url. Entries without URL are skipped. Returns: number of written lines. """ path = manifests_root / f"{split_name}.txt" path.parent.mkdir(parents=True, exist_ok=True) lines: List[str] = [] for it in used_items: im = images_by_id.get(it["id"]) if images_by_id is not None else None url = _candidate_url_from_image(im) if im else None if not url: continue lines.append(f"{Path(it['file_name']).name}\t{url}") if lines: with path.open("w", encoding="utf-8") as f: f.write("\n".join(lines)) return len(lines) else: with path.open("w", encoding="utf-8") as f: f.write("") return 0 # --------------------------- # Pipeline: single annotations (stratified split) # --------------------------- def prepare_single_annotations( images_dir: Path, annotations_path: Path, label_map_path: Path, out_root: Path, train_ratio: float, val_ratio: float, seed: int, dedupe: bool, images_urls_csv: Optional[Path], min_poly_area_px: float = 0.0 ): label_map = load_label_map(label_map_path) target_classes: List[str] = label_map["target_classes"] class_to_index = {c: i for i, c in enumerate(target_classes)} coco = read_json(annotations_path) images_by_id, cats_by_id, anns_by_image = index_coco(coco) # 1) Annotation-based download (if directory is empty/not present) if not images_dir.exists() or not any(images_dir.glob("*")): print("[INFO] images_dir empty/not found -> downloading from annotation URLs (flickr_640_url/flickr_url/etc.)") ensure_images_from_ann_urls(images_by_id, images_dir) # 2) CSV fallback (optional) if still empty if images_urls_csv and (not any(images_dir.glob("*"))): print(f"[INFO] images_dir still empty -> downloading from CSV {images_urls_csv}") ensure_images_from_urls(images_urls_csv, images_dir, images_by_id) # Build items items, raw_stats, unmapped = items_from_coco(coco, images_dir, cats_by_id, anns_by_image, label_map) # Filter images without annotations items = [x for x in items if x["anns"]] # Dedupe if dedupe: seen = {} deduped = [] for it in items: cand = images_dir / it["file_name"] if not cand.exists(): cand = images_dir / Path(it["file_name"]).name ph = image_md5(cand) if cand.exists() else None key = ph or it["file_name"] if key in seen: continue seen[key] = True deduped.append(it) items = deduped # Stratified split y = [it["primary_label"] or "other" for it in items] train_idx, val_idx, test_idx = stratified_split(items, y, train_ratio, val_ratio, seed) idx_set = {"train": set(train_idx), "val": set(val_idx), "test": set(test_idx)} # Write splits (one item per call -> logic unchanged) per_split_counts = {"train": 0, "val": 0, "test": 0} per_split_class = {s: Counter() for s in ("train", "val", "test")} per_split_used: Dict[str, List[dict]] = {"train": [], "val": [], "test": []} mapping_report = [] for i, it in enumerate(items): split = "train" if i in idx_set["train"] else "val" if i in idx_set["val"] else "test" n_used_before = per_split_counts[split] used, class_presence, used_items = write_split( split, [it], out_root, images_dir, cats_by_id, label_map, class_to_index, min_poly_area_px=min_poly_area_px ) per_split_counts[split] += used for k, v in class_presence.items(): per_split_class[split][k] += v if used_items: per_split_used[split].extend(used_items) if per_split_counts[split] > n_used_before: mapped_classes = [] for a in it["anns"]: cat = cats_by_id.get(a["category_id"]) if cat: mapped_classes.append(map_class(cat["name"], label_map)) mapping_report.append({ "file_name": it["file_name"], "primary": it["primary_label"] or "other", "classes_in_image": list(sorted(set(mapped_classes))) }) # dataset.yaml names = target_classes dataset_yaml = { "path": str(out_root.resolve()), "train": "images/train", "val": "images/val", "test": "images/test", "names": names } write_json(out_root / "dataset.yolov8.json", dataset_yaml) write_text(out_root / "dataset.yaml", "path: {}\ntrain: {}\nval: {}\ntest: {}\nnames:\n".format( dataset_yaml["path"], dataset_yaml["train"], dataset_yaml["val"], dataset_yaml["test"] ) + "".join([f" {i}: {n}\n" for i, n in enumerate(names)])) # Stats stats = { "total_images_after_filter": sum(per_split_counts.values()), "per_split_counts": per_split_counts, "per_split_class_presence": {k: dict(v) for k, v in per_split_class.items()}, "raw_class_counts": dict(raw_stats), "unmapped_buckets_seen": dict(unmapped) } write_json(out_root / "class_stats.json", stats) write_json(out_root / "mapping_report.jsonl", mapping_report) # Write manifests per split manif_root = out_root / "manifests" for split in ("train", "val", "test"): write_split_manifest(manif_root, split, per_split_used[split], images_by_id) print("=== DONE (single annotations) ===") print(json.dumps(stats, indent=2, ensure_ascii=False)) print(f"dataset.yaml -> {out_root / 'dataset.yaml'}") # --------------------------- # Pipeline: pre-split annotations (train/val/test provided) # --------------------------- def prepare_presplit_annotations( images_dir: Path, annotations_train: Path, annotations_val: Path, annotations_test: Path, label_map_path: Path, out_root: Path, dedupe: bool, images_urls_csv: Optional[Path], min_poly_area_px: float = 0.0 ): label_map = load_label_map(label_map_path) target_classes: List[str] = label_map["target_classes"] class_to_index = {c: i for i, c in enumerate(target_classes)} # load three COCOs coco_train = read_json(annotations_train) coco_val = read_json(annotations_val) coco_test = read_json(annotations_test) # index imgs_tr, cats_tr, anns_tr = index_coco(coco_train) imgs_vl, cats_vl, anns_vl = index_coco(coco_val) imgs_te, cats_te, anns_te = index_coco(coco_test) # Ensure images (first: from annotation URLs) if not images_dir.exists() or not any(images_dir.glob("*")): print("[INFO] images_dir empty/not found -> downloading from annotation URLs (train/val/test merged)") merged = {**imgs_tr, **imgs_vl, **imgs_te} ensure_images_from_ann_urls(merged, images_dir) # CSV fallback if images_urls_csv and (not any(images_dir.glob("*"))): print(f"[INFO] images_dir still empty -> downloading from CSV {images_urls_csv}") merged = {**imgs_tr, **imgs_vl, **imgs_te} ensure_images_from_urls(images_urls_csv, images_dir, merged) # build items per split items_tr, raw_tr, unm_tr = items_from_coco(coco_train, images_dir, cats_tr, anns_tr, label_map) items_vl, raw_vl, unm_vl = items_from_coco(coco_val, images_dir, cats_vl, anns_vl, label_map) items_te, raw_te, unm_te = items_from_coco(coco_test, images_dir, cats_te, anns_te, label_map) # Filter empties items_tr = [x for x in items_tr if x["anns"]] items_vl = [x for x in items_vl if x["anns"]] items_te = [x for x in items_te if x["anns"]] # Dedupe (within each split) if dedupe: def _dedupe(items: List[dict]) -> List[dict]: seen = {} out = [] for it in items: cand = images_dir / it["file_name"] if not cand.exists(): cand = images_dir / Path(it["file_name"]).name ph = image_md5(cand) if cand.exists() else None key = ph or it["file_name"] if key in seen: continue seen[key] = True out.append(it) return out items_tr = _dedupe(items_tr) items_vl = _dedupe(items_vl) items_te = _dedupe(items_te) # Write each split per_split_counts = {"train": 0, "val": 0, "test": 0} per_split_class = {s: Counter() for s in ("train", "val", "test")} used, cpres, used_items_tr = write_split("train", items_tr, out_root, images_dir, cats_tr, label_map, class_to_index, min_poly_area_px=min_poly_area_px) per_split_counts["train"] += used per_split_class["train"].update(cpres) used, cpres, used_items_vl = write_split("val", items_vl, out_root, images_dir, cats_vl, label_map, class_to_index, min_poly_area_px=min_poly_area_px) per_split_counts["val"] += used per_split_class["val"].update(cpres) used, cpres, used_items_te = write_split("test", items_te, out_root, images_dir, cats_te, label_map, class_to_index, min_poly_area_px=min_poly_area_px) per_split_counts["test"] += used per_split_class["test"].update(cpres) # dataset.yaml names = target_classes dataset_yaml = { "path": str(out_root.resolve()), "train": "images/train", "val": "images/val", "test": "images/test", "names": names } write_json(out_root / "dataset.yolov8.json", dataset_yaml) write_text(out_root / "dataset.yaml", "path: {}\ntrain: {}\nval: {}\ntest: {}\nnames:\n".format( dataset_yaml["path"], dataset_yaml["train"], dataset_yaml["val"], dataset_yaml["test"] ) + "".join([f" {i}: {n}\n" for i, n in enumerate(names)])) # Stats raw_total = Counter() raw_total.update(raw_tr); raw_total.update(raw_vl); raw_total.update(raw_te) unm_total = Counter() unm_total.update(unm_tr); unm_total.update(unm_vl); unm_total.update(unm_te) stats = { "total_images_after_filter": sum(per_split_counts.values()), "per_split_counts": per_split_counts, "per_split_class_presence": {k: dict(v) for k, v in per_split_class.items()}, "raw_class_counts": dict(raw_total), "unmapped_buckets_seen": dict(unm_total) } write_json(out_root / "class_stats.json", stats) # Manifests per split manif_root = out_root / "manifests" write_split_manifest(manif_root, "train", used_items_tr, imgs_tr) write_split_manifest(manif_root, "val", used_items_vl, imgs_vl) write_split_manifest(manif_root, "test", used_items_te, imgs_te) print("=== DONE (pre-split annotations) ===") print(json.dumps(stats, indent=2, ensure_ascii=False)) print(f"dataset.yaml -> {out_root / 'dataset.yaml'}") # --------------------------- # CLI # --------------------------- def parse_args(): p = argparse.ArgumentParser(description="Prepare TACO for YOLOv8-seg with Alami label mapping.") # Image source p.add_argument("--images_dir", required=True, type=Path, help="Path to target image directory (created if missing). " "If empty -> download from annotation URLs; optional CSV fallback.") p.add_argument("--images_urls_csv", type=Path, default=None, help="CSV with image URLs (one per line). Optional fallback.") # Mode A (single annotations file; we produce the split) p.add_argument("--annotations", type=Path, help="Path to COCO annotations.json (single-file mode with internal split).") # Mode B (pre-split annotations) p.add_argument("--annotations_train", type=Path, help="COCO annotations train.json") p.add_argument("--annotations_val", type=Path, help="COCO annotations val.json") p.add_argument("--annotations_test", type=Path, help="COCO annotations test.json") # General p.add_argument("--label_map", required=True, type=Path, help="Path to ml/configs/label_map.json") p.add_argument("--out", required=True, type=Path, help="Output root, e.g., ml/datasets/taco") # Only for Mode A (we create the split) p.add_argument("--train_ratio", type=float, default=0.8) p.add_argument("--val_ratio", type=float, default=0.1) p.add_argument("--seed", type=int, default=42) # Optional p.add_argument("--no_dedupe", action="store_true", help="Disable duplicate filtering") # NEW p.add_argument("--min_poly_area_px", type=float, default=0.0, help="Minimum polygon area in pixels; polygons below are discarded (default 0 = off)") args = p.parse_args() have_single = args.annotations is not None have_presplit = all([args.annotations_train, args.annotations_val, args.annotations_test]) if not have_single and not have_presplit: p.error("Specify either --annotations OR (--annotations_train, --annotations_val, --annotations_test).") if have_single: if not (0 < args.train_ratio < 1): p.error("--train_ratio must be in (0,1)") if not (0 <= args.val_ratio < 1): p.error("--val_ratio must be in [0,1)") if args.train_ratio + args.val_ratio >= 1: p.error("train_ratio + val_ratio must be < 1 (remainder is test)") return args def main(): args = parse_args() out_root = args.out out_root.mkdir(parents=True, exist_ok=True) if args.annotations: prepare_single_annotations( images_dir=args.images_dir, annotations_path=args.annotations, label_map_path=args.label_map, out_root=out_root, train_ratio=args.train_ratio, val_ratio=args.val_ratio, seed=args.seed, dedupe=(not args.no_dedupe), images_urls_csv=args.images_urls_csv, min_poly_area_px=float(args.min_poly_area_px) ) else: prepare_presplit_annotations( images_dir=args.images_dir, annotations_train=args.annotations_train, annotations_val=args.annotations_val, annotations_test=args.annotations_test, label_map_path=args.label_map, out_root=out_root, dedupe=(not args.no_dedupe), images_urls_csv=args.images_urls_csv, min_poly_area_px=float(args.min_poly_area_px) ) if __name__ == "__main__": main()