face_recognition1/face_detect/test.py

import os
import sys
sys.path.append(os.path.dirname(__file__))

import cv2
import math
import torch
import torch.backends.cudnn as cudnn
import numpy as np
from data.config import cfg
from layers.functions.prior_box import PriorBox
from utils.nms_wrapper import nms
from models.faceboxes import FaceBoxes
from utils.box_utils import decode
from utils.timer import Timer

trained_model = os.path.join(os.path.dirname(__file__), './checkpoints/FaceBoxesProd.pth')
save_folder = 'eval'
dataset = 'Custom'
confidence_threshold = 0.2
top_k = 5000
nms_threshold = 0.3
keep_top_k = 750
show_image = True
vis_thres = 0.5


def check_keys(model, pretrained_state_dict):
    ckpt_keys = set(pretrained_state_dict.keys())
    model_keys = set(model.state_dict().keys())
    used_pretrained_keys = model_keys & ckpt_keys
    unused_pretrained_keys = ckpt_keys - model_keys
    missing_keys = model_keys - ckpt_keys
    print('Missing keys:{}'.format(len(missing_keys)))
    print('Unused checkpoint keys:{}'.format(len(unused_pretrained_keys)))
    print('Used keys:{}'.format(len(used_pretrained_keys)))
    assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
    return True


def remove_prefix(state_dict, prefix):
    """ Old style model is stored with all names of parameters sharing common prefix 'module.' """
    print('remove prefix \'{}\''.format(prefix))
    f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
    return {f(key): value for key, value in state_dict.items()}


def load_model(model, pretrained_path, device):
    print('Loading pretrained model from {}'.format(pretrained_path))
    pretrained_dict = torch.load(pretrained_path, map_location=device)

    if "state_dict" in pretrained_dict.keys():
        pretrained_dict = remove_prefix(pretrained_dict['state_dict'], 'module.')
    else:
        pretrained_dict = remove_prefix(pretrained_dict, 'module.')
    check_keys(model, pretrained_dict)
    model.load_state_dict(pretrained_dict, strict=False)
    return model


torch.set_grad_enabled(False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = FaceBoxes(phase='test', size=None, num_classes=2)
net = load_model(net, trained_model, device)
net.eval()
cudnn.benchmark = True
net = net.to(device)


def get_bbox(orig_image):
    # testing scale
    resize = 0.5

    _t = {'forward_pass': Timer(), 'misc': Timer()}

    img_raw = orig_image
    img = np.float32(img_raw)
    if resize != 1:
        img = cv2.resize(img, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)
    im_height, im_width, _ = img.shape
    scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
    img -= (104, 117, 123)
    img = img.transpose(2, 0, 1)
    img = torch.from_numpy(img).unsqueeze(0)
    img = img.to(device)
    scale = scale.to(device)

    _t['forward_pass'].tic()
    loc, conf = net(img)  # forward pass
    _t['forward_pass'].toc()
    _t['misc'].tic()
    priorbox = PriorBox(cfg, image_size=(im_height, im_width))
    priors = priorbox.forward()
    priors = priors.to(device)
    prior_data = priors.data
    boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])
    boxes = boxes * scale / resize
    boxes = boxes.cpu().numpy()
    scores = conf.squeeze(0).data.cpu().numpy()[:, 1]

    # ignore low scores
    inds = np.where(scores > confidence_threshold)[0]
    boxes = boxes[inds]
    scores = scores[inds]

    # keep top-K before NMS
    order = scores.argsort()[::-1][:top_k]
    boxes = boxes[order]
    scores = scores[order]

    # do NMS
    dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
    #keep = py_cpu_nms(dets, nms_threshold)
    keep = nms(dets, nms_threshold, force_cpu=True)
    dets = dets[keep, :]

    # keep top-K faster NMS
    dets = dets[:keep_top_k, :]
    _t['misc'].toc()

    boxes, scores = [], []
    for k in range(dets.shape[0]):
        xmin = dets[k, 0]
        ymin = dets[k, 1]
        xmax = dets[k, 2]
        ymax = dets[k, 3]
        ymin += 0.2 * (ymax - ymin + 1)
        score = dets[k, 4]
        boxes.append([int(xmin), int(ymin), int(xmax - xmin), int(ymax - ymin)])
        scores.append(score)

    max_score = 0.0
    final_box = None
    for i, score in enumerate(scores):
        if max_score < score:
            max_score = score
            final_box = boxes[i]

    return final_box


class Detection:
    def __init__(self):
        src_dir = os.path.dirname(__file__)
        if not os.path.exists(os.path.join(src_dir, "checkpoints")):
            os.makedirs(os.path.join(src_dir, "checkpoints"))

        caffemodel = os.path.join(src_dir, "checkpoints/Widerface-RetinaFace.caffemodel")
        deploy = os.path.join(src_dir, "checkpoints/deploy.prototxt")

        self.detector = cv2.dnn.readNetFromCaffe(deploy, caffemodel)
        self.detector_confidence = 0.6

    def get_bbox(self, img):
        height, width = img.shape[0], img.shape[1]
        aspect_ratio = width / height
        if img.shape[1] * img.shape[0] >= 192 * 192:
            img = cv2.resize(img,
                             (int(192 * math.sqrt(aspect_ratio)),
                              int(192 / math.sqrt(aspect_ratio))), interpolation=cv2.INTER_LINEAR)

        blob = cv2.dnn.blobFromImage(img, 1, mean=(104, 117, 123))
        self.detector.setInput(blob, 'data')
        out = self.detector.forward('detection_out').squeeze()
        max_conf_index = np.argmax(out[:, 2])
        left, top, right, bottom = out[max_conf_index, 3]*width, out[max_conf_index, 4]*height, \
                                   out[max_conf_index, 5]*width, out[max_conf_index, 6]*height

        if right == left or bottom == top:
            return None

        bbox = [int(left), int(top), int(right-left+1), int(bottom-top+1)]
        return bbox

    def check_face(self):
        pass


if __name__ == '__main__':

    # image = cv2.imread('arun_2.jpg')

    # box = get_bbox(image)
    # cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 2)
    #
    src_dir = 'D:/19.Database/office_angled_db'
    dst_dir = 'D:/19.Database/office_angled_db_result'
    detector = Detection()

    for file in os.listdir(src_dir):
        image1 = cv2.imread(os.path.join(src_dir, file))
        box = detector.get_bbox(image1)
        if box:
            cv2.rectangle(image1, (box[0], box[1]), (box[0] + box[2], box[1] + box[3]), (0, 0, 255), 5)

        cv2.imwrite(os.path.join(dst_dir, file), image1)
        # cv2.waitKey(0)