-
Notifications
You must be signed in to change notification settings - Fork 0
/
generate_yolo_structure.py
128 lines (98 loc) · 4.31 KB
/
generate_yolo_structure.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import argparse
import pandas as pd
import shutil
import os
import csv
import re
COCO_FOLDER = 'myfiles/yolo-export'
IMAGES_FOLDER = COCO_FOLDER + '/images'
TRAIN_IMAGES = IMAGES_FOLDER + '/train'
TEST_IMAGES = IMAGES_FOLDER + '/val'
LABELS_FOLDER = COCO_FOLDER + '/labels'
TRAIN_LABELS_FOLDER = LABELS_FOLDER + '/train'
TEST_LABELS_FOLDER = LABELS_FOLDER + '/val'
def copy_file(file_path_from: str, file_path_to: str) -> None:
"""Function which copy file path to images"""
if os.path.exists(file_path_to):
return
shutil.copyfile(file_path_from, file_path_to)
def write_annotation(file: str, label, box) -> None:
"""Function which write one annotation into file in txt"""
with open(file, mode='a', encoding='utf-8') as file:
file.write(
"{label} {x} {y} {width} {height}\n".format(label=label, x=box[0], y=box[1], width=box[2], height=box[3]))
def string_to_list(s: str) -> list[float]:
"""Function from string row return array of float"""
# Use regular expression to find all numbers in the string
numbers = re.findall(r"[-+]?\d*\.\d+|\d+", s)
return [float(num) for num in numbers[:4]]
def create_output_structure():
"""Function which create output yolo folder"""
if os.path.exists(COCO_FOLDER):
shutil.rmtree(COCO_FOLDER)
os.mkdir(COCO_FOLDER)
os.mkdir(IMAGES_FOLDER)
os.mkdir(TRAIN_IMAGES)
os.mkdir(TEST_IMAGES)
os.mkdir(LABELS_FOLDER)
os.mkdir(TRAIN_LABELS_FOLDER)
os.mkdir(TEST_LABELS_FOLDER)
def generate_yolo_yaml(labels: list[str]):
"""Function generate yolo yaml"""
with open(COCO_FOLDER + '/' + 'yolo.yaml', mode='w', encoding='utf-8') as file:
dirname = os.path.dirname(__file__)
file.write(
f'path: {dirname}/myfiles/yolo-export\n'
'\n'
'train: images/train\n'
'val: images/val\n'
'\n'
'names:\n'
)
for index, label in enumerate(labels):
file.write(f" {index}: {label}\n")
def generate_yolo(label_file_path: str, detection_file_path: str):
"""Function which generate yolo"""
create_output_structure()
labels = []
with open(label_file_path, mode='r', encoding='utf-8') as label_file:
reader = csv.DictReader(label_file)
for row in reader:
labels.append(re.sub('a photo of an? ', '', row['0']))
generate_yolo_yaml(labels=labels)
image_empty_bed = set()
images_router = {}
df = pd.read_csv(detection_file_path)
df.sort_values("labels", ascending=True, inplace=True)
for index, row in df.iterrows():
file_name = 'myfiles/images/' + row['video_name'] + '-' + row['file_name']
coco_file_name_image = row['video_name'] + '-' + row['file_name']
coco_file_name_txt = coco_file_name_image.replace('.jpg', '.txt')
label = int(row['labels'])
box = string_to_list(row['boxes'])
if label == 0:
image_empty_bed.add(file_name)
if label == 2 and (file_name in image_empty_bed): # if is people lying in the bed and
continue
# split train and test data to 80% / 20%
if not file_name in images_router:
images_router[file_name] = 'test' if len(images_router.keys()) % 5 == 0 else 'train'
current_folder = images_router[file_name]
if current_folder == 'test':
copy_file(file_path_from=file_name, file_path_to=TEST_IMAGES + '/' + coco_file_name_image)
write_annotation(file=TEST_LABELS_FOLDER + '/' + coco_file_name_txt, label=label, box=box)
else:
copy_file(file_path_from=file_name, file_path_to=TRAIN_IMAGES + '/' + coco_file_name_image)
write_annotation(file=TRAIN_LABELS_FOLDER + '/' + coco_file_name_txt, label=label, box=box)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
prog='Yolo structure creator',
description='Program generate yolo creating structure'
)
parser.add_argument('-l', '--label-file', help='Set the label file', default="myfiles/pre-annotation/labels.csv")
parser.add_argument('-d', '--detection-file', help='Set the detection file in csv format', default="myfiles/pre-annotation/detections.csv")
args = parser.parse_args()
generate_yolo(
label_file_path=args.label_file,
detection_file_path=args.detection_file,
)