# -*- coding: utf-8 -*- """ rectlabel.py -- Tools to implement a labeling UI for bounding boxes in images Copyright 2020 Monterey Bay Aquarium Research Institute Distributed under MIT license. See license.txt for more information. """ import cv2 import numpy as np import time import os from lxml import etree from lxml.etree import Element, SubElement from xml.etree import ElementTree from xml.dom import minidom import xml.dom.minidom def prettify(elem): """Return a pretty-printed XML string for the Element. """ return etree.tostring(elem, pretty_print=True).decode() def ppxml(xml): parser = etree.XMLParser(remove_blank_text=True) tree = etree.parse(xml, parser) tree.write(xml, encoding='utf-8', pretty_print=True, xml_declaration=True) class Annotation: def __init__(self, xml_path=None, image=None, filename=None, folder='LCM', lcmlog='lcm-log-file', camera_name='cam', utime=0): if xml_path is not None: # load the existing annotation self.xml_path = xml_path parser = etree.XMLParser(remove_blank_text=True) self.xml_tree = etree.parse(self.xml_path, parser) self.root = self.xml_tree.getroot() self.filename = self.xml_tree.find('filename').text self.folder = self.xml_tree.find('folder').text self.bboxes = [] else: self.utime = utime self.camera_name = camera_name self.image_width = image.shape[1] self.image_height = image.shape[0] self.image_depth = image.shape[2] self.lcmlog = lcmlog self.filename = filename self.folder = folder self.image = image self.bboxes = [] self.root = Element('annotation') self.folder_tag = self.add_tag(self.root, 'folder', self.folder) self.filename_tag = self.add_tag(self.root, 'filename', self.filename) self.source = SubElement(self.root, 'source') self.add_tag(self.source, 'database', 'Unknown') self.add_tag(self.source, 'lcmlog', lcmlog) self.add_tag(self.source, 'utime', utime) self.add_tag(self.source, 'camera-name', self.camera_name) self.size = SubElement(self.root, 'size') self.add_tag(self.size, 'width', self.image_width) self.add_tag(self.size, 'height', self.image_height) self.add_tag(self.size, 'depth', self.image_depth) def add_tag(self, parent, child_name, value=None): tag = SubElement(parent, child_name) if value is not None: tag.text = str(value) return tag def get_box(self, ind): objs = self.xml_tree.findall('object') if ind >= 0 and ind < len(objs): xmin = int(float(objs[ind].find('bndbox').find('xmin').text)) xmax = int(float(objs[ind].find('bndbox').find('xmax').text)) ymin = int(float(objs[ind].find('bndbox').find('ymin').text)) ymax = int(float(objs[ind].find('bndbox').find('ymax').text)) return xmin, ymin, xmax, ymax else: return [] def remove_boxes(self): for obj in self.xml_tree.findall('object'): self.root.remove(obj) def set_box(self, ind, bbox, label): objs = self.xml_tree.findall('object') if ind >= 0 and ind < len(objs): objs[ind].find('bndbox').find('xmin').text = str(int(bbox[0])) objs[ind].find('bndbox').find('ymin').text = str(int(bbox[1])) objs[ind].find('bndbox').find('xmax').text = str(int(bbox[2])) objs[ind].find('bndbox').find('ymax').text = str(int(bbox[3])) objs[ind].find('name').text = label def add_box(self, bbox): pt0, pt1, utime, class_name, scores = bbox obj = SubElement(self.root, 'object') self.add_tag(obj, 'name', class_name) self.add_tag(obj, 'pose', 'Unspecified') self.add_tag(obj, 'truncated', '0') self.add_tag(obj, 'occluded', '0') self.add_tag(obj, 'difficult', '0') bndbox = SubElement(obj, 'bndbox') self.add_tag(bndbox, 'xmin', int(pt0[0])) self.add_tag(bndbox, 'ymin', int(pt0[1])) self.add_tag(bndbox, 'xmax', int(pt1[0])) self.add_tag(bndbox, 'ymax', int(pt1[1])) self.bboxes.append(bbox) def save_xml(self, filepath=None): if self.xml_path is not None: with open(self.xml_path, 'w') as f: f.write(prettify(self.root)) return if filepath is None: filepath = os.path.join(self.folder, 'xml', self.filename + '.xml') with open(filepath, 'w') as f: self.xml_tree.write(filepath, encoding='utf-8', pretty_print=True, xml_declaration=True) #f.write(prettify(self.root)) def save_boxed_image(self, filepath=None): # save out the full image with boxes drawn marked_image = self.image.copy() for bbox in self.bboxes: pt0, pt1, utime, class_name, scores = bbox cv2.rectangle(marked_image, tuple(pt0), tuple(pt1), (0, 255, 0), 2) cv2.putText(marked_image, class_name, (pt0[0], pt1[1]), cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 255)) if filepath is None: filepath = os.path.join(self.folder, 'boxed_images', self.filename + ".jpg") print(filepath) cv2.imwrite(filepath, marked_image) def save_image(self, filepath=None): if filepath is None: filepath = os.path.join(self.folder, 'images', self.filename + ".jpg") cv2.imwrite(filepath, self.image) def save_rois(self, filepath=None): # for each annotation, extract the ROI from the raw image and save if filepath is None: fileprefix = self.filename else: fileprefix = filepath for bbox in self.bboxes: pt0, pt1, utime, class_name, scores = bbox roi = self.image[pt0[1]:pt1[1], pt0[0]:pt1[0], :] coord_string = str(pt0[0]) + "-" + str(pt0[1]) + "-" + str(pt1[0]) + "-" + str(pt1[1]) filepath = os.path.join(self.folder, 'rois', fileprefix + "-" + coord_string + ".jpg") cv2.imwrite(filepath, roi) if __name__=="__main__": # A test bbox = ((100, 200), (130, 260), 1567897687000, 'test_class1', 0.25) bbox2 = ((800, 400), (870, 420), 1567899687000, 'test_class2', 0.25) img = 255*np.random.random((540, 960, 3)) filename = 'a_test_file' ann = Annotation(image=img, filename=filename, folder='.', utime=int(time.time()*1000)) ann.add_box(bbox) ann.add_box(bbox2) ann.save_xml() ann.save_boxed_image() ann.save_rois() ann2 = Annotation(xml_path=filename+'.xml')