From f97f4640a925771644b4a96af1f62a81815d7abe Mon Sep 17 00:00:00 2001 From: uvos Date: Fri, 7 Jun 2024 14:04:07 +0200 Subject: [PATCH] PersonDatasetAssembler: add the option to mach images that do NOT contain the specified person --- PersonDatasetAssembler/PersonDatasetAssembler.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/PersonDatasetAssembler/PersonDatasetAssembler.py b/PersonDatasetAssembler/PersonDatasetAssembler.py index d7086c9..c720547 100755 --- a/PersonDatasetAssembler/PersonDatasetAssembler.py +++ b/PersonDatasetAssembler/PersonDatasetAssembler.py @@ -48,10 +48,11 @@ def extract_video_images(video: cv2.VideoCapture, interval: int = 0): ret = True frame_counter = 0 while ret: + video.set(cv2.CAP_PROP_POS_FRAMES, frame_counter) ret, frame = video.read() - if ret and frame_counter % interval == 0: + if ret: yield frame - frame_counter += 1 + frame_counter += interval def contains_face_match(detector: cv2.FaceDetectorYN, recognizer: cv2.FaceRecognizerSF, image: numpy.ndarray, referance_features: list(), thresh: float) -> bool: @@ -88,6 +89,9 @@ def process_referance(detector: cv2.FaceDetectorYN, recognizer: cv2.FaceRecogniz for image in images: detector.setInputSize([image.shape[1], image.shape[0]]) faces = detector.detect(image)[1] + if faces is None: + print("unable to find face in referance image") + exit(1) image = recognizer.alignCrop(image, faces[0]) features = recognizer.feature(image) out.append(features) @@ -103,6 +107,7 @@ if __name__ == "__main__": parser.add_argument('--match_model', '-m', required=True, help="Path to the onnx recognition model to be used") parser.add_argument('--detect_model', '-d', required=True, help="Path to the onnx detection model to be used") parser.add_argument('--threshold', '-t', default=0.362, type=float, help="match threshold to use") + parser.add_argument('--invert', '-n', action='store_true', help="output files that DONT match") args = parser.parse_args() recognizer = cv2.FaceRecognizerSF.create(model=args.match_model, config="", backend_id=cv2.dnn.DNN_BACKEND_DEFAULT , target_id=cv2.dnn.DNN_TARGET_CPU) @@ -140,7 +145,7 @@ if __name__ == "__main__": else: resized = image score, match = contains_face_match(detector, recognizer, resized, referance_features, args.threshold) - if match: + if match and not args.invert or not match and args.invert: filename = f"{counter:04}.png" cv2.imwrite(os.path.join(args.out, filename), image) counter += 1