動画を使ったYOLO用の学習データを作成します。汎用性を持たせるために、動画からはPascalVOC形式で学習データを作成し、その後、YOLO用の学習データに変換します。ただし、プログラムでは、検出する物体は固定されている物体とします。
フォルダ構成
フォルダ構成を次に示します。「result」フォルダ直下に置かれるファイルについては「Google Colaboratoryを使ってYOLOv4-tinyで学習」を参照してください。
video
動画ファイル
result
＞tree /f result
\RESULT
│ classes.txt
│ obj.data
│ test.txt
│ train.txt
│ yolov4-tiny-my.cfg
│
├─annotation
│ video_0000.xml
│ video_0001.xml
│ video_0002.xml
│ video_0003.xml
│ video_0004.xml
│ ・・・l
│ video_0094.xml
│ video_0095.xml
│ video_0096.xml
│ video_0097.xml
│ video_0098.xml
│ video_0099.xml
│
├─image
│ video_0000.jpg
│ video_0000.txt
│ video_0001.jpg
│ video_0001.txt
│ video_0002.jpg
│ video_0002.txt
│ video_0003.jpg
│ video_0003.txt
│ video_0004.jpg
│ video_0004.txt
│ ・・・・・
│
│ video_0096.jpg
│ video_0096.txt
│ video_0097.jpg
│ video_0097.txt
│ video_0098.jpg
│ video_0098.txt
│ video_0099.jpg
│ video_0099.txt
│
└─train
train.txt
動画からPascalVOC形式の学習データの作成
動画からPascalVOC形式の学習データを作成するプログラムを示します。
- bounding boxを表す座標系には、画像の左上を原点とし、右側にx座標正方向、下側にy座標正方向としたものが利用されています。bounding boxは、boxの左上の点のx座標(xmin)、y座標(ymin)、boxの右下の点のx座標(xmax)、y座標(ymax)として与えられています。
- 12行目と13行目のコメントを外すと、3か所のbounding boxが作成されます。
- 14行目でbounding boxの左上の原点を指定し、16行目でbounding boxののサイズを指定します。
- 15行目でクラスの名前を指定し、[0, 0]は名前の終了を示します。
- 19行目で動画のパスを設定します。
- 37行目で動画から切り取った画像に名前を付けて保存します。
- 38行目から90行目で作成したbounding boxの情報をPascalVOC形式で作成して保存します。
- 118行目で学習データを個数を判断します。
vocFormat/main.py
import cv2
import os
import xml.etree.cElementTree as ET
import random
# bounding boxを表す座標系には、画像の左上を原点とし、右側にx座標正方向、下側にy座標正方向としたものが利用されています。
# bounding boxは、boxの左上の点のx座標(xmin)、y座標(ymin)、boxの右下の点のx座標(xmax)、y座標(ymax)として与えられています。
# box_m = [234, 246], [769, 61], [1287, 171] # bounding boxは3箇所
# name = ["m1", "m2", "m3" ]
box_m = [234, 246], [0, 0] # bounding boxは3箇所
name = ["m1"]
img_size = 416
# video_path = 'video/WIN_20220118_17_19_13_Pro.mp4'
video_path = 'video/WIN_20220210_18_25_55_Pro.mp4'
class TrafficLights:
annotation_dir = 'result/annotation/'
images_dir = 'result/image/'
train_dir = 'result/train'
ext = 'jpg'
basename = "video"
jpeg_filenames_list = []
def __init__(self, name, box):
self.name = name
self.box = box
def setimage(self, num, frame):
cv2.imwrite('{}_{}.{}'.format(TrafficLights.images_dir + TrafficLights.basename, num,
TrafficLights.ext), frame)
filename = '{}_{}'.format(TrafficLights.basename, num)
jpeg_filename = filename + TrafficLights.ext
# テキストファイルの作成
TrafficLights.jpeg_filenames_list.append(filename)
# XMLファイルの保存
xml_filename = filename + '.xml'
print(xml_filename)
new_root = ET.Element('annotation')
new_filename = ET.SubElement(new_root, 'filename')
new_filename.text = jpeg_filename
Size = ET.SubElement(new_root, 'size')
Width = ET.SubElement(Size, 'width')
Height = ET.SubElement(Size, 'height')
Depth = ET.SubElement(Size, 'depth')
Width.text = str(frame.shape[1])
Height.text = str(frame.shape[0])
Depth.text = str(frame.shape[2])
i = 0
for index in self.box:
if index[0] == 0: break
Object = ET.SubElement(new_root, 'object')
Name = ET.SubElement(Object, 'name')
Name.text = self.name[i]
Difficult = ET.SubElement(Object, 'difficult')
Difficult.text = '0'
Bndbox = ET.SubElement(Object, 'bndbox')
Xmin = ET.SubElement(Bndbox, 'xmin')
Ymin = ET.SubElement(Bndbox, 'ymin')
Xmax = ET.SubElement(Bndbox, 'xmax')
Ymax = ET.SubElement(Bndbox, 'ymax')
x = random.randint(0, 32) - 16
y = random.randint(0, 32) - 16
Xmin.text = str(index[0] + x)
Ymin.text = str(index[1] + y)
Xmax.text = str(index[0] + x + img_size)
Ymax.text = str(index[1] + y + img_size)
i += 1
new_tree = ET.ElementTree(new_root)
new_tree.write(os.path.join(TrafficLights.annotation_dir, xml_filename))
@staticmethod
def cleatetrain():
# テキストファイルの保存
text = "\n".join(TrafficLights.jpeg_filenames_list)
with open(os.path.join(TrafficLights.train_dir, 'train.txt'), "w") as f:
f.write(text)
def main():
trafficlights_m = TrafficLights(name, box_m)
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return
digit = len(str(int(cap.get(cv2.CAP_PROP_FRAME_COUNT))))
n = 0
while True:
ret, frame = cap.read()
if ret:
num = str(n).zfill(digit)
trafficlights_m.setimage(num, frame)
n += 1
if n == 100:
break
else:
break
TrafficLights.cleatetrain()
if __name__ == '__main__':
main()
実行すると次のようなPascalVOC形式のファイルが「RESULT/annotation」フォルダに作成されます。
PascalVOC形式学習データをYOLO用学習データに変換
PascalVOC形式学習データをYOLO用学習データに変換するプログラムを示します。
- 34行目のvoc2yoloメソッドで作成したPascalVOC形式学習データの「annotation」フォルダからYOLO用学習データに変換します。
- 77行目のimglist2fileメソッドで作成したリストをシャッフルしてから、学習用のデータリストを「train.txt」ファイルに、評価用のデータリストを「test.txt」に保存します。
vocFormat/pascalVOC2yolov3.py
# coding:utf-8
from __future__ import print_function
import os
import random
import glob
import xml.etree.ElementTree as ET
base_dir = 'result/'
annotation_dir = base_dir + 'annotation/'
images_dir = base_dir + 'image/'
yolo_dir = base_dir + 'image'
def xml_reader(filename):
""" Parse a PASCAL VOC xml file """
tree = ET.parse(filename)
size = tree.find('size')
width = int(size.find('width').text)
height = int(size.find('height').text)
objects = []
for obj in tree.findall('object'):
obj_struct = {}
obj_struct['name'] = obj.find('name').text
bbox = obj.find('bndbox')
obj_struct['bbox'] = [int(bbox.find('xmin').text),
int(bbox.find('ymin').text),
int(bbox.find('xmax').text),
int(bbox.find('ymax').text)]
objects.append(obj_struct)
return width, height, objects
def voc2yolo(filename):
classes_dict = {}
with open("result/classes.txt") as f:
for idx, line in enumerate(f.readlines()):
class_name = line.strip()
classes_dict[class_name] = idx
width, height, objects = xml_reader(filename)
lines = []
for obj in objects:
x, y, x2, y2 = obj['bbox']
class_name = obj['name']
label = classes_dict[class_name]
cx = (x2 + x) * 0.5 / width
cy = (y2 + y) * 0.5 / height
w = (x2 - x) * 1. / width
h = (y2 - y) * 1. / height
line = "%s %.6f %.6f %.6f %.6f\n" % (label, cx, cy, w, h)
lines.append(line)
txt_name = filename.replace(".xml", ".txt").replace("annotation", "image")
with open(txt_name, "w") as f:
f.writelines(lines)
def get_image_list(image_dir, suffix=['jpg', 'jpeg', 'JPG', 'JPEG', 'png']):
'''get all image path ends with suffix'''
if not os.path.exists(image_dir):
print("PATH:%s not exists" % image_dir)
return []
imglist = []
for root, sdirs, files in os.walk(image_dir):
if not files:
continue
for filename in files:
filepath = os.path.join(root, filename) + "\n"
if filename.split('.')[-1] in suffix:
print(filepath)
imglist.append(filepath)
return imglist
def imglist2file(imglist):
random.shuffle(imglist)
# size = 100
size = 20
train_list = imglist[:-size]
valid_list = imglist[-size:]
with open(base_dir + "train.txt", "w") as f:
f.writelines(train_list)
with open(base_dir + "test.txt", "w") as f:
f.writelines(valid_list)
if __name__ == "__main__":
xml_path_list = glob.glob(annotation_dir + "*.xml")
for xml_path in xml_path_list:
voc2yolo(xml_path)
imglist = get_image_list(images_dir)
imglist2file(imglist)
実行すると次のように変換されたYOLO用学習データのファイルが、画像データと組み合わせて「RESULT/image」フォルダに作成されます。
video_0000.txt
0 0.226562 0.418519 0.216667 0.385185