def _main():
annotation_path = 'train.txt'
data_path = 'train.npz'
output_path = 'cfg/my_yolo.h5'
log_dir = 'logs/000/'
classes_path = 'cfg/voc_classes.txt'
anchors_path = 'cfg/yolo_anchors.txt'
class_names = get_classes(classes_path)
anchors = get_anchors(anchors_path)
input_shape = (416, 416) # multiple of 32
image_data, box_data = get_training_data(annotation_path,
data_path,
input_shape,
max_boxes=100,
load_previous=True)
y_true = preprocess_true_boxes(box_data, input_shape, anchors,
len(class_names))
infer_model, model = create_model(input_shape,
anchors,
len(class_names),
load_pretrained=True,
freeze_body=True)
train(model, image_data / 255., y_true, log_dir=log_dir)
infer_model.save(output_path)
def data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes):
'''data generator for fit_generator'''
n = len(annotation_lines)
i = 0
while True:
image_data = []
box_data = []
for b in range(batch_size):
if i==0:
np.random.shuffle(annotation_lines)
image, box = get_random_data(annotation_lines[i], input_shape, random=True)
image_data.append(image)
box_data.append(box)
i = (i+1) % n
image_data = np.array(image_data)
box_data = np.array(box_data)
y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)
yield [image_data, *y_true], np.zeros(batch_size)