Update model.py

mrcnn/model.py

@@ -2866,3 +2866,68 @@ def denorm_boxes_graph(boxes, shape):
     scale = tf.concat([h, w, h, w], axis=-1) - tf.constant(1.0)
     shift = tf.constant([0., 0., 1., 1.])
     return tf.cast(tf.round(tf.multiply(boxes, scale) + shift), tf.int32)
+
+############################################################
+#  Custom Callbacks
+############################################################
+
+class MeanAveragePrecisionCallback(Callback):
+    def __init__(self, train_model: MaskRCNN, inference_model: MaskRCNN, dataset: Dataset,
+                 calculate_map_at_every_X_epoch=5, dataset_limit=None,
+                 verbose=1):
+        super().__init__()
+        self.train_model = train_model
+        self.inference_model = inference_model
+        self.dataset = dataset
+        self.calculate_map_at_every_X_epoch = calculate_map_at_every_X_epoch
+        self.dataset_limit = len(self.dataset.image_ids)
+        if dataset_limit is not None:
+            self.dataset_limit = dataset_limit
+        self.dataset_image_ids = self.dataset.image_ids.copy()
+
+        if inference_model.config.BATCH_SIZE != 1:
+            raise ValueError("This callback only works with the bacth size of 1")
+
+        self._verbose_print = print if verbose > 0 else lambda *a, **k: None
+
+    def on_epoch_end(self, epoch, logs=None):
+
+        if epoch > 2 and (epoch+1)%self.calculate_map_at_every_X_epoch == 0:
+            self._verbose_print("Calculating mAP...")
+            self._load_weights_for_model()
+
+            mAPs = self._calculate_mean_average_precision()
+            mAP = np.mean(mAPs)
+
+            if logs is not None:
+                logs["val_mean_average_precision"] = mAP
+
+            self._verbose_print("mAP at epoch {0} is: {1}".format(epoch+1, mAP))
+
+        super().on_epoch_end(epoch, logs)
+
+    def _load_weights_for_model(self):
+        last_weights_path = self.train_model.find_last()
+        self._verbose_print("Loaded weights for the inference model (last checkpoint of the train model): {0}".format(
+            last_weights_path))
+        self.inference_model.load_weights(last_weights_path,
+                                          by_name=True)
+
+    def _calculate_mean_average_precision(self):
+        mAPs = []
+
+        # Use a random subset of the data when a limit is defined
+        np.random.shuffle(self.dataset_image_ids)
+
+        for image_id in self.dataset_image_ids[:self.dataset_limit]:
+            image, image_meta, gt_class_id, gt_bbox, gt_mask = load_image_gt(self.dataset, self.inference_model.config,
+                                                                             image_id, use_mini_mask=False)
+            molded_images = np.expand_dims(mold_image(image, self.inference_model.config), 0)
+            results = self.inference_model.detect(molded_images, verbose=0)
+            r = results[0]
+            # Compute mAP - VOC uses IoU 0.5
+            AP, _, _, _ = utils.compute_ap(gt_bbox, gt_class_id, gt_mask, r["rois"],
+                                           r["class_ids"], r["scores"], r['masks'])
+            mAPs.append(AP)
+
+        return np.array(mAPs)