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)