[POT] Update face detection sample (#10471)

* support cascade model for sw api

* update mtcnnengine

* delete empty line
This commit is contained in:
Indira Salyahova
2022-02-22 17:46:08 +03:00
committed by GitHub
parent dab1a34aa2
commit e2df6d149b

View File

@@ -6,12 +6,18 @@ from argparse import ArgumentParser
from functools import partial
from time import time
import copy
import cv2
import numpy as np
from openvino.runtime import PartialShape # pylint: disable=E0611,E0401
from openvino.tools.pot import Metric, DataLoader, IEEngine, \
load_model, compress_model_weights, create_pipeline
from openvino.tools.pot.graph.model_utils import add_outputs
from openvino.tools.pot.samplers.batch_sampler import BatchSampler
from openvino.tools.pot.engines.utils import process_accumulated_stats, \
restore_original_node_names, align_stat_names_with_results, \
add_tensor_names, collect_model_outputs
from openvino.tools.pot.utils.logger import init_logger, get_logger
from openvino.tools.pot.api.samples.face_detection import utils
@@ -28,7 +34,7 @@ class WiderFaceLoader(DataLoader):
def __init__(self, config):
super().__init__(config)
self._min_height_ann = 60
self._img_ids, self._annotations = self._read_image_ids_annotations(config.annotation_file)
self._img_ids, self._annotations = self._read_image_ids_annotations(self.config.annotation_file)
def __getitem__(self, index):
"""
@@ -98,14 +104,68 @@ class MTCNNEngine(IEEngine):
self._model = self._set_model(model)
self._output_layers = {}
stage_names = ['pnet', 'rnet', 'onet']
for stage, model_dict in enumerate(model.models):
for stage, _ in enumerate(model.models):
self._output_layers[stage_names[stage]] = {
'probabilities': model_dict['name'] + '_' + self.config['outputs']['probabilities'][stage],
'regions': model_dict['name'] + '_' + self.config['outputs']['regions'][stage],
'probabilities': self.config['outputs']['probabilities'][stage],
'regions': self.config['outputs']['regions'][stage],
}
def _add_outputs(self, nodes_name):
add_outputs(self._model, nodes_name)
return add_outputs(self._model, nodes_name)
def predict(self, stats_layout=None, sampler=None, stat_aliases=None,
metric_per_sample=False, print_progress=False):
stat_names_aliases = None
if sampler is None:
sampler = BatchSampler(self.data_loader)
if stats_layout:
model_with_stat_op, nodes_names_map, output_to_node_names = self._statistic_graph_builder. \
insert_statistic(copy.deepcopy(self._nx_model),
stats_layout, stat_aliases)
self.set_model(model_with_stat_op)
nodes_name = []
for names_map in nodes_names_map.values():
nodes_name.extend(list(names_map.keys()))
outputs = self._add_outputs(nodes_names_map)
for model_name, outputs_data in outputs.items():
add_tensor_names(outputs_data, nodes_names_map[model_name].keys())
model_output_names = []
for model in self._model:
model_output_names.extend(collect_model_outputs(model['model']))
align_stat_names_with_results(model_output_names,
nodes_name,
output_to_node_names,
stats_layout,
stat_aliases)
# Creating statistics layout with IE-like names
stats_layout, stat_names_aliases = self._convert_stats_names(stats_layout)
self._predict(stats_layout=stats_layout,
sampler=sampler,
print_progress=print_progress,
need_metrics_per_sample=metric_per_sample)
accumulated_stats = \
process_accumulated_stats(stat_names_aliases=stat_names_aliases,
accumulated_stats=self._accumulated_layer_stats)
if stats_layout:
restore_original_node_names(output_to_node_names, accumulated_stats, stats_layout, stat_aliases)
metrics = None
if self._metric:
metrics = self._metric.avg_value
if metric_per_sample:
metrics = (sorted(self._per_sample_metrics, key=lambda i: i['sample_id']), metrics)
self._reset()
return metrics, accumulated_stats
def _predict(self, stats_layout, sampler, print_progress=False,
need_metrics_per_sample=False):
@@ -142,12 +202,12 @@ class MTCNNEngine(IEEngine):
progress_log_fn('Inference finished')
def _infer(self, data, ie_network, stats_collect_callback=None):
ie_network.reshape(PartialShape(data.shape))
filled_input = self._fill_input(ie_network, data)
input_shapes = {layer_name: data.shape for layer_name, data in filled_input.items()}
ie_network.reshape(input_shapes)
exec_model = self._ie.load_network(network=ie_network,
device_name=self.config.device)
result = exec_model.infer(filled_input)
compiled_model = self._ie.compile_model(model=ie_network,
device_name=self.config.device)
infer_request = compiled_model.create_infer_request()
result = infer_request.infer(filled_input)
# Collect statistics
if stats_collect_callback:
stats_collect_callback(self._transform_for_callback(result))
@@ -173,9 +233,12 @@ class MTCNNEngine(IEEngine):
total_boxes = np.zeros((0, 9), np.float)
for idx, outputs in enumerate(output):
scales = input_meta['scales'][idx]
mapping = outputs[[i for i, _ in outputs.items()
if i.any_name == self._output_layers['pnet']['probabilities']][0]][0, 1]
regions = outputs[[i for i, _ in outputs.items()
if i.any_name == self._output_layers['pnet']['regions']][0]][0]
mapping = outputs[self._output_layers['pnet']['probabilities']][0, 1]
regions = outputs[self._output_layers['pnet']['regions']][0]
boxes = utils.generate_bounding_box(mapping, regions, scales, 0.6)
if len(boxes) != 0:
pick = utils.nms(boxes, 0.5)
@@ -207,8 +270,9 @@ class MTCNNEngine(IEEngine):
return np.transpose(img, [0, 3, 2, 1])
def postprocess(output):
score = output[self._output_layers['rnet']['probabilities']][:, 1]
regions = output[self._output_layers['rnet']['regions']]
score = output[[i for i, _ in output.items()
if i.any_name == self._output_layers['rnet']['probabilities']][0]][:, 1]
regions = output[[i for i, _ in output.items() if i.any_name == self._output_layers['rnet']['regions']][0]]
return utils.calibrate_bboxes(prev_stage_output, score, regions, nms_type='union')
ie_network = self._model[1]['model']
@@ -223,8 +287,9 @@ class MTCNNEngine(IEEngine):
return np.transpose(img, [0, 3, 2, 1])
def postprocess(output):
score = output[self._output_layers['onet']['probabilities']][:, 1]
regions = output[self._output_layers['onet']['regions']]
score = output[[i for i, _ in output.items()
if i.any_name == self._output_layers['onet']['probabilities']][0]][:, 1]
regions = output[[i for i, _ in output.items() if i.any_name == self._output_layers['onet']['regions']][0]]
bboxes = utils.calibrate_bboxes(prev_stage_output, score, regions)
pick = utils.nms(bboxes, 0.7, 'min')
bboxes_to_remove = np.setdiff1d(np.arange(len(bboxes)), pick)