File size: 3,775 Bytes
48e3e9e | 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 | """Clean-room VisDrone2019-DET -> YOLO builder.
Written from the public VisDrone annotation spec (no third-party converter
code copied). Each annotation line is:
bbox_left,bbox_top,bbox_width,bbox_height,score,category,truncation,occlusion
Categories: 0=ignored-regions 1=pedestrian 2=people 3=bicycle 4=car 5=van
6=truck 7=tricycle 8=awning-tricycle 9=bus 10=motor 11=others. We keep 1..10
(remapped to 0..9) and drop 0/11 and score==0 (ignored) boxes.
Builds: <datasets>/visdrone/{images,labels}/{train,val}/ + train.txt/val.txt
+ visdrone.yaml
from the extracted <datasets>/visdrone_raw/VisDrone2019-DET-{train,val}/.
Usage:
python build_visdrone.py --datasets /path/to/datasets
"""
import argparse
import shutil
from pathlib import Path
from PIL import Image
SPLITS = {"train": "VisDrone2019-DET-train", "val": "VisDrone2019-DET-val"}
NAMES = ["pedestrian", "people", "bicycle", "car", "van", "truck",
"tricycle", "awning-tricycle", "bus", "motor"] # VisDrone 1..10 -> 0..9
def convert_split(raw: Path, out: Path, split: str, raw_name: str) -> int:
src = raw / raw_name
img_src, ann_src = src / "images", src / "annotations"
img_dst = out / "images" / split
lbl_dst = out / "labels" / split
img_dst.mkdir(parents=True, exist_ok=True)
lbl_dst.mkdir(parents=True, exist_ok=True)
listing = []
n_imgs = n_boxes = n_dropped = 0
for img_path in sorted(img_src.glob("*.jpg")):
with Image.open(img_path) as im:
W, H = im.size
ann = ann_src / (img_path.stem + ".txt")
lines = []
if ann.exists():
for raw in ann.read_text().splitlines():
raw = raw.strip().rstrip(",")
if not raw:
continue
p = raw.split(",")
if len(p) < 6:
continue
left, top, w, h, score, cat = (int(float(x)) for x in p[:6])
if score == 0 or cat < 1 or cat > 10 or w <= 0 or h <= 0:
n_dropped += 1
continue
cx = (left + w / 2) / W
cy = (top + h / 2) / H
nw, nh = w / W, h / H
# clip to [0,1] (a few VisDrone boxes bleed past the edge)
cx, cy = min(max(cx, 0), 1), min(max(cy, 0), 1)
nw, nh = min(nw, 1), min(nh, 1)
lines.append(f"{cat - 1} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}")
n_boxes += 1
# copy image into the YOLO tree (move would empty the raw dir; copy is safe)
shutil.copy2(img_path, img_dst / img_path.name)
(lbl_dst / (img_path.stem + ".txt")).write_text("\n".join(lines))
listing.append(f"images/{split}/{img_path.name}")
n_imgs += 1
(out / f"{split}.txt").write_text("\n".join(listing))
print(f"{split}: {n_imgs} imgs, {n_boxes} boxes kept, {n_dropped} dropped")
return n_imgs
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--datasets",
required=True,
help="Directory containing visdrone_raw/ with the extracted official zips",
)
args = parser.parse_args()
datasets = Path(args.datasets)
raw = datasets / "visdrone_raw"
out = datasets / "visdrone"
for split, raw_name in SPLITS.items():
convert_split(raw, out, split, raw_name)
names_block = "\n".join(f" {i}: {n}" for i, n in enumerate(NAMES))
(out / "visdrone.yaml").write_text(
f"# VisDrone2019-DET (research/non-commercial license)\n"
f"path: visdrone\ntrain: train.txt\nval: val.txt\n"
f"nc: {len(NAMES)}\nnames:\n{names_block}\n"
)
print("WROTE", out / "visdrone.yaml")
if __name__ == "__main__":
main()
|