bbox_utils.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Library with methods for handling bounding boxes in images

License_info:
# ==============================================================================
# ISC License (ISC)
# Copyright 2020 Christian Doppler Laboratory for Embedded Machine Learning
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
# REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND
# FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
# INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
# LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
# OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
# PERFORMANCE OF THIS SOFTWARE.

# The following script uses several method fragments from Tensorflow
https://github.com/tensorflow/models/blob/master/research/object_detection/dataset_tools/create_pascal_tf_record.py

Tensorflow has the following licence:
# ==============================================================================
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""

# Futures
from __future__ import print_function

# Built-in/Generic Imports
import os
import time

# Libs
import argparse
import numpy as np
import glob
import xml.etree.ElementTree as ET
from multiprocessing import Pool
import matplotlib
from six import BytesIO
import re
import pandas as pd
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageFont
import tensorflow as tf
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as viz_utils
import tkinter

# Own modules
import image_utils as im

__author__ = 'Alexander Wendt'
__copyright__ = 'Copyright 2020, Christian Doppler Laboratory for ' \
                'Embedded Machine Learning'
__credits__ = ['']
__license__ = 'ISC'
__version__ = '0.2.0'
__maintainer__ = 'Alexander Wendt'
__email__ = 'alexander.wendt@tuwien.ac.at'
__status__ = 'Experiental'


def xml_to_csv(path, filter=None):
    """Iterates through all .xml files (generated by labelImg) in a given directory and combines
    them in a single Pandas dataframe.

    Parameters:
    ----------
    path : str
        The path containing the .xml files
    filter: list of image file names. Default None. If no filter is given, all xml files are used

    Returns
    -------
    Pandas DataFrame
        The produced dataframe
    """
    xml_file_list=[]

    if filter is not None:
        print("Filter available. Using only xml files with corresponding image files")
        #xml_filename = os.path.join(xml_source, os.path.splitext(filename)[0] + '.xml')
        xml_file_list = [os.path.join(path, os.path.splitext(image_name)[0] + '.xml') for image_name in filter]
    else:
        print("Filter not used. Select all xml files of the folder")
        xml_file_list = glob.glob(path + '/*.xml')

    xml_list = []
    for xml_file in xml_file_list: #glob.glob(path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        for member in root.findall('object'):
            #Check if detection or ground truth
            if isinstance(member.find("score"), ET.Element):
                score = float(member.find("score").text)
            else:
                score = None

            value = (root.find('filename').text,
                     int(root.find('size')[0].text),
                     int(root.find('size')[1].text),
                     member.find("name").text,
                     int(member.find("bndbox")[0].text),
                     int(member.find("bndbox")[1].text),
                     int(member.find("bndbox")[2].text),
                     int(member.find("bndbox")[3].text),
                     score
                     )
            xml_list.append(value)
    column_name = ['filename', 'width', 'height',
                   'class', 'xmin', 'ymin', 'xmax', 'ymax', 'score']
    xml_df = pd.DataFrame(xml_list, columns=column_name)
    return xml_df

def extract_info_from_annotations(annotation, category_index, color_gt=False):

    boxes = np.zeros([annotation.shape[0], 4])
    classes = np.zeros([annotation.shape[0]]).astype('int')

    if 'score' in annotation.columns and annotation['score'][0] is not None:
        scores = np.zeros([annotation.shape[0]])
    elif color_gt:
        scores = np.ones([annotation.shape[0]])
    else:
        scores = None

    for i in range(annotation.shape[0]):
        boxes[i][0] = annotation['ymin'][i] / annotation['height'][i]
        boxes[i][1] = annotation['xmin'][i] / annotation['width'][i]
        boxes[i][2] = annotation['ymax'][i] / annotation['height'][i]
        boxes[i][3] = annotation['xmax'][i] / annotation['width'][i]
            #[annotation['ymin'][i] / annotation['height'][i],
            # annotation['xmin'][i] / annotation['width'][i],
            # annotation['ymax'][i] / annotation['height'][i],
            ## annotation['xmax'][i] / annotation['width'][i]]).reshape(i, 4)
        print("Boxes: ", boxes)

        class_index = [category_index[j + 1].get('id') for j in range(len(category_index)) if
                       category_index[j + 1].get('name') == annotation['class'][i]][0]

        classes[i] = class_index
        print("Class index: ", class_index)

        if 'score' in annotation.columns and annotation['score'][i] is not None:
            scores[i] = annotation['score'][i]


    return boxes, classes, scores

def visualize_image(image_name, image_np, scores, boxes, classes, category_index, min_score=0.4):
    # Get objects
    #image_np, boxes, classes, scores = value

    if scores is None or (max(scores) >= min_score):
        #print("Visualize image")
        print(image_name)
        # print(value)
        # print(classes)
        # print(scores)
        # print(boxes)
        # input()
        if image_np.shape[0] < 1000:
            line_thickness = 2
        elif image_np.shape[0] < 2000:
            line_thickness = 4
        else:
            line_thickness = 8

        print("Image width: {}->Setting line thinckness: {}".format(image_np.shape[0], line_thickness))

        plt.rcParams['figure.figsize'] = [42, 21]
        label_id_offset = 1
        image_np_with_detections = image_np.copy()
        viz_utils.visualize_boxes_and_labels_on_image_array(
            image_np_with_detections,
            boxes,
            classes,
            scores,
            category_index,
            use_normalized_coordinates=True,
            max_boxes_to_draw=200,
            min_score_thresh=min_score,
            line_thickness=line_thickness,
            agnostic_mode=False)
        # plt.show()
        # plt.subplot(5, 1, 1)
    else:
        image_np_with_detections = image_np

    return image_np_with_detections

def visualize_image_with_boundingbox(annotation_dir, category_index, image_name, image_path, min_score=0.4, color_gt=False):
    image_np = im.load_image_into_numpy_array(image_path)# load_image(image_path)
    filter = []
    filter.append(image_name)
    annotation_df = xml_to_csv(annotation_dir, filter)
    print("Annotation: ", annotation_df)
    boxes, classes, scores = extract_info_from_annotations(annotation_df, category_index, color_gt=color_gt)
    fig1 = visualize_image(image_name, image_np, scores, boxes, classes, category_index, min_score)

    return fig1