3. 检查下转换YOLO数据集是否成功

import os
import cv2

img_formats = ['.bmp', '.jpg', '.jpeg', '.png', '.tif']

if __name__ == '__main__':
    classes = ['person']
    path = '/home/wangsong/data/val2017.txt'
    with open(path, 'r') as f:
        img_files = [x for x in f.read().splitlines() if os.path.splitext(x)[-1].lower() in img_formats]

    for img_path in img_files:
        img = cv2.imread(img_path)

        im_h = img.shape[0]
        im_w = img.shape[1]

        label_path = img_path.replace('images', 'labels').replace(os.path.splitext(img_path)[-1], '.txt')

        with open(label_path) as file:
            line = file.readline()
            while line:
                cls = int(line.split(' ')[0])
                x = float(line.split(' ')[1])
                y = float(line.split(' ')[2])
                w = float(line.split(' ')[3])
                h = float(line.split(' ')[4].split('\n')[0])

                # print(cls,x,y,w,h)
                real_w = im_w * w
                real_h = im_h * h
                xmin = ((x * im_w + 1.0) * 2.0 - real_w) / 2.0
                xmax = ((x * im_w + 1.0) * 2.0 + real_w) / 2.0
                ymin = ((y * im_h + 1.0) * 2.0 - real_h) / 2.0
                ymax = ((y * im_h + 1.0) * 2.0 + real_h) / 2.0

                img = cv2.rectangle(img, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (0, 255, 0), 2)
                line = file.readline()
            cv2.imshow('result', img)
            cv2.waitKey(0)

发布了16 篇原创文章 · 获赞 37 · 访问量 8465

猜你喜欢

转载自blog.csdn.net/songwsx/article/details/102698186
今日推荐