import os.path as path
from copy import deepcopy
from typing import Optional
from functools import partial
import cv2
import numpy as np
import emloop as el
from ..datasets import RectangleRPNDataset
from ..utils.rpn import decode_rpn_predictions
[docs]class ComputeRegionProposals(el.AbstractHook):
[docs] def __init__(self,
dataset: RectangleRPNDataset,
non_maximum_suppression: Optional[float]=0.3,
min_probability: Optional[float]=None,
top_k: Optional[int]=None,
classifier_probabilities_name: str='classifier_probabilities',
regression_predictions_name: str='regression_predictions',
proposed_regions_name: str='rectangle_proposals',
**kwargs):
"""
Create new ``ComputeRectangleProposal`` hook.
:param dataset: RPN dataset which created the anchors
:param non_maximum_suppression: if specified, use nms with the given threshold
:param min_probability: if specified, filter out regions with probability lower than the specified threshold
:param top_k: if specified, output only top k regions
:param classifier_probabilities_name: name of the classifier probabilities variable
:param regression_predictions_name: name of the regression predictions variable
:param proposed_regions_name: name of the output proposed regions variable
"""
self._classifier_probabilities_name = classifier_probabilities_name
self._regression_predictions_name = regression_predictions_name
self._proposed_regions_name = proposed_regions_name
self._decode_fn = partial(decode_rpn_predictions,
dataset=dataset,
non_maximum_suppression=non_maximum_suppression,
min_probability=min_probability,
top_k=top_k)
super().__init__(**kwargs)
[docs] def after_batch(self, stream_name: str, batch_data: el.Batch) -> None:
"""Decode RPN predictions and save the proposed regions to batch data."""
regions = []
for classifier_probabilities, regression_predictions in \
zip(batch_data[self._classifier_probabilities_name], batch_data[self._regression_predictions_name]):
regions.append(self._decode_fn(classifier_probabilities=classifier_probabilities,
regression_predictions=regression_predictions))
batch_data[self._proposed_regions_name] = regions
[docs]class VisualizeRectangleProposals(el.AbstractHook):
"""Visualize proposed rectangle regions."""
[docs] def __init__(self, output_dir: str,
images_name: str='images',
rectangles_name: str='rectangle_proposals',
expected_rectangles_name: str='target_regions',
errors_only: bool=False,
**kwargs):
"""
Create new ``VisualizeRectangleProposals`` hook.
:param output_dir: output directory
:param images_name: images variable name
:param rectangles_name: rectangles variable name
:param expected_rectangles_name: expected rectangles variable name
:param errors_only: visualize only images with recall < 1.0
"""
self._output_dir = output_dir
self._images_name = images_name
self._rectangles_name = rectangles_name
self._expected_rectangles_name = expected_rectangles_name
self._errors_only = errors_only
self._tot = 0
super().__init__(**kwargs)
[docs] def after_batch(self, stream_name: str, batch_data: el.Batch) -> None:
"""Visualize rectangle proposals and save the results to the ``output_dir``."""
for i, (image, proposals, targets) in enumerate(zip(batch_data[self._images_name],
batch_data[self._rectangles_name],
batch_data[self._expected_rectangles_name])):
if self._errors_only and batch_data['region_recall'][i] == 1.:
continue
canvas = deepcopy(image)
# draw targets in blue, proposals in green
for color, rectangles in (((255, 0, 0), targets), ((0, 255, 0), proposals)):
for rectangle in rectangles:
rectangle = np.int32(rectangle)
cv2.rectangle(canvas, (rectangle[0] - rectangle[2] // 2, rectangle[1] - rectangle[3] // 2),
(rectangle[0] + rectangle[2] // 2, rectangle[1] + rectangle[3] // 2), color, 1)
self._tot += 1
cv2.imwrite(path.join(self._output_dir, 'visual_{}_{}.jpg'.format(stream_name, self._tot)), canvas)
[docs]class ComputeRegionStats(el.AbstractHook):
"""Compute rectangle RPN stats (recall and overlap ratio)."""
[docs] def __init__(self,
dataset: RectangleRPNDataset,
expected_rectangles_name: str='target_regions',
predicted_rectangles_name: str='rectangle_proposals',
**kwargs):
"""
Create new ``ComputeRectangleRPNStats`` hook.
:param dataset: RectangleRPNDataset which created the anchors
:param expected_rectangles_name: expected rectangles variable name
:param predicted_rectangles_name: predicted rectangles variable name
"""
self._dataset = dataset
self._expected_rectangles_name = expected_rectangles_name
self._predicted_rectangles_name = predicted_rectangles_name
super().__init__(**kwargs)
[docs] def after_batch(self, stream_name: str, batch_data: el.Batch) -> None:
"""Compute recall and overlap ratio for the expected and predicted rectangles."""
overall_recalls = []
overall_overlaps = []
for expected_rectangles, predicted_rectangles in zip(batch_data[self._expected_rectangles_name],
batch_data[self._predicted_rectangles_name]):
overlaps = []
for expected_rectangle in expected_rectangles:
for predicted_rectangle in predicted_rectangles:
overlap = self._dataset.anchor_region_overlap(expected_rectangle, predicted_rectangle)
if overlap > 0.5:
overlaps.append(overlap)
break
overall_recalls.append(len(overlaps)/len(expected_rectangles))
overall_overlaps.append(np.mean(overlaps))
batch_data['region_recall'] = overall_recalls
batch_data['region_overlap'] = overall_overlaps