80 lines
2.9 KiB
Python
80 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""标注预览图生成: 动物红色实线框 + 生境 cover 黄色虚线框
|
|
|
|
用法:
|
|
venv/bin/python visualize_labels.py --input datasets/images --labels datasets/labels --output datasets/pheasant_label
|
|
"""
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
from PIL import Image, ImageDraw
|
|
|
|
CLASSES = ["pheasant", "cover"]
|
|
LABELS = {"pheasant": "野鸡", "cover": "疑似区域"}
|
|
|
|
|
|
def draw_dashed(draw, box, outline, width, dash=12, gap=8):
|
|
x1, y1, x2, y2 = box
|
|
for (ax, ay, bx, by) in [(x1, y1, x2, y1), (x2, y1, x2, y2),
|
|
(x2, y2, x1, y2), (x1, y2, x1, y1)]:
|
|
length = max(abs(bx - ax), abs(by - ay))
|
|
steps = max(int(length / (dash + gap)), 1)
|
|
for i in range(steps):
|
|
s = i / steps
|
|
e = min((i * (dash + gap) + dash) / length, 1.0)
|
|
if e <= s:
|
|
continue
|
|
draw.line([ax + (bx - ax) * s, ay + (by - ay) * s,
|
|
ax + (bx - ax) * e, ay + (by - ay) * e],
|
|
fill=outline, width=width)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="标注预览图生成")
|
|
ap.add_argument("--input", default="datasets/images")
|
|
ap.add_argument("--labels", default="datasets/labels")
|
|
ap.add_argument("--output", default="datasets/pheasant_label")
|
|
args = ap.parse_args()
|
|
|
|
img_dir = Path(args.input)
|
|
lab_dir = Path(args.labels)
|
|
out_dir = Path(args.output)
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
count = 0
|
|
for img_path in sorted(img_dir.rglob("*")):
|
|
if img_path.suffix.lower() not in (".jpg", ".jpeg", ".png", ".webp", ".bmp"):
|
|
continue
|
|
rel = img_path.relative_to(img_dir)
|
|
lab_path = lab_dir / rel.with_suffix(".txt")
|
|
if not lab_path.exists():
|
|
continue
|
|
im = Image.open(img_path).convert("RGB")
|
|
W, H = im.size
|
|
d = ImageDraw.Draw(im)
|
|
for line in lab_path.read_text().splitlines():
|
|
parts = line.split()
|
|
if not parts:
|
|
continue
|
|
cid, cx, cy, w, h = map(float, parts)
|
|
box = ((cx - w / 2) * W, (cy - h / 2) * H,
|
|
(cx + w / 2) * W, (cy + h / 2) * H)
|
|
if int(cid) == 1: # cover 生境: 黄色虚线
|
|
draw_dashed(d, box, (255, 200, 0), width=5)
|
|
d.text((box[0] + 6, max(box[1] - 28, 4)), "疑似区域",
|
|
fill=(255, 200, 0))
|
|
else: # 动物: 红色实线
|
|
d.rectangle(box, outline=(255, 0, 0), width=6)
|
|
label = LABELS.get(CLASSES[int(cid)], CLASSES[int(cid)])
|
|
d.text((box[0] + 6, max(box[1] - 28, 4)), label,
|
|
fill=(255, 0, 0))
|
|
out = out_dir / (img_path.stem + ".jpg")
|
|
im.save(out, quality=92)
|
|
count += 1
|
|
print(out.name)
|
|
print(f"共生成 {count} 张")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|