src/process_images.py

import os
from tqdm import tqdm

from src.landmark_detector import MediaPipeLandmarkDetector, DlibLandmarkDetector


def get_images_and_landmarks(image_list, is_dlib):
     
    # Get the list of images in the directory
    image_paths = []
    for file in image_list:
        if file.endswith(".jpg") or file.endswith(".png"):
            image_paths.append(file)
        else:
            print(f"Skipping file: {file}. Not a supported image format. (jpg or png)")

    if len(image_paths) < 2:
        raise ValueError("At least two images are required for morphing.")
        # exit()


    landmarks_list = [] # List of landmarks for each image
    images_list = []    # List of images

    # Initialize the landmark detector
    landmark_detector = DlibLandmarkDetector() if is_dlib else MediaPipeLandmarkDetector()

    print("Generating landmarks for the images...")

    # Detect landmarks for each image
    for image_path in tqdm(image_paths):
        try:
            landmarks, image = landmark_detector.get_landmarks(image_path)
            landmarks_list.append(landmarks)
            images_list.append(image)
        except Exception as e:
            print(f"{e} \nSkipping image: {image_path}\n")
            continue

    if len(landmarks_list) < 2:
        raise ValueError("At least two faces are required for morphing.")
        # exit()

    # if images dont have the same dimensions raise an error
    if len(set([image.shape for image in images_list])) > 1:
        raise ValueError("Images must have the same dimensions for morphing.")
        # exit()

    return images_list, landmarks_list