Spaces:
No application file
No application file
| """ | |
| 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() | |