def __init__(self, config_file): self.config_file_path = config_file # ----- CharNet configuration and initialization ----- cfg.merge_from_file(config_file) cfg.freeze() print(cfg) self.charnet = CharNet() self.charnet.load_state_dict(torch.load(cfg.WEIGHT)) self.charnet.eval() if torch.cuda.is_available(): print("[*] Using cuda!") self.charnet.cuda()
(word_bbox[0], word_bbox[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1) return img_word_ins if __name__ == '__main__': parser = argparse.ArgumentParser(description="Test") parser.add_argument("config_file", help="path to config file", type=str) parser.add_argument("image_dir", type=str) parser.add_argument("results_dir", type=str) args = parser.parse_args() cfg.merge_from_file(args.config_file) cfg.freeze() print(cfg) charnet = CharNet() charnet.load_state_dict(torch.load(cfg.WEIGHT)) charnet.eval() charnet.cuda() for im_name in tqdm(sorted(os.listdir(args.image_dir))): #print("Processing {}...".format(im_name)) im_file = os.path.join(args.image_dir, im_name) im_original = cv2.imread(im_file) im, scale_w, scale_h, original_w, original_h = resize( im_original, size=cfg.INPUT_SIZE) with torch.no_grad():