ysir.py

from deepface import DeepFace
import os
from deepface.basemodels import VGGFace
import pandas as pd
import csv


model = VGGFace.load_model()

# Specify the directory path to your folder
folder_path = 'train_small\\train_small\\'

# List to hold image names
image_names = []

# Loop through all files in the folder
for filename in os.listdir(folder_path):
    if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
        image_names.append(filename)

filename = 'output.csv'

with open(filename, mode='w', newline='') as file:
    writer = csv.writer(file)
    
    # Write the header row
    writer.writerow(['Id', 'Category'])
    for i in range(4977):
        try:
            dfs = DeepFace.find(img_path = f'test\\test\\{i}.jpg', db_path = "train_small\\train_small\\", enforce_detection= False, model_name= "VGG-Face", distance_metric= "cosine")
            find = dfs[0].iloc[0].identity
            find = find.split("\\")[2].split("_")[0]
        except (ValueError,IndexError):
            find = "error"
        row = [i, find]
        writer.writerow(row)