Spaces:
Build error
Build error
File size: 2,289 Bytes
9067b6a |
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 |
# coding=utf-8
# Description: visualize yolo label image.
import argparse
import os
import cv2
import numpy as np
IMG_FORMATS = ["bmp", "jpg", "jpeg", "png", "tif", "tiff", "dng", "webp", "mpo"]
def main(args):
img_dir, label_dir, class_names = args.img_dir, args.label_dir, args.class_names
label_map = dict()
for class_id, classname in enumerate(class_names):
label_map[class_id] = classname
for file in os.listdir(img_dir):
if file.split('.')[-1] not in IMG_FORMATS:
print(f'[Warning]: Non-image file {file}')
continue
img_path = os.path.join(img_dir, file)
label_path = os.path.join(label_dir, file[: file.rindex('.')] + '.txt')
try:
img_data = cv2.imread(img_path)
height, width, _ = img_data.shape
color = [tuple(np.random.choice(range(256), size=3)) for i in class_names]
thickness = 2
with open(label_path, 'r') as f:
for bbox in f:
cls, x_c, y_c, w, h = [float(v) if i > 0 else int(v) for i, v in enumerate(bbox.split('\n')[0].split(' '))]
x_tl = int((x_c - w / 2) * width)
y_tl = int((y_c - h / 2) * height)
cv2.rectangle(img_data, (x_tl, y_tl), (x_tl + int(w * width), y_tl + int(h * height)), tuple([int(x) for x in color[cls]]), thickness)
cv2.putText(img_data, label_map[cls], (x_tl, y_tl - 10), cv2.FONT_HERSHEY_COMPLEX, 1, tuple([int(x) for x in color[cls]]), thickness)
cv2.imshow('image', img_data)
cv2.waitKey(0)
except Exception as e:
print(f'[Error]: {e} {img_path}')
print('======All Done!======')
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--img_dir', default='VOCdevkit/voc_07_12/images')
parser.add_argument('--label_dir', default='VOCdevkit/voc_07_12/labels')
parser.add_argument('--class_names', default=['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'])
args = parser.parse_args()
print(args)
main(args)
|