print("[Training HTR] Model: SFR | Dataset: ", args.dataset) data_folder = os.getenv("DATA_FOLDER") if os.getenv("DATA_FOLDER") else "data" target_folder = os.path.join(data_folder, "sfr", args.dataset) pages_folder = os.path.join(target_folder, "pages") char_set_path = os.path.join(pages_folder, "character_set.json") with open(char_set_path) as f: char_set = json.load(f) idx_to_char = {} for k, v in char_set['idx_to_char'].items(): idx_to_char[int(k)] = v training_set_list_path = os.path.join(pages_folder, "training.json") training_set_list = load_file_list_direct(training_set_list_path) train_dataset = HwDataset(training_set_list, char_set['char_to_idx'], augmentation=True, img_height=args.input_height) train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0, drop_last=True, collate_fn=hw_dataset.collate) batches_per_epoch = int(args.images_per_epoch / args.batch_size) train_dataloader = DatasetWrapper(train_dataloader, batches_per_epoch) testing_set_list_path = os.path.join(pages_folder, "testing.json")
args_filename = os.path.join(args.output, args.name, 'args.json') create_folders(args_filename) with open(args_filename, 'w') as fp: json.dump(args.__dict__, fp, indent=4) print("[Training Line-Outliner] Model: SFRS | Dataset: ", args.dataset) data_folder = os.getenv("DATA_FOLDER") if os.getenv("DATA_FOLDER") else "data" target_folder = os.path.join(data_folder, "sfrs", args.dataset) pages_folder = os.path.join(target_folder, "pages") char_set_path = os.path.join(pages_folder, "character_set.json") training_set_list_path = os.path.join(pages_folder, "training.json") training_set_list = load_file_list_direct(training_set_list_path) train_dataset = LolDataset(training_set_list, augmentation=True) train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0, collate_fn=lol_dataset.collate) batches_per_epoch = int(int(args.images_per_epoch) / args.batch_size) train_dataloader = DatasetWrapper(train_dataloader, batches_per_epoch) test_set_list_path = os.path.join(pages_folder, "testing.json") test_set_list = load_file_list_direct(test_set_list_path) test_dataset = LolDataset(test_set_list) test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False,
parser.add_argument("--stop_after_no_improvement", default=20) parser.add_argument("--learning_rate", default=0.0002) parser.add_argument("--tsa_size", default=3) parser.add_argument("--patch_ratio", default=5) parser.add_argument("--output", default="scripts/original/snapshots/training") parser.add_argument("--model", default="scripts/new/snapshots/training2/lol-last.pt") args = parser.parse_args() data_folder = os.getenv("DATA_FOLDER") if os.getenv("DATA_FOLDER") else "data" target_folder = os.path.join(data_folder, "sfrs", args.dataset) pages_folder = os.path.join(target_folder, "pages") char_set_path = os.path.join(pages_folder, "character_set.json") test_set_list_path = os.path.join(pages_folder, "validation.json") test_set_list = load_file_list_direct(test_set_list_path) test_dataset = LolDataset(test_set_list[0:1]) test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=0, collate_fn=lol_dataset.collate) count = 0 while True: for t in ["training", "training2"]: model_path = "scripts/new/snapshots/" + t + "/last.pt" paint_model_run(model_path, test_dataloader, destination=os.path.join("screenshots", t,