118 lines
4.1 KiB
Python
118 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
import random
|
|
import shutil
|
|
from pathlib import Path
|
|
from typing import Dict
|
|
|
|
import pandas as pd
|
|
from PIL import Image
|
|
|
|
from .utils import ensure_dir
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
|
|
|
|
STYLE_TO_CLASS = {
|
|
"unknown": 0,
|
|
"filled": 1,
|
|
"hollow": 2,
|
|
"border": 3,
|
|
}
|
|
|
|
|
|
def _crop_image_and_labels(img: Image.Image, boxes: pd.DataFrame, x0: int, y0: int, size: int):
|
|
crop = img.crop((x0, y0, x0 + size, y0 + size)).convert("RGB")
|
|
labels = []
|
|
for _, r in boxes.iterrows():
|
|
bx1, by1, bx2, by2 = float(r.x), float(r.y), float(r.x + r.w), float(r.y + r.h)
|
|
ix1, iy1 = max(bx1, x0), max(by1, y0)
|
|
ix2, iy2 = min(bx2, x0 + size), min(by2, y0 + size)
|
|
if ix2 <= ix1 or iy2 <= iy1:
|
|
continue
|
|
visible_area = (ix2 - ix1) * (iy2 - iy1)
|
|
box_area = max((bx2 - bx1) * (by2 - by1), 1)
|
|
if visible_area / box_area < 0.35:
|
|
continue
|
|
cx = ((ix1 + ix2) / 2 - x0) / size
|
|
cy = ((iy1 + iy2) / 2 - y0) / size
|
|
w = (ix2 - ix1) / size
|
|
h = (iy2 - iy1) / size
|
|
cls = STYLE_TO_CLASS.get(str(r.get("fill_style", "unknown")), 0)
|
|
labels.append(f"{cls} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}")
|
|
return crop, labels
|
|
|
|
|
|
def export_candidates_to_yolo(
|
|
tif_path: str | Path,
|
|
candidates_csv: str | Path,
|
|
out_dir: str | Path,
|
|
cfg: Dict,
|
|
sheet_id: str,
|
|
tile_size: int = 1024,
|
|
overlap: int = 128,
|
|
val_fraction: float = 0.20,
|
|
include_empty_tiles: bool = True,
|
|
max_empty_tiles: int = 250,
|
|
) -> Path:
|
|
out_dir = Path(out_dir)
|
|
for split in ["train", "val"]:
|
|
ensure_dir(out_dir / "images" / split)
|
|
ensure_dir(out_dir / "labels" / split)
|
|
|
|
img = Image.open(tif_path).convert("RGB")
|
|
boxes = pd.read_csv(candidates_csv)
|
|
step = max(1, tile_size - overlap)
|
|
random.seed(42)
|
|
empty_written = 0
|
|
total_written = 0
|
|
|
|
for y0 in range(0, max(1, img.height - tile_size + 1), step):
|
|
for x0 in range(0, max(1, img.width - tile_size + 1), step):
|
|
in_tile = boxes[
|
|
(boxes.cx >= x0) & (boxes.cx < x0 + tile_size) &
|
|
(boxes.cy >= y0) & (boxes.cy < y0 + tile_size)
|
|
]
|
|
if in_tile.empty:
|
|
if not include_empty_tiles or empty_written >= max_empty_tiles:
|
|
continue
|
|
# Keep some empty/hard-negative tiles to stop the model from detecting all blue map details.
|
|
if random.random() > 0.08:
|
|
continue
|
|
empty_written += 1
|
|
crop, labels = _crop_image_and_labels(img, boxes, x0, y0, tile_size)
|
|
split = "val" if random.random() < val_fraction else "train"
|
|
stem = f"{sheet_id}_{x0}_{y0}"
|
|
crop.save(out_dir / "images" / split / f"{stem}.jpg", quality=92)
|
|
with open(out_dir / "labels" / split / f"{stem}.txt", "w", encoding="utf-8") as f:
|
|
f.write("\n".join(labels))
|
|
total_written += 1
|
|
|
|
data_yaml = out_dir / "data.yaml"
|
|
names = cfg.get("export", {}).get("yolo_class_names", ["blue_rect_unknown", "blue_rect_filled", "blue_rect_hollow", "blue_rect_border"])
|
|
with open(data_yaml, "w", encoding="utf-8") as f:
|
|
f.write(f"path: {out_dir.resolve()}\n")
|
|
f.write("train: images/train\n")
|
|
f.write("val: images/val\n")
|
|
f.write("names:\n")
|
|
for i, name in enumerate(names):
|
|
f.write(f" {i}: {name}\n")
|
|
LOG.info("YOLO export complete: %s (%d tiles)", data_yaml, total_written)
|
|
return data_yaml
|
|
|
|
|
|
def merge_yolo_datasets(src_dirs: list[str | Path], out_dir: str | Path) -> Path:
|
|
out_dir = Path(out_dir)
|
|
for split in ["train", "val"]:
|
|
ensure_dir(out_dir / "images" / split)
|
|
ensure_dir(out_dir / "labels" / split)
|
|
for src in src_dirs:
|
|
src = Path(src)
|
|
for split in ["train", "val"]:
|
|
for img in (src / "images" / split).glob("*.jpg"):
|
|
shutil.copy2(img, out_dir / "images" / split / img.name)
|
|
for lab in (src / "labels" / split).glob("*.txt"):
|
|
shutil.copy2(lab, out_dir / "labels" / split / lab.name)
|
|
return out_dir
|