import os
import shutil
import string
import random
from retinaface import RetinaFace
import tensorflow as tf
from PIL import Image

# Set the GPU device
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

print(tf.test.gpu_device_name())

#pip install tensorflow-gpu

original_images_folder = r"G:\woman dataset\raw_unprocessed"  # org images to process folder
target_folder_post_processed = r"G:\woman dataset\a1_post_processed"  # processed images here
target_folder_multiple = r"G:\woman dataset\a1_more_faces"  # more than 1 face found pictures
target_folder_0 = r"G:\woman dataset\a1_0_faces"  # no face found pictures
target_folder_low_colors = r"G:\woman dataset\a1_low_colors"  # low colors found pictures
target_downscaled = r"G:\woman dataset\downscaled"  # low colors found pictures
target_corrupted= r"G:\woman dataset\corrupted"

folders_to_create = [original_images_folder, target_folder_post_processed, target_folder_multiple, target_folder_0, target_folder_low_colors,target_corrupted,target_downscaled]

for folder in folders_to_create:
    if not os.path.exists(folder):
        os.makedirs(folder)

def move_image(org_image_path, target_folder):
    if not os.path.exists(org_image_path):
        return False
    image_filename = os.path.basename(org_image_path)
    new_image_path = os.path.join(target_folder, image_filename)
    shutil.move(org_image_path, new_image_path)

def generate_random_name(length=20):
    letters = string.ascii_lowercase
    return ''.join(random.choice(letters) for _ in range(length))

from PIL import Image

def resize_image(image_path, max_dimension, move_folder):
    try:
        img = Image.open(image_path)
        width, height = img.size
        if width > max_dimension or height > max_dimension:
            ratio = min(max_dimension / width, max_dimension / height)
            new_width = int(width * ratio)
            new_height = int(height * ratio)
            img = img.resize((new_width, new_height), resample=Image.LANCZOS)
        return img
    except Exception as e:
        print(f"Error resizing image: {e}")
        try:
            # Create the move folder if it doesn't exist
            if not os.path.exists(move_folder):
                os.makedirs(move_folder)
            # Move the original image to the move folder
            image_filename = os.path.basename(image_path)
            new_path = os.path.join(move_folder, image_filename)
            move_image(image_path,new_path)
            print(f"Original image moved to {new_path}")
        except Exception as move_error:
            print(f"Error moving image: {move_error}")
        return None


def detect_and_move_image(org_image_path, target_folder_post_processed, target_folder_multiple, target_folder_low_colors):
    random_name = generate_random_name()
    resized_image_path = f"temp_{random_name}.jpg"
    resized_image_path = os.path.join(target_folder_post_processed, resized_image_path)

    resized_image = resize_image(org_image_path, 1000, target_folder_low_colors)
    if resized_image is None:
        return None
    resized_image.save(resized_image_path)

    faces = RetinaFace.detect_faces(resized_image_path)
    num_faces = len(faces)

    if num_faces > 1:
        move_image(org_image_path, target_folder_multiple)
    elif num_faces < 1:
        move_image(org_image_path, target_folder_0)
    elif count_colors(resized_image_path) < 20000:
        move_image(org_image_path, target_folder_low_colors)

    move_image(org_image_path, target_folder_post_processed)

    # Delete the temporary resized image
    #os.remove(resized_image_path)

def count_colors(image_path):
    if not os.path.exists(image_path):
        return 0

    image = Image.open(image_path)
    # Convert the image to RGB mode if it's not already
    image = image.convert("RGB")

    # Get the size of the image
    width, height = image.size

    # Create an empty set to store unique colors
    colors = set()

    # Iterate over each pixel in the image
    for x in range(width):
        for y in range(height):
            # Get the RGB values of the pixel
            r, g, b = image.getpixel((x, y))

            # Add the RGB values as a tuple to the set
            colors.add((r, g, b))

    # Return the number of unique colors
    return len(colors)

# Counter variable to keep track of the number of images processed
image_counter = 0

# Iterate through each image in the original images folder
for image_filename in os.listdir(original_images_folder):
    image_path = os.path.join(original_images_folder, image_filename)

    if os.path.isfile(image_path):
        detect_and_move_image(image_path, target_folder_post_processed, target_folder_multiple, target_folder_low_colors)
        image_counter += 1
        print("Total images processed:", image_counter)

# Print the total number of images processed
print("Total number of images processed:", image_counter)