File size: 2,256 Bytes
c8c00f0 | 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 | import os
from PIL import Image
def preprocess_dataset(input_dir, output_dir, target_size=(512, 512)):
"""
Center crops images to a square, resizes them to 512x512,
and converts them to PNG format.
"""
os.makedirs(output_dir, exist_ok=True)
valid_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff')
for root, _, files in os.walk(input_dir):
for file in files:
if file.lower().endswith(valid_extensions):
# Setup paths
rel_path = os.path.relpath(root, input_dir)
out_folder = os.path.join(output_dir, rel_path)
os.makedirs(out_folder, exist_ok=True)
img_path = os.path.join(root, file)
filename_without_ext = os.path.splitext(file)[0]
save_path = os.path.join(out_folder, f"{filename_without_ext}.png")
with Image.open(img_path) as img:
w, h = img.size
# 1. Calculate center crop box
min_dim = min(w, h)
left = (w - min_dim) // 2
top = (h - min_dim) // 2
right = left + min_dim
bottom = top + min_dim
# 2. Crop to square
img_cropped = img.crop((left, top, right, bottom))
# 3. Resize to target resolution (512x512)
# For masks (binary), use NEAREST; for images, use LANCZOS
if "ground_truth" in root.lower() or "mask" in root.lower():
img_resized = img_cropped.resize(target_size, Image.Resampling.NEAREST)
else:
img_resized = img_cropped.resize(target_size, Image.Resampling.LANCZOS)
# 4. Save as PNG
img_resized.save(save_path, "PNG")
print(f"Processed: {file} -> {save_path}")
# Example Usage:
preprocess_dataset(
input_dir="./engine/DefectDiffu/few-shot-training/vscel/img/tiger-strip",
output_dir="./engine/DefectDiffu/few-shot-training/vcsel_dataset_512"
) |