from ultralytics import YOLO
import cv2
import os
import threading
import time
from PIL import Image

Image.MAX_IMAGE_PIXELS = None

NUM_THREADS = 1

ASPECT_RATIOS = [(1536, 1536), (1280, 1536), (1536, 1280), (1024, 1536), (1536, 1024), (1024, 768), (768, 1024), (1024, 1280), (1280, 1024),(768,1368),(1368,768),(768,1344),(1344,768),(896,1152),(1152,896),(1536,640),(640,1536),(1216,832),(832,1216)]

#ASPECT_RATIOS = [(1536, 1536)]

# Constants
INPUT_FOLDER = r"C:\Users\King\Pictures\big_expression_dataset_raw"
OUTPUT_FOLDER = r"C:\Users\King\Pictures\big_expression_dataset_raw_Cropped"
YOLO_DETECTED_FOLDER = r"C:\Users\King\Pictures\YoloDetected"
SAVE_TO_YOLO_DETECTED_FOLDER = True

# Thread-safe variables and lock
counter_lock = threading.Lock()
image_processed_counter = 0
write_lock = threading.Lock()
model_lock = threading.Lock()
start_time = time.time()

def resize_bbox_to_dimensions(bbox, target_width, target_height, img_width, img_height):
    x1, y1, x2, y2 = bbox
    current_width = x2 - x1
    current_height = y2 - y1
    desired_aspect_ratio = target_width / target_height
    current_aspect_ratio = current_width / current_height
    
    if current_aspect_ratio < desired_aspect_ratio:
        new_width = desired_aspect_ratio * current_height
        x1 -= (new_width - current_width) / 2
        x2 += (new_width - current_width) / 2
    elif current_aspect_ratio > desired_aspect_ratio:
        new_height = current_width / desired_aspect_ratio
        y1 -= (new_height - current_height) / 2
        y2 += (new_height - current_height) / 2

    x1 = max(x1, 0)
    y1 = max(y1, 0)
    x2 = min(x2, img_width)
    y2 = min(y2, img_height)

    return [int(x1), int(y1), int(x2), int(y2)]

created_folders = set()

def process_files(filelist):
    global image_processed_counter, start_time
    with counter_lock:
        model = YOLO("yolov10x.pt")
    for filename in filelist:
        try:
            img_path = os.path.join(INPUT_FOLDER, filename)
            image = cv2.imread(img_path)
            if image is None:
                raise ValueError(f"Could not read image {filename}")
            img_width, img_height = image.shape[1], image.shape[0]

            with model_lock:
                results = model(img_path)
            
            detections = results[0].boxes

            person_detected = [box for box in detections if model.names[int(box.cls)] == 'person']
            if person_detected:
                x1, y1, x2, y2 = person_detected[0].xyxy[0].tolist()
                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)

                for target_width, target_height in ASPECT_RATIOS:
                    new_x1, new_y1, new_x2, new_y2 = resize_bbox_to_dimensions([x1, y1, x2, y2], target_width, target_height, img_width, img_height)
                    new_x1, new_y1 = max(new_x1, 0), max(new_y1, 0)
                    new_x2, new_y2 = min(new_x2, img_width), min(new_y2, img_height)
                    cropped_img = image[new_y1:new_y2, new_x1:new_x2]

                    aspect_folder = os.path.join(OUTPUT_FOLDER, f"{target_width}x{target_height}")
                    with counter_lock:
                        if aspect_folder not in created_folders:
                            if not os.path.exists(aspect_folder):
                                os.makedirs(aspect_folder)
                            created_folders.add(aspect_folder)

                    save_path = os.path.join(aspect_folder, os.path.splitext(filename)[0] + '.png')
                    cv2.imwrite(save_path, cropped_img)

                if SAVE_TO_YOLO_DETECTED_FOLDER:
                    if not os.path.exists(YOLO_DETECTED_FOLDER):
                        os.makedirs(YOLO_DETECTED_FOLDER)
                    yolo_detected_img_path = os.path.join(YOLO_DETECTED_FOLDER, os.path.splitext(filename)[0] + '.png')
                    cv2.imwrite(yolo_detected_img_path, results[0].plot())
                    
        except Exception as e:
            print(f"An error occurred while processing {filename}: {e}")
        finally:
            with counter_lock:
                image_processed_counter += 1
                current_time = time.time()
                elapsed_time = current_time - start_time
                average_speed_ms = (elapsed_time / image_processed_counter) * 1000
                print(f"Processed {image_processed_counter}/{len(all_files)} images. Remaining: {len(all_files) - image_processed_counter}. Average processing speed: {average_speed_ms:.2f} ms/image")

def split_list(lst, n):
    avg = len(lst) // n
    remainder = len(lst) % n
    start = 0
    results = []
    for i in range(n):
        end = start + avg + (1 if i < remainder else 0)
        results.append(lst[start:end])
        start = end
    return results

print("Reading all image files may take a while, please wait...")
all_files = os.listdir(INPUT_FOLDER)

split_files = split_list(all_files, NUM_THREADS)
print("All files read")
threads = []

for sublist in split_files:
    thread = threading.Thread(target=process_files, args=(sublist,))
    thread.start()
    threads.append(thread)

for thread in threads:
    thread.join()

print("Processing complete!")