AnnotationLocationLoader.py 6.67 KB
Newer Older
sjjsmuel's avatar
sjjsmuel committed
1
import json
sjjsmuel's avatar
sjjsmuel committed
2
from pathlib import Path, PosixPath
sjjsmuel's avatar
sjjsmuel committed
3 4


sjjsmuel's avatar
sjjsmuel committed
5
class AnnotationLocationLoader:    
sjjsmuel's avatar
sjjsmuel committed
6

sjjsmuel's avatar
sjjsmuel committed
7
    def __init__(self, annotation_file='input/caries_dataset_annotation.json', images_base_folder=Path('input/test_data/'), mouth_annotations_folder=None):
sjjsmuel's avatar
sjjsmuel committed
8
        self._annotation_file = annotation_file
sjjsmuel's avatar
sjjsmuel committed
9 10 11 12 13
        self._img_path = None
        self._annotated_images = set()
        self._available_annotations = set()
        self._available_images = None
        self._data = {} 
sjjsmuel's avatar
sjjsmuel committed
14 15 16

        if not type(images_base_folder) == PosixPath:
            images_base_folder = Path(images_base_folder)
sjjsmuel's avatar
sjjsmuel committed
17 18 19 20 21
        self._img_path = images_base_folder

        # get the names of the images witch are available as files
        self._available_images = self._get_names_from_available_images()
        self._load_annotations()
sjjsmuel's avatar
sjjsmuel committed
22 23 24 25 26 27 28

        if not mouth_annotations_folder == None:
            if not type(mouth_annotations_folder) == PosixPath:
                mouth_annotations_folder = Path(mouth_annotations_folder)
            self._mouth_annotations_folder = mouth_annotations_folder

            self._load_mouth_annotation_additon()
sjjsmuel's avatar
sjjsmuel committed
29 30 31

        self._annotated_images = list(self._annotated_images)
        self._available_annotations = list(self._available_annotations)
sjjsmuel's avatar
sjjsmuel committed
32 33

    def _get_names_from_available_images(self):
sjjsmuel's avatar
sjjsmuel committed
34 35 36 37
        """
        Finds the names for all pictures available in the image base folder.
        :return:
        """
sjjsmuel's avatar
sjjsmuel committed
38
        names_from_available_images = []
sjjsmuel's avatar
sjjsmuel committed
39
        for path in [path for path in self._img_path.iterdir() if path.is_dir()]:
sjjsmuel's avatar
sjjsmuel committed
40 41 42 43 44
            names_from_available_images.extend([filename.name for filename in path.iterdir() if filename.is_file() and not filename.name.startswith('.')])
        return names_from_available_images


    def _load_annotations(self):
sjjsmuel's avatar
sjjsmuel committed
45 46 47
        """
        Loads all annotations for the available images from the json file to an internal structure.
        """
sjjsmuel's avatar
sjjsmuel committed
48 49 50 51 52 53 54 55 56 57
        with open(self._annotation_file) as file:
            json_data = json.load(file)

            for picture in json_data:
                picture_filename = picture['External ID']

                if not picture_filename in self._available_images:
                    #print('File ”{}” not found.'.format(picture_filename))
                    continue

sjjsmuel's avatar
sjjsmuel committed
58 59 60
                if not picture_filename in self._data.keys():
                    self._data[picture_filename] = []

sjjsmuel's avatar
sjjsmuel committed
61 62 63 64 65 66 67 68 69 70 71 72 73 74 75
                # Skip the 'Skip' entries in the file
                if not type(picture['Label']) == dict:
                    continue

                for annotation_type in picture['Label'].keys():
                    self._available_annotations.add(annotation_type)
                    self._annotated_images.add(picture_filename)
                    for box in picture['Label'][annotation_type]:
                        x_all = []
                        y_all = []
                        for point in box['geometry']:
                            x_all.append(point['x'])
                            y_all.append(point['y'])
                        box_coord = [(min(x_all), min(y_all)), (max(x_all), max(y_all))]
                        self._data[picture_filename].append((annotation_type.lower(), box_coord))
