diff --git a/src/face.py b/src/face.py new file mode 100644 index 0000000..8fef10f --- /dev/null +++ b/src/face.py @@ -0,0 +1,25 @@ +import cv2 + +class Face(): + + def __init__(self, data, img, name): + self.data = data + self.img = img + self.name = name + + self.feature = None + + self.l2 = None + self.cosine = None + + def set_feature(self, fetaure): + self.feature = feature + + def set_cosine(self, cosine:float): + self.cosine = cosine + + def set_l2(self, l2:float): + self.l2 = l2 + + def save(self, path:str, extension:str = "png"): + cv2.imwrite(f"{path}/{self.name}.{extension}", self.img) \ No newline at end of file diff --git a/src/face_recognition.py b/src/face_recognition.py index 93c53f2..e8e9df3 100644 --- a/src/face_recognition.py +++ b/src/face_recognition.py @@ -1,16 +1,21 @@ import cv2 from recognizer import Recognizer +from face import * -def extract_faces(faces:list, img:list, filename:str): +def extract_faces(faces_data:list, img:list, filename:str): recognizer = Recognizer.get() + res = [] + for i in range(1, len(faces_data[1])): + face_data = faces_data[1][i] - for i in range(1, len(faces)): - face = faces[i] - - croped_img = recognizer.alignCrop(img, face[0]) - cv2.imwrite(f"./data/extracted/{filename}_{i}.png", croped_img) + croped_img = recognizer.alignCrop(img, face_data) + face = Face(face_data, croped_img, f"{filename}_{i}") + + res.append(face) + + return res def is_same_face( diff --git a/src/main.py b/src/main.py index 0464b81..8963a37 100644 --- a/src/main.py +++ b/src/main.py @@ -9,6 +9,7 @@ from face_recognition import * files = list_all_files("./data/raw", ["jpg", "webp"]) # process each files +faces = [] for file_path in files: filename = file_path.split("/")[-1].split(".")[0] @@ -18,12 +19,16 @@ for file_path in files: if img is None: print(f"ERROR: Could not read image at {file_path}") continue - cv2.imshow("image1", img) # detect faces file - faces = detect_face_in_image(img, scale=1) + faces_data = detect_face_in_image(img, scale=1) + print("faces_data: ", faces_data[1]) # extract faces from file - extract_faces(faces, img, filename) + faces += extract_faces(faces_data, img, filename) + +# classify each faces +for face in faces: + face.save("./data/extracted") cv2.destroyAllWindows() \ No newline at end of file