def crnnSource(): alphabet = keys_crnn.alphabet converter = util.strLabelConverter(alphabet) if torch.cuda.is_available() and GPU: print("GPU") model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda() else: print("CPU") model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cpu() #model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda() path = './CRNN/crnn/samples/newmodel.pth' # 将GUP训练模型的权重转换为单个CPU可用 trainWeights = torch.load(path, map_location=lambda storage, loc: storage) modelWeights = OrderedDict() for k, v in trainWeights.items(): name = k.replace('module.', '') modelWeights[name] = v model.eval() model.load_state_dict(modelWeights) return model, converter
def crnnSource(): alphabet = keys.alphabet converter = util.strLabelConverter(alphabet) if torch.cuda.is_available() and GPU: model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda() else: model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cpu() path = './crnn/samples/model_acc97.pth' model.eval() model.load_state_dict(torch.load(path)) return model, converter
def new_model(nclass, preModel): # 定义你自己的模型 if torch.cuda.is_available() and opt.ngpu: model = crnn.CRNN(32, 1, nclass + 1, 256, 1).cuda() else: model = crnn.CRNN(32, 1, nclass + 1, 256, 1).cpu() modelDict = model.state_dict() ## preModelDict = preModel.state_dict() ## preModelDict = {k: v for k, v in preModelDict.items() if 'rnn.1' not in k} modelDict.update(preModelDict) ##更新权重 model.load_state_dict(modelDict) ##加载预训练模型权重 return model
def crnnSource(self, dir_model): alphabet = '0123456789abcdefghijklmnopqrstuvwxyz' # keys.alphabet converter = utils.strLabelConverter(alphabet) if self.gpuid == '-1': print('cpu-version') model = crnn.CRNN(32, 1, len(alphabet) + 1, 256) else: print('gpu-version') model = crnn.CRNN(32, 1, len(alphabet) + 1, 256).cuda() path = dir_model + '/bib_number/CRNN/crnn.pth' # './crnn/samples/netCRNN63.pth' model.load_state_dict(torch.load(path)) print('Loaded network {:s}'.format(path)) return model, converter
def __init__(self, model_path): #def crnnSource(model_path, use_gpu=True): alphabet = keys.alphabet # Chinese words self.converter = crnn_utils.strLabelConverter(alphabet) # note that in https://github.com/bear63/sceneReco support multi GPU. # model = crnn.CRNN(32, 1, len(alphabet)+1, 256, 1).cuda() self.model = crnn.CRNN(32, 1, len(alphabet) + 1, 256) self.cpu_model = crnn.CRNN(32, 1, len(alphabet) + 1, 256) if torch.cuda.is_available(): self.model = self.model.cuda() print('loading pretrained model from %s' % model_path) #model_path = './crnn/samples/netCRNN63.pth' model_state_dict = torch.load(model_path) self.model.load_state_dict(model_state_dict) self.cpu_model.load_state_dict(model_state_dict)
def __init__(self, model_path, gpu_id=None): ''' 初始化pytorch模型 :param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件) :param gpu_id: 在哪一块gpu上运行 ''' self.gpu_id = gpu_id self.converter = utils.strLabelConverter(alphabet) if self.gpu_id is not None and isinstance( self.gpu_id, int) and torch.cuda.is_available(): checkpoint = torch.load(model_path) self.device = torch.device("cuda:%s" % self.gpu_id) else: checkpoint = torch.load(model_path, map_location='cpu') self.device = torch.device("cpu") print('text recognition running on device:', self.device) self.net = crnn.CRNN(CNN=CNN, RNN=RNN, nIn=n_in, n_class=new_class, nHidden=n_hidden, nLayer=n_layer, dropout=dp) self.net.load_state_dict(checkpoint['state_dict']) self.net.to(self.device) self.net.eval() self.transform = transforms.Compose([transforms.ToTensor()])
def recognize(image_path, alphabet, snapshot, gpu): model = crnn.CRNN(32, 1, 37, 256) if torch.cuda.is_available(): model = model.cuda() print('loading pretrained model from %s' % snapshot) model.load_state_dict(torch.load(snapshot)) converter = utils.strLabelConverter(alphabet) transformer = dataset.resizeNormalize((100, 32)) image = Image.open(image_path).convert('L') image = transformer(image) if torch.cuda.is_available(): image = image.cuda() image = image.view(1, *image.size()) image = Variable(image) model.eval() preds = model(image) _, preds = preds.max(2) preds = preds.transpose(1, 0).contiguous().view(-1) preds_size = Variable(torch.IntTensor([preds.size(0)])) raw_pred = converter.decode(preds.data, preds_size.data, raw=True) sim_pred = converter.decode(preds.data, preds_size.data, raw=False) print('%-20s => %-20s' % (raw_pred, sim_pred)) return sim_pred
def crnnSource(): alphabet = keys.alphabet converter = util.strLabelConverter(alphabet) model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda() path = './crnn/samples/netCRNN63.pth' model.load_state_dict(torch.load(path)) return model, converter
def batch_test(dirpath): alphabet = keys_crnn.alphabet #print(len(alphabet)) #input('\ninput:') converter = util.strLabelConverter(alphabet) # model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda() model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1) path = './samples/model_acc97.pth' model.load_state_dict(torch.load(path)) #print(model) paths=glob.glob(os.path.join(dirpath,'*.[jp][pn]g')) for i in paths: print(i) image = Image.open(i).convert('L') #print(image.size) scale = image.size[1] * 1.0 / 32 w = image.size[0] / scale w = int(w) #print("width:" + str(w)) transformer = dataset.resizeNormalize((w, 32)) # image = transformer(image).cuda() image = transformer(image) image = image.view(1, *image.size()) image = Variable(image) model.eval() preds = model(image) #print(preds.shape) _, preds = preds.max(2) #print(preds.shape) preds = preds.squeeze(1) preds = preds.transpose(-1, 0).contiguous().view(-1) preds_size = Variable(torch.IntTensor([preds.size(0)])) raw_pred = converter.decode(preds.data, preds_size.data, raw=True) sim_pred = converter.decode(preds.data, preds_size.data, raw=False) print(sim_pred)
def load_images_to_predict(): # load model model_path = './expr/netCRNN_199_423.pth' alphabet = '0123456789,.:(%$!^&-/);<~|`>?+=_[]{}"\'@#*ABCDEFGHIJKLMNOPQRSTUVWXYZ\ ' imgH = 32 # should be 32 nclass = len(alphabet) + 1 nhiddenstate = 256 model = crnn.CRNN(imgH, 1, nclass, nhiddenstate) if torch.cuda.is_available(): model = model.cuda() print('loading pretrained model from %s' % model_path) model.load_state_dict({k.replace('module.',''):v for k,v in torch.load(model_path).items()}) # load image #img = [os.path.basename(f) for f in glob.glob(IMAGE_FOLDER)] #filenames, extns = [os.path.splitext(f)[0:1] for f in glob.glob(IMAGE_FOLDER)] #exts = [os.path.splitext(f)[1] for f in glob.glob(IMAGE_FOLDER)] #img_files = [s + e for s , e in (filenames, extns)] for img in glob.glob(IMAGE_FOLDER): image = Image.open(img).convert('L') words_list = [] #with open(BB_FOLDER + img.split('/')[1].split('.')[0]+'.txt', 'r') as boxes: print('Ïmage name:', os.path.basename(img)) with open(BB_INT_FOLDER + os.path.basename(img).split('.')[0]+'.txt', 'r') as boxes: for line in csv.reader(boxes): box = [int(string, 10) for string in line[0:8]] boxImg = image.crop((box[0], box[1], box[4], box[5])) words = predict_this_box(boxImg, model, alphabet) words_list.append(words) with open( TXT_RESULT_FLD + os.path.basename(img).split('.')[0] +'.txt', 'w+') as resultfile: for line in words_list: resultfile.writelines(line+'\n')
def load_images_to_predict(): # load model model_path = './expr/netCRNN_99_500.pth' alphabet = '0123456789,.:(%$!^&-/);<~|`>?+=_[]{}"\'@#*ABCDEFGHIJKLMNOPQRSTUVWXYZ\ ' imgH = 32 # should be 32 nclass = len(alphabet) + 1 nhiddenstate = 256 model = crnn.CRNN(imgH, 1, nclass, nhiddenstate) if torch.cuda.is_available(): model = model.cuda() print('loading pretrained model from %s' % model_path) model.load_state_dict({ k.replace('module.', ''): v for k, v in torch.load(model_path).items() }) # load image filenames = [os.path.splitext(f)[0] for f in glob.glob("data_test/*.jpg")] jpg_files = [s + ".jpg" for s in filenames] for jpg in jpg_files: image = Image.open(jpg).convert('L') words_list = [] with open('boundingbox/' + jpg.split('/')[1].split('.')[0] + '.txt', 'r') as boxes: for line in csv.reader(boxes): box = [int(string, 10) for string in line[0:8]] boxImg = image.crop((box[0], box[1], box[4], box[5])) words = predict_this_box(boxImg, model, alphabet) words_list.append(words) with open('test_result/' + jpg.split('/')[1].split('.')[0] + '.txt', 'w+') as resultfile: for line in words_list: resultfile.writelines(line + '\n')
def crnnModel(): alphabet = keys.alphabet # define a converter for convert rnn results to string converter = util.strLabelConverter(alphabet) model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda() path = './crnn/samples/CRNN.pth' model.load_state_dict(torch.load(path)) return model, converter
def __init__(self, weight): super(Ocr, self).__init__() # 读取字母表 with open('../data/lol.alphabet', encoding='utf-8') as f: self.alphabet = f.read().strip() self.net = crnn.CRNN(32, 1, len(self.alphabet) + 1, 256).to('cuda') self.net.load_state_dict(torch.load(weight)) self.converter = utils.StrLabelConverter(self.alphabet)
def __init__(self, model_path='./crnn.pytorch-master/data/crnn.pth'): self.model = crnn.CRNN(32, 1, 37, 256) if torch.cuda.is_available(): self.model = self.model.cuda() print('loading pretrained self.model from %s' % model_path) self.model.load_state_dict(torch.load(model_path)) self.alphabet = '0123456789abcdefghijklmnopqrstuvwxyz' self.converter = utils.strLabelConverter(self.alphabet)
def __init__(self, model_path): self.model = crnn.CRNN(32, 1, 37, 256) if torch.cuda.is_available(): self.model = self.model.cuda() self.model.load_state_dict(torch.load(model_path)) self.converter = utils.strLabelConverter( '0123456789abcdefghijklmnopqrstuvwxyz') self.model.eval()
def main(): if not os.path.exists(opt.output): os.makedirs(opt.output) converter = utils.strLabelConverter(opt.alphabet) collate = dataset.AlignCollate() train_dataset = dataset.TextLineDataset(text_file=opt.train_list, transform=dataset.ResizeNormalize(100, 32), converter=converter) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=opt.batchsize, shuffle=True, num_workers=opt.num_workers, collate_fn=collate) test_dataset = dataset.TextLineDataset(text_file=opt.train_list, transform=dataset.ResizeNormalize(100, 32), converter=converter) test_loader = torch.utils.data.DataLoader(test_dataset, shuffle=False, batch_size=opt.batchsize, num_workers=opt.num_workers, collate_fn=collate) criterion = nn.CTCLoss() import models.crnn as crnn crnn = crnn.CRNN(opt.imgH, opt.nc, opt.num_classes, opt.nh) crnn.apply(utils.weights_init) if opt.pretrained != '': print('loading pretrained model from %s' % opt.pretrained) crnn.load_state_dict(torch.load(opt.pretrained), strict=False) print(crnn) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") crnn = crnn.to(device) criterion = criterion.to(device) # setup optimizer optimizer = optim.Adam(crnn.parameters(), lr=opt.lr) for epoch in range(opt.num_epochs): loss_avg = 0.0 i = 0 while i < len(train_loader): time0 = time.time() # 训练 train_iter = iter(train_loader) cost = trainBatch(crnn, train_iter, criterion, optimizer, device) # 一个批次,一个批次训练 loss_avg += cost i += 1 if i % opt.interval == 0: print('[%d/%d][%d/%d] Loss: %f Time: %f s' % (epoch, opt.num_epochs, i, len(train_loader), loss_avg, time.time() - time0)) loss_avg = 0.0 if (epoch + 1) % opt.valinterval == 0: val(crnn, test_loader, criterion, converter=converter, device=device, max_iter=100)
def main(): model_path = './data/crnn.pth' img_path = './data/demo.png' alphabet = '0123456789abcdefghijklmnopqrstuvwxyz' model = crnn.CRNN(32, 1, 37, 256) if torch.cuda.is_available(): model = model.cuda() print('loading pretrained model from %s' % model_path) model.load_state_dict(torch.load(model_path)) converter = utils.strLabelConverter(alphabet) #transformer = dataset.resizeNormalize((100, 32)) dir = Path("images/").glob( "**/*.jpg") # Modifiable. Takes all .jpg files in the given directory for image_path in dir: # Applies a Face - Detect for every image in the above directory image = cv2.imread(str(image_path)) cv2.imshow('Original Image', image) cv2.waitKey(0) plate_contour = plate_detection(image) if plate_contour is None: print("No license plates detected") continue img = image.copy() try: plate = crop_contour(plate_contour, image, img) except: print("Plate could not be processed") continue cv2.destroyAllWindows() plate_recognition(plate) image = plate.convert('L') #image = transformer(image) cv2.resize(image, (100, 32)) if torch.cuda.is_available(): image = image.cuda() image = image.view(1, *image.size()) image = Variable(image) model.eval() preds = model(image) _, preds = preds.max(2) preds = preds.transpose(1, 0).contiguous().view(-1) preds_size = Variable(torch.IntTensor([preds.size(0)])) raw_pred = converter.decode(preds.data, preds_size.data, raw=True) sim_pred = converter.decode(preds.data, preds_size.data, raw=False) print('%-20s => %-20s' % (raw_pred, sim_pred)) exit()
def pre_model(nclass, ocrModelPath): # @@parm nclass:字符总数 # @@预训练模型文件 if torch.cuda.is_available() and opt.ngpu: model = crnn.CRNN(32, 1, nclass + 1, 256, 1).cuda() else: model = crnn.CRNN(32, 1, nclass + 1, 256, 1).cpu() state_dict = torch.load(ocrModelPath, map_location=lambda storage, loc: storage) new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k.replace('module.', '') # remove `module.` new_state_dict[name] = v model.load_state_dict(new_state_dict) model.eval() return model
def test(): config = config_args() transform_val=\ transforms.Compose([ transforms.Resize(size=(config.DATASET.IMAGE_SIZE.H,config.DATASET.IMAGE_SIZE.W)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5],std=[0.5, 0.5, 0.5])]) criterion = nn.CTCLoss() label_tool = LabelTool(char_std_path=config.DATASET.CHAR_FILE) model = crnn.CRNN(config).to(torch.device("cuda:" + config.CUDNN.GPU)) model.load_state_dict( torch.load(config.TRAIN.RESUME.MODEL_SAVE)['state_dict']) images_val, labels_val = label_tool.get_labels( labels_path=config.DATASET.LABELS_FILE.VAL) dataset_val = Dataset_OCR(images_root_dir=config.DATASET.IMAGE_ROOT, images_name=images_val, transform=transform_val) dataloader_val = DataLoader(dataset=dataset_val, batch_size=config.TEST.BATCH, shuffle=config.TEST.SHUFFLE, num_workers=config.TEST.WORKERS) loss_all_val = 0 step_val = 0 nums_all = 0 nums_all_correct = 0 for images_val, indexs_val in dataloader_val: images_val = images_val.to(torch.device("cuda:" + config.CUDNN.GPU)) with torch.no_grad(): output_val = model(images_val) sequence_len_val = output_val.shape[0] target_val, input_lengths_val, target_lengths_val = label_tool.convert_ctcloss_labels( indexs_val, labels_val, sequence_len_val) preds_val = output_val.permute(1, 0, 2).argmax(2).cpu().numpy() preds_str_val, preds_str_val_blank = label_tool.decode_batch(preds_val) correct_nums = label_tool.cal_correct_nums(preds_str_val, preds_str_val_blank, indexs_val, labels_val, step_val) nums_all_correct += correct_nums nums_all += output_val.shape[1] print('nums_all_correct{},nums_all{}'.format(nums_all_correct, nums_all)) loss_val = criterion(output_val, target_val, input_lengths_val, target_lengths_val) loss_all_val += loss_val step_val += 1 acc = nums_all_correct / nums_all print("val_loss_avarage={},val_accuracy={}".format( loss_all_val / step_val, nums_all_correct / nums_all))
def crnnSource(): # Choose alphabet from arg or key.py(default) if opt.alphabet is None: alphabet = keys.alphabet else: alphabet = opt.alphabet converter = utils.strLabelConverter(alphabet) if torch.cuda.is_available() and opt.gpu: model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1) model = model.cuda() else: model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1) model = model.cpu() path = opt.model print('loading pretrained model from %s' % path) model.eval() model.load_state_dict(torch.load(path)) return model, converter
def net_init(): nclass = len(params.alphabet) + 1 crnn = net.CRNN(params.imgH, params.nc, nclass, params.nh) crnn.apply(weights_init) if params.pretrained != '': print('loading pretrained model from %s' % params.pretrained) if params.multi_gpu: crnn = torch.nn.DataParallel(crnn) crnn.load_state_dict(torch.load(params.pretrained)) return crnn
def __init__(self, weights, char, height=32, cuda=None): alphabet = open(char).read().rstrip() nclass = len(alphabet) + 1 self.height = height self.device = 'cuda:{}'.format(cuda) if cuda != None else 'cpu' self.model = crnn.CRNN(32, 3, nclass, 256) self.model.load_state_dict( torch.load(weights, map_location=self.device)) if cuda != None: self.model.cuda(self.device) self.converter = utils.strLabelConverter(alphabet, ignore_case=False) self.model.eval()
def crnnSource(): alphabet = keys_crnn.alphabet converter = util.strLabelConverter(alphabet) if torch.cuda.is_available() and GPU: print("GPU") model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda() else: print("CPU") model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cpu() print(os.getcwd()) path = '../ModelSet/Chinese-OCR/crnn/samples/model_acc97.pth' # 模型识别种类:5530 # path = '../ModelSet/Chinese-OCR/crnn/samples/mixed_second_finetune_acc97p7.pth' # 对比效果:6736,需要同步修改keys_crnn.py model.eval() # 带参数:不使用GPU model.load_state_dict(torch.load(path, map_location='cpu')) return model, converter
def get_recognition_model(): model_path = './weights/crnn.pth' #Intial model model = crnn.CRNN(32, 1, 37, 256) try: model = model.cuda() except: #there is some issue with model loading, #and loading it twice solves the problem model = model.cuda() print('loading pretrained model from %s' % model_path) model.load_state_dict(torch.load(model_path)) return model
def load_model(model_path): # net init global transformer, model, converter print('loading pretrained model from %s' % model_path) nclass = len(params.alphabet) + 1 model = crnn.CRNN(params.imgH, params.nc, nclass, params.nh) if torch.cuda.is_available(): model = model.cuda() model = torch.nn.DataParallel(model) else: model.load_state_dict(torch.load(model_path, map_location='cpu')) model.eval() converter = utils.strLabelConverter(params.alphabet) transformer = dataset.resizeNormalize((100, 32))
def __init__(self, imgp, labp, bs): self.crnn_net = None nclass = len(params.alphabet) + 1 self.crnn_net = crnn.CRNN(32, 1, nclass, 256, n_rnn=2, leakyRelu=False) crnn_model_path = './crnn.pth' self.crnn_handle = CRNNHandle.CRNNHandle(crnn_model_path, self.crnn_net, gpu_id=1.1) self.imgp = imgp self.labp = labp self.bs = bs
def crnnloader(): #model_path = '/home/ahmed/crnn/data/crnn.pth' model_path = './CRNN/data/crnn.pth' #img_path = '/home/ahmed/crnn/data/demo.png' #img_path='/home/ahmed/Pictures/cogedis/2-total.png' img_path = './CRNN/data/demo.png' alphabet = '0123456789abcdefghijklmnopqrstuvwxyz' model = crnn.CRNN(32, 1, 37, 256, 1) print('loading pretrained model from %s' % model_path) model.load_state_dict(torch.load(model_path)) converter = utils.strLabelConverter(alphabet) return model, converter
def net_init(): nclass = len(params.alphabet) + 1 crnn = net.CRNN(params.imgH, params.nc, nclass, params.nh) crnn.apply(weights_init) if params.pretrained != '': print('loading pretrained model from %s' % params.pretrained) if params.multi_gpu: crnn = torch.nn.DataParallel(crnn) std = torch.load(params.pretrained) # # Remove the last FC layer std.popitem(last=True) std.popitem(last=True) crnn.load_state_dict(std, strict=False) return crnn
def load_images_to_predict(): print('load_images_to_predict') # load model model_path = './save/netCRNN_225_1250.pth' alphabet = '0123456789,.:(%$!^&-/);<~|`>?+=_[]{}"\'@#*abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ\ ' imgH = 32 # should be 32 nclass = len(alphabet) + 1 nhiddenstate = 256 model = crnn.CRNN(imgH, 1, nclass, nhiddenstate) if torch.cuda.is_available(): model = model.cuda() print('loading pretrained model from %s' % model_path) model.load_state_dict({ k.replace('module.', ''): v for k, v in torch.load(model_path, map_location='cpu').items() }) # load image main_path = os.path.abspath(os.path.join(os.getcwd(), "..")) image_path = os.path.join(main_path, 'result/step2/image2/*.jpg') label_path = os.path.join(main_path, 'result/step2/label/') image_list = [os.path.splitext(f)[0] for f in glob.glob(image_path)] jpg_files = [name + ".jpg" for name in image_list] print('total images: ', len(jpg_files)) count = 1 for jpg in jpg_files: image = Image.open(jpg).convert('L') words_list = [] label = label_path + jpg.split('/')[-1][:-3] + 'txt' txt = sort_lines(label) for line in txt: box = line.split(',') crop_image = image.crop( (int(box[0]), int(box[1]), int(box[4]), int(box[5]))) words = predict_this_box(crop_image, model, alphabet) words_list.append(words) result_path = os.path.join(main_path, 'result/step3/') save_path = result_path + jpg.split('/')[-1][:-3] + 'txt' with open(save_path, 'w+') as result: for line in words_list: result.writelines(line + '\n') if count % 1000 == 0: print(str(count), 'images finished.', 'total images: ', len(jpg_files)) count += 1
def __init__(self, model_path): self.model_path = model_path nclass = len(params.alphabet) + 1 self.model = crnn.CRNN(params.imgH, params.nc, nclass, params.nh) if torch.cuda.is_available(): self.model = self.model.cuda() if params.multi_gpu: self.model = torch.nn.DataParallel(self.model) self.converter = utils.strLabelConverter(params.alphabet) self.transformer = dataset.resizeNormalize((100, 32)) print('loading pretrained model from %s' % model_path) self.model.load_state_dict(torch.load(model_path))