Python+OpenCV实现Faster R-CNN样本查看器

Python+OpenCV实现Faster R-CNN样本查看器

一、上代码

import cv2
import os


def get_samples(dir):
    datasets = []
    files = os.listdir(dir)
    for file in files:
        ext_sp = os.path.splitext(file)
        if ext_sp[1] not in ['.jpg', '.bmp', '.png']:
            continue
        if not os.path.exists(os.path.join(dir, ext_sp[0] + '.txt')):
            continue
        lines = []
        with open(os.path.join(dir, ext_sp[0] + '.txt'), mode='r', encoding='utf-8') as fs:
            lines = fs.readlines()
        boxes = []
        for line in lines:
            sp_arr = line.split(',')
            boxes.append({'label': sp_arr[0], 'left': int(sp_arr[1]), 'top': int(sp_arr[2]), 'width': int(sp_arr[3]),
                          'height': int(sp_arr[4])})
        datasets.append({'dir': dir, 'file': ext_sp, 'txt': ext_sp[0] + '.txt', 'boxes': boxes})
    return datasets


if __name__ == '__main__':
    dir = 'E:\\BaiduNetdiskDownload\\MeterPhotos\\20171218152456'
    datasets = get_samples(dir)

    for ds in datasets:
        imgpath = os.path.join(ds['dir'], ds['file'][0] + ds['file'][1])
        img = cv2.imread(imgpath, cv2.IMREAD_COLOR)
        for box in ds['boxes']:
            cv2.rectangle(img, (box['left'], box['top']), (box['left'] + box['width'], box['top'] + box['height']),
                          (255, 0, 0), 2)
            cv2.putText(img, box['label'], (box['left'], box['top']), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
        cv2.imshow('win', img)
        print('Current show:%s' %imgpath)
        k = cv2.waitKey(0)
        if  k==27:
            break

二、效果（按ESC键退出）