diff --git a/face/ArcFace.py b/face/ArcFace.py new file mode 100644 index 0000000..e69de29 diff --git a/face/face_embedding.py b/face/face_embedding.py new file mode 100644 index 0000000..e69de29 diff --git a/face/main.py b/face/main.py new file mode 100644 index 0000000..9f2db69 --- /dev/null +++ b/face/main.py @@ -0,0 +1,121 @@ + +from yoloface import face_detector +from util.distance import findEuclideanDistance +from itertools import combinations +from util.VideoCapture import VideoCapture +from util.plot import * +import torch +from similarImageFinder_functions import * +from verifier import FaceVerifier, VGGFace, ArcFace + + +class SimilarImageFinder: + def __init__(self, detector, verifier, video_name, video_path, period=30, video_capture=True): + self.video_name = video_name + if video_capture: + # 영상 -> 이미지 + VideoCapture(period, video_name, video_path) + # 같은 프레임의 이미지를 중에서 적어도 2번 이상 사람이 감지된 프레임만 추출해서 data frame 생성 + df = self._make_df_detection_by_detector(detector, video_name) + else: + df = pd.read_parquet(f'./dataset/{video_name}/detection_data/yoloface_data.parquet') + # 유사한 이미지 탐색 + similarity_list = self.find_similarity_images(df, verifier) + print(similarity_list) + # 이미지 출력 + plot_by_cv2(similarity_list, video_name, 30) + + + def _make_df_detection_by_detector(self,detector, video_name): + """ + Face Detection 데이터 프레임을 만드는 함수입니다. + :param: detector ['opencv', 'ssd', 'mtcnn', 'retinaface', 'yoloface] + :return: + """ + if detector == 'yoloface': + try: + torch.tensor(1).cuda() + # de + yolo_detector = face_detector.YoloDetector(target_size=720, gpu=0, min_face=90) + print('gpu 사용') + df_detection = detect_images_by_gpu(yolo_detector, video_name, batch_size=16) + except: + yolo_detector = face_detector.YoloDetector(target_size=720, gpu=-1, min_face=90) + print('cpu 사용, 시간이 오래걸리기 때문에 gpu를 이용하는 것을 권장합니다.') + df_detection = detect_images_by_cpu(yolo_detector, video_name) + else: + print('아직 미구현') + # frame_num 순으로 정렬(시간 순으로 정렬됨) + df_detection = df_detection.sort_values(by='frame_num') + # filter 된 dataframe + df_filtered = filter_df(df_detection) + # 데이터프레임 저장 + df_filtered.to_parquet(f'./dataset/{video_name}/detection_data/yoloface_data.parquet',engine='pyarrow') + + return df_filtered + + + def find_similarity_images(self,df_filtered, verifier='VGGFace'): + similarity_list = [] + if verifier == 'VGGFace': # Face Verification 을 위한 VGGFace Model 불러오기 + model = VGGFace.loadModel() + elif verifier == 'ArcFace': + model = ArcFace.loadModel() + + for frame in df_filtered.index.unique(): + df_temp = df_filtered.loc[frame] + # 감지된 사람수 + detect_person_num = df_temp['detect_person_num'].iloc[0] + # 영상 번호 리스트 + video_num_list = list(df_temp['video_num']) + # 선택된 영상 번호 + selected_video = None + # 거리 + dis_min = 100000000 + + for selected_video_nums in list(combinations(video_num_list, 2)): + videonum1, videonum2 = selected_video_nums[0], selected_video_nums[1] + boxes = list(df_temp.loc[(df_temp['video_num'] == videonum1) | (df_temp['video_num'] == videonum2)]['boxes']) + imgs_path = [f'./dataset/{self.video_name}/frame/{frame}/{videonum1}.jpg', + f'./dataset/{self.video_name}/frame/{frame}/{videonum2}.jpg'] + + # 2개의 이미지에서 가장 큰 얼굴 간의 비율, face verification 을 위한 crop face image + area_fraction, crop_face1, crop_face2 = get_max_area_fraction_and_crop_faces(boxes, imgs_path) + + # #Verified Test : 검증 확인 + # verified = FaceVerifier.verify(crop_face1, crop_face2, model)['verified'] + # #check not verified face + # if not verified: plot_crop_face(crop_face1, crop_face2) + # # check verified face + # if verified: plot_crop_face(crop_face1, crop_face2) + + # compare area_fraction and face recognition + if 0.8 < area_fraction < 1.2 and FaceVerifier.verify(crop_face1, crop_face2, model)['verified']: + landmarks = list( + df_temp.loc[(df_temp['video_num'] == videonum1) | (df_temp['video_num'] == videonum2)]['landmarks']) + + dis = findEuclideanDistance(landmarks[0], landmarks[1], detect_person_num) + + if dis < dis_min: + dis_min = dis + selected_video = selected_video_nums + + if selected_video is not None: + similarity_list.append((frame, selected_video, dis_min)) + + similarity_list = sorted(similarity_list, key=lambda x: x[2]) + + return similarity_list + + +if __name__ == '__main__': + video_path = '/content/drive/MyDrive/HAI/HAI_Project/dataset' + # video capture 생략하고 싶으면 False + # SimilarImageFinder('yoloface', 'idle_tomboy3', video_path, period=30, video_capture=True) + SimilarImageFinder('yoloface', 'ArcFace', 'tomboy1', video_path, period=30, video_capture=False) + +if __name__ == '__main__': + video_path = '/content/drive/MyDrive/HAI/HAI_Project/dataset' + # video capture 생략하고 싶으면 False + # SimilarImageFinder('yoloface', 'idle_tomboy3', video_path, period=30, video_capture=True) + SimilarImageFinder('yoloface', 'VGGFace', 'tomboy2', video_path, period=30, video_capture=False) \ No newline at end of file