add app files

notebooks/app/FastMtcnn.py

@@ -0,0 +1,79 @@
+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)]