PyCS/pycs/pipeline/tf1/detection/detector.py

#!/usr/bin/env python

"""Detector:  Face detection implementation."""

import logging

import numpy as np

from ..utils import TFModel


# import tensorflow.contrib.slim as slim


class Detector(TFModel):
    def __init__(self, config):
        TFModel.__init__(self, config)

        try:
            # (1) Find feature tensor
            self.tf_image_tensor = self.tf_graph.get_tensor_by_name("import/image_tensor:0")

            self.tf_detection_boxes = self.tf_graph.get_tensor_by_name('import/detection_boxes:0')
            self.tf_detection_scores = self.tf_graph.get_tensor_by_name('import/detection_scores:0')
            self.tf_detection_classes = self.tf_graph.get_tensor_by_name('import/detection_classes:0')
            self.tf_num_detections = self.tf_graph.get_tensor_by_name('import/num_detections:0')

            self.input_shape = self.tf_image_tensor.shape[1:].as_list()
            if "downscale-to" in config.keys():
                for i in range(len(self.input_shape)):
                    if self.input_shape[i] is None:
                        self.input_shape[i] = config["downscale-to"]
            logging.debug("Input shape: %s" % self.input_shape)

        except:
            self._report_error("Could not access tensors by name")

    def detect_faces(self, image):
        if None not in self.input_shape:
            resized_image = image.resize(size=self.input_shape[0:2])
        else:
            resized_image = image

        (boxes, scores, classes, num) = self.tf_session.run(
            [self.tf_detection_boxes, self.tf_detection_scores, self.tf_detection_classes, self.tf_num_detections],
            feed_dict={self.tf_image_tensor: np.expand_dims(resized_image, axis=0)})

        sample_num = int(num[0])
        sample_scores = scores[0][0:sample_num]
        sample_boxes = boxes[0][0:sample_num]

        filtered_boxes = sample_boxes[sample_scores > 0.5]
        filtered_scores = sample_scores[sample_scores > 0.5]

        ret_boxes = []
        for index, box in enumerate(filtered_boxes):
            score = sample_scores[index]
            ymin, xmin, ymax, xmax = box
            ret_box = {'x': float(xmin), 'y': float(ymin),
                       'w': float(xmax - xmin), 'h': float(ymax - ymin),
                       'score': float(score)}
            ret_boxes += [ret_box]

        return ret_boxes