从摄像头批量生成训练图片

创建以label为名称的文件夹，生成图片放入label目录下

import cv2
import os
import sys
import time

# 从摄像头中获取图像，按w则保存一张到label目录下
if __name__ == '__main__':
    #capture = cv2.VideoCapture(0) #摄像头/dev/video0
    capture = cv2.VideoCapture(1) #摄像头/dev/video1
    if (False == capture.isOpened()):
        print("capture open failed")
        sys.exit(-1)

    if (len(sys.argv) != 2):
        print("error argc != 2")
        sys.exit(-1)

    label = sys.argv[1] # label
    curDir = os.path.dirname(os.path.realpath(__file__))
    labelDir = os.path.join(curDir, label)

    if (False == os.path.exists(labelDir)):
        os.makedirs(labelDir) #创建label目录
        
    namei = 0

    while (1):
        key = cv2.waitKey(1)
        if (key == ord('q')):
            print('quit')
            sys.exit(-1)
        elif (key == ord('w')): #按w保存
            namei += 1
            fileName = os.path.join(labelDir, str(namei))
            fileName += ".jpg" 
            print(fileName)
            cv2.imwrite(fileName, img)

        ret, img = capture.read()
        if (False == ret):
            print("read failed")
            time.sleep(1)
            continue

        cv2.imshow("test", img)

        #time.sleep(1)

    capture.release()
    cv2.destroyAllWindows()

读取代码如下，遍历文件夹，得到路径及对应label

def GetLabelAndFileName(rootdir, level, label):
    if not os.path.isdir(rootdir):
        print("not dir")
        return "",""

    level += 1
    #print("level:", level)

    labelNames = []
    list = os.listdir(rootdir)
    for i in range(0, len(list)):
        path = os.path.join(rootdir, list[i])
        if os.path.isfile(path):
            if (level == 1):
                #print("ignore:", path)
                continue
            else:
                #print(path, label)
                labelNames.append({"path":path, "label":label})
        elif os.path.isdir(path):
            if (level == 1):
                label = list[i]
                #print("set label:", label)
            labelNames.extend(GetLabelAndFileName(path, level, label))

    return labelNames

if __name__ == "__main__":
    labelNames = GetLabelAndFileName(sys.argv[1], 0, "")
    ##print(labelNames)
    for ln in labelNames:
        print(ln["path"], ln["label"])

imgs/4/155.jpg 4
imgs/4/156.jpg 4
imgs/4/157.jpg 4
imgs/6/1.jpg 6
imgs/6/2.jpg 6
imgs/6/3.jpg 6

作者：帅得不敢出门