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

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

点击(此处)折叠或打开

1. import cv2
   
2. import os
   
3. import sys
   
4. import time
   
5.  
   
6. # 从摄像头中获取图像，按w则保存一张到label目录下
   
7. if __name__ == '__main__':
   
8.     #capture = cv2.VideoCapture(0) #摄像头/dev/video0
   
9.     capture = cv2.VideoCapture(1) #摄像头/dev/video1
   
10.     if (False == capture.isOpened()):
    
11.         print("capture open failed")
    
12.         sys.exit(-1)
    
13.  
    
14.     if (len(sys.argv) != 2):
    
15.         print("error argc != 2")
    
16.         sys.exit(-1)
    
17.  
    
18.     label = sys.argv[1] # label
    
19.     curDir = os.path.dirname(os.path.realpath(__file__))
    
20.     labelDir = os.path.join(curDir, label)
    
21.  
    
22.     if (False == os.path.exists(labelDir)):
    
23.         os.makedirs(labelDir) #创建label目录
    
24.         
    
25.     namei = 0
    
26.  
    
27.     while (1):
    
28.         key = cv2.waitKey(1)
    
29.         if (key == ord('q')):
    
30.             print('quit')
    
31.             sys.exit(-1)
    
32.         elif (key == ord('w')): #按w保存
    
33.             namei += 1
    
34.             fileName = os.path.join(labelDir, str(namei))
    
35.             fileName += ".jpg"
    
36.             print(fileName)
    
37.             cv2.imwrite(fileName, img)
    
38.  
    
39.         ret, img = capture.read()
    
40.         if (False == ret):
    
41.             print("read failed")
    
42.             time.sleep(1)
    
43.             continue
    
44.  
    
45.         cv2.imshow("test", img)
    
46.  
    
47.         #time.sleep(1)
    
48.  
    
49.     capture.release()
    
50.     cv2.destroyAllWindows()

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

点击(此处)折叠或打开

1. def GetLabelAndFileName(rootdir, level, label):
   
2.     if not os.path.isdir(rootdir):
   
3.         print("not dir")
   
4.         return "",""
   
5.  
   
6.     level += 1
   
7.     #print("level:", level)
   
8.  
   
9.     labelNames = []
   
10.     list = os.listdir(rootdir)
    
11.     for i in range(0, len(list)):
    
12.         path = os.path.join(rootdir, list[i])
    
13.         if os.path.isfile(path):
    
14.             if (level == 1):
    
15.                 #print("ignore:", path)
    
16.                 continue
    
17.             else:
    
18.                 #print(path, label)
    
19.                 labelNames.append({"path":path, "label":label})
    
20.         elif os.path.isdir(path):
    
21.             if (level == 1):
    
22.                 label = list[i]
    
23.                 #print("set label:", label)
    
24.             labelNames.extend(GetLabelAndFileName(path, level, label))
    
25.  
    
26.     return labelNames
    
27.  
    
28. if __name__ == "__main__":
    
29.     labelNames = GetLabelAndFileName(sys.argv[1], 0, "")
    
30.     ##print(labelNames)
    
31.     for ln in labelNames:
    
32.         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

作者：帅得不敢出门