Source code for akida_models.detection.processing

#!/usr/bin/env python
# ******************************************************************************
# Copyright 2020 Brainchip Holdings Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ******************************************************************************
"""
Processing tools for YOLO data handling.
"""

import os

import xml.etree.ElementTree as et
import numpy as np
import tensorflow as tf


[docs]class BoundingBox: """ Utility class to represent a bounding box. The box is defined by its top left corner (x1, y1), bottom right corner (x2, y2), label, score and classes. """ def __init__(self, x1, y1, x2, y2, score=-1, classes=None): self.x1 = x1 self.y1 = y1 self.x2 = x2 self.y2 = y2 self.label = -1 self.score = score self.classes = classes def __repr__(self): return "<BoundingBox({}, {}, {}, {}, {}, {}, {})>\n".format( self.x1, self.x2, self.y1, self.y2, self.get_label(), self.get_score(), self.classes)
[docs] def get_label(self): """ Returns the label for this bounding box. Returns: Index of the label as an integer. """ if self.label == -1: self.label = np.argmax(self.classes) return self.label
[docs] def get_score(self): """ Returns the score for this bounding box. Returns: Confidence as a float. """ if self.score == -1: self.score = self.classes[self.get_label()] return self.score
[docs] def iou(self, other): """ Computes intersection over union ratio between this bounding box and another one. Args: other (BoundingBox): the other bounding box for IOU computation Returns: IOU value as a float """ def _interval_overlap(interval_1, interval_2): x1, x2 = interval_1 x3, x4 = interval_2 if x3 < x1: if x4 < x1: return 0 return min(x2, x4) - x1 if x2 < x3: return 0 return min(x2, x4) - x3 intersect_w = _interval_overlap([self.x1, self.x2], [other.x1, other.x2]) intersect_h = _interval_overlap([self.y1, self.y2], [other.y1, other.y2]) intersect = intersect_w * intersect_h w1, h1 = self.x2 - self.x1, self.y2 - self.y1 w2, h2 = other.x2 - other.x1, other.y2 - other.y1 union = w1 * h1 + w2 * h2 - intersect return float(intersect) / union
[docs]def load_image(image_path): """ Loads an image from a path. Args: image_path (string): full path of the image to load Returns: a Tensorflow image Tensor """ raw_image = tf.io.read_file(image_path) return tf.image.decode_jpeg(raw_image, channels=3)
[docs]def preprocess_image(image_buffer, output_size): """ Preprocess an image for YOLO inference. Args: image_buffer (tf.Tensor): image to preprocess output_size (tuple): shape of the image after preprocessing Returns: A resized and normalized image as a Numpy array. """ # Resize width = tf.constant(output_size[0]) height = tf.constant(output_size[1]) image = tf.compat.v1.image.resize(image_buffer, [height, width], method=tf.image.ResizeMethod.BILINEAR, align_corners=False) return image.numpy()
[docs]def decode_output(output, anchors, nb_classes, obj_threshold=0.5, nms_threshold=0.5): """ Decodes a YOLO model output. Args: output (tf.Tensor): model output to decode anchors (list): list of anchors boxes nb_classes (int): number of classes obj_threshold (float): confidence threshold for a box nms_threshold (float): non-maximal supression threshold Returns: List of `BoundingBox` objects """ def _sigmoid(x): return 1. / (1. + np.exp(-x)) def _softmax(x, axis=-1, t=-100.): x = x - np.max(x) if np.min(x) < t: x = x / np.min(x) * t e_x = np.exp(x) return e_x / e_x.sum(axis, keepdims=True) grid_h, grid_w, nb_box = output.shape[:3] boxes = [] # decode the output by the network output[..., 4] = _sigmoid(output[..., 4]) output[..., 5:] = output[..., 4][..., np.newaxis] * _softmax(output[..., 5:]) output[..., 5:] *= output[..., 5:] > obj_threshold for row in range(grid_h): for col in range(grid_w): for b in range(nb_box): # from 5th element onwards are class classes classes = output[row, col, b, 5:] if np.sum(classes) > 0: # first 4 elements are x, y, w, and h x, y, w, h = output[row, col, b, :4] x = (col + _sigmoid(x) ) / grid_w # center position, unit: image width y = (row + _sigmoid(y) ) / grid_h # center position, unit: image height w = anchors[b][0] * np.exp(w) / grid_w # unit: image width h = anchors[b][1] * np.exp(h) / grid_h # unit: image height confidence = output[row, col, b, 4] x1 = max(x - w / 2, 0) y1 = max(y - h / 2, 0) x2 = min(x + w / 2, grid_w) y2 = min(y + h / 2, grid_h) box = BoundingBox(x1, y1, x2, y2, confidence, classes) boxes.append(box) # suppress non-maximal boxes for c in range(nb_classes): sorted_indices = list( reversed(np.argsort([box.classes[c] for box in boxes]))) for ind, index_i in enumerate(sorted_indices): if boxes[index_i].classes[c] == 0: continue for j in range(ind + 1, len(sorted_indices)): index_j = sorted_indices[j] # filter out redundant boxes (same class and overlapping too # much) if (boxes[index_i].iou(boxes[index_j]) >= nms_threshold) and ( c == boxes[index_i].get_label()) and ( c == boxes[index_j].get_label()): boxes[index_j].score = 0 # remove the boxes which are less likely than a obj_threshold boxes = [box for box in boxes if box.get_score() > obj_threshold] return boxes
[docs]def parse_voc_annotations(gt_folder, image_folder, file_path, labels): """ Loads PASCAL-VOC data. Data is loaded using the groundtruth informations and stored in a dictionary. Args: gt_folder (str): path to the folder containing ground truth files image_folder (str): path to the folder containing the images file_path (str): file containing the list of files to parse labels (list): list of labels of interest Returns: dict: a dictionnary containing all data present in the ground truth file """ def _build_list(file_path): """ Builds a list of file names from a text file path. Args: file_path (string): path of the file containing the names Returns: list: list of file names """ file_list = [] with open(file_path, "r", encoding='UTF-8') as f: for line in f: file_list.append(line.rstrip()) return file_list file_list = _build_list(file_path) all_data = [] for f in file_list: tree = et.parse(os.path.join(gt_folder, f) + ".xml") data = {"boxes": []} for elem in tree.iter(): if 'filename' in elem.tag: data['image_path'] = os.path.join(image_folder, elem.text) if 'size' in elem.tag: for attr in list(elem): if 'width' in attr.tag: width = int(attr.text) if 'height' in attr.tag: height = int(attr.text) if 'depth' in attr.tag: depth = int(attr.text) data['image_shape'] = (width, height, depth) if 'object' in elem.tag: box = {} for attr in list(elem): if 'difficult' in attr.tag: if int(attr.text) != 0: box.clear() break if 'name' in attr.tag: box['label'] = attr.text if box['label'] not in labels: box.clear() break if 'bndbox' in attr.tag: for dim in list(attr): if 'xmin' in dim.tag: box['x1'] = int(round(float(dim.text))) if 'ymin' in dim.tag: box['y1'] = int(round(float(dim.text))) if 'xmax' in dim.tag: box['x2'] = int(round(float(dim.text))) if 'ymax' in dim.tag: box['y2'] = int(round(float(dim.text))) if len(box) == 5: data["boxes"].append(box) if len(data["boxes"]) != 0: all_data.append(data) return all_data
[docs]def parse_widerface_annotations(gt_file, image_folder): """ Loads WiderFace data. Data is loaded using the groundtruth informations and stored in a dictionary. Args: gt_file (str): path to the ground truth file image_folder (str): path to the directory containing the images Returns: dict: a dictionnary containing all data present in the ground truth file """ def _is_valid_box(w_box, h_box, w_img, h_img): box_area = w_box * h_box img_area = w_img * h_img return box_area >= img_area / 60.0 all_data = [] with open(gt_file, "r", encoding='UTF-8') as f: for line in f: file_path = line.rstrip() full_file_path = os.path.join(image_folder, str(file_path)) if not os.path.isfile(full_file_path): raise RuntimeError(f"Not a file path : {full_file_path}") image = load_image(full_file_path) shape = (image.shape[1], image.shape[0], image.shape[2]) data = { "image_path": full_file_path, "image_shape": shape, "boxes": [] } nb_boxes = int(f.readline()) if nb_boxes == 0: f.readline() for _ in range(nb_boxes): annotation = f.readline() values = annotation.split(' ') x1 = int(values[0]) y1 = int(values[1]) w = int(values[2]) h = int(values[3]) if _is_valid_box(w, h, shape[0], shape[1]): box = { "x1": x1, "y1": y1, "x2": x1 + w, "y2": y1 + h, "label": "face" } data["boxes"].append(box) if len(data["boxes"]) != 0: all_data.append(data) return all_data