{"cells":[{"metadata":{},"cell_type":"markdown","source":"This script contains tool to detect faces using MTCNN  based on ROI, that were driven from previous MTCNN detection."},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install facenet_pytorch\n!pip install opencv-contrib-python\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"from functools import wraps\n\nfrom functools import wraps\n\nimport cv2\nimport numpy as np\nfrom facenet_pytorch.models.mtcnn import MTCNN\n\n\ndef roi_supported_detector(dict_of_id_to_bbs, bbs_width):\n    def mydecorator(detector_function):\n        @wraps(detector_function)\n        def wrapper(*args, **kwargs):\n            person_id = kwargs['person_id']\n            image = kwargs['img']\n            if person_id in dict_of_id_to_bbs.keys():\n                bb = dict_of_id_to_bbs[person_id][0]\n                margin = 0.5\n                width = abs(bb[0] - bb[2])\n                roi_coords = np.asarray([max(bb[0] - width * margin, 0),\n                                         max(bb[1] - width * margin, 0),\n                                         min(bb[2] + width * margin, image.shape[1]),\n                                         min(bb[3] + width * margin, image.shape[0])]).astype(int)\n                roi = image[roi_coords[1]:roi_coords[3], roi_coords[0]:roi_coords[2], :]\n                kwargs['img'] = roi\n                cv2.imshow(\"img_roi\", roi)\n                result = detector_function(*args, **kwargs)\n                result = list(result)\n                if result[0] is not None:\n                    result[0] = (result[0].reshape(-1, 2) + np.asarray(roi_coords)[:2]).reshape(1, -1)\n                    if len(result) > 2:\n                        result[2] = result[2] + np.asarray(roi_coords)[:2]\n                    new_bb = result[0]\n                    dict_of_id_to_bbs[person_id] = new_bb\n                    bbs_width[person_id] = width\n\n                result = tuple(result)\n            else:\n                result = detector_function(*args, **kwargs)\n                if result[0] is not None:\n                    new_bb = result[0]\n                    dict_of_id_to_bbs[person_id] = new_bb\n                    bb = dict_of_id_to_bbs[person_id][0]\n                    width = abs(bb[0] - bb[2])\n                    bbs_width[person_id] = width\n\n            return result\n\n        return wrapper\n\n    return mydecorator\n\n\nclass Extended_MTCNN(MTCNN):\n    dict_of_id_to_bbs = {}\n    bbs_width = {}\n\n    def __init__(self, *args, **kwargs):\n        super().__init__(args, kwargs)\n\n    @roi_supported_detector(dict_of_id_to_bbs, bbs_width)\n    def detect_with_roi(self, person_id=1, *args, **kwargs):\n        return self.detect(*args, **kwargs)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## And then use this like:"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"cell_type":"code","source":"mtcnn = Extended_MTCNN(device=device,min_face_size=20)\ncap = cv2.VideoCapture(\"path_to_video\")\nwhile ret:\n\n    ret, frame = cap.read()\n    person_id = 1\n    \n    if person_id in mtcnn.dict_of_id_to_bbs.keys():\n        mtcnn.min_face_size = int(mtcnn.bbs_width[person_id] * 0.8)\n        print(mtcnn.min_face_size)\n    boxes, probs, landmarks = mtcnn.detect_with_roi(person_id=person_id,\n                                                    img=cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), landmarks=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"##### Note , that you should know person id in order to select roi correctly"},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}