File size: 2,399 Bytes
11f89a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
# display image with masks and bounding boxes
from os import listdir
import json
import cv2
import easyocr
reader = easyocr.Reader(['en'])
# fit a bounding box cnn on the EID dataset
from numpy import zeros
from numpy import asarray
from numpy import expand_dims
from matplotlib import pyplot
from matplotlib.patches import Rectangle
from bboxcnn.config import Config
from bboxcnn.model import BBoxCNN
from bboxcnn.model import mold_image
from bboxcnn.utils import Dataset
class PredictionConfig(Config):
# define the name of the configuration
NAME = "eid_cfg"
# number of classes (background + EID Field classes(10))
NUM_CLASSES = 1 + 10
# simplify GPU config(Here the GPU and CPU Config are same: It works if Dex sys doesnt have GPU)
GPU_COUNT = 1
IMAGES_PER_GPU = 1
# create config
cfg = PredictionConfig()
# define the model
model = BBoxCNN(mode='inference', model_dir='./', config=cfg)
# load model weights
model_path = 'model/bboxcnn_eid_cfg_0033.h5'
model.load_weights(model_path, by_name=True)
class_ids_to_class_name = {1: 'Sex', 2: 'DOB', 3: 'Country',
4: 'DOE', 5: 'Card Number', 6: 'Document Type',
7: 'Id Number', 8: 'MRZ', 9: 'Name', 10: 'Nationality'}
def load_image(path):
source_image = cv2.imread(path, cv2.IMREAD_COLOR)
scaled_image = mold_image(source_image, cfg)
# convert image into one sample
sample = expand_dims(scaled_image, 0)
return source_image, sample
def extract_text(image_path):
# piece images using bboxes
dict_ = {}
class_ids_to_class_name = {1: 'Sex', 2: 'DOB', 3: 'Country',
4: 'DOE', 5: 'Card Number', 6: 'Document Type',
7: 'Id Number', 8: 'MRZ', 9: 'Name', 10: 'Nationality'}
source_image, scaled_samp_image = load_image(image_path)
yhat = model.detect(scaled_samp_image, verbose=0)[0]
bboxes, class_ids, mask = yhat['rois'], yhat['class_ids'], yhat['masks']
for bbox, class_id in zip(bboxes, class_ids):
xmin, ymin, xmax, ymax = bbox
piece_image = source_image[xmin:xmax, ymin:ymax]
classname = class_ids_to_class_name[class_id]
text = reader.readtext(piece_image, detail=0)
dict_.update({class_ids_to_class_name[class_id]: text})
return dict_
def process(im_path):
eid_data = extract_text(im_path)
return eid_data |