import os import sys import json import argparse import random import logging import fnmatch from operator import xor import numpy as np import cv2 work_dir = os.path.dirname(os.path.abspath(__file__)) cvat_dir = os.path.join(work_dir, '..', '..') sys.path.insert(0, cvat_dir) from cvat.apps.auto_annotation.inference import run_inference_engine_annotation def _get_kwargs(): parser = argparse.ArgumentParser() parser.add_argument('--py', help='Path to the python interpt file') parser.add_argument('--xml', help='Path to the xml file') parser.add_argument('--bin', help='Path to the bin file') parser.add_argument('--json', help='Path to the JSON mapping file') parser.add_argument('--model-name', help='Name of the model in the Model Manager') parser.add_argument('--task-id', type=int, help='ID task used to test the model') parser.add_argument('--restricted', dest='restricted', action='store_true') parser.add_argument('--unrestricted', dest='restricted', action='store_false') parser.add_argument('--image-files', nargs='*', help='Paths to image files you want to test') parser.add_argument('--show-images', action='store_true', help='Show the results of the annotation in a window') parser.add_argument('--show-image-delay', default=0, type=int, help='Displays the images for a set duration in milliseconds, default is until a key is pressed') parser.add_argument('--serialize', default=False, action='store_true', help='Try to serialize the result') parser.add_argument('--show-labels', action='store_true', help='Show the labels on the window') return vars(parser.parse_args()) def _init_django(settings): import django os.environ['DJANGO_SETTINGS_MODULE'] = settings django.setup() def random_color(): rgbl=[255,0,0] random.shuffle(rgbl) return tuple(rgbl) def pairwise(iterable): result = [] for i in range(0, len(iterable) - 1, 2): result.append((iterable[i], iterable[i+1])) return np.array(result, dtype=np.int32) def find_min_y(array): min_ = sys.maxsize index = None for i, pair in enumerate(array): if pair[1] < min_: min_ = pair[1] index = i return array[index] def _get_docker_files(model_name: str, task_id: int): _init_django('cvat.settings.development') from cvat.apps.auto_annotation.models import AnnotationModel from cvat.apps.engine.models import Task as TaskModel task = TaskModel(pk=task_id) model = AnnotationModel.objects.get(name=model_name) images_dir = task.data.get_data_dirname() py_file = model.interpretation_file.name mapping_file = model.labelmap_file.name xml_file = model.model_file.name bin_file = model.weights_file.name image_files = [] images_dir = os.path.abspath(images_dir) for root, _, filenames in os.walk(images_dir): for filename in fnmatch.filter(filenames, '*.jpg'): image_files.append(os.path.join(root, filename)) return py_file, mapping_file, bin_file, xml_file, image_files def main(): kwargs = _get_kwargs() py_file = kwargs.get('py') bin_file = kwargs.get('bin') mapping_file = os.path.abspath(kwargs.get('json')) xml_file = kwargs.get('xml') model_name = kwargs.get('model_name') task_id = kwargs.get('task_id') is_docker = model_name and task_id # xor is `exclusive or`. English is: if one or the other but not both if xor(bool(model_name), bool(task_id)): logging.critical('Must provide both `--model-name` and `--task-id` together!') return if is_docker: files = _get_docker_files(model_name, task_id) py_file = files[0] mapping_file = files[1] bin_file = files[2] xml_file = files[3] image_files = files[4] else: return_ = False if not py_file: logging.critical('Must provide --py file!') return_ = True if not bin_file: logging.critical('Must provide --bin file!') return_ = True if not xml_file: logging.critical('Must provide --xml file!') return_ = True if not mapping_file: logging.critical('Must provide --json file!') return_ = True if return_: return if not os.path.isfile(py_file): logging.critical('Py file not found! Check the path') return if not os.path.isfile(bin_file): logging.critical('Bin file is not found! Check path!') return if not os.path.isfile(xml_file): logging.critical('XML File not found! Check path!') return if not os.path.isfile(mapping_file): logging.critical('JSON file is not found! Check path!') return with open(mapping_file) as json_file: try: mapping = json.load(json_file) except json.decoder.JSONDecodeError: logging.critical('JSON file not able to be parsed! Check file') return try: mapping = mapping['label_map'] except KeyError: logging.critical("JSON Mapping file must contain key `label_map`!") logging.critical("Exiting") return mapping = {int(k): v for k, v in mapping.items()} restricted = kwargs['restricted'] if not is_docker: image_files = kwargs.get('image_files') if image_files: image_data = [cv2.imread(f) for f in image_files] else: test_image = np.ones((1024, 1980, 3), np.uint8) * 255 image_data = [test_image,] attribute_spec = {} results = run_inference_engine_annotation(image_data, xml_file, bin_file, mapping, attribute_spec, py_file, restricted=restricted) logging.warning('Inference didn\'t have any errors.') show_images = kwargs.get('show_images', False) if show_images: if image_files is None: logging.critical("Warning, no images provided!") logging.critical('Exiting without presenting results') return if not results['shapes']: logging.warning(str(results)) logging.critical("No objects detected!") return show_image_delay = kwargs['show_image_delay'] show_labels = kwargs.get('show_labels') for index, data in enumerate(image_data): for detection in results['shapes']: if not detection['frame'] == index: continue points = detection['points'] label_str = detection['label_id'] # Cv2 doesn't like floats for drawing points = [int(p) for p in points] color = random_color() if detection['type'] == 'rectangle': cv2.rectangle(data, (points[0], points[1]), (points[2], points[3]), color, 3) if show_labels: cv2.putText(data, label_str, (points[0], points[1] - 7), cv2.FONT_HERSHEY_COMPLEX, 0.6, color, 1) elif detection['type'] in ('polygon', 'polyline'): # polylines is picky about datatypes points = pairwise(points) cv2.polylines(data, [points], 1, color) if show_labels: min_point = find_min_y(points) cv2.putText(data, label_str, (min_point[0], min_point[1] - 7), cv2.FONT_HERSHEY_COMPLEX, 0.6, color, 1) cv2.imshow(str(index), data) cv2.waitKey(show_image_delay) cv2.destroyWindow(str(index)) if kwargs['serialize']: _init_django('cvat.settings.production') from cvat.apps.engine.serializers import LabeledDataSerializer # NOTE: We're actually using `run_inference_engine_annotation` # incorrectly here. The `mapping` dict is supposed to be a mapping # of integers -> integers and represents the transition from model # integers to the labels in the database. We're using a mapping of # integers -> strings. For testing purposes, this shortcut is fine. # We just want to make sure everything works. Until, that is.... # we want to test using the label serializer. Then we have to transition # back to integers, otherwise the serializer complains about have a string # where an integer is expected. We'll just brute force that. for shape in results['shapes']: # Change the english label to an integer for serialization validation shape['label_id'] = 1 serializer = LabeledDataSerializer(data=results) if not serializer.is_valid(): logging.critical('Data unable to be serialized correctly!') serializer.is_valid(raise_exception=True) if __name__ == '__main__': main()