preprocess_final.py

import cv2
import pandas as pd
from PIL import Image
import os
from mtcnn_cv2 import MTCNN

def is_rgba(image_path):
    """
    Checks if image is_rgba, if unable to open, attempts to convert video to first frame
    Parameters:
        * image_path: path to image file
    Returns:
        * boolean: True if image is RGBA format
    """
    try:
        img = Image.open(image_path)
        return img.mode == 'RGBA'
    except IOError:
        # If can't open the image: (possibly because it is a video or some other format)
        return False

def convert_to_RGB(data_folder):
    """
    Parameters:
        * csv_path: e.g. "./purdue-face-recognition/train.csv"
        * data_folder: e.g. "./train_big/train_big/"
    Converts images in data_folder file to RGBA
    """

    # Convert all rgba images to rgb
    curr = 1
    for filename in os.listdir(data_folder):
        print(f"Converting... {filename} {curr}/{len(os.listdir(data_folder))}")
        curr += 1
        filename = data_folder + filename

        if is_rgba(filename):
            rgba_image = Image.open(filename)
            rgb_image = Image.new("RGB", rgba_image.size, (255, 255, 255))
            rgb_image.paste(rgba_image, mask=rgba_image.split()[3])
            rgb_image.save(filename) # Saves in same location as was found

def convert_first_face_mtcnn(img_filename, old_folder, new_folder):
    """
    Parameters:
        * str:img_filename: e.g. 1.jpg
        * str:old_folder: e.g. ./train_big/train_big/
        * str:new_folder: e.g. ./train_mtcnn/

    Returns:
        * bool: True if face found
    """
    file_path = old_folder + img_filename
    new_file_path = new_folder + img_filename
    detector = MTCNN()

    try:
        img = cv2.imread(file_path)
        rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    except:
        cap = cv2.VideoCapture(file_path)
        ret, first_frame = cap.read()
        cap.release()
        if ret:
            cv2.imwrite(file_path, first_frame)
            img = cv2.imread(file_path)
            rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        else:
            return False
            print(f'ERROR: {img_filename}')

    faces = detector.detect_faces(rgb)
    if len(faces) > 0:
        if len(faces) > 1:
            print(f"{img_filename} = multiple faces")
        (x,y,w,h) = faces[0]['box']
        roi_color = img[y:y+h, x:x+w]
        cv2.imwrite(new_file_path, roi_color)
        return True
    else:
        cv2.imwrite(new_file_path, rgb)
        return False

def convert_train_mtcnn_to_readable_cropped():
    """
    Converts ./train_big/train_big/ to ./train_mtcnn/ and saves image files that have a detected face in new_train_mtcnn.csv
    """
    old_folder = "./train_big/train_big/"
    new_folder = "./train_mtcnn/"
    num_files = 10#len(os.listdir(old_folder))
    not_found_list = []

    #convert_to_RGB(old_folder)

    # Convert images to MTCNN
    for file_num in range(num_files):
        filename = str(file_num) + ".jpg"
        print(f"Converting... {filename}")
        is_found = convert_first_face_mtcnn(filename, old_folder, new_folder)
        if not is_found:
            not_found_list.append(filename)
    
    # Keep track of converted images in new_train_mtcnn.csv: this will be used in training
    with open("./purdue-face-recognition-challenge-2024/train.csv", "r") as csv_file:
        df = pd.read_csv(csv_file, delimiter=',')

    filtered_df = df[~df['File Name'].isin(not_found_list)]
    filtered_df.to_csv('new_train_mtcnn.csv', index=False)

def convert_test_mtcnn():
    """
    Converts ./test/test/ to ./test_mtcnn/ cropping to faces when able to
    """
    old_folder = "./test/test/"
    new_folder = "./test_mtcnn/"
    num_files = 10#len(os.listdir(old_folder))

    convert_to_RGB(old_folder)

    for file_num in range(num_files):
        filename = str(file_num) + ".jpg"
        print(f"Converting ... {filename}")
        is_found = convert_first_face_mtcnn(filename, old_folder, new_folder)


def main():
    convert_train_mtcnn_to_readable_cropped()
    convert_test_mtcnn()

if __name__ == "__main__":
    main()