본문 바로가기

Python

Resize and crop images in a specific directory

반응형
from tqdm import tqdm
import logging
import traceback
import argparse
from PIL import Image
import matplotlib.pyplot as plt
import cv2
import random
from pathlib import Path
from utils.functions import read_config, read_json
from IPython.display import display

def resize_and_crop_image(img, resize_dims, crop):
    img = img.resize(resize_dims, resample=Image.BILINEAR)
    img = img.crop(crop)
    return img

cwd = Path.cwd()
cfg_path = cwd.parent / 'config/CVTpro/data.json'
cfg = read_json(path=str(cfg_path))


# Image resize and crop
ori_dims = (cfg['original_image']['w'], cfg['original_image']['h'])
resize_dims = (cfg['image']['w'], cfg['image']['h'] + cfg['image']['top_crop'])
crop = (0, cfg['image']['top_crop'], resize_dims[0], resize_dims[1])
img_aug_params = {'scale_width': resize_dims[0] / ori_dims[0],
                       'scale_height': resize_dims[1] / ori_dims[1],
                       'resize_dims': resize_dims,
                       'crop': crop}

root = Path('/home/dooseop/DATASET/nuscenes/samples')
directories = list(root.glob('*'))

# create a new directory
new_path = root.parent / 'samples_new'
if (new_path.exists() is not True):
    new_path.mkdir(parents=True, exist_ok=True)

# 'nuscenes/samples/FOLDER_NAME
for _, dir in enumerate(directories):

    # if (_ > 1):
    #     continue

    FOLDER_NAME = dir.parts[-1]
    if (FOLDER_NAME == 'LIDAR_TOP'):
        continue

    # check files in current dir
    files = list(dir.glob('*'))

    # create target folder
    new_dir = new_path / FOLDER_NAME
    if (new_dir.exists() is not True):
        new_dir.mkdir(parents=True, exist_ok=True)

    # copy from dir to new_dir
    for _, file in enumerate(tqdm(files, desc=f'{FOLDER_NAME}')):
        file_name = file.name
        file_suffix = file.suffix
        file_stem = file.stem

        # Surround Image
        img = Image.open(file)
        img = resize_and_crop_image(img, resize_dims=img_aug_params['resize_dims'],
                                    crop=img_aug_params['crop'])

        new_file_path = new_dir / file_stem
        img.save(new_file_path.with_suffix('.png'))