Commit for merge

+import torch
+from torch import nn
+import numpy as np
+import os
+
+from .utils.detect_face import detect_face, extract_face
+
+
+class PNet(nn.Module):
+    """MTCNN PNet.
+    
+    Keyword Arguments:
+        pretrained {bool} -- Whether or not to load saved pretrained weights (default: {True})
+    """
+
+    def __init__(self, pretrained=True):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(3, 10, kernel_size=3)
+        self.prelu1 = nn.PReLU(10)
+        self.pool1 = nn.MaxPool2d(2, 2, ceil_mode=True)
+        self.conv2 = nn.Conv2d(10, 16, kernel_size=3)
+        self.prelu2 = nn.PReLU(16)
+        self.conv3 = nn.Conv2d(16, 32, kernel_size=3)
+        self.prelu3 = nn.PReLU(32)
+        self.conv4_1 = nn.Conv2d(32, 2, kernel_size=1)
+        self.softmax4_1 = nn.Softmax(dim=1)
+        self.conv4_2 = nn.Conv2d(32, 4, kernel_size=1)
+
+        self.training = False
+
+        if pretrained:
+            state_dict_path = os.path.join(os.path.dirname(__file__), 'data/pnet.pt')
+            state_dict = torch.load(state_dict_path)
+            self.load_state_dict(state_dict)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.prelu1(x)
+        x = self.pool1(x)
+        x = self.conv2(x)
+        x = self.prelu2(x)
+        x = self.conv3(x)
+        x = self.prelu3(x)
+        a = self.conv4_1(x)
+        a = self.softmax4_1(a)
+        b = self.conv4_2(x)
+        return b, a
+
+
+class RNet(nn.Module):
+    """MTCNN RNet.
+    
+    Keyword Arguments:
+        pretrained {bool} -- Whether or not to load saved pretrained weights (default: {True})
+    """
+
+    def __init__(self, pretrained=True):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(3, 28, kernel_size=3)
+        self.prelu1 = nn.PReLU(28)
+        self.pool1 = nn.MaxPool2d(3, 2, ceil_mode=True)
+        self.conv2 = nn.Conv2d(28, 48, kernel_size=3)
+        self.prelu2 = nn.PReLU(48)
+        self.pool2 = nn.MaxPool2d(3, 2, ceil_mode=True)
+        self.conv3 = nn.Conv2d(48, 64, kernel_size=2)
+        self.prelu3 = nn.PReLU(64)
+        self.dense4 = nn.Linear(576, 128)
+        self.prelu4 = nn.PReLU(128)
+        self.dense5_1 = nn.Linear(128, 2)
+        self.softmax5_1 = nn.Softmax(dim=1)
+        self.dense5_2 = nn.Linear(128, 4)
+
+        self.training = False
+
+        if pretrained:
+            state_dict_path = os.path.join(os.path.dirname(__file__), 'data/rnet.pt')
+            state_dict = torch.load(state_dict_path)
+            self.load_state_dict(state_dict)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.prelu1(x)
+        x = self.pool1(x)
+        x = self.conv2(x)
+        x = self.prelu2(x)
+        x = self.pool2(x)
+        x = self.conv3(x)
+        x = self.prelu3(x)
+        x = x.permute(0, 3, 2, 1).contiguous()
+        x = self.dense4(x.view(x.shape[0], -1))
+        x = self.prelu4(x)
+        a = self.dense5_1(x)
+        a = self.softmax5_1(a)
+        b = self.dense5_2(x)
+        return b, a
+
+
+class ONet(nn.Module):
+    """MTCNN ONet.
+    
+    Keyword Arguments:
+        pretrained {bool} -- Whether or not to load saved pretrained weights (default: {True})
+    """
+
+    def __init__(self, pretrained=True):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=3)
+        self.prelu1 = nn.PReLU(32)
+        self.pool1 = nn.MaxPool2d(3, 2, ceil_mode=True)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
+        self.prelu2 = nn.PReLU(64)
+        self.pool2 = nn.MaxPool2d(3, 2, ceil_mode=True)
+        self.conv3 = nn.Conv2d(64, 64, kernel_size=3)
+        self.prelu3 = nn.PReLU(64)
+        self.pool3 = nn.MaxPool2d(2, 2, ceil_mode=True)
+        self.conv4 = nn.Conv2d(64, 128, kernel_size=2)
+        self.prelu4 = nn.PReLU(128)
+        self.dense5 = nn.Linear(1152, 256)
+        self.prelu5 = nn.PReLU(256)
+        self.dense6_1 = nn.Linear(256, 2)
+        self.softmax6_1 = nn.Softmax(dim=1)
+        self.dense6_2 = nn.Linear(256, 4)
+        self.dense6_3 = nn.Linear(256, 10)
+
+        self.training = False
+
+        if pretrained:
+            state_dict_path = os.path.join(os.path.dirname(__file__), 'data/onet.pt')
+            state_dict = torch.load(state_dict_path)
+            self.load_state_dict(state_dict)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.prelu1(x)
+        x = self.pool1(x)
+        x = self.conv2(x)
+        x = self.prelu2(x)
+        x = self.pool2(x)
+        x = self.conv3(x)
+        x = self.prelu3(x)
+        x = self.pool3(x)
+        x = self.conv4(x)
+        x = self.prelu4(x)
+        x = x.permute(0, 3, 2, 1).contiguous()
+        x = self.dense5(x.view(x.shape[0], -1))
+        x = self.prelu5(x)
+        a = self.dense6_1(x)
+        a = self.softmax6_1(a)
+        b = self.dense6_2(x)
+        c = self.dense6_3(x)
+        return b, c, a
+
+
+class MTCNN(nn.Module):
+    """MTCNN face detection module.
+
+    This class loads pretrained P-, R-, and O-nets and returns images cropped to include the face
+    only, given raw input images of one of the following types:
+        - PIL image or list of PIL images
+        - numpy.ndarray (uint8) representing either a single image (3D) or a batch of images (4D).
+    Cropped faces can optionally be saved to file
+    also.
+    
+    Keyword Arguments:
+        image_size {int} -- Output image size in pixels. The image will be square. (default: {160})
+        margin {int} -- Margin to add to bounding box, in terms of pixels in the final image. 
+            Note that the application of the margin differs slightly from the davidsandberg/facenet
+            repo, which applies the margin to the original image before resizing, making the margin
+            dependent on the original image size (this is a bug in davidsandberg/facenet).
+            (default: {0})
+        min_face_size {int} -- Minimum face size to search for. (default: {20})
+        thresholds {list} -- MTCNN face detection thresholds (default: {[0.6, 0.7, 0.7]})
+        factor {float} -- Factor used to create a scaling pyramid of face sizes. (default: {0.709})
+        post_process {bool} -- Whether or not to post process images tensors before returning.
+            (default: {True})
+        select_largest {bool} -- If True, if multiple faces are detected, the largest is returned.
+            If False, the face with the highest detection probability is returned.
+            (default: {True})
+        keep_all {bool} -- If True, all detected faces are returned, in the order dictated by the
+            select_largest parameter. If a save_path is specified, the first face is saved to that
+            path and the remaining faces are saved to <save_path>1, <save_path>2 etc.
+        device {torch.device} -- The device on which to run neural net passes. Image tensors and
+            models are copied to this device before running forward passes. (default: {None})
+    """
+
+    def __init__(
+        self, image_size=160, margin=0, min_face_size=20,
+        thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,
+        select_largest=True, keep_all=False, device=None
+    ):
+        super().__init__()
+
+        self.image_size = image_size
+        self.margin = margin
+        self.min_face_size = min_face_size
+        self.thresholds = thresholds
+        self.factor = factor
+        self.post_process = post_process
+        self.select_largest = select_largest
+        self.keep_all = keep_all
+
+        self.pnet = PNet()
+        self.rnet = RNet()
+        self.onet = ONet()
+
+        self.device = torch.device('cpu')
+        if device is not None:
+            self.device = device
+            self.to(device)
+
+    def forward(self, img, save_path=None, return_prob=False):
+        """Run MTCNN face detection on a PIL image or numpy array. This method performs both
+        detection and extraction of faces, returning tensors representing detected faces rather
+        than the bounding boxes. To access bounding boxes, see the MTCNN.detect() method below.
+        
+        Arguments:
+            img {PIL.Image, np.ndarray, or list} -- A PIL image, np.ndarray, or list.
+        
+        Keyword Arguments:
+            save_path {str} -- An optional save path for the cropped image. Note that when
+                self.post_process=True, although the returned tensor is post processed, the saved
+                face image is not, so it is a true representation of the face in the input image.
+                If `img` is a list of images, `save_path` should be a list of equal length.
+                (default: {None})
+            return_prob {bool} -- Whether or not to return the detection probability.
+                (default: {False})
+        
+        Returns:
+            Union[torch.Tensor, tuple(torch.tensor, float)] -- If detected, cropped image of a face
+                with dimensions 3 x image_size x image_size. Optionally, the probability that a
+                face was detected. If self.keep_all is True, n detected faces are returned in an
+                n x 3 x image_size x image_size tensor with an optional list of detection
+                probabilities. If `img` is a list of images, the item(s) returned have an extra 
+                dimension (batch) as the first dimension.
+
+        Example:
+        >>> from facenet_pytorch import MTCNN
+        >>> mtcnn = MTCNN()
+        >>> face_tensor, prob = mtcnn(img, save_path='face.png', return_prob=True)
+        """
+
+        # Detect faces
+        with torch.no_grad():
+            batch_boxes, batch_probs = self.detect(img)
+
+        # Determine if a batch or single image was passed
+        batch_mode = True
+        if not isinstance(img, (list, tuple)) and not (isinstance(img, np.ndarray) and len(img.shape) == 4):
+            img = [img]
+            batch_boxes = [batch_boxes]
+            batch_probs = [batch_probs]
+            batch_mode = False
+
+        # Parse save path(s)
+        if save_path is not None:
+            if isinstance(save_path, str):
+                save_path = [save_path]
+        else:
+            save_path = [None for _ in range(len(img))]
+        
+        # Process all bounding boxes and probabilities
+        faces, probs = [], []
+        for im, box_im, prob_im, path_im in zip(img, batch_boxes, batch_probs, save_path):
+            if box_im is None:
+                faces.append(None)
+                probs.append([None] if self.keep_all else None)
+                continue
+
+            if not self.keep_all:
+                box_im = box_im[[0]]
+
+            faces_im = []
+            for i, box in enumerate(box_im):
+                face_path = path_im
+                if path_im is not None and i > 0:
+                    save_name, ext = os.path.splitext(path_im)
+                    face_path = save_name + '_' + str(i + 1) + ext
+
+                face = extract_face(im, box, self.image_size, self.margin, face_path)
+                if self.post_process:
+                    face = fixed_image_standardization(face)
+                faces_im.append(face)
+
+            if self.keep_all:
+                faces_im = torch.stack(faces_im)
+            else:
+                faces_im = faces_im[0]
+                prob_im = prob_im[0]
+            
+            faces.append(faces_im)
+            probs.append(prob_im)
+    
+        if not batch_mode:
+            faces = faces[0]
+            probs = probs[0]
+
+        if return_prob:
+            return faces, probs
+        else:
+            return faces
+
+    def detect(self, img, landmarks=False):
+        """Detect all faces in PIL image and return bounding boxes and optional facial landmarks.
+
+        This method is used by the forward method and is also useful for face detection tasks
+        that require lower-level handling of bounding boxes and facial landmarks (e.g., face
+        tracking). The functionality of the forward function can be emulated by using this method
+        followed by the extract_face() function.
+        
+        Arguments:
+            img {PIL.Image, np.ndarray, or list} -- A PIL image or a list of PIL images.
+
+        Keyword Arguments:
+            landmarks {bool} -- Whether to return facial landmarks in addition to bounding boxes.
+                (default: {False})
+        
+        Returns:
+            tuple(numpy.ndarray, list) -- For N detected faces, a tuple containing an
+                Nx4 array of bounding boxes and a length N list of detection probabilities.
+                Returned boxes will be sorted in descending order by detection probability if
+                self.select_largest=False, otherwise the largest face will be returned first.
+                If `img` is a list of images, the items returned have an extra dimension
+                (batch) as the first dimension. Optionally, a third item, the facial landmarks,
+                are returned if `landmarks=True`.
+
+        Example:
+        >>> from PIL import Image, ImageDraw
+        >>> from facenet_pytorch import MTCNN, extract_face
+        >>> mtcnn = MTCNN(keep_all=True)
+        >>> boxes, probs, points = mtcnn.detect(img, landmarks=True)
+        >>> # Draw boxes and save faces
+        >>> img_draw = img.copy()
+        >>> draw = ImageDraw.Draw(img_draw)
+        >>> for i, (box, point) in enumerate(zip(boxes, points)):
+        ...     draw.rectangle(box.tolist(), width=5)
+        ...     for p in point:
+        ...         draw.rectangle((p - 10).tolist() + (p + 10).tolist(), width=10)
+        ...     extract_face(img, box, save_path='detected_face_{}.png'.format(i))
+        >>> img_draw.save('annotated_faces.png')
+        """
+
+        with torch.no_grad():
+            batch_boxes, batch_points = detect_face(
+                img, self.min_face_size,
+                self.pnet, self.rnet, self.onet,
+                self.thresholds, self.factor,
+                self.device
+            )
+
+        boxes, probs, points = [], [], []
+        for box, point in zip(batch_boxes, batch_points):
+            box = np.array(box)
+            point = np.array(point)
+            if len(box) == 0:
+                boxes.append(None)
+                probs.append([None])
+                points.append(None)
+            elif self.select_largest:
+                box_order = np.argsort((box[:, 2] - box[:, 0]) * (box[:, 3] - box[:, 1]))[::-1]
+                box = box[box_order]
+                point = point[box_order]
+                boxes.append(box[:, :4])
+                probs.append(box[:, 4])
+                points.append(point)
+            else:
+                boxes.append(box[:, :4])
+                probs.append(box[:, 4])
+                points.append(point)
+        boxes = np.array(boxes)
+        probs = np.array(probs)
+        points = np.array(points)
+
+        if not isinstance(img, (list, tuple)) and not (isinstance(img, np.ndarray) and len(img.shape) == 4):
+            boxes = boxes[0]
+            probs = probs[0]
+            points = points[0]
+
+        if landmarks:
+            return boxes, probs, points
+
+        return boxes, probs
+
+
+def fixed_image_standardization(image_tensor):
+    processed_tensor = (image_tensor - 127.5) / 128.0
+    return processed_tensor
+
+def prewhiten(x):
+    mean = x.mean()
+    std = x.std()
+    std_adj = std.clamp(min=1.0/(float(x.numel())**0.5))
+    y = (x - mean) / std_adj
+    return y

