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/DefectFill/data/xray_PCB", # 
    # "./engine/DefectFill/data/xray_PCB/train/defective_masks/xray_die"
    output_dir="./engine/DefectFill/data/xray_PCB_dataset_512"
)