Logo_watermark_detection / scripts /verify_labels.py
Amitkumar001's picture
Upload folder using huggingface_hub (part 8)
a3a407d verified
Raw History Blame Contribute Delete
1.24 kB
import cv2
import os
from pathlib import Path
def verify_labels(img_dir="dataset_subset/images/train", lbl_dir="dataset_subset/labels/train", out_dir="verify_labels"):
img_dir = Path(img_dir)
lbl_dir = Path(lbl_dir)
out_dir = Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
img_files = list(img_dir.glob("*.jpg"))[:10]
for img_p in img_files:
img = cv2.imread(str(img_p))
h, w = img.shape[:2]
lbl_p = lbl_dir / (img_p.stem + ".txt")
if not lbl_p.exists():
continue
with open(lbl_p, "r") as f:
for line in f:
cls, xc, yc, nw, nh = map(float, line.split())
x1 = int((xc - nw/2) * w)
y1 = int((yc - nh/2) * h)
x2 = int((xc + nw/2) * w)
y2 = int((yc + nh/2) * h)
color = (0, 255, 0) if cls == 0 else (0, 0, 255)
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
cv2.putText(img, str(int(cls)), (x1, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
cv2.imwrite(str(out_dir / img_p.name), img)
print(f"Verified 10 labels in {out_dir}")
if __name__ == "__main__":
verify_labels()