Create prediction.py

prediction.py

@@ -0,0 +1,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