megadetector 10.0.15__py3-none-any.whl

megadetector/data_management/resize_coco_dataset.py

@@ -0,0 +1,665 @@
+"""
+
+resize_coco_dataset.py
+
+Given a COCO-formatted dataset, resizes all the images to a target size,
+scaling bounding boxes accordingly.
+
+"""
+
+#%% Imports and constants
+
+import os
+import json
+import shutil
+import argparse
+import sys
+
+from collections import defaultdict
+from multiprocessing.pool import Pool, ThreadPool
+from functools import partial
+
+from PIL import Image
+from tqdm import tqdm
+
+from megadetector.utils.path_utils import insert_before_extension
+from megadetector.visualization.visualization_utils import \
+    open_image, resize_image, exif_preserving_save
+from megadetector.utils.ct_utils import make_test_folder
+from megadetector.utils.ct_utils import write_json
+
+
+#%% Functions
+
+def _process_single_image_for_resize(image_data,
+                                     input_folder,
+                                     output_folder,
+                                     target_size,
+                                     correct_size_image_handling,
+                                     unavailable_image_handling,
+                                     no_enlarge_width,
+                                     verbose):
+    """
+    Processes a single image: loads, resizes/copies, updates metadata, and scales annotations.
+
+    [image_data] is a tuple of [im,annotations]
+    """
+
+    assert unavailable_image_handling in ('error','omit'), \
+        f'Illegal unavailable_image_handling {unavailable_image_handling}'
+
+    assert isinstance(image_data,tuple) and len(image_data) == 2
+    assert isinstance(image_data[0],dict)
+    assert isinstance(image_data[1],list)
+    im = image_data[0].copy()
+    annotations_this_image = [ann.copy() for ann in image_data[1]]
+
+    input_fn_relative = im['file_name']
+    input_fn_abs = os.path.join(input_folder, input_fn_relative)
+
+    if not os.path.isfile(input_fn_abs):
+        if unavailable_image_handling == 'error':
+            raise FileNotFoundError('Could not find file {}'.format(input_fn_abs))
+        else:
+            print("Can't find image {}, skipping".format(input_fn_relative))
+            return None, None
+
+    output_fn_abs = os.path.join(output_folder, input_fn_relative)
+    output_dir = os.path.dirname(output_fn_abs)
+    if len(output_dir) > 0:
+        os.makedirs(output_dir, exist_ok=True)
+
+    if verbose:
+        print('Resizing {} to {}'.format(input_fn_abs,output_fn_abs))
+
+    try:
+        pil_im = open_image(input_fn_abs)
+        input_w = pil_im.width
+        input_h = pil_im.height
+    except Exception as e:
+        if unavailable_image_handling == 'error':
+            raise Exception('Could not open image {}: {}'.format(
+                input_fn_relative, str(e)))
+        else:
+            print("Can't open image {}, skipping".format(input_fn_relative))
+            return None, None
+
+    image_is_already_target_size = \
+        (input_w == target_size[0]) and (input_h == target_size[1])
+    if no_enlarge_width and (input_w < target_size[0]):
+        image_is_already_target_size = True
+    preserve_original_size = \
+        (target_size[0] == -1) and (target_size[1] == -1)
+
+    # Do we need to resize, or can we try to get away with a copy?
+    if image_is_already_target_size or preserve_original_size:
+        output_w = input_w
+        output_h = input_h
+        if correct_size_image_handling == 'copy':
+            if input_fn_abs != output_fn_abs: # only copy if src and dst are different
+                 shutil.copyfile(input_fn_abs, output_fn_abs)
+        elif correct_size_image_handling == 'rewrite':
+            exif_preserving_save(pil_im, output_fn_abs)
+        else:
+            raise ValueError(
+                f'Unrecognized value {correct_size_image_handling} for correct_size_image_handling')
+    else:
+        try:
+            pil_im = resize_image(pil_im, target_size[0], target_size[1],
+                                  no_enlarge_width=no_enlarge_width)
+            output_w = pil_im.width
+            output_h = pil_im.height
+            # We've already applied the rotation at the time we loaded the image, so don't
+            # write the orientation tag.
+            exif_preserving_save(pil_im, output_fn_abs, tags_to_exclude=('Orientation',))
+        except Exception as e:
+            if unavailable_image_handling == 'error':
+                raise Exception('Could not resize image {}: {}'.format(
+                    input_fn_relative, str(e)))
+            else:
+                print("Can't resize image {}, skipping".format(input_fn_relative))
+                return None,None
+
+    im['width'] = output_w
+    im['height'] = output_h
+
+    for ann in annotations_this_image:
+
+        if 'bbox' in ann:
+            bbox = ann['bbox']
+            if (output_w != input_w) or (output_h != input_h):
+                width_scale = output_w / input_w
+                height_scale = output_h / input_h
+                bbox = [
+                    bbox[0] * width_scale,
+                    bbox[1] * height_scale,
+                    bbox[2] * width_scale,
+                    bbox[3] * height_scale
+                ]
+            ann['bbox'] = bbox
+
+    # ...for each annotation associated with this image
+
+    return im, annotations_this_image
+
+# ...def _process_single_image_for_resize(...)
+
+
+def resize_coco_dataset(input_folder,
+                        input_filename,
+                        output_folder,
+                        output_filename=None,
+                        target_size=(-1,-1),
+                        correct_size_image_handling='copy',
+                        unavailable_image_handling='error',
+                        n_workers=1,
+                        pool_type='thread',
+                        no_enlarge_width=True,
+                        verbose=False):
+    """
+    Given a COCO-formatted dataset (images in input_folder, data in input_filename), resizes
+    all the images to a target size (in output_folder) and scales bounding boxes accordingly.
+
+    Args:
+        input_folder (str): the folder where images live; filenames in [input_filename] should
+            be relative to [input_folder]
+        input_filename (str): the (input) COCO-formatted .json file containing annotations
+        output_folder (str): the folder to which we should write resized images; can be the
+            same as [input_folder], in which case images are over-written
+        output_filename (str, optional): the COCO-formatted .json file we should generate that refers
+            to the resized images
+        target_size (list or tuple of ints, optional): this should be tuple/list of ints, with length 2 (w,h).
+            If either dimension is -1, aspect ratio will be preserved.  If both dimensions are -1, this means
+            "keep the original size".  If  both dimensions are -1 and correct_size_image_handling is copy, this
+            function is basically a no-op.
+        correct_size_image_handling (str, optional): what to do in the case where the original size
+            already matches the target size.  Can be 'copy' (in which case the original image is just copied
+            to the output folder) or 'rewrite' (in which case the image is opened via PIL and re-written,
+            attempting to preserve the same quality).  The only reason to do use 'rewrite' 'is the case where
+            you're superstitious about biases coming from images in a training set being written by different
+            image encoders.
+        unavailable_image_handling (str, optional): what to do when a file can't be opened.  Can be
+            'error' or 'omit'.
+        n_workers (int, optional): number of workers to use for parallel processing.
+            Defaults to 1 (no parallelization). If <= 1, processing is sequential.
+        pool_type (str, optional): type of multiprocessing pool to use ('thread' or 'process').
+            Defaults to 'thread'. Only used if n_workers > 1.
+        no_enlarge_width (bool, optional): if [no_enlarge_width] is True, and
+            [target width] is larger than the original image width, does not modify the image,
+            but still writes it
+        verbose (bool, optional): enable additional debug output
+
+    Returns:
+        dict: the COCO database with resized images, identical to the content of [output_filename]
+    """
+
+    # Validate arguments
+
+    assert unavailable_image_handling in ('error','omit'), \
+        f'Illegal unavailable_image_handling {unavailable_image_handling}'
+
+    # Read input data
+    with open(input_filename,'r') as f:
+        d = json.load(f)
+
+    # Map image IDs to annotations
+    image_id_to_annotations = defaultdict(list)
+    for ann in d['annotations']:
+        image_id_to_annotations[ann['image_id']].append(ann)
+
+    original_images = d['images']
+
+    # Our worker function will take tuples of images and their
+    # associated annotations
+    image_annotation_tuples = []
+    for im in original_images:
+        if im['id'] not in image_id_to_annotations:
+            annotations_this_image = []
+        else:
+            annotations_this_image = image_id_to_annotations[im['id']]
+        image_annotation_tuple = (im,annotations_this_image)
+        image_annotation_tuples.append(image_annotation_tuple)
+
+    processed_results = []
+
+    if n_workers <= 1:
+
+        for image_annotation_tuple in tqdm(image_annotation_tuples,
+                                           desc="Resizing images sequentially"):
+            result = _process_single_image_for_resize(
+                image_data=image_annotation_tuple,
+                input_folder=input_folder,
+                output_folder=output_folder,
+                target_size=target_size,
+                correct_size_image_handling=correct_size_image_handling,
+                unavailable_image_handling=unavailable_image_handling,
+                no_enlarge_width=no_enlarge_width,
+                verbose=verbose
+            )
+            processed_results.append(result)
+
+    else:
+
+        pool = None
+
+        try:
+
+            assert pool_type in ('process', 'thread'), f'Illegal pool type {pool_type}'
+            selected_pool = ThreadPool if (pool_type == 'thread') else Pool
+
+            print(f'Starting a {pool_type} pool of {n_workers} workers for image resizing')
+            pool = selected_pool(n_workers)
+
+            p_process_image = partial(_process_single_image_for_resize,
+                                       input_folder=input_folder,
+                                       output_folder=output_folder,
+                                       target_size=target_size,
+                                       correct_size_image_handling=correct_size_image_handling,
+                                       unavailable_image_handling=unavailable_image_handling,
+                                       no_enlarge_width=no_enlarge_width,
+                                       verbose=verbose)
+
+            processed_results = list(tqdm(pool.imap(p_process_image, image_annotation_tuples),
+                                        total=len(image_annotation_tuples),
+                                        desc=f"Resizing images with {pool_type} pool"))
+
+        finally:
+            if pool is not None:
+                pool.close()
+                pool.join()
+                print('Pool closed and joined for COCO dataset resizing')
+
+    new_images_list = []
+    new_annotations_list = []
+    for res_im_data, res_annotations in processed_results:
+        if res_im_data is None or res_annotations is None:
+            assert res_annotations is None and res_im_data is None
+            assert unavailable_image_handling == 'omit'
+            continue
+        new_images_list.append(res_im_data)
+        new_annotations_list.extend(res_annotations)
+
+    d['images'] = new_images_list
+    d['annotations'] = new_annotations_list
+
+    if output_filename is not None:
+        write_json(output_filename,d)
+
+    return d
+
+# ...def resize_coco_dataset(...)
+
+
+#%% Interactive driver
+
+if False:
+
+    pass
+
+    #%% Test resizing
+
+    input_folder = 'i:/data/lila/ena24'
+    # input_filename = 'i:/data/lila/ena24.json'
+    input_filename = 'i:/data/lila/ena24-mini.json'
+
+    output_folder = 'i:/data/lila/ena24-resized'
+    output_filename = insert_before_extension(input_filename,'resized')
+
+    target_size = (640,-1)
+
+    correct_size_image_handling = 'rewrite'
+
+    _ = resize_coco_dataset(input_folder=input_folder,
+                            input_filename=input_filename,
+                            output_folder=output_folder,
+                            output_filename=output_filename,
+                            target_size=target_size,
+                            correct_size_image_handling=correct_size_image_handling,
+                            unavailable_image_handling='omit',
+                            n_workers=10,
+                            pool_type='process')
+
+
+    #%% Preview
+
+    from megadetector.visualization import visualize_db
+    options = visualize_db.DbVizOptions()
+    options.parallelize_rendering = True
+    options.viz_size = (640, -1)
+    options.num_to_visualize = 100
+
+    preview_folder = 'i:/data/lila/ena24-resized-preview'
+    html_file,_ = visualize_db.visualize_db(output_filename,
+                                            preview_folder,
+                                            output_folder,options)
+
+
+    from megadetector.utils import path_utils # noqa
+    path_utils.open_file(html_file)
+
+
+#%% Command-line driver
+
+def main():
+    """
+    Command-line driver for resize_coco_dataset
+    """
+
+    parser = argparse.ArgumentParser(
+        description='Resize images in a COCO dataset and scale annotations'
+    )
+    parser.add_argument(
+        'input_folder',
+        type=str,
+        help='Path to the folder containing original images'
+    )
+    parser.add_argument(
+        'input_filename',
+        type=str,
+        help='Path to the input COCO .json file'
+    )
+    parser.add_argument(
+        'output_folder',
+        type=str,
+        help='Path to the folder where resized images will be saved'
+    )
+    parser.add_argument(
+        'output_filename',
+        type=str,
+        help='Path to the output COCO .json file for resized data'
+    )
+    parser.add_argument(
+        '--target_size',
+        type=str,
+        default='-1,-1',
+        help='Target size as "width,height". Use -1 to preserve aspect ratio for a dimension. ' + \
+             'E.g., "800,600" or "1024,-1".'
+    )
+    parser.add_argument(
+        '--correct_size_image_handling',
+        type=str,
+        default='copy',
+        choices=['copy', 'rewrite'],
+        help='How to handle images already at target size'
+    )
+    parser.add_argument(
+        '--n_workers',
+        type=int,
+        default=1,
+        help='Number of workers for parallel processing. <=1 for sequential'
+    )
+    parser.add_argument(
+        '--pool_type',
+        type=str,
+        default='thread',
+        choices=['thread', 'process'],
+        help='Type of multiprocessing pool if n_workers > 1'
+    )
+
+    if len(sys.argv[1:]) == 0:
+        parser.print_help()
+        parser.exit()
+
+    args = parser.parse_args()
+
+    try:
+        target_size_parts = args.target_size.split(',')
+        if len(target_size_parts) != 2:
+            raise ValueError("target_size must have two comma-separated parts (width,height).")
+        parsed_target_size = (int(target_size_parts[0]), int(target_size_parts[1]))
+    except ValueError as e:
+        print(f"Error parsing target_size: {e}")
+        parser.print_help()
+        parser.exit()
+
+    resize_coco_dataset(
+        args.input_folder,
+        args.input_filename,
+        args.output_folder,
+        args.output_filename,
+        target_size=parsed_target_size,
+        correct_size_image_handling=args.correct_size_image_handling,
+        n_workers=args.n_workers,
+        pool_type=args.pool_type
+    )
+    print("Dataset resizing complete")
+
+if __name__ == '__main__':
+    main()
+
+
+#%% Tests
+
+class TestResizeCocoDataset:
+    """
+    Test class for the resize_coco_dataset function.
+    """
+
+    def set_up(self): # noqa
+        self.test_dir = make_test_folder(subfolder='resize_coco_tests')
+
+        self.input_images_dir_seq = os.path.join(self.test_dir, 'input_images_seq')
+        os.makedirs(self.input_images_dir_seq, exist_ok=True)
+
+        self.input_images_dir_par = os.path.join(self.test_dir, 'input_images_par')
+        os.makedirs(self.input_images_dir_par, exist_ok=True)
+
+        self.output_images_dir_seq = os.path.join(self.test_dir, 'output_images_seq')
+        os.makedirs(self.output_images_dir_seq, exist_ok=True)
+
+        self.output_images_dir_par = os.path.join(self.test_dir, 'output_images_par')
+        os.makedirs(self.output_images_dir_par, exist_ok=True)
+
+    def tear_down(self): # noqa
+
+        # Ensure shutil is imported if not already globally in the file
+        # (it is, under '#%% Imports and constants')
+        if hasattr(self, 'test_dir') and os.path.exists(self.test_dir):
+            shutil.rmtree(self.test_dir)
+
+    def _create_dummy_image_and_coco_json(self,
+                                          image_dir,
+                                          json_filename_base="input_coco.json",
+                                          num_images=2,
+                                          original_size=(100, 100),
+                                          num_annotations_per_image=2):
+        coco_data = {
+            "images": [],
+            "annotations": [],
+            "categories": [{"id": 1, "name": "test_category"}]
+        }
+
+        annotation_id_counter = 1
+
+        for i in range(num_images):
+            image_name = f"image_{i}.png"
+            image_path = os.path.join(image_dir, image_name)
+
+            # Create a dummy image
+            try:
+                img = Image.new('RGB', original_size, color='red')
+                img.save(image_path)
+            except Exception as e:
+                # In some environments, font loading for default PIL text might fail.
+                # For a simple color image, this shouldn't be an issue.
+                # If it is, consider a simpler save or pre-creating a tiny PNG.
+                print(f"Warning: Could not create dummy image {image_path}: {e}")
+                # Fallback: create an empty file, though this will fail later steps
+                # open(image_path, 'a').close()
+
+            image_entry = {
+                "id": i + 1,
+                "file_name": image_name, # Filename only, not path
+                "width": original_size[0],
+                "height": original_size[1]
+            }
+            coco_data["images"].append(image_entry)
+
+            for j in range(num_annotations_per_image):
+                annotation_entry = {
+                    "id": annotation_id_counter,
+                    "image_id": image_entry["id"],
+                    "category_id": 1, # Corresponds to "test_category"
+                    # Simple, non-overlapping bbox for testing scaling
+                    "bbox": [10 + j*30, 10 + j*5, 20, 15]
+                }
+                coco_data["annotations"].append(annotation_entry)
+                annotation_id_counter += 1
+
+        json_file_path = os.path.join(self.test_dir, json_filename_base)
+        with open(json_file_path, 'w') as f:
+            json.dump(coco_data, f, indent=1)
+
+        return json_file_path, coco_data
+
+    def test_resize_sequential_vs_parallel(self):
+        """
+        Test driver for sequence vs. parallel COCO dataset resizing.
+        """
+
+        self.set_up()
+
+        try:
+            num_images_to_test = 3
+            original_w, original_h = 120, 80
+            target_w, target_h = 60, 40
+            target_size_test = (target_w, target_h)
+
+            # Sequential run
+            input_json_path_seq, _ = self._create_dummy_image_and_coco_json(
+                image_dir=self.input_images_dir_seq,
+                json_filename_base="input_coco_seq.json",
+                num_images=num_images_to_test,
+                original_size=(original_w, original_h)
+            )
+            output_json_path_seq = os.path.join(self.test_dir, 'output_coco_seq.json')
+
+            print("Test: starting sequential resize (1 worker)...")
+            resize_coco_dataset(
+                input_folder=self.input_images_dir_seq,
+                input_filename=input_json_path_seq,
+                output_folder=self.output_images_dir_seq,
+                output_filename=output_json_path_seq,
+                target_size=target_size_test,
+                n_workers=1
+            )
+            print(f"Test: Sequential resize complete. Output: {output_json_path_seq}")
+
+            # Parallel run
+            # For the parallel run, we use different input/output directories but can reuse the same logic
+            # for creating the dummy dataset structure. The image files will be new.
+            input_json_path_par, _ = self._create_dummy_image_and_coco_json(
+                image_dir=self.input_images_dir_par,
+                json_filename_base="input_coco_par.json",
+                num_images=num_images_to_test,
+                original_size=(original_w, original_h)
+            )
+            output_json_path_par = os.path.join(self.test_dir, 'output_coco_par.json')
+
+            print("Test: Starting parallel resize (2 workers, thread pool)...")
+            resize_coco_dataset(
+                input_folder=self.input_images_dir_par,
+                input_filename=input_json_path_par,
+                output_folder=self.output_images_dir_par,
+                output_filename=output_json_path_par,
+                target_size=target_size_test,
+                n_workers=2, # Using 2 workers for testing parallelism
+                pool_type='thread'
+            )
+            print(f"Test: Parallel resize complete. Output: {output_json_path_par}")
+
+            # Load results
+            with open(output_json_path_seq, 'r') as f:
+                data_seq = json.load(f)
+            with open(output_json_path_par, 'r') as f:
+                data_par = json.load(f)
+
+            # Compare COCO JSON data
+            # Compare images
+            assert len(data_seq['images']) == num_images_to_test
+            assert len(data_seq['images']) == len(data_par['images']), "Number of images differs"
+
+            sorted_images_seq = sorted(data_seq['images'], key=lambda x: x['id'])
+            sorted_images_par = sorted(data_par['images'], key=lambda x: x['id'])
+
+            for img_s, img_p in zip(sorted_images_seq, sorted_images_par, strict=True):
+                assert img_s['id'] == img_p['id'], \
+                    f"Image IDs differ: {img_s['id']} vs {img_p['id']}"
+                # Filenames are generated independently, so we only check structure, not exact name matching
+                # across seq/par runs' inputs, but output structure should be consistent if input
+                # names were e.g. image_0, image_1
+                assert img_s['file_name'] == img_p['file_name']
+                assert img_s['width'] == target_w, \
+                    f"Seq image {img_s['id']} width incorrect"
+                assert img_s['height'] == target_h, \
+                    f"Seq image {img_s['id']} height incorrect"
+                assert img_p['width'] == target_w, \
+                    f"Par image {img_p['id']} width incorrect"
+                assert img_p['height'] == target_h, \
+                    f"Par image {img_p['id']} height incorrect"
+
+            # Compare annotations
+            assert len(data_seq['annotations']) == len(data_par['annotations']), \
+                "Number of annotations differs"
+            # Assuming _create_dummy_image_and_coco_json creates the same number of annotations for each test run
+
+            sorted_anns_seq = sorted(data_seq['annotations'], key=lambda x: x['id'])
+            sorted_anns_par = sorted(data_par['annotations'], key=lambda x: x['id'])
+
+            for ann_s, ann_p in zip(sorted_anns_seq, sorted_anns_par, strict=True):
+                assert ann_s['id'] == ann_p['id'], \
+                    f"Annotation IDs differ: {ann_s['id']} vs {ann_p['id']}"
+                assert ann_s['image_id'] == ann_p['image_id'], \
+                    f"Annotation image_ids differ for ann_id {ann_s['id']}"
+                assert ann_s['category_id'] == ann_p['category_id'], \
+                    f"Annotation category_ids differ for ann_id {ann_s['id']}"
+
+                # Check bbox scaling (example: original width 120, target 60 -> scale 0.5)
+                # Original bbox: [10, 10, 20, 15] -> Scaled: [5, 5, 10, 7.5] (Floats possible)
+                # Need to compare with tolerance or ensure rounding is handled if expecting ints
+                # For this test, let's assume direct comparison works due to simple scaling.
+                # If PIL's resize causes slight pixel shifts affecting precise sub-pixel bbox calculations,
+                # then a tolerance (pytest.approx) would be better.
+                # Given the current resize_coco_dataset logic, it's direct multiplication.
+                for i in range(4):
+                    assert abs(ann_s['bbox'][i] - ann_p['bbox'][i]) < 1e-5, \
+                        f"Bbox element {i} differs for ann_id {ann_s['id']}: {ann_s['bbox']} vs {ann_p['bbox']}"
+
+            # Compare actual image files
+            seq_files = sorted(os.listdir(self.output_images_dir_seq))
+            par_files = sorted(os.listdir(self.output_images_dir_par))
+
+            assert len(seq_files) == num_images_to_test, "Incorrect number of output images (sequential)"
+            assert len(seq_files) == len(par_files), "Number of output image files differs"
+
+            for fname_s, fname_p in zip(seq_files, par_files, strict=True):
+                assert fname_s == fname_p, "Output image filenames differ between seq and par runs"
+                img_s_path = os.path.join(self.output_images_dir_seq, fname_s)
+                img_p_path = os.path.join(self.output_images_dir_par, fname_p)
+
+                with Image.open(img_s_path) as img_s_pil:
+                    assert img_s_pil.size == target_size_test, \
+                        f"Image {fname_s} (seq) has wrong dimensions: {img_s_pil.size}"
+                with Image.open(img_p_path) as img_p_pil:
+                    assert img_p_pil.size == target_size_test, \
+                        f"Image {fname_p} (par) has wrong dimensions: {img_p_pil.size}"
+
+            print("Test test_resize_sequential_vs_parallel PASSED")
+
+        finally:
+            self.tear_down()
+
+    # ...def test_resize_sequential_vs_parallel(...)
+
+# ...class TestResizeCocoDataset
+
+
+def test_resize_coco_dataset_main():
+    """
+    Driver for the TestResizeCocoDataset() class.
+    """
+
+    print("Starting TestResizeCocoDataset main runner...")
+    test_runner = TestResizeCocoDataset()
+    test_runner.test_resize_sequential_vs_parallel()
+    print("TestResizeCocoDataset main runner finished.")