+import torch
+from torch.nn.functional import interpolate
+from torchvision.transforms import functional as F
+from torchvision.ops.boxes import batched_nms
+import cv2
+from PIL import Image
+import numpy as np
+import os
+
+
+def detect_face(imgs, minsize, pnet, rnet, onet, threshold, factor, device):
+    if isinstance(imgs, (np.ndarray, torch.Tensor)):
+        imgs = torch.as_tensor(imgs, device=device)
+        if len(imgs.shape) == 3:
+            imgs = imgs.unsqueeze(0)
+    else:
+        if not isinstance(imgs, (list, tuple)):
+            imgs = [imgs]
+        if any(img.size != imgs[0].size for img in imgs):
+            raise Exception("MTCNN batch processing only compatible with equal-dimension images.")
+        imgs = np.stack([np.uint8(img) for img in imgs])
+
+    imgs = torch.as_tensor(imgs, device=device)
+
+    model_dtype = next(pnet.parameters()).dtype
+    imgs = imgs.permute(0, 3, 1, 2).type(model_dtype)
+
+    batch_size = len(imgs)
+    h, w = imgs.shape[2:4]
+    m = 12.0 / minsize
+    minl = min(h, w)
+    minl = minl * m
+
+    # Create scale pyramid
+    scale_i = m
+    scales = []
+    while minl >= 12:
+        scales.append(scale_i)
+        scale_i = scale_i * factor
+        minl = minl * factor
+
+    # First stage
+    boxes = []
+    image_inds = []
+    all_inds = []
+    all_i = 0
+    for scale in scales:
+        im_data = imresample(imgs, (int(h * scale + 1), int(w * scale + 1)))
+        im_data = (im_data - 127.5) * 0.0078125
+        reg, probs = pnet(im_data)
+    
+        boxes_scale, image_inds_scale = generateBoundingBox(reg, probs[:, 1], scale, threshold[0])
+        boxes.append(boxes_scale)
+        image_inds.append(image_inds_scale)
+        all_inds.append(all_i + image_inds_scale)
+        all_i += batch_size
+
+    boxes = torch.cat(boxes, dim=0)
+    image_inds = torch.cat(image_inds, dim=0).cpu()
+    all_inds = torch.cat(all_inds, dim=0)
+
+    # NMS within each scale + image
+    pick = batched_nms(boxes[:, :4], boxes[:, 4], all_inds, 0.5)
+    boxes, image_inds = boxes[pick], image_inds[pick]
+    
+    # NMS within each image
+    pick = batched_nms(boxes[:, :4], boxes[:, 4], image_inds, 0.7)
+    boxes, image_inds = boxes[pick], image_inds[pick]
+
+    regw = boxes[:, 2] - boxes[:, 0]
+    regh = boxes[:, 3] - boxes[:, 1]
+    qq1 = boxes[:, 0] + boxes[:, 5] * regw
+    qq2 = boxes[:, 1] + boxes[:, 6] * regh
+    qq3 = boxes[:, 2] + boxes[:, 7] * regw
+    qq4 = boxes[:, 3] + boxes[:, 8] * regh
+    boxes = torch.stack([qq1, qq2, qq3, qq4, boxes[:, 4]]).permute(1, 0)
+    boxes = rerec(boxes)
+    y, ey, x, ex = pad(boxes, w, h)
+    
+    # Second stage
+    if len(boxes) > 0:
+        im_data = []
+        for k in range(len(y)):
+            if ey[k] > (y[k] - 1) and ex[k] > (x[k] - 1):
+                img_k = imgs[image_inds[k], :, (y[k] - 1):ey[k], (x[k] - 1):ex[k]].unsqueeze(0)
+                im_data.append(imresample(img_k, (24, 24)))
+        im_data = torch.cat(im_data, dim=0)
+        im_data = (im_data - 127.5) * 0.0078125
+        out = rnet(im_data)
+
+        out0 = out[0].permute(1, 0)
+        out1 = out[1].permute(1, 0)
+        score = out1[1, :]
+        ipass = score > threshold[1]
+        boxes = torch.cat((boxes[ipass, :4], score[ipass].unsqueeze(1)), dim=1)
+        image_inds = image_inds[ipass]
+        mv = out0[:, ipass].permute(1, 0)
+
+        # NMS within each image
+        pick = batched_nms(boxes[:, :4], boxes[:, 4], image_inds, 0.7)
+        boxes, image_inds, mv = boxes[pick], image_inds[pick], mv[pick]
+        boxes = bbreg(boxes, mv)
+        boxes = rerec(boxes)
+
+    # Third stage
+    points = torch.zeros(0, 5, 2, device=device)
+    if len(boxes) > 0:
+        y, ey, x, ex = pad(boxes, w, h)
+        im_data = []
+        for k in range(len(y)):
+            if ey[k] > (y[k] - 1) and ex[k] > (x[k] - 1):
+                img_k = imgs[image_inds[k], :, (y[k] - 1):ey[k], (x[k] - 1):ex[k]].unsqueeze(0)
+                im_data.append(imresample(img_k, (48, 48)))
+        im_data = torch.cat(im_data, dim=0)
+        im_data = (im_data - 127.5) * 0.0078125
+        out = onet(im_data)
+
+        out0 = out[0].permute(1, 0)
+        out1 = out[1].permute(1, 0)
+        out2 = out[2].permute(1, 0)
+        score = out2[1, :]
+        points = out1
+        ipass = score > threshold[2]
+        points = points[:, ipass]
+        boxes = torch.cat((boxes[ipass, :4], score[ipass].unsqueeze(1)), dim=1)
+        image_inds = image_inds[ipass]
+        mv = out0[:, ipass].permute(1, 0)
+
+        w_i = boxes[:, 2] - boxes[:, 0] + 1
+        h_i = boxes[:, 3] - boxes[:, 1] + 1
+        points_x = w_i.repeat(5, 1) * points[:5, :] + boxes[:, 0].repeat(5, 1) - 1
+        points_y = h_i.repeat(5, 1) * points[5:10, :] + boxes[:, 1].repeat(5, 1) - 1
+        points = torch.stack((points_x, points_y)).permute(2, 1, 0)
+        boxes = bbreg(boxes, mv)
+
+        # NMS within each image using "Min" strategy
+        # pick = batched_nms(boxes[:, :4], boxes[:, 4], image_inds, 0.7)
+        pick = batched_nms_numpy(boxes[:, :4], boxes[:, 4], image_inds, 0.7, 'Min')
+        boxes, image_inds, points = boxes[pick], image_inds[pick], points[pick]
+
+    boxes = boxes.cpu().numpy()
+    points = points.cpu().numpy()
+
+    batch_boxes = []
+    batch_points = []
+    for b_i in range(batch_size):
+        b_i_inds = np.where(image_inds == b_i)
+        batch_boxes.append(boxes[b_i_inds].copy())
+        batch_points.append(points[b_i_inds].copy())
+
+    batch_boxes, batch_points = np.array(batch_boxes), np.array(batch_points)
+
+    return batch_boxes, batch_points
+
+
+def bbreg(boundingbox, reg):
+    if reg.shape[1] == 1:
+        reg = torch.reshape(reg, (reg.shape[2], reg.shape[3]))
+
+    w = boundingbox[:, 2] - boundingbox[:, 0] + 1
+    h = boundingbox[:, 3] - boundingbox[:, 1] + 1
+    b1 = boundingbox[:, 0] + reg[:, 0] * w
+    b2 = boundingbox[:, 1] + reg[:, 1] * h
+    b3 = boundingbox[:, 2] + reg[:, 2] * w
+    b4 = boundingbox[:, 3] + reg[:, 3] * h
+    boundingbox[:, :4] = torch.stack([b1, b2, b3, b4]).permute(1, 0)
+
+    return boundingbox
+
+
+def generateBoundingBox(reg, probs, scale, thresh):
+    stride = 2
+    cellsize = 12
+
+    reg = reg.permute(1, 0, 2, 3)
+
+    mask = probs >= thresh
+    mask_inds = mask.nonzero()
+    image_inds = mask_inds[:, 0]
+    score = probs[mask]
+    reg = reg[:, mask].permute(1, 0)
+    bb = mask_inds[:, 1:].type(reg.dtype).flip(1)
+    q1 = ((stride * bb + 1) / scale).floor()
+    q2 = ((stride * bb + cellsize - 1 + 1) / scale).floor()
+    boundingbox = torch.cat([q1, q2, score.unsqueeze(1), reg], dim=1)
+    return boundingbox, image_inds
+
+
+def nms_numpy(boxes, scores, threshold, method):
+    if boxes.size == 0:
+        return np.empty((0, 3))
+
+    x1 = boxes[:, 0].copy()
+    y1 = boxes[:, 1].copy()
+    x2 = boxes[:, 2].copy()
+    y2 = boxes[:, 3].copy()
+    s = scores
+    area = (x2 - x1 + 1) * (y2 - y1 + 1)
+
+    I = np.argsort(s)
+    pick = np.zeros_like(s, dtype=np.int16)
+    counter = 0
+    while I.size > 0:
+        i = I[-1]
+        pick[counter] = i
+        counter += 1
+        idx = I[0:-1]
+
+        xx1 = np.maximum(x1[i], x1[idx]).copy()
+        yy1 = np.maximum(y1[i], y1[idx]).copy()
+        xx2 = np.minimum(x2[i], x2[idx]).copy()
+        yy2 = np.minimum(y2[i], y2[idx]).copy()
+
+        w = np.maximum(0.0, xx2 - xx1 + 1).copy()
+        h = np.maximum(0.0, yy2 - yy1 + 1).copy()
+
+        inter = w * h
+        if method is "Min":
+            o = inter / np.minimum(area[i], area[idx])
+        else:
+            o = inter / (area[i] + area[idx] - inter)
+        I = I[np.where(o <= threshold)]
+
+    pick = pick[:counter].copy()
+    return pick
+
+
+def batched_nms_numpy(boxes, scores, idxs, threshold, method):
+    device = boxes.device
+    if boxes.numel() == 0:
+        return torch.empty((0,), dtype=torch.int64, device=device)
+    # strategy: in order to perform NMS independently per class.
+    # we add an offset to all the boxes. The offset is dependent
+    # only on the class idx, and is large enough so that boxes
+    # from different classes do not overlap
+    max_coordinate = boxes.max()
+    offsets = idxs.to(boxes) * (max_coordinate + 1)
+    boxes_for_nms = boxes + offsets[:, None]
+    boxes_for_nms = boxes_for_nms.cpu().numpy()
+    scores = scores.cpu().numpy()
+    keep = nms_numpy(boxes_for_nms, scores, threshold, method)
+    return torch.as_tensor(keep, dtype=torch.long, device=device)
+
+
+def pad(boxes, w, h):
+    boxes = boxes.trunc().int().cpu().numpy()
+    x = boxes[:, 0]
+    y = boxes[:, 1]
+    ex = boxes[:, 2]
+    ey = boxes[:, 3]
+
+    x[x < 1] = 1
+    y[y < 1] = 1
+    ex[ex > w] = w
+    ey[ey > h] = h
+
+    return y, ey, x, ex
+
+
+def rerec(bboxA):
+    h = bboxA[:, 3] - bboxA[:, 1]
+    w = bboxA[:, 2] - bboxA[:, 0]
+    
+    l = torch.max(w, h)
+    bboxA[:, 0] = bboxA[:, 0] + w * 0.5 - l * 0.5
+    bboxA[:, 1] = bboxA[:, 1] + h * 0.5 - l * 0.5
+    bboxA[:, 2:4] = bboxA[:, :2] + l.repeat(2, 1).permute(1, 0)
+
+    return bboxA
+
+
+def imresample(img, sz):
+    im_data = interpolate(img, size=sz, mode="area")
+    return im_data
+
+
+def crop_resize(img, box, image_size):
+    if isinstance(img, np.ndarray):
+        out = cv2.resize(
+            img[box[1]:box[3], box[0]:box[2]],
+            (image_size, image_size),
+            interpolation=cv2.INTER_AREA
+        ).copy()
+    else:
+        out = img.crop(box).copy().resize((image_size, image_size), Image.BILINEAR)
+    return out
+
+
+def save_img(img, path):
+    if isinstance(img, np.ndarray):
+        cv2.imwrite(path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+    else:
+        img.save(path)
+
+
+def get_size(img):
+    if isinstance(img, np.ndarray):
+        return img.shape[1::-1]
+    else:
+        return img.size
+
+
+def extract_face(img, box, image_size=160, margin=0, save_path=None):
+    """Extract face + margin from PIL Image given bounding box.
+    
+    Arguments:
+        img {PIL.Image} -- A PIL Image.
+        box {numpy.ndarray} -- Four-element bounding box.
+        image_size {int} -- Output image size in pixels. The image will be square.
+        margin {int} -- Margin to add to bounding box, in terms of pixels in the final image. 
+            Note that the application of the margin differs slightly from the davidsandberg/facenet
+            repo, which applies the margin to the original image before resizing, making the margin
+            dependent on the original image size.
+        save_path {str} -- Save path for extracted face image. (default: {None})
+    
+    Returns:
+        torch.tensor -- tensor representing the extracted face.
+    """
+    margin = [
+        margin * (box[2] - box[0]) / (image_size - margin),
+        margin * (box[3] - box[1]) / (image_size - margin),
+    ]
+    raw_image_size = get_size(img)
+    box = [
+        int(max(box[0] - margin[0] / 2, 0)),
+        int(max(box[1] - margin[1] / 2, 0)),
+        int(min(box[2] + margin[0] / 2, raw_image_size[0])),
+        int(min(box[3] + margin[1] / 2, raw_image_size[1])),
+    ]
+
+    face = crop_resize(img, box, image_size)
+
+    if save_path is not None:
+        os.makedirs(os.path.dirname(save_path) + "/", exist_ok=True)
+        save_img(face, save_path)
+
+    face = F.to_tensor(np.float32(face))
+
+    return face

+import tensorflow as tf
+import torch
+import json
+import os, sys
+
+from dependencies.facenet.src import facenet
+from dependencies.facenet.src.models import inception_resnet_v1 as tf_mdl
+from dependencies.facenet.src.align import detect_face
+
+from models.inception_resnet_v1 import InceptionResnetV1
+from models.mtcnn import PNet, RNet, ONet
+
+
+def import_tf_params(tf_mdl_dir, sess):
+    """Import tensorflow model from save directory.
+    
+    Arguments:
+        tf_mdl_dir {str} -- Location of protobuf, checkpoint, meta files.
+        sess {tensorflow.Session} -- Tensorflow session object.
+    
+    Returns:
+        (list, list, list) -- Tuple of lists containing the layer names,
+            parameter arrays as numpy ndarrays, parameter shapes.
+    """
+    print('\nLoading tensorflow model\n')
+    if callable(tf_mdl_dir):
+        tf_mdl_dir(sess)
+    else:
+        facenet.load_model(tf_mdl_dir)
+
+    print('\nGetting model weights\n')
+    tf_layers = tf.trainable_variables()
+    tf_params = sess.run(tf_layers)
+
+    tf_shapes = [p.shape for p in tf_params]
+    tf_layers = [l.name for l in tf_layers]
+
+    if not callable(tf_mdl_dir):
+        path = os.path.join(tf_mdl_dir, 'layer_description.json')
+    else:
+        path = 'data/layer_description.json'
+    with open(path, 'w') as f:
+        json.dump({l: s for l, s in zip(tf_layers, tf_shapes)}, f)
+
+    return tf_layers, tf_params, tf_shapes
+
+
+def get_layer_indices(layer_lookup, tf_layers):
+    """Giving a lookup of model layer attribute names and tensorflow variable names,
+    find matching parameters.
+    
+    Arguments:
+        layer_lookup {dict} -- Dictionary mapping pytorch attribute names to (partial)
+            tensorflow variable names. Expects dict of the form {'attr': ['tf_name', ...]}
+            where the '...'s are ignored.
+        tf_layers {list} -- List of tensorflow variable names.
+    
+    Returns:
+        list -- The input dictionary with the list of matching inds appended to each item.
+    """
+    layer_inds = {}
+    for name, value in layer_lookup.items():
+        layer_inds[name] = value + [[i for i, n in enumerate(tf_layers) if value[0] in n]]
+    return layer_inds
+
+
+def load_tf_batchNorm(weights, layer):
+    """Load tensorflow weights into nn.BatchNorm object.
+    
+    Arguments:
+        weights {list} -- Tensorflow parameters.
+        layer {torch.nn.Module} -- nn.BatchNorm.
+    """
+    layer.bias.data = torch.tensor(weights[0]).view(layer.bias.data.shape)
+    layer.weight.data = torch.ones_like(layer.weight.data)
+    layer.running_mean = torch.tensor(weights[1]).view(layer.running_mean.shape)
+    layer.running_var = torch.tensor(weights[2]).view(layer.running_var.shape)
+
+
+def load_tf_conv2d(weights, layer, transpose=False):
+    """Load tensorflow weights into nn.Conv2d object.
+    
+    Arguments:
+        weights {list} -- Tensorflow parameters.
+        layer {torch.nn.Module} -- nn.Conv2d.
+    """
+    if isinstance(weights, list):
+        if len(weights) == 2:
+            layer.bias.data = (
+                torch.tensor(weights[1])
+                    .view(layer.bias.data.shape)
+            )
+        weights = weights[0]
+    
+    if transpose:
+        dim_order = (3, 2, 1, 0)
+    else:
+        dim_order = (3, 2, 0, 1)
+
+    layer.weight.data = (
+        torch.tensor(weights)
+            .permute(dim_order)
+            .view(layer.weight.data.shape)
+    )
+
+
+def load_tf_conv2d_trans(weights, layer):
+    return load_tf_conv2d(weights, layer, transpose=True)
+
+
+def load_tf_basicConv2d(weights, layer):
+    """Load tensorflow weights into grouped Conv2d+BatchNorm object.
+    
+    Arguments:
+        weights {list} -- Tensorflow parameters.
+        layer {torch.nn.Module} -- Object containing Conv2d+BatchNorm.
+    """
+    load_tf_conv2d(weights[0], layer.conv)
+    load_tf_batchNorm(weights[1:], layer.bn)
+
+
+def load_tf_linear(weights, layer):
+    """Load tensorflow weights into nn.Linear object.
+    
+    Arguments:
+        weights {list} -- Tensorflow parameters.
+        layer {torch.nn.Module} -- nn.Linear.
+    """
+    if isinstance(weights, list):
+        if len(weights) == 2:
+            layer.bias.data = (
+                torch.tensor(weights[1])
+                    .view(layer.bias.data.shape)
+            )
+        weights = weights[0]
+    layer.weight.data = (
+        torch.tensor(weights)
+            .transpose(-1, 0)
+            .view(layer.weight.data.shape)
+    )
+
+
+# High-level parameter-loading functions:
+
+def load_tf_block35(weights, layer):
+    load_tf_basicConv2d(weights[:4], layer.branch0)
+    load_tf_basicConv2d(weights[4:8], layer.branch1[0])
+    load_tf_basicConv2d(weights[8:12], layer.branch1[1])
+    load_tf_basicConv2d(weights[12:16], layer.branch2[0])
+    load_tf_basicConv2d(weights[16:20], layer.branch2[1])
+    load_tf_basicConv2d(weights[20:24], layer.branch2[2])
+    load_tf_conv2d(weights[24:26], layer.conv2d)
+
+
+def load_tf_block17_8(weights, layer):
+    load_tf_basicConv2d(weights[:4], layer.branch0)
+    load_tf_basicConv2d(weights[4:8], layer.branch1[0])
+    load_tf_basicConv2d(weights[8:12], layer.branch1[1])
+    load_tf_basicConv2d(weights[12:16], layer.branch1[2])
+    load_tf_conv2d(weights[16:18], layer.conv2d)
+
+
+def load_tf_mixed6a(weights, layer):
+    if len(weights) != 16:
+        raise ValueError(f'Number of weight arrays ({len(weights)}) not equal to 16')
+    load_tf_basicConv2d(weights[:4], layer.branch0)
+    load_tf_basicConv2d(weights[4:8], layer.branch1[0])
+    load_tf_basicConv2d(weights[8:12], layer.branch1[1])
+    load_tf_basicConv2d(weights[12:16], layer.branch1[2])
+
+
+def load_tf_mixed7a(weights, layer):
+    if len(weights) != 28:
+        raise ValueError(f'Number of weight arrays ({len(weights)}) not equal to 28')
+    load_tf_basicConv2d(weights[:4], layer.branch0[0])
+    load_tf_basicConv2d(weights[4:8], layer.branch0[1])
+    load_tf_basicConv2d(weights[8:12], layer.branch1[0])
+    load_tf_basicConv2d(weights[12:16], layer.branch1[1])
+    load_tf_basicConv2d(weights[16:20], layer.branch2[0])
+    load_tf_basicConv2d(weights[20:24], layer.branch2[1])
+    load_tf_basicConv2d(weights[24:28], layer.branch2[2])
+
+
+def load_tf_repeats(weights, layer, rptlen, subfun):
+    if len(weights) % rptlen != 0:
+        raise ValueError(f'Number of weight arrays ({len(weights)}) not divisible by {rptlen}')
+    weights_split = [weights[i:i+rptlen] for i in range(0, len(weights), rptlen)]
+    for i, w in enumerate(weights_split):
+        subfun(w, getattr(layer, str(i)))
+
+
+def load_tf_repeat_1(weights, layer):
+    load_tf_repeats(weights, layer, 26, load_tf_block35)
+
+
+def load_tf_repeat_2(weights, layer):
+    load_tf_repeats(weights, layer, 18, load_tf_block17_8)
+
+
+def load_tf_repeat_3(weights, layer):
+    load_tf_repeats(weights, layer, 18, load_tf_block17_8)
+
+
+def test_loaded_params(mdl, tf_params, tf_layers):
+    """Check each parameter in a pytorch model for an equivalent parameter
+    in a list of tensorflow variables.
+    
+    Arguments:
+        mdl {torch.nn.Module} -- Pytorch model.
+        tf_params {list} -- List of ndarrays representing tensorflow variables.
+        tf_layers {list} -- Corresponding list of tensorflow variable names.
+    """
+    tf_means = torch.stack([torch.tensor(p).mean() for p in tf_params])
+    for name, param in mdl.named_parameters():
+        pt_mean = param.data.mean()
+        matching_inds = ((tf_means - pt_mean).abs() < 1e-8).nonzero()
+        print(f'{name} equivalent to {[tf_layers[i] for i in matching_inds]}')
+
+
+def compare_model_outputs(pt_mdl, sess, test_data):
+    """Given some testing data, compare the output of pytorch and tensorflow models.
+    
+    Arguments:
+        pt_mdl {torch.nn.Module} -- Pytorch model.
+        sess {tensorflow.Session} -- Tensorflow session object.
+        test_data {torch.Tensor} -- Pytorch tensor.
+    """
+    print('\nPassing test data through TF model\n')
+    if isinstance(sess, tf.Session):
+        images_placeholder = tf.get_default_graph().get_tensor_by_name("input:0")
+        phase_train_placeholder = tf.get_default_graph().get_tensor_by_name("phase_train:0")
+        embeddings = tf.get_default_graph().get_tensor_by_name("embeddings:0")
+        feed_dict = {images_placeholder: test_data.numpy(), phase_train_placeholder: False}
+        tf_output = torch.tensor(sess.run(embeddings, feed_dict=feed_dict))
+    else:
+        tf_output = sess(test_data)
+
+    print(tf_output)
+
+    print('\nPassing test data through PT model\n')
+    pt_output = pt_mdl(test_data.permute(0, 3, 1, 2))
+    print(pt_output)
+
+    distance = (tf_output - pt_output).norm()
+    print(f'\nDistance {distance}\n')
+
+
+def compare_mtcnn(pt_mdl, tf_fun, sess, ind, test_data):
+    tf_mdls = tf_fun(sess)
+    tf_mdl = tf_mdls[ind]
+
+    print('\nPassing test data through TF model\n')
+    tf_output = tf_mdl(test_data.numpy())
+    tf_output = [torch.tensor(out) for out in tf_output]
+    print('\n'.join([str(o.view(-1)[:10]) for o in tf_output]))
+
+    print('\nPassing test data through PT model\n')
+    with torch.no_grad():
+        pt_output = pt_mdl(test_data.permute(0, 3, 2, 1))
+    pt_output = [torch.tensor(out) for out in pt_output]
+    for i in range(len(pt_output)):
+        if len(pt_output[i].shape) == 4:
+            pt_output[i] = pt_output[i].permute(0, 3, 2, 1).contiguous()
+    print('\n'.join([str(o.view(-1)[:10]) for o in pt_output]))
+
+    distance = [(tf_o - pt_o).norm() for tf_o, pt_o in zip(tf_output, pt_output)]
+    print(f'\nDistance {distance}\n')
+
+
+def load_tf_model_weights(mdl, layer_lookup, tf_mdl_dir, is_resnet=True, arg_num=None):
+    """Load tensorflow parameters into a pytorch model.
+    
+    Arguments:
+        mdl {torch.nn.Module} -- Pytorch model.
+        layer_lookup {[type]} -- Dictionary mapping pytorch attribute names to (partial)
+            tensorflow variable names, and a function suitable for loading weights.
+            Expects dict of the form {'attr': ['tf_name', function]}. 
+        tf_mdl_dir {str} -- Location of protobuf, checkpoint, meta files.
+    """
+    tf.reset_default_graph()
+    with tf.Session() as sess:
+        tf_layers, tf_params, tf_shapes = import_tf_params(tf_mdl_dir, sess)
+        layer_info = get_layer_indices(layer_lookup, tf_layers)
+
+        for layer_name, info in layer_info.items():
+            print(f'Loading {info[0]}/* into {layer_name}')
+            weights = [tf_params[i] for i in info[2]]
+            layer = getattr(mdl, layer_name)
+            info[1](weights, layer)
+
+        test_loaded_params(mdl, tf_params, tf_layers)
+
+        if is_resnet:
+            compare_model_outputs(mdl, sess, torch.randn(5, 160, 160, 3).detach())
+
+
+def tensorflow2pytorch():
+    lookup_inception_resnet_v1 = {
+        'conv2d_1a': ['InceptionResnetV1/Conv2d_1a_3x3', load_tf_basicConv2d],
+        'conv2d_2a': ['InceptionResnetV1/Conv2d_2a_3x3', load_tf_basicConv2d],
+        'conv2d_2b': ['InceptionResnetV1/Conv2d_2b_3x3', load_tf_basicConv2d],
+        'conv2d_3b': ['InceptionResnetV1/Conv2d_3b_1x1', load_tf_basicConv2d],
+        'conv2d_4a': ['InceptionResnetV1/Conv2d_4a_3x3', load_tf_basicConv2d],
+        'conv2d_4b': ['InceptionResnetV1/Conv2d_4b_3x3', load_tf_basicConv2d],
+        'repeat_1': ['InceptionResnetV1/Repeat/block35', load_tf_repeat_1],
+        'mixed_6a': ['InceptionResnetV1/Mixed_6a', load_tf_mixed6a],
+        'repeat_2': ['InceptionResnetV1/Repeat_1/block17', load_tf_repeat_2],
+        'mixed_7a': ['InceptionResnetV1/Mixed_7a', load_tf_mixed7a],
+        'repeat_3': ['InceptionResnetV1/Repeat_2/block8', load_tf_repeat_3],
+        'block8': ['InceptionResnetV1/Block8', load_tf_block17_8],
+        'last_linear': ['InceptionResnetV1/Bottleneck/weights', load_tf_linear],
+        'last_bn': ['InceptionResnetV1/Bottleneck/BatchNorm', load_tf_batchNorm],
+        'logits': ['Logits', load_tf_linear],
+    }
+
+    print('\nLoad VGGFace2-trained weights and save\n')
+    mdl = InceptionResnetV1(num_classes=8631).eval()
+    tf_mdl_dir = 'data/20180402-114759'
+    data_name = 'vggface2'
+    load_tf_model_weights(mdl, lookup_inception_resnet_v1, tf_mdl_dir)
+    state_dict = mdl.state_dict()
+    torch.save(state_dict, f'{tf_mdl_dir}-{data_name}.pt')    
+    torch.save(
+        {
+            'logits.weight': state_dict['logits.weight'],
+            'logits.bias': state_dict['logits.bias'],
+        },
+        f'{tf_mdl_dir}-{data_name}-logits.pt'
+    )
+    state_dict.pop('logits.weight')
+    state_dict.pop('logits.bias')
+    torch.save(state_dict, f'{tf_mdl_dir}-{data_name}-features.pt')
+    
+    print('\nLoad CASIA-Webface-trained weights and save\n')
+    mdl = InceptionResnetV1(num_classes=10575).eval()
+    tf_mdl_dir = 'data/20180408-102900'
+    data_name = 'casia-webface'
+    load_tf_model_weights(mdl, lookup_inception_resnet_v1, tf_mdl_dir)
+    state_dict = mdl.state_dict()
+    torch.save(state_dict, f'{tf_mdl_dir}-{data_name}.pt')    
+    torch.save(
+        {
+            'logits.weight': state_dict['logits.weight'],
+            'logits.bias': state_dict['logits.bias'],
+        },
+        f'{tf_mdl_dir}-{data_name}-logits.pt'
+    )
+    state_dict.pop('logits.weight')
+    state_dict.pop('logits.bias')
+    torch.save(state_dict, f'{tf_mdl_dir}-{data_name}-features.pt')
+    
+    lookup_pnet = {
+        'conv1': ['pnet/conv1', load_tf_conv2d_trans],
+        'prelu1': ['pnet/PReLU1', load_tf_linear],
+        'conv2': ['pnet/conv2', load_tf_conv2d_trans],
+        'prelu2': ['pnet/PReLU2', load_tf_linear],
+        'conv3': ['pnet/conv3', load_tf_conv2d_trans],
+        'prelu3': ['pnet/PReLU3', load_tf_linear],
+        'conv4_1': ['pnet/conv4-1', load_tf_conv2d_trans],
+        'conv4_2': ['pnet/conv4-2', load_tf_conv2d_trans],
+    }
+    lookup_rnet = {
+        'conv1': ['rnet/conv1', load_tf_conv2d_trans],
+        'prelu1': ['rnet/prelu1', load_tf_linear],
+        'conv2': ['rnet/conv2', load_tf_conv2d_trans],
+        'prelu2': ['rnet/prelu2', load_tf_linear],
+        'conv3': ['rnet/conv3', load_tf_conv2d_trans],
+        'prelu3': ['rnet/prelu3', load_tf_linear],
+        'dense4': ['rnet/conv4', load_tf_linear],
+        'prelu4': ['rnet/prelu4', load_tf_linear],
+        'dense5_1': ['rnet/conv5-1', load_tf_linear],
+        'dense5_2': ['rnet/conv5-2', load_tf_linear],
+    }
+    lookup_onet = {
+        'conv1': ['onet/conv1', load_tf_conv2d_trans],
+        'prelu1': ['onet/prelu1', load_tf_linear],
+        'conv2': ['onet/conv2', load_tf_conv2d_trans],
+        'prelu2': ['onet/prelu2', load_tf_linear],
+        'conv3': ['onet/conv3', load_tf_conv2d_trans],
+        'prelu3': ['onet/prelu3', load_tf_linear],
+        'conv4': ['onet/conv4', load_tf_conv2d_trans],
+        'prelu4': ['onet/prelu4', load_tf_linear],
+        'dense5': ['onet/conv5', load_tf_linear],
+        'prelu5': ['onet/prelu5', load_tf_linear],
+        'dense6_1': ['onet/conv6-1', load_tf_linear],
+        'dense6_2': ['onet/conv6-2', load_tf_linear],
+        'dense6_3': ['onet/conv6-3', load_tf_linear],
+    }
+
+    print('\nLoad PNet weights and save\n')
+    tf_mdl_dir = lambda sess: detect_face.create_mtcnn(sess, None)
+    mdl = PNet()
+    data_name = 'pnet'
+    load_tf_model_weights(mdl, lookup_pnet, tf_mdl_dir, is_resnet=False, arg_num=0)
+    torch.save(mdl.state_dict(), f'data/{data_name}.pt')
+    tf.reset_default_graph()
+    with tf.Session() as sess:
+        compare_mtcnn(mdl, tf_mdl_dir, sess, 0, torch.randn(1, 256, 256, 3).detach())
+
+    print('\nLoad RNet weights and save\n')
+    mdl = RNet()
+    data_name = 'rnet'
+    load_tf_model_weights(mdl, lookup_rnet, tf_mdl_dir, is_resnet=False, arg_num=1)
+    torch.save(mdl.state_dict(), f'data/{data_name}.pt')
+    tf.reset_default_graph()
+    with tf.Session() as sess:
+        compare_mtcnn(mdl, tf_mdl_dir, sess, 1, torch.randn(1, 24, 24, 3).detach())
+
+    print('\nLoad ONet weights and save\n')
+    mdl = ONet()
+    data_name = 'onet'
+    load_tf_model_weights(mdl, lookup_onet, tf_mdl_dir, is_resnet=False, arg_num=2)
+    torch.save(mdl.state_dict(), f'data/{data_name}.pt')
+    tf.reset_default_graph()
+    with tf.Session() as sess:
+        compare_mtcnn(mdl, tf_mdl_dir, sess, 2, torch.randn(1, 48, 48, 3).detach())

+import torch
+import numpy as np
+import time
+
+
+class Logger(object):
+
+    def __init__(self, mode, length, calculate_mean=False):
+        self.mode = mode
+        self.length = length
+        self.calculate_mean = calculate_mean
+        if self.calculate_mean:
+            self.fn = lambda x, i: x / (i + 1)
+        else:
+            self.fn = lambda x, i: x
+
+    def __call__(self, loss, metrics, i):
+        track_str = '\r{} | {:5d}/{:<5d}| '.format(self.mode, i + 1, self.length)
+        loss_str = 'loss: {:9.4f} | '.format(self.fn(loss, i))
+        metric_str = ' | '.join('{}: {:9.4f}'.format(k, self.fn(v, i)) for k, v in metrics.items())
+        print(track_str + loss_str + metric_str + '   ', end='')
+        if i + 1 == self.length:
+            print('')
+
+
+class BatchTimer(object):
+    """Batch timing class.
+    Use this class for tracking training and testing time/rate per batch or per sample.
+    
+    Keyword Arguments:
+        rate {bool} -- Whether to report a rate (batches or samples per second) or a time (seconds
+            per batch or sample). (default: {True})
+        per_sample {bool} -- Whether to report times or rates per sample or per batch.
+            (default: {True})
+    """
+
+    def __init__(self, rate=True, per_sample=True):
+        self.start = time.time()
+        self.end = None
+        self.rate = rate
+        self.per_sample = per_sample
+
+    def __call__(self, y_pred, y):
+        self.end = time.time()
+        elapsed = self.end - self.start
+        self.start = self.end
+        self.end = None
+
+        if self.per_sample:
+            elapsed /= len(y_pred)
+        if self.rate:
+            elapsed = 1 / elapsed
+
+        return torch.tensor(elapsed)
+
+
+def accuracy(logits, y):
+    _, preds = torch.max(logits, 1)
+    return (preds == y).float().mean()
+
+
+def pass_epoch(
+    model, loss_fn, loader, optimizer=None, scheduler=None,
+    batch_metrics={'time': BatchTimer()}, show_running=True,
+    device='cpu', writer=None
+):
+    """Train or evaluate over a data epoch.
+    
+    Arguments:
+        model {torch.nn.Module} -- Pytorch model.
+        loss_fn {callable} -- A function to compute (scalar) loss.
+        loader {torch.utils.data.DataLoader} -- A pytorch data loader.
+    
+    Keyword Arguments:
+        optimizer {torch.optim.Optimizer} -- A pytorch optimizer.
+        scheduler {torch.optim.lr_scheduler._LRScheduler} -- LR scheduler (default: {None})
+        batch_metrics {dict} -- Dictionary of metric functions to call on each batch. The default
+            is a simple timer. A progressive average of these metrics, along with the average
+            loss, is printed every batch. (default: {{'time': iter_timer()}})
+        show_running {bool} -- Whether or not to print losses and metrics for the current batch
+            or rolling averages. (default: {False})
+        device {str or torch.device} -- Device for pytorch to use. (default: {'cpu'})
+        writer {torch.utils.tensorboard.SummaryWriter} -- Tensorboard SummaryWriter. (default: {None})
+    
+    Returns:
+        tuple(torch.Tensor, dict) -- A tuple of the average loss and a dictionary of average
+            metric values across the epoch.
+    """
+    
+    mode = 'Train' if model.training else 'Valid'
+    logger = Logger(mode, length=len(loader), calculate_mean=show_running)
+    loss = 0
+    metrics = {}
+
+    for i_batch, (x, y) in enumerate(loader):
+        x = x.to(device)
+        y = y.to(device)
+        y_pred = model(x)
+        loss_batch = loss_fn(y_pred, y)
+
+        if model.training:
+            loss_batch.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        metrics_batch = {}
+        for metric_name, metric_fn in batch_metrics.items():
+            metrics_batch[metric_name] = metric_fn(y_pred, y).detach().cpu()
+            metrics[metric_name] = metrics.get(metric_name, 0) + metrics_batch[metric_name]
+            
+        if writer is not None and model.training:
+            if writer.iteration % writer.interval == 0:
+                writer.add_scalars('loss', {mode: loss_batch.detach().cpu()}, writer.iteration)
+                for metric_name, metric_batch in metrics_batch.items():
+                    writer.add_scalars(metric_name, {mode: metric_batch}, writer.iteration)
+            writer.iteration += 1
+        
+        loss_batch = loss_batch.detach().cpu()
+        loss += loss_batch
+        if show_running:
+            logger(loss, metrics, i_batch)
+        else:
+            logger(loss_batch, metrics_batch, i_batch)
+    
+    if model.training and scheduler is not None:
+        scheduler.step()
+
+    loss = loss / (i_batch + 1)
+    metrics = {k: v / (i_batch + 1) for k, v in metrics.items()}
+            
+    if writer is not None and not model.training:
+        writer.add_scalars('loss', {mode: loss.detach()}, writer.iteration)
+        for metric_name, metric in metrics.items():
+            writer.add_scalars(metric_name, {mode: metric})
+
+    return loss, metrics
+
+
+def collate_pil(x): 
+    out_x, out_y = [], [] 
+    for xx, yy in x: 
+        out_x.append(xx) 
+        out_y.append(yy) 
+    return out_x, out_y 

+##################################################
+#1. webcam에서 얼굴을 인식합니다
+#2. 인식한 얼굴을 등록합니다
+##################################################
+import torch
+import numpy as np
+import cv2
+import asyncio
+import websockets
+import json
+import os
+import timeit
+import base64
+
+from PIL import Image
+from io import BytesIO
+import requests
+
+from models.mtcnn import MTCNN
+
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+print('Running on device: {}'.format(device))
+
+mtcnn = MTCNN(keep_all=True, device=device)
+
+uri = 'ws://169.56.95.131:8765'
+
+async def send_face(face_list, image_list):
+    global uri
+    async with websockets.connect(uri) as websocket:
+        for face, image in zip(face_list, image_list):
+            #type: np.float32
+            send = json.dumps({'action': 'register', 'student_id':'2014104149', 'student_name':'정해갑', 'MTCNN': face.tolist()})
+            await websocket.send(send)
+            recv = await websocket.recv()
+            data = json.loads(recv)
+            if data['status'] == 'success':
+                # 성공
+                print(data['student_id'], 'is registered')
+
+def detect_face(frame):
+    # If required, create a face detection pipeline using MTCNN:
+    global mtcnn
+    results = mtcnn.detect(frame)
+    image_list = []
+    if results[1][0] == None:
+        return []
+    for box, prob in zip(results[0], results[1]):
+        if prob < 0.95:
+            continue
+        print('face detected. prob:', prob)
+        x1, y1, x2, y2 = box
+        image = frame[int(y1-10):int(y2+10), int(x1-10):int(x2+10)]
+        image_list.append(image)
+    return image_list
+
+def detect_face(frame):
+    results = mtcnn.detect(frame)
+    faces = mtcnn(frame, return_prob = False)
+    image_list = []
+    face_list = []
+    if results[1][0] == None:
+        return [], []
+    for box, face, prob in zip(results[0], faces, results[1]):
+        if prob < 0.97:
+            continue
+        print('face detected. prob:', prob)
+        x1, y1, x2, y2 = box
+        if (x2-x1) * (y2-y1) < 15000:
+            # 얼굴 해상도가 너무 낮으면 무시
+            continue
+        # 얼굴 주변 ±3 영역  저장
+        image = frame[int(y1-3):int(y2+3), int(x1-3):int(x2+3)]
+        image_list.append(image)
+        # MTCNN 데이터 저장
+        face_list.append(face.numpy())
+    return image_list, face_list
+
+cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
+cap.set(3, 720)
+cap.set(4, 480)
+ret, frame = cap.read()
+frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+image_list, face_list = detect_face(frame)
+if face_list:
+    asyncio.get_event_loop().run_until_complete(send_face(face_list, image_list))
\ No newline at end of file