sjjsmuel's avatar
sjjsmuel committed
76 77 78


    def _load_mouth_annotation_additon(self):
sjjsmuel's avatar
sjjsmuel committed
79 80 81 82
        """
        Loads the additional annotations for the mouth for all available images if there is a corresponding folder given
        This represents the local constrains for the advanced training.
        """
sjjsmuel's avatar
sjjsmuel committed
83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104
        annotation_files = [file for file in self._mouth_annotations_folder.iterdir() if file.is_file() and not file.name.startswith('.')]
        counter_number_of_annotated_but_missing_files = 0
        for annotation_file in annotation_files:
            with open(annotation_file) as file:
                mouth_annotation = json.load(file)
                picture_filename = mouth_annotation['asset']['name']

                if not picture_filename in self._available_images:
                    #print('File ”{}” not found.'.format(picture_filename))
                    counter_number_of_annotated_but_missing_files += 1
                    continue

                if not picture_filename in self._data.keys():
                    self._data[picture_filename] = []
                self._annotated_images.add(picture_filename)

                bb = mouth_annotation['regions'][0]['boundingBox']
                top_left = (round(bb['left']), round(bb['top']))
                bottom_right = (round(bb['left']+bb['width']), round(bb['top'] + bb['height']))
                img_width = mouth_annotation['asset']['size']['width']
                img_height = mouth_annotation['asset']['size']['height']
                self._data[picture_filename].append(('mouth', [(img_width, img_height), top_left, bottom_right]))
sjjsmuel's avatar
sjjsmuel committed
105 106
        if counter_number_of_annotated_but_missing_files > 0:
            print('[INFO] {} mouth annotations were skipped during loading. This was done due to missing corresponding files in the assigned folder.'.format(counter_number_of_annotated_but_missing_files))
sjjsmuel's avatar
sjjsmuel committed
107

sjjsmuel's avatar
sjjsmuel committed
108 109 110 111 112 113 114 115 116 117 118 119 120 121


    def get_all_types_of_annotations(self):
        """
        :return: list of all the types of annotations witch appeared at least once in the annotation_file
        """
        return self._available_annotations

    def get_all_annotated_images(self):
        """
        :return: list of the names of all images witch have at least one annotation
        """
        return self._annotated_images

sjjsmuel's avatar
sjjsmuel committed
122 123 124 125 126 127
    def get_all_available_images(self):
        """
        :return: list of the names of all images witch have at least one annotation
        """
        return self._available_images

sjjsmuel's avatar
sjjsmuel committed
128 129 130 131 132 133 134 135 136 137 138 139 140
    def is_annotated(self, image_name):
        """
        Should check weather for the given filename an annotation exists
        :param image_name: complete name of the file including the filetype as a string
        :return: boolean weather there is an annotation for the image
        """
        return image_name in self._annotated_images


    def get_annotations(self, image_name, filter=None):
        """
            Returns a list of annotations for the given image_name
            e.g. [ ('caries', [(x1,y1), (x2,y2)]),  _more_entries_ ]
sjjsmuel's avatar
sjjsmuel committed
141
            coords in form: [(top left), (bottom right)]
sjjsmuel's avatar
sjjsmuel committed
142 143
            :param filter: a list of strings representing the types of annotations the user wants to derive
        """
sjjsmuel's avatar
sjjsmuel committed
144 145
        if not filter:
            filter = self._available_annotations # return anything but the annotation for the mouth
sjjsmuel's avatar
sjjsmuel committed
146 147 148 149 150 151 152 153

        if self.is_annotated(image_name):
            if filter and len(filter)>0:
                filter = [category.lower() for category in filter]
                return [annotation for annotation in self._data[image_name] if annotation[0] in filter]
            return self._data[image_name]
        else:
            return []