test.py

import argparse
import sys

import cv2
from face_detector import FaceDetector
from face_recognition import FaceRecognizer
from face_recognition import FaceClustering

# The test module of the face recognition system. This comprises the following workflow:
#   1) Capturing new video frame.
#   2) Run face detection / tracking.
#   3) Extract face embedding and perform face identification (mode "indent") or re-identification (mode "cluster").
#   4) Display face detection / tracking along with the prediction of face identification.

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # The training mode ("ident" to train face identification, "cluster" for face clustering)

    parser.add_argument('--mode', type=str, default="ident")
    # parser.add_argument('--mode', type=str, default="cluster")
    # The video capture input. In case of "None" the default video capture (webcam) is used. Use a filename(s) to read
    # video data from image file (see VideoCapture documentation)
    #parser.add_argument('--video', type=str, default="data/Nancy_Sinatra_test/Nancy_Sinatra_Video.mp4")

    parser.add_argument('--video', type=str, default=None)

    # those that are not in training
    #parser.add_argument('--video', type=str, default="data/Al_Pacino/Al_Pacino_Video.mp4")



    args = parser.parse_args()
    print(args.video)
    if not hasattr(args, "video"):
        args.video = None
    #additional to check accuracy by archit
    if args.video is None:
        expected_label = 'Unknown cam'
    elif args.video.split('/')[-2] == 'Al_Pacino' or args.video.split('/')[-2] == 'Maria_Shriver':
        expected_label = 'Unknown'
    else:
        expected_label = args.video.split('/')[-2]
        expected_label = expected_label.replace('_test','')
        expected_label = expected_label.replace('_train', '')
    print("Expected Label: ", expected_label)

    # Setup OpenCV video capture.
    try:
        if args.video is None:
            try:
                camera = cv2.VideoCapture(0)
                if not camera.isOpened():
                    raise IOError("Cannot open webcam")
                wait_for_frame = 1
            except Exception as e:
                print(f"Error accessing webcam: {e}")
                camera = None
                wait_for_frame = 0
        else:
            try:
                camera = cv2.VideoCapture(args.video)
                if not camera.isOpened():
                    raise IOError(f"Cannot open video file: {args.video}")
                wait_for_frame = 10
            except Exception as e:
                print(f"Error accessing video file: {e}")
                camera = None
                wait_for_frame = 0
    except Exception as e:
        print(f"Unexpected error in video setup: {e}")
        camera = None
        wait_for_frame = 0
        sys.exit(1)
    camera.set(3, 640)
    camera.set(4, 480)

    # Image display
    cv2.namedWindow("Camera")
    cv2.moveWindow("Camera", 0, 0)

    # Prepare face detection, identification, and clustering.
    detector = FaceDetector()
    recognizer = FaceRecognizer(expected_label=expected_label)
    clustering = FaceClustering()

    # The video capturing loop.
    on_track = False
    while True:
        label_str = 'NA'
        confidence_str = 'NA'
        char = cv2.waitKey(wait_for_frame) & 0xFF
        if char == 27:
            # Stop capturing using ESC.
            break

        # Capture new video frame.
        _, frame = camera.read()
        if frame is None:
            print("End of stream")
            break
        # Resize the frame.
        height, width, channels = frame.shape
        if width < 640:
            s = 640.0 / width
            frame = cv2.resize(frame, (int(s * width), int(s * height)))
        # Flip frame if it is live video.
        if args.video is None:
            frame = cv2.flip(frame, 1)

        # Track (or initially detect if required) a face in the current frame.
        face = detector.track_face(frame)

        if face is not None and not on_track:
            # We found a new face that we can track over time.
            on_track = True

            if args.mode == "ident":
                # Face identification: predict identity for the current frame.
                predicted_label, prob, dist_to_prediction = recognizer.predict(face["aligned"])
                label_str = "{}".format(predicted_label)
                confidence_str = "Prob.: {:1.2f}, Dist.: {:1.2f}".format(prob, dist_to_prediction)
            elif args.mode == "cluster":
                # Face clustering: determine cluster for the current frame.
                predicted_label, distances_to_clusters = clustering.predict(face["aligned"])
                label_str = "Cluster {}".format(predicted_label)
                confidence_str = "Dist.: {}".format(distances_to_clusters)
            else:
                print('Please give --mode as ident or cluster')
                sys.exit(1)

        if face is None or face["response"] < detector.tm_threshold:
            # We lost the track of the face visible in the previous frame.
            on_track = False

        # Display annotations for face tracking, identification, and clustering.
        if face is not None:
            face_rect = face["rect"]

            cv2.rectangle(frame, (face_rect[0], face_rect[1]),
                          (face_rect[0] + face_rect[2] - 1, face_rect[1] + face_rect[3] - 1),
                          (0, 255, 0), 2)
            cv2.putText(frame, label_str, (face_rect[0] + face_rect[2] + 10, face_rect[1] + face_rect[3] + 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
            cv2.putText(frame, confidence_str, (face_rect[0] + face_rect[2] + 10, face_rect[1] + face_rect[3] + 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

        cv2.imshow("Camera", frame)