79 lines
3.2 KiB
Python
79 lines
3.2 KiB
Python
import cv2
|
|
from facenet_pytorch import MTCNN
|
|
from PIL import Image
|
|
import torch
|
|
from imutils.video import FileVideoStream
|
|
import cv2
|
|
import time
|
|
from catboost import CatBoostClassifier
|
|
from tqdm.notebook import tqdm
|
|
from deepface import DeepFace
|
|
import pandas as pd
|
|
import pickle
|
|
import numpy as np
|
|
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
|
|
|
|
class FastMTCNN(object):
|
|
"""Fast MTCNN implementation."""
|
|
|
|
def __init__(self, stride, resize=1, *args, **kwargs):
|
|
"""Constructor for FastMTCNN class.
|
|
|
|
Arguments:
|
|
stride (int): The detection stride. Faces will be detected every `stride` frames
|
|
and remembered for `stride-1` frames.
|
|
|
|
Keyword arguments:
|
|
resize (float): Fractional frame scaling. [default: {1}]
|
|
*args: Arguments to pass to the MTCNN constructor. See help(MTCNN).
|
|
**kwargs: Keyword arguments to pass to the MTCNN constructor. See help(MTCNN).
|
|
"""
|
|
self.stride = stride
|
|
self.resize = resize
|
|
self.mtcnn = MTCNN(*args, **kwargs)
|
|
self.catboost_model_usa = CatBoostClassifier()
|
|
self.catboost_model_usa.load_model("../catboost_usa.cbm")
|
|
with open('../model/saved_dictionary.pkl', 'rb') as f:
|
|
self.name_usa = pickle.load(f)
|
|
|
|
def __call__(self, frames):
|
|
"""Detect faces in frames using strided MTCNN."""
|
|
if self.resize != 1:
|
|
frames = [
|
|
cv2.resize(f, (int(f.shape[1] * self.resize), int(f.shape[0] * self.resize)))
|
|
for f in frames
|
|
]
|
|
|
|
boxes, probs = self.mtcnn.detect(frames[::self.stride])
|
|
dicter3 = {}
|
|
faces = []
|
|
names = {}
|
|
all_x = pd.DataFrame()
|
|
for i, frame in enumerate(frames[::self.stride]):
|
|
box_ind = int(i / self.stride)
|
|
if boxes[box_ind] is None:
|
|
continue
|
|
for box in boxes[box_ind]:
|
|
box = [int(b) for b in box]
|
|
faces.append(frame[box[1]:box[3], box[0]:box[2]])
|
|
image = frame[box[1]:box[3], box[0]:box[2]]
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
embedding = DeepFace.represent(image, model_name='Facenet', enforce_detection=False)
|
|
try:
|
|
ebd = embedding[0]["embedding"]
|
|
except:
|
|
continue
|
|
dicter3[1] = ebd
|
|
data_usa = pd.DataFrame.from_dict(dicter3.items())
|
|
data_usa.rename(columns={0: "id", 1: "embd"}, inplace=True, errors="ignore")
|
|
new_cols = pd.DataFrame(data_usa['embd'].apply(pd.Series))
|
|
df_usa = pd.concat([data_usa, new_cols], axis=1)
|
|
df_usa.drop(["embd"], axis=1, inplace=True, errors="ignore")
|
|
df_usa["id"] = df_usa["id"].apply(lambda x: str(x)[:str(x).find("_")])
|
|
y, X = df_usa["id"], df_usa.drop(["id"], axis=1)
|
|
all_x = pd.concat([all_x, X], axis=0)
|
|
result = self.catboost_model_usa.predict(all_x)
|
|
vals, counts = np.unique(result, return_counts=True)
|
|
return vals[np.argmax(counts)] |