class Classifier_Vietocr:
def __init__(self, ckpt_path=None, gpu='0'):
print('Classifier_Vietocr. Init')
self.config = Cfg.load_config(cls_base_config_path, cls_config_path)
if ckpt_path is not None:
self.config['weights'] = ckpt_path
self.config['cnn']['pretrained'] = False
if gpu is not None:
self.config['device'] = 'cuda:' + str(gpu)
else:
self.config['device'] = 'cpu'
self.config['predictor']['beamsearch'] = False
self.model = Predictor(self.config)
def inference(self, numpy_list, debug=False):
print('Classifier_Vietocr. Inference', len(numpy_list), 'boxes')
text_values = []
prob_value = []
for idx, f in enumerate(numpy_list):
img = Image.fromarray(f)
s, prob = self.model.predict(img, True)
if debug:
print(round(prob, 3), s)
cv2.imshow('sample', f)
cv2.waitKey(0)
text_values.append(s)
prob_value.append(prob)
return text_values, prob_value
class Classifier_Vietocr:
def __init__(self,
ckpt_path=None,
gpu='0',
config_name='vgg_seq2seq',
write_file=False,
debug=False):
print('Classifier_Vietocr. Init')
self.config = Cfg.load_config_from_name(config_name)
# config['weights'] = './weights/transformerocr.pth'
if ckpt_path is not None:
self.config['weights'] = ckpt_path
self.config['cnn']['pretrained'] = False
if gpu is not None:
self.config['device'] = 'cuda:' + str(gpu)
self.config['predictor']['beamsearch'] = False
self.model = Predictor(self.config)
def inference(self, numpy_list, debug=False):
print('Classifier_Vietocr. Inference', len(numpy_list), 'boxes')
text_values = []
prob_value = []
# t = tqdm(iter(val_loader), total=len(val_loader), desc='Classifier_CRNN. Inference...')
for idx, f in enumerate(numpy_list):
img = Image.fromarray(f)
s, prob = self.model.predict(img, True)
if debug:
print(round(prob, 3), s)
cv2.imshow('sample', f)
cv2.waitKey(0)
text_values.append(s)
prob_value.append(prob)
return text_values, prob_value
class TEXT_IMAGES(object):
def __init__(self, cmnd_detect_config_path='./center/config/cmnd.yml', line_detect_weight_path='weights/line_detect_weight.pth', reg_model='vgg_seq2seq', ocr_weight_path='weights/vgg-seq2seq.pth'):
print("Loading TEXT_MODEL...")
cmnd_detect_config = Cfg.load_config_from_file(cmnd_detect_config_path)
self.cmnd_detect_module = CENTER_MODEL(cmnd_detect_config)
self.line_detect_module = LineDetection(line_detect_weight_path)
config = Cfg_reg.load_config_from_name(reg_model)
config['weights'] = ocr_weight_path
config['device'] = 'cpu'
config['predictor']['beamsearch'] = False
self.recognition_text_module = Predictor(config)
def get_content_image(self, image, show_line=False):
# cv image
# return image_drawed, texts, boxes
img_detected, have_cmnd = self.cmnd_detect_module.detect_obj(image)
if not have_cmnd:
print("Không phát hiện CMND!!!")
return None, None
result_line_img, img_draw_box = self.line_detect_module.predict_box(img_detected, show_line)
result_ocr = {}
for key, values in result_line_img.items():
label = key
imgs = values
result_ocr[label] = []
for img in imgs:
res_str = self.recognition_text_module.predict(img)
result_ocr[label].append(res_str)
print(result_ocr)
return result_ocr, img_draw_box
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config',
type=str,
default='./logs/hw_word_seq2seq/config.yml')
parser.add_argument('--weight',
type=str,
default='./logs/hw_word_seq2seq/best.pt')
parser.add_argument('--img', type=str, default=None, required=True)
args = parser.parse_args()
config = Cfg.load_config_from_file(args.config, download_base=False)
config['weights'] = args.weight
print(config.pretty_text())
detector = Predictor(config)
if os.path.isdir(args.img):
img_paths = os.listdir(args.img)
for img_path in img_paths:
try:
img = Image.open(args.img + '/' + img_path)
except:
continue
t1 = time.time()
s, prob = detector.predict(img, return_prob=True)
print('Text in {} is:\t {} | prob: {:.2f} | times: {:.2f}'.format(
img_path, s, prob,
time.time() - t1))
else:
t1 = time.time()
img = Image.open(args.img)
s, prob = detector.predict(img, return_prob=True)
print('Text in {} is:\t {} | prob: {:.2f} | times: {:.2f}'.format(
args.img, s, prob,
time.time() - t1))
class OCR():
def __init__(self):
config = Cfg.load_config_from_name('vgg_transformer')
config['weights'] = './model/transformerocr.pth'
# config['weights'] = 'https://drive.google.com/uc?id=13327Y1tz1ohsm5YZMyXVMPIOjoOA0OaA'
# config['device'] = ''
config['device'] = 'cuda'
config['predictor']['beamsearch'] = False
self.detector = Predictor(config)
def recognize(self, img):
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = Image.fromarray(img)
return self.detector.predict(img)
def create_text_annotation_ocr(imgs, dest):
config = Cfg.load_config_from_name('vgg_transformer')
config['export'] = 'transformerocr_checkpoint.pth'
config['device'] = 'cuda'
config['predictor']['beamsearch'] = False
detector = Predictor(config)
f = io.open(os.path.join(dest, "annotation.txt"), "a", encoding="utf-8")
for idx, image in enumerate(imgs):
text = detector.predict(image)
if idx + 1 == len(imgs):
f.write('crop_img/{:06d}.jpg\t{}'.format(idx + 1, text))
else:
f.write('crop_img/{:06d}.jpg\t{}\n'.format(idx+1, text))
f.close()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--img', required=True, help='foo help')
parser.add_argument('--config', required=True, help='foo help')
args = parser.parse_args()
config = Cfg.load_config_from_file(args.config)
detector = Predictor(config)
img = Image.open(args.img)
s = detector.predict(img)
print(s)
class ocr:
def __init__(
self,
config,
) -> None:
super(ocr, self).__init__()
self.config = config
config_base = Cfg.load_config_from_file("config/base.yml")
config = Cfg.load_config_from_file(self.config)
config_base.update(config)
config = config_base
config['vocab'] = character
self.text_r = Predictor(config)
def run(self, im: Image):
s = self.text_r.predict(im)
return index_decode(s)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--img', required=True, help='foo help')
parser.add_argument('--config', required=True, help='foo help')
args = parser.parse_args()
config_base = Cfg.load_config_from_file("config/base.yml")
config = Cfg.load_config_from_file(args.config)
config_base.update(config)
config = config_base
config['vocab'] = character
detector = Predictor(config)
img = Image.open(args.img)
s = detector.predict(img)
print(s)
class PredictorImage(object):
def __init__(self):
self.yolo = YOLOv4()
self.yolo.classes = './coco.names'
self.yolo.make_model()
self.yolo.load_weights("./model/yolov4-custom_last.weights",
weights_type="yolo")
self.config = Cfg.load_config()
self.config['weights'] = './model/transformerocr.pth'
self.config['predictor']['beamsearch'] = False
self.config['device'] = 'cpu'
self.detector = Predictor(self.config)
self.classes = ['id', 'name', 'dmy', 'add1', 'add2']
self.res = dict.fromkeys(self.classes, '')
# self.address_correction = AddressCorrection()
def predict(self, img):
return self.yolo.predict(img)
def inference(self, img):
res = self.predict(img)
img_h, img_w = img.shape[:2]
for r in res:
x_center_p, y_center_p, w_p, h_p, cl, pro = r[0], r[1], r[2], r[
3], r[4], r[5]
w = int(img_w * w_p)
h = int(img_h * h_p)
x = int(x_center_p * img_w - w / 2)
y = int(y_center_p * img_h - h / 2)
resize_img = img[y:y + h, x:x + w, :]
self.predict_ocr(resize_img, self.classes[int(cl)])
cv2.rectangle(img, (int(x), int(y)), (int(x + w), int(y + h)),
(0, 255, 0), 1)
print(self.res)
cv2.imshow('predict', img)
cv2.waitKey(0)
def predict_ocr(self, img, label):
img = Image.fromarray(img)
s = self.detector.predict(img)
self.res[label] += s + ' '
def img_to_text(list_img):
results = []
for img in list_img:
# sử dụng config mặc định của mô hình
config = Cfg.load_config_from_name('vgg_transformer')
# đường dẫn đến trọng số đã huấn luyện hoặc comment để sử dụng #pretrained model mặc định
config['weights'] = 'checkpoints/transformerocr.pth'
config['device'] = 'cpu' # device chạy 'cuda:0', 'cuda:1', 'cpu'
detector = Predictor(config)
img = Image.fromarray(img.astype(np.uint8))
# img = Image.fromarray((img * 255).astype(np.uint8))
# img.show()
# dự đoán
# muốn trả về xác suất của câu dự đoán thì đổi return_prob=True
text = detector.predict(img)
if len(text) > 0:
results.append(text)
return results
def img_to_text(list_img):
results = []
config = Cfg.load_config_from_name("vgg_transformer")
# đường dẫn đến trọng số đã huấn luyện hoặc comment để sử dụng #pretrained model mặc định
config[
"weights"] = "https://drive.google.com/uc?id=13327Y1tz1ohsm5YZMyXVMPIOjoOA0OaA"
# config['weights'] = 'transformerocr.pth'
config["device"] = "cpu" # device chạy 'cuda:0', 'cuda:1', 'cpu'
detector = Predictor(config)
for i in range(len(list_img)):
if i == 0:
continue
# sử dụng config mặc định của mô hình
img = Image.fromarray(list_img[i].astype(np.uint8))
# dự đoán
# muốn trả về xác suất của câu dự đoán thì đổi return_prob=True
text = detector.predict(img)
if len(text) > 0:
results.append(text)
return results
from PIL import Image
import matplotlib.pyplot as plt
import time
import glob
config = Cfg.load_config_from_name('vgg_transformer')
config[
'weights'] = 'https://drive.google.com/uc?id=13327Y1tz1ohsm5YZMyXVMPIOjoOA0OaA'
config['device'] = 'cuda:0'
config['predictor']['beamsearch'] = False
detector = Predictor(config)
# img = "image/20.JPG"
# img = cv2.imread(img)
# pil_img = Image.fromarray(img)
# start = time.time()
# s = detector.predict(pil_img)
# print(s, "predicted: ", time.time()-start)
# cv2.imshow("image", img)
# cv2.waitKey(0)
for f in glob.glob("image/*.JPG"):
img = cv2.imread(f)
pil_img = Image.fromarray(img)
start = time.time()
s = detector.predict(pil_img)
print(s, "predicted: ", time.time() - start)
cv2.imshow("image", img)
cv2.waitKey(0)
from vietocr.tool.config import Cfg
from vietocr.tool.predictor import Predictor
import cv2
config = Cfg.load_config_from_name('vgg_transformer')
# load pretrained weight
config['weights'] = './transformerocr.pth'
# set device to use cpu
config['device'] = 'cpu'
config['cnn']['pretrained'] = False
config['predictor']['beamsearch'] = False
detector = Predictor(config)
img = cv2.imread('img_check.png')
result = detector.predict(img)
print(result)
from PIL import Image
from vietocr.tool.predictor import Predictor
from vietocr.tool.config import Cfg
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path
import cv2
import time
config = Cfg.load_config_from_name('vgg_transformer')
config['export'] = 'transformerocr_checkpoint.pth'
config['device'] = 'cpu'
config['predictor']['beamsearch'] = False
start1 = time.time()
detector = Predictor(config)
end1 = time.time()
img = Image.open('./ANH_1321.jpeg')
print("Load image: ", end1 - start1)
start = time.time()
print(detector.predict(img))
end = time.time()
print('Required time: ', end - start)
cv2.imshow('image', np.array(img))
cv2.waitKey(0)
else:
list_files = [img_path]
total_cer = 0
total_inference_time = 0
print('Total files:', len(list_files))
for idx, f in enumerate(list_files):
img_path = os.path.join(src_dir, f)
label_path = img_path.replace('.jpg', '.txt').replace('.png', '.txt').replace('.PNG', '.txt').replace('.JPG',
'.txt')
with open(label_path, 'r', encoding='utf-8') as f:
label = f.readline()
img = Image.open(img_path)
begin = time.time()
s = detector.predict(img)
end = time.time()
cer = cer_loss_one_image(label, s)
total_cer += cer
total_inference_time += (end - begin)
if debug:
print('Predict:', s)
img_cv = cv2.imread(img_path)
cv2.imshow('img', img_cv)
cv2.waitKey(0)
else:
print(idx, 'pred:', s, ', gt:', label, ', cer:', round(cer, 2), ', time:', round(end - begin, 4))
print('avg cer: ', total_cer / len(list_files))
print('avg infer time: ', total_inference_time / len(list_files), ', fps:', len(list_files) / total_inference_time)
class Reader:
def __init__(self, weights):
# list tinh & tp co dau
self.tinh_list = [
'An Giang', 'Bà Rịa - Vũng Tàu', 'Bắc Giang', 'Bắc Kạn',
'Bạc Liêu', 'Bắc Ninh', 'Bến Tre', 'Bình Định', 'Bình Dương',
'Bình Phước', 'Bình Thuận', 'Cà Mau', 'Cao Bằng', 'Đắk Lắk',
'Đắk Nông', 'Điện Biên', 'Đồng Nai', 'Đồng Tháp', 'Gia Lai',
'Hà Giang', 'Hà Nam', 'Hà Tĩnh', 'Hải Dương', 'Hậu Giang',
'Hòa Bình', 'Hưng Yên', 'Khánh Hòa', 'Kiên Giang', 'Kon Tum',
'Lai Châu', 'Lâm Đồng', 'Lạng Sơn', 'Lào Cai', 'Long An',
'Nam Định', 'Nghệ An', 'Ninh Bình', 'Ninh Thuận', 'Phú Thọ',
'Quảng Bình', 'Quảng Nam', 'Quảng Ngãi', 'Quảng Ninh', 'Quảng Trị',
'Sóc Trăng', 'Sơn La', 'Tây Ninh', 'Thái Bình', 'Thái Nguyên',
'Thanh Hóa', 'Thừa Thiên Huế', 'Tiền Giang', 'Trà Vinh',
'Tuyên Quang', 'Vĩnh Long', 'Vĩnh Phúc', 'Yên Bái', 'Phú Yên',
'Cần Thơ', 'Đà Nẵng', 'Hải Phòng', 'Hà Nội', 'TP Hồ Chí Minh'
]
# list tinh & tp khong co dau
self.provinces = [
self.remove_accent(tinh).lower() for tinh in self.tinh_list
]
self.config = Cfg.load_config_from_name('vgg_transformer')
self.config['weights'] = weights
self.config['cnn']['pretrained'] = False
self.config['device'] = 'cpu'
self.config['predictor']['beamsearch'] = False
self.reader = Predictor(self.config)
def read(self, image):
"""
Recognise text from image
:param image: ndarray of image
:return: text
"""
text = self.reader.predict(image)
return text
def remove_accent(self, text):
return unidecode.unidecode(text)
def postprocess_address(self, original_text, thresold):
# preprocess text
text = self.remove_accent(original_text)
text = text.lower()
# calculate editance between text with each of address in provinces list
edits = [
levenshtein_distance(text, address) for address in self.provinces
]
edits = np.array(edits)
arg_min = np.argmin(edits)
if edits[arg_min] < thresold:
return self.tinh_list[arg_min]
else:
return original_text
