"""
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(...)
[docs]
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
[docs]
class TestResizeCocoDataset:
"""
Test class for the resize_coco_dataset function.
"""
[docs]
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)
[docs]
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
[docs]
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
[docs]
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.")