diff --git a/mrcnn/model.py b/mrcnn/model.py index 62cb2b0951..7adcd8f6e2 100644 --- a/mrcnn/model.py +++ b/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)