Spaces:
No application file
No application file
File size: 5,306 Bytes
b26e93d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 |
"""
Copyright (c) 2024 The D-FINE Authors. All Rights Reserved.
"""
import argparse
import json
import os
from concurrent.futures import ThreadPoolExecutor
from PIL import Image
def resize_image_and_update_annotations(image_path, annotations, max_size=640):
print(f"Processing image: {image_path}")
try:
with Image.open(image_path) as img:
w, h = img.size
if max(w, h) <= max_size:
return annotations, w, h, False # No need to resize
scale = max_size / max(w, h)
new_w = int(w * scale)
new_h = int(h * scale)
print(f"Resizing image to width={new_w}, height={new_h}")
img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
new_image_path = image_path.replace(".jpg", "_resized.jpg")
img.save(new_image_path)
print(f"Resized image saved: {new_image_path}")
print(f"Original size: ({w}, {h}), New size: ({new_w}, {new_h})")
# Update annotations
for ann in annotations:
ann["area"] = ann["area"] * (scale**2)
ann["bbox"] = [coord * scale for coord in ann["bbox"]]
if "orig_size" in ann:
ann["orig_size"] = (new_w, new_h)
if "size" in ann:
ann["size"] = (new_w, new_h)
except Exception as e:
print(f"Error processing {image_path}: {e}")
return None
return annotations, new_w, new_h, True
def resize_images_and_update_annotations(base_dir, subset, max_size=640, num_workers=4):
print(f"Starting to resize images and update annotations for subset: {subset}")
json_file = os.path.join(base_dir, subset, "new_zhiyuan_objv2_{}.json".format(subset))
if not os.path.isfile(json_file):
print(f"Error: JSON file not found at {json_file}")
return
print(f"Loading JSON file: {json_file}")
with open(json_file, "r") as f:
data = json.load(f)
print("JSON file loaded.")
print("Preparing image annotations mapping...")
image_annotations = {img["id"]: [] for img in data["images"]}
for ann in data["annotations"]:
image_annotations[ann["image_id"]].append(ann)
print("Image annotations mapping prepared.")
def process_image(image_info):
image_path = os.path.join(base_dir, subset, image_info["file_name"])
results = resize_image_and_update_annotations(
image_path, image_annotations[image_info["id"]], max_size
)
if results is None:
updated_annotations, new_w, new_h, resized = None, None, None, None
else:
updated_annotations, new_w, new_h, resized = results
return image_info, updated_annotations, new_w, new_h, resized
print(f"Processing images with {num_workers} worker threads...")
with ThreadPoolExecutor(max_workers=num_workers) as executor:
results = list(executor.map(process_image, data["images"]))
print("Image processing completed.")
new_images = []
new_annotations = []
print("Updating image and annotation data...")
for image_info, updated_annotations, new_w, new_h, resized in results:
if updated_annotations is not None:
image_info["width"] = new_w
image_info["height"] = new_h
image_annotations[image_info["id"]] = updated_annotations
if resized:
image_info["file_name"] = image_info["file_name"].replace(".jpg", "_resized.jpg")
new_images.append(image_info)
new_annotations.extend(updated_annotations)
print(f"Total images processed: {len(new_images)}")
print(f"Total annotations updated: {len(new_annotations)}")
new_data = {
"images": new_images,
"annotations": new_annotations,
"categories": data["categories"],
}
new_json_file = json_file.replace(".json", "_resized.json")
print("Saving new training annotations...")
with open(new_json_file, "w") as f:
json.dump(new_data, f)
print(f"New JSON file saved to {new_json_file}")
def parse_arguments():
parser = argparse.ArgumentParser(
description="Resize images and update dataset annotations for both train and val sets."
)
parser.add_argument(
"--base_dir",
type=str,
required=True,
help="Base directory of the dataset, e.g., /data/Objects365/data",
)
parser.add_argument(
"--max_size",
type=int,
default=640,
help="Maximum size for the longer side of the image (default: 640)",
)
parser.add_argument(
"--num_workers",
type=int,
default=4,
help="Number of worker threads for parallel processing (default: 4)",
)
args = parser.parse_args()
return args
def main():
args = parse_arguments()
base_dir = args.base_dir
max_size = args.max_size
num_workers = args.num_workers
subsets = ["train", "val"]
for subset in subsets:
print(f"Processing subset: {subset}")
resize_images_and_update_annotations(
base_dir=base_dir, subset=subset, max_size=max_size, num_workers=num_workers
)
print("All subsets processed.")
if __name__ == "__main__":
main()
|