common.py 1.91 KB
Newer Older
1 2
import os

3 4
import cv2

5 6 7 8 9 10 11 12 13 14
DATA_DIR = "./data"
DATABASE_LOC = os.path.join(DATA_DIR, "faces.pickle")
CASCADE_DIR = "./cascades"
# Output location for pickle database files
OUTPUT_DIR = "./output"

RES_DIRECTORY = "./res"
# Directory for the face detection model.
FACE_DETECTION_MODEL_DIR = os.path.join(RES_DIRECTORY, "face_detection_model")
EMBEDDINGS_PROCESSOR_LOC = os.path.join(RES_DIRECTORY, "openface_nn4.small2.v1.t7")
15 16 17 18 19 20 21 22 23 24


def display_frame(frame):
    """Displays the frame to the user."""
    cv2.imshow("Frame", frame)


def start_video_stream(camera: int):
    """Starts the video stream and returns the created stream.
    Also waits for the video stream to open before returning it."""
Kyle Anderson's avatar
Kyle Anderson committed
25
    video_stream = cv2.VideoCapture(0)
26
    return video_stream
27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64


def load_cascade(cascade_loc: str) -> cv2.CascadeClassifier:
    """
    Opens the cascade classifier at the given path.
    :param cascade_loc: The file location of the cascade.
    :return:The CascadeClassifier class.
    """
    return cv2.CascadeClassifier(cascade_loc)


def load_detector(proto_path: str, model_path: str):
    """
    Loads the caffe detector with the given proto text file and the model file.
    :param proto_path: The path location of the prototext file.
    :param model_path: The path to the caffe model.
    :return: The detector.
    """
    return cv2.dnn.readNetFromCaffe(proto_path, model_path)


def load_embedding_model(model_path: str):
    """
    Loads the torch embedding model at the given location.
    :param model_path: The path to the model.
    :return: The embedding model
    """
    return cv2.dnn.readNetFromTorch(model_path)


CAFFE_MODEL_NAME = "res10_300x300_ssd_iter_140000.caffemodel"
PROTOTXT_NAME = "deploy.prototxt"


def load_detector_from_dir(detector_dir: str):
    prototxt: str = os.path.join(detector_dir, PROTOTXT_NAME)
    caffe_model: str = os.path.join(detector_dir, CAFFE_MODEL_NAME)
    return load_detector(prototxt, caffe_model)