crop.py

import os
from PIL import Image
from tqdm import tqdm
import torch
import torchvision.transforms as transforms


def center_crop_images(input_folder, output_folder, crop_size=(512, 512)):
    # Create output folder if it doesn't exist
    os.makedirs(output_folder, exist_ok=True)

    # Define the center crop transform
    crop_transform = transforms.CenterCrop(crop_size)

    # Get list of all image files in the input folder
    image_files = [
        f
        for f in os.listdir(input_folder)
        if f.endswith((".png", ".jpg", ".jpeg", ".bmp", ".gif"))
    ]

    # Loop over each image file with progress bar
    for image_file in tqdm(image_files, desc="Center cropping images"):
        img_path = os.path.join(input_folder, image_file)
        try:
            # Open the image
            img = Image.open(img_path)

            # Apply the center crop transform
            img_cropped = crop_transform(img)

            # Save the cropped image to the output folder
            img_cropped.save(os.path.join(output_folder, image_file))

        except Exception as e:
            print(f"Error processing {img_path}: {e}")


# Usage
input_folder = "/cm/shared/ninhnq3/datasets/cyclegan-rain/testB"
output_folder = input_folder
center_crop_images(input_folder, output_folder)