import os
import cv2
from PIL import Image

def is_image_corrupted(path):
    """
    Returns True if the image cannot be opened by cv2 or PIL.
    """
    # ---- Try reading with OpenCV ----
    img = cv2.imread(path)
    if img is None:
        return True  # cv2 failed to decode

    # ---- Try reading with PIL ----
    try:
        with Image.open(path) as im:
            im.verify()  # Validate without fully loading
    except Exception:
        return True  # PIL failed to open / verify

    return False  # Image is valid


def delete_corrupted_images(directory):
    """
    Deletes all unreadable / corrupted images in the directory.
    """
    valid_extensions = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"}

    deleted_count = 0

    for filename in os.listdir(directory):
        path = os.path.join(directory, filename)

        # Skip folders
        if not os.path.isfile(path):
            continue

        # Check valid image extensions
        ext = os.path.splitext(filename)[1].lower()
        if ext not in valid_extensions:
            continue

        # Check corruption
        if is_image_corrupted(path):
            print(f"Deleting corrupted image: {filename}")
            os.remove(path)
            deleted_count += 1

    print(f"\nDone. Deleted {deleted_count} corrupted images.")


# --------- RUN ----------
directory = "/media/NAS/DATASET/wacv_safe_challenge_2025_dataset/physical_fakes/PS-Battles/originals/"   # <-- change this
delete_corrupted_images(directory)