File size: 3,557 Bytes
d176ecd | 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 | #!/usr/bin/env python3
"""Convert WIDER FACE (HF parquet, CUHK-CSE/wider_face @ refs/convert/parquet) to single-class YOLO format.
- Skips boxes flagged `invalid` and degenerate boxes (<1px).
- Keeps small faces (no min-size filter) since blurring needs high recall.
- Drops images whose only boxes are invalid (they would teach the model to ignore real faces).
- Downscales images so the long side is <= --max-side (matches training imgsz; boxes are normalized).
Usage:
python convert_wider_to_yolo.py --raw raw/default --out dataset
"""
import argparse
import io
import os
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
import pyarrow.parquet as pq
from PIL import Image
SPLITS = {"train": "train", "validation": "val"}
def convert_row(args):
name, img_bytes, bboxes, invalid, out_img, out_lbl, max_side = args
valid = [b for b, inv in zip(bboxes, invalid) if not inv and b[2] >= 1 and b[3] >= 1]
if not valid:
return "dropped", 0, 0
im = Image.open(io.BytesIO(img_bytes)).convert("RGB")
w, h = im.size
scale = min(1.0, max_side / max(w, h))
if scale < 1.0:
im = im.resize((max(1, round(w * scale)), max(1, round(h * scale))), Image.LANCZOS)
stem = Path(name).stem
im.save(out_img / f"{stem}.jpg", quality=95)
lines = []
for x, y, bw, bh in valid:
x1, y1 = max(0.0, x), max(0.0, y)
x2, y2 = min(float(w), x + bw), min(float(h), y + bh)
if x2 - x1 < 1 or y2 - y1 < 1:
continue
lines.append(f"0 {(x1 + x2) / 2 / w:.6f} {(y1 + y2) / 2 / h:.6f} {(x2 - x1) / w:.6f} {(y2 - y1) / h:.6f}")
(out_lbl / f"{stem}.txt").write_text("\n".join(lines) + "\n")
return "ok", len(lines), len(bboxes) - len(valid)
def rows(files, out_img, out_lbl, max_side):
for f in files:
pf = pq.ParquetFile(f)
for batch in pf.iter_batches(batch_size=64):
for r in batch.to_pylist():
yield (r["image"]["path"], r["image"]["bytes"], r["faces"]["bbox"],
r["faces"]["invalid"], out_img, out_lbl, max_side)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--raw", default="raw/default")
ap.add_argument("--out", default="dataset")
ap.add_argument("--max-side", type=int, default=640)
ap.add_argument("--limit", type=int, default=0, help="max images per split (0 = all; for testing)")
ap.add_argument("--workers", type=int, default=os.cpu_count())
a = ap.parse_args()
out = Path(a.out)
for src, dst in SPLITS.items():
files = sorted((Path(a.raw) / src).glob("*.parquet"))
out_img, out_lbl = out / "images" / dst, out / "labels" / dst
out_img.mkdir(parents=True, exist_ok=True)
out_lbl.mkdir(parents=True, exist_ok=True)
gen = rows(files, out_img, out_lbl, a.max_side)
if a.limit:
import itertools
gen = itertools.islice(gen, a.limit)
ok = dropped = faces = 0
with ProcessPoolExecutor(a.workers) as ex:
for status, n, _ in ex.map(convert_row, gen, chunksize=16):
if status == "ok":
ok += 1
faces += n
else:
dropped += 1
print(f"{dst}: {ok} images, {faces} valid faces, {dropped} images dropped (no valid boxes)")
(out / "data.yaml").write_text(
f"path: {out.resolve()}\ntrain: images/train\nval: images/val\nnames:\n 0: face\n")
print("wrote", out / "data.yaml")
if __name__ == "__main__":
main()
|