megadetector 10.0.15__py3-none-any.whl

megadetector/data_management/cct_json_utils.py

@@ -0,0 +1,566 @@
+"""
+
+cct_json_utils.py
+
+Utilities for working with COCO Camera Traps .json databases:
+
+https://github.com/agentmorris/MegaDetector/blob/main/megadetector/data_management/README.md#coco-cameratraps-format
+
+"""
+
+#%% Constants and imports
+
+import json
+import os
+import datetime
+import dateutil
+
+from tqdm import tqdm
+from collections import defaultdict, OrderedDict
+
+
+#%% Classes
+
+class CameraTrapJsonUtils:
+    """
+    Miscellaneous utility functions for working with COCO Camera Traps databases
+    """
+
+    @staticmethod
+    def annotations_to_string(annotations, cat_id_to_name):
+        """
+        Given a list of annotations and a mapping from class IDs to names, produces
+        a comma-delimited string containing a list of class names, sorted alphabetically.
+
+        Args:
+            annotations (list): a list of annotation dicts
+            cat_id_to_name (dict): a dict mapping category IDs to category names
+
+        Returns:
+            str: a comma-delimited list of class names
+        """
+
+        class_names = CameraTrapJsonUtils.annotations_to_class_names(annotations, cat_id_to_name)
+        return ','.join(class_names)
+
+
+    @staticmethod
+    def annotations_to_class_names(annotations, cat_id_to_name):
+        """
+        Given a list of annotations and a mapping from class IDs to names, produces
+        a list of class names, sorted alphabetically.
+
+        Args:
+            annotations (list): a list of annotation dicts
+            cat_id_to_name (dict): a dict mapping category IDs to category names
+
+        Returns:
+            list: a list of class names present in [annotations]
+        """
+
+        # Collect all names
+        class_names = [cat_id_to_name[ann['category_id']] for ann in annotations]
+        # Make names unique and sort
+        class_names = sorted(set(class_names))
+        return class_names
+
+
+    @staticmethod
+    def order_db_keys(db):
+        """
+        Given a dict representing a JSON database in the COCO Camera Trap
+        format, returns an OrderedDict with keys in the order of 'info',
+        'categories', 'annotations' and 'images'. When this OrderedDict is
+        serialized with json.dump(), the order of the keys are preserved.
+
+        Args:
+            db (dict): a JSON database in the COCO Camera Trap format
+
+        Returns:
+            dict: the same content as [db] but as an OrderedDict with keys ordered for
+            readability
+        """
+
+        ordered = OrderedDict([
+            ('info', db['info']),
+            ('categories', db['categories']),
+            ('annotations', db['annotations']),
+            ('images', db['images'])])
+        return ordered
+
+
+    @staticmethod
+    def group_annotations_by_image_field(db_indexed, image_field='seq_id'):
+        """
+        Given an instance of IndexedJsonDb, group annotation entries by a field in the
+        image entry.  Typically used to find all the annotations associated with a sequence.
+
+        Args:
+            db_indexed (IndexedJsonDb): an initialized IndexedJsonDb, typically loaded from a
+                COCO Camera Traps .json file
+            image_field (str, optional): a field by which to group annotations (defaults
+                to 'seq_id')
+
+        Returns:
+            dict: a dict mapping objects (typically strings, in fact typically sequence IDs) to
+            lists of annotations
+        """
+
+        image_id_to_image_field = {}
+        for image_id, image_entry in db_indexed.image_id_to_image.items():
+            image_id_to_image_field[image_id] = image_entry[image_field]
+
+        res = defaultdict(list)
+        for annotations in db_indexed.image_id_to_annotations.values():
+            for annotation_entry in annotations:
+                field_value = image_id_to_image_field[annotation_entry['image_id']]
+                res[field_value].append(annotation_entry)
+        return res
+
+
+    @staticmethod
+    def get_entries_for_locations(db, locations):
+        """
+        Given a dict representing a JSON database in the COCO Camera Trap format, returns a dict
+        with the 'images' and 'annotations' fields in the CCT format, each is an array that only
+        includes entries in the original [db] that are in the [locations] set.
+
+        Args:
+            db (dict): a dict representing a JSON database in the COCO Camera Trap format
+            locations (set): a set or list of locations to include; each item is a string
+
+        Returns:
+            dict: a dict with the 'images' and 'annotations' fields in the CCT format
+        """
+
+        locations = set(locations)
+        print('Original DB has {} image and {} annotation entries.'.format(
+            len(db['images']), len(db['annotations'])))
+        new_db = { 'images': [], 'annotations': [] }
+        new_images = set()
+        for i in db['images']:
+            # cast location to string as the entries in locations are strings
+            if str(i['location']) in locations:
+                new_db['images'].append(i)
+                new_images.add(i['id'])
+        for a in db['annotations']:
+            if a['image_id'] in new_images:
+                new_db['annotations'].append(a)
+        print(
+            'New DB has {} image and {} annotation entries.'.format(
+                len(new_db['images']), len(new_db['annotations'])))
+        return new_db
+
+
+class IndexedJsonDb:
+    """
+    Wrapper for a COCO Camera Traps database.
+
+    Handles boilerplate dictionary creation that we do almost every time we load
+    a .json database.
+    """
+
+    def __init__(self,
+                 json_filename,
+                 b_normalize_paths=False,
+                 filename_replacements=None,
+                 b_convert_classes_to_lower=True,
+                 b_force_forward_slashes=True):
+        """
+        Constructor for IndexedJsonDb that loads from a .json file or CCT-formatted dict.
+
+        Args:
+            json_filename (str): filename to load, or an already-loaded dict
+            b_normalize_paths (bool, optional): whether to invoke os.path.normpath on
+                all filenames.  Not relevant if b_force_forward_slashes is True.
+            filename_replacements (dict, optional): a set of string --> string mappings
+                that will trigger replacements in all filenames, typically used to remove
+                leading folders
+            b_convert_classes_to_lower (bool, optional): whether to convert all class
+                names to lowercase
+            b_force_forward_slashes (bool, optional): whether to convert backslashes to
+                forward slashes in all path names
+        """
+
+        if isinstance(json_filename, str):
+            with open(json_filename) as f:
+                self.db = json.load(f)
+        else:
+            self.db = json_filename
+
+        assert 'images' in self.db, (
+            f'Could not find image list in file {json_filename}, are you sure '
+            'this is a COCO camera traps file?')
+
+        if b_convert_classes_to_lower:
+            # Convert classnames to lowercase to simplify comparisons later
+            for c in self.db['categories']:
+                c['name'] = c['name'].lower()
+
+        # Normalize paths to simplify comparisons later
+        if b_normalize_paths:
+            for im in self.db['images']:
+                im['file_name'] = os.path.normpath(im['file_name'])
+
+        if b_force_forward_slashes:
+            for im in self.db['images']:
+                im['file_name'] = im['file_name'].replace('\\','/')
+
+        if filename_replacements is not None:
+            for s in filename_replacements:
+                # Make custom replacements in filenames, typically used to
+                # accommodate changes in root paths after DB construction
+                r = filename_replacements[s]
+                for im in self.db['images']:
+                    im['file_name'] = im['file_name'].replace(s, r)
+
+        ### Build useful mappings to facilitate working with the DB
+
+        # Category ID <--> name
+        self.cat_id_to_name = {
+            cat['id']: cat['name'] for cat in self.db['categories']}
+        self.cat_name_to_id = {
+            cat['name']: cat['id'] for cat in self.db['categories']}
+
+        # Image filename --> ID
+        self.filename_to_id = {
+            im['file_name']: im['id'] for im in self.db['images']}
+
+        # Image ID --> image object
+        self.image_id_to_image = {im['id']: im for im in self.db['images']}
+
+        # Image ID --> annotations
+        # Each image can potentially multiple annotations, hence using lists
+        self.image_id_to_annotations = defaultdict(list)
+        for ann in self.db['annotations']:
+            self.image_id_to_annotations[ann['image_id']].append(ann)
+
+    # ...__init__
+
+
+    def get_annotations_for_image(self, image):
+        """
+        Finds all the annnotations associated with the image dict [image].
+
+        Args:
+            image (dict): an image dict loaded from a CCT .json file.  Only the 'id' field
+                is used.
+
+        Returns:
+            list: list of annotations associated with this image.  Returns None if the db
+            has not been loaded, or [] if no annotations are available for this image.
+        """
+
+        if self.db is None:
+            return None
+
+        if image['id'] not in self.image_id_to_annotations:
+            return []
+
+        image_annotations = self.image_id_to_annotations[image['id']]
+        return image_annotations
+
+
+    def get_classes_for_image(self, image):
+        """
+        Returns a list of class names associated with [image].
+
+        Args:
+            image (dict): an image dict loaded from a CCT .json file.  Only the 'id' field
+                is used.
+
+        Returns:
+            list: list of class names associated with this image.  Returns None if the db
+            has not been loaded, or [] if no annotations are available for this image.
+        """
+
+        if self.db is None:
+            return None
+
+        if image['id'] not in self.image_id_to_annotations:
+            return []
+
+        class_ids = []
+        image_annotations = self.image_id_to_annotations[image['id']]
+        for ann in image_annotations:
+            class_ids.append(ann['category_id'])
+        class_ids = sorted(set(class_ids))
+        class_names = [self.cat_id_to_name[x] for x in class_ids]
+
+        return class_names
+
+# ...class IndexedJsonDb
+
+class SequenceOptions:
+    """
+    Options parameterizing the grouping of images into sequences by time.
+    """
+
+    def __init__(self):
+
+        #: Images separated by <= this duration will be grouped into the same sequence.
+        self.episode_interval_seconds = 60.0
+
+        #: How to handle invalid datetimes: 'error' or 'none'
+        self.datetime_conversion_failure_behavior = 'none'
+
+        #: Enable additional debug output
+        self.verbose = False
+
+
+#%% Functions
+
+def write_object_with_serialized_datetimes(d,json_fn):
+    """
+    Writes the object [d] to the .json file [json_fn] with a standard approach
+    to serializing Python datetime objects.
+
+    Args:
+        d (obj): the object to write, typically a dict
+        json_fn (str): the output filename
+    """
+
+    # This writes datetimes as:
+    #
+    # 2022-12-31T09:52:50
+    def json_serialize_datetime(obj):
+        if isinstance(obj, (datetime.datetime, datetime.date)):
+            return obj.isoformat()
+        raise TypeError('Object {} (type {}) not serializable'.format(
+            str(obj),type(obj)))
+
+    with open(json_fn,'w') as f:
+        json.dump(d,f,indent=1,default=json_serialize_datetime)
+
+
+def parse_datetimes_from_cct_image_list(images,
+                                        conversion_failure_behavior='error',
+                                        verbose=False):
+    """
+    Given the "images" field from a COCO camera traps dictionary, converts all
+    string-formatted datetime fields to Python datetimes, making reasonable assumptions
+    about datetime representations.  Modifies [images] in place.
+
+    Args:
+        images (list): a list of dicts in CCT images format
+        conversion_failure_behavior (str, optional): determines what happens on a failed
+            conversion; can be "error" (raise an error), "str" (leave as a string), or
+            "none" (convert to None)
+        verbose (bool, optional): enable additional debug output
+
+    Returns:
+        images: the input list, with datetimes converted (after modifying in place)
+    """
+
+    assert isinstance(images,list)
+
+    print('Parsing datetimes from CCT image list...')
+
+    for im in tqdm(images):
+
+        if 'datetime' not in im:
+            continue
+        if isinstance(im['datetime'],datetime.datetime):
+            continue
+        try:
+            dt = dateutil.parser.parse(im['datetime'])
+            im['datetime'] = dt
+        except Exception as e:
+            s = 'could not parse datetime {} from {}: {}'.format(
+                str(im['datetime']),im['file_name'],str(e))
+            if conversion_failure_behavior == 'error':
+                raise ValueError(s)
+            elif conversion_failure_behavior == 'str':
+                if verbose:
+                    print('Warning: {}'.format(s))
+                pass
+            elif conversion_failure_behavior == 'none':
+                if verbose:
+                    print('Warning: {}'.format(s))
+                im['datetime'] = None
+
+    # ...for each image
+
+    return images
+
+# ...def parse_datetimes_from_cct_image_list(...)
+
+
+def parse_datetimes_from_cct_dict(d,conversion_failure_behavior='error'):
+    """
+    Given a COCO camera traps dictionary that may just have been loaded from file,
+    converts all string-formatted datetime fields to Python datetimes, making
+    reasonable assumptions about datetime representations.  Modifies [d] in place
+    if [d] is supplied as a dict
+
+    Args:
+        d (dict or str): a dict in CCT format or a filename pointing to a CCT .json file
+        conversion_failure_behavior (str, optional): determines what happens on a failed
+            conversion; can be "error" (raise an error), "str" (leave as a string), or
+            "none" (convert to None)
+
+    Returns:
+        dict: the CCT dict with converted datetimes.
+    """
+
+    if isinstance(d,str):
+        assert os.path.isfile(d), 'Could not find .json file {}'.format(d)
+        with open(d,'r') as f:
+            d = json.load(f)
+
+    images = d['images']
+
+    # Modifies in place
+    _ = parse_datetimes_from_cct_image_list(images)
+
+    return d
+
+# ...def parse_datetimes_from_cct_dict(...)
+
+
+def create_sequences(image_info,options=None):
+    """
+    Synthesizes episodes/sequences/bursts for the images in [image_info].
+
+    Modifies [image_info] in place, populating the 'seq_id', 'seq_num_frames', and 'frame_num'
+    fields for each image.
+
+    Args:
+        image_info (str, dict, or list): a dict in CCT format, a CCT .json file, or just the
+            'images' component of a CCT dataset (a list of dicts with  fields 'file_name' (str),
+            'datetime' (datetime), and 'location' (str)).
+        options (SequenceOptions, optional): options parameterizing the assembly of images into
+            sequences; see the SequenceOptions class for details.
+
+    Returns:
+        image_info: if [image_info] is passed as a list, returns the list, otherwise returns
+        a CCT-formatted dict.
+    """
+
+    if options is None:
+        options = SequenceOptions()
+
+    to_return = None
+
+    if isinstance(image_info,list):
+        to_return = image_info
+
+    elif isinstance(image_info,str):
+        print('Reading image information from {}'.format(image_info))
+        with open(image_info,'r') as f:
+            d = json.load(f)
+            to_return = d
+            image_info = d['images']
+
+    elif isinstance(image_info,dict):
+        to_return = image_info
+        image_info = image_info['images']
+
+    else:
+        raise ValueError('Unrecognized type for [image_info]')
+
+    # Modifies the images in place
+    _ = parse_datetimes_from_cct_image_list(image_info,
+            conversion_failure_behavior=options.datetime_conversion_failure_behavior,
+            verbose=options.verbose)
+
+    n_invalid_datetimes = 0
+    for im in image_info:
+        if not isinstance(im['datetime'],datetime.datetime):
+            assert im['datetime'] is None, 'At this point, datetimes should be valid or None'
+            n_invalid_datetimes += 1
+    if n_invalid_datetimes > 0:
+        print('Warning: {} of {} images have invalid datetimes'.format(
+            n_invalid_datetimes,len(image_info)))
+
+    # Find all unique locations
+    locations = set()
+    for im in image_info:
+        locations.add(im['location'])
+
+    print('Found {} locations'.format(len(locations)))
+    locations = list(locations)
+    locations.sort()
+
+    all_sequences = set()
+
+    # i_location = 0; location = locations[i_location]
+    for i_location,location in tqdm(enumerate(locations),total=len(locations)):
+
+        images_this_location = [im for im in image_info if im['location'] == location]
+
+        # Sorting datetimes fails when there are None's in the list.  So instead of sorting datetimes
+        # directly, sort tuples with a boolean for none-ness, then the datetime itself.
+        #
+        # https://stackoverflow.com/questions/18411560/sort-list-while-pushing-none-values-to-the-end
+        sorted_images_this_location = sorted(images_this_location,
+                                             key = lambda im: (im['datetime'] is None,im['datetime']))
+
+        sequence_id_to_images_this_location = defaultdict(list)
+
+        current_sequence_id = None
+        next_frame_number = 0
+        next_sequence_number = 0
+        previous_datetime = None
+
+        # previous_datetime = sorted_images_this_location[0]['datetime']
+        # im = sorted_images_this_location[1]
+        for im in sorted_images_this_location:
+
+            invalid_datetime = False
+
+            if previous_datetime is None:
+                delta = None
+            elif im['datetime'] is None:
+                invalid_datetime = True
+            else:
+                delta = (im['datetime'] - previous_datetime).total_seconds()
+
+            # Start a new sequence if necessary, including the case where this datetime is invalid
+            if (delta is None) or (delta > options.episode_interval_seconds) or (invalid_datetime):
+                next_frame_number = 0
+                current_sequence_id = 'location_{}_sequence_index_{}'.format(
+                    location,str(next_sequence_number).zfill(5))
+                next_sequence_number = next_sequence_number + 1
+                assert current_sequence_id not in all_sequences
+                all_sequences.add(current_sequence_id)
+
+            im['seq_id'] = current_sequence_id
+            im['seq_num_frames'] = None
+            im['frame_num'] = next_frame_number
+            sequence_id_to_images_this_location[current_sequence_id].append(im)
+            next_frame_number = next_frame_number + 1
+
+            # If this was an invalid datetime, this will record the previous datetime
+            # as None, which will force the next image to start a new sequence.
+            previous_datetime = im['datetime']
+
+        # ...for each image in this location
+
+        # Fill in seq_num_frames
+        for seq_id in sequence_id_to_images_this_location.keys():
+            assert seq_id in sequence_id_to_images_this_location
+            images_this_sequence = sequence_id_to_images_this_location[seq_id]
+            assert len(images_this_sequence) > 0
+            for im in images_this_sequence:
+                im['seq_num_frames'] = len(images_this_sequence)
+
+    # ...for each location
+
+    print('Created {} sequences from {} images'.format(len(all_sequences),len(image_info)))
+
+    return to_return
+
+# ...def create_sequences(...)
+
+
+#%% Test drivers
+
+if False:
+
+    pass
+
+    #%%
+
+    fn = r'g:\temp\test.json'
+    d = parse_datetimes_from_cct_dict(fn,conversion_failure_behavior='error')