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