Upload train.py

train.py

@@ -0,0 +1,131 @@
+# display image with masks and bounding boxes
+from os import listdir
+from xml.etree import ElementTree
+import json
+from numpy import zeros
+from numpy import asarray
+from bboxcnn.utils import Dataset
+from bboxcnn.config import Config
+from bboxcnn.model import BBoxCNN
+
+class PASSPORT_Dataset(Dataset):
+    # load the dataset definitions
+    def load_dataset(self, dataset_dir, is_train=True):
+        # define one class
+        
+        self.add_class("dataset", 1, "Country Name")
+        self.add_class("dataset", 2, "Document Type")
+        self.add_class("dataset", 3, "Country Code")
+        self.add_class("dataset", 4, "Passport Number")
+        self.add_class("dataset", 5, "Surname")
+        self.add_class("dataset", 6, "Given Name")
+        self.add_class("dataset", 7, "Nationality")
+        self.add_class("dataset", 8, "Sex")
+        self.add_class("dataset", 9, "DOB")
+        self.add_class("dataset", 10, "Place Of Birth")
+        self.add_class("dataset", 11, "Place Of Issue")
+        self.add_class("dataset", 12, "DOI")
+        self.add_class("dataset", 13, "DOE")
+        self.add_class("dataset", 14, "MRZ")
+        self.add_class("dataset", 15, "Name Of Father")
+        self.add_class("dataset", 16, "Name Of Mother")
+        self.add_class("dataset", 17, "Name Of Spouse")
+        self.add_class("dataset", 18, "Address")
+        self.add_class("dataset", 19, "Old Passport Information")
+        self.add_class("dataset", 20, "File Number")
+        # define data locations
+        images_dir = dataset_dir + '/images/'
+        annotations_dir = dataset_dir + '/annots/'
+        # find all images
+        for filename in listdir(images_dir):
+            # extract image id
+            image_id = filename[:-4]
+            # skip bad images
+            if image_id in ['017']:
+                continue
+            # skip all images after 150 if we are building the train set
+            if is_train and int(image_id) >= 79:
+                continue
+            # skip all images before 150 if we are building the test/val set
+            if not is_train and int(image_id) < 79:
+                continue
+            img_path = images_dir + filename
+            ann_path = annotations_dir + image_id + '.json'
+            # add to dataset
+            self.add_image('dataset', image_id=image_id, path=img_path, annotation=ann_path)
+
+    # extract bounding boxes from an annotation file
+    def extract_boxes(self, filename):
+        # load and parse the file
+        with open(filename, 'r') as f:
+            data = json.load(f)
+        boxes = list()
+        bndboxes = [i['bndbox'] for i in data['object']]
+        class_names = [i['name'] for i in data['object']]
+        for box in bndboxes:
+            xmin = int(box['xmin'])
+            ymin = int(box['ymin'])
+            xmax = int(box['xmax'])
+            ymax = int(box['ymax'])
+            coors = [xmin, ymin, xmax, ymax]
+            boxes.append(coors)
+        # extract image dimensions
+        width = int(data['size']['width'])
+        height = int(data['size']['height'])
+        return boxes, class_names, width, height
+
+    # load the masks for an image
+    def load_mask(self, image_id):
+        # get details of image
+        info = self.image_info[image_id]
+        # define box file location
+        path = info['annotation']
+        # load XML
+        boxes, class_names, w, h = self.extract_boxes(path)
+        # create one array for all masks, each on a different channel
+        masks = zeros([h, w, len(boxes)], dtype='uint8')
+        # create masks
+        class_ids = list()
+        for i, entity in enumerate(zip(boxes, class_names)):
+            box, class_name = entity
+            row_s, row_e = box[1], box[3]
+            col_s, col_e = box[0], box[2]
+            masks[row_s:row_e, col_s:col_e, i] = i+1
+            class_ids.append(self.class_names.index(class_name))
+        return masks, asarray(class_ids, dtype='int32')
+
+    # load an image reference
+    def image_reference(self, image_id):
+        info = self.image_info[image_id]
+        return info['path']
+
+# define a configuration for the model
+class PASSPORT_Config(Config):
+    # define the name of the configuration
+    NAME = "passport_cfg"
+    # number of classes (background + Object Classes)
+    NUM_CLASSES = 1 + 20
+    # number of training steps per epoch
+    STEPS_PER_EPOCH = 81
+    
+
+# train set
+train_set = PASSPORT_Dataset()
+train_set.load_dataset('passport_data', is_train=True)
+train_set.prepare()
+print('Train: %d' % len(train_set.image_ids))
+# prepare test/val set
+test_set = PASSPORT_Dataset()
+test_set.load_dataset('passport_data', is_train=False)
+test_set.prepare()
+print('Test: %d' % len(test_set.image_ids))
+# prepare config
+config = PASSPORT_Config()
+config.display()
+# define the model
+model = BBoxCNN(mode='training', model_dir='./', config=config)
+# load weights (mscoco) and exclude the output layers
+# model.load_weights('bboxcnn_base.h5', by_name=True, exclude=["bboxcnn_class_logits", "bboxcnn_bbox_fc",  "bboxcnn_bbox", "bboxcnn_mask"])
+model.load_weights('passport_cfg20220520T2226/bboxcnn_passport_cfg_0090.h5', by_name=True, exclude=["bboxcnn_class_logits", "bboxcnn_bbox_fc",  "bboxcnn_bbox", "bboxcnn_mask"])
+# train weights (output layers or 'heads')
+model.train(train_set, test_set, learning_rate=config.LEARNING_RATE, epochs=90, layers='heads')