Ejemplo n.º 1
0
def gen_landmark48_data(data_dir,
                        anno_file,
                        pnet_model_file,
                        rnet_model_file,
                        prefix_path='',
                        use_cuda=True,
                        vis=False):

    pnet, rnet, _ = create_mtcnn_net(p_model_path=pnet_model_file,
                                     r_model_path=rnet_model_file,
                                     use_cuda=use_cuda)
    mtcnn_detector = MtcnnDetector(pnet=pnet, rnet=rnet, min_face_size=12)

    imagedb = ImageDB(anno_file, mode="test", prefix_path=prefix_path)
    imdb = imagedb.load_imdb()
    image_reader = TestImageLoader(imdb, 1, False)

    all_boxes = list()
    batch_idx = 0

    for databatch in image_reader:
        if batch_idx % 100 == 0:
            print("%d images done" % batch_idx)
        im = databatch

        if im.shape[0] >= 1200 or im.shape[1] >= 1200:
            all_boxes.append(np.array([]))
            batch_idx += 1
            continue

        t = time.time()

        p_boxes, p_boxes_align = mtcnn_detector.detect_pnet(im=im)

        boxes, boxes_align = mtcnn_detector.detect_rnet(im=im,
                                                        dets=p_boxes_align)

        if boxes_align is None:
            all_boxes.append(np.array([]))
            batch_idx += 1
            continue
        if vis:
            rgb_im = cv2.cvtColor(np.asarray(im), cv2.COLOR_BGR2RGB)
            vision.vis_two(rgb_im, boxes, boxes_align)

        t1 = time.time() - t
        t = time.time()
        all_boxes.append(boxes_align)
        batch_idx += 1

    save_path = config.MODEL_STORE_DIR

    if not os.path.exists(save_path):
        os.mkdir(save_path)

    save_file = os.path.join(save_path, "detections_%d.pkl" % int(time.time()))
    with open(save_file, 'wb') as f:
        cPickle.dump(all_boxes, f, cPickle.HIGHEST_PROTOCOL)

    gen_sample_data(data_dir, anno_file, save_file, prefix_path)
Ejemplo n.º 2
0
import cv2
from core.detect import create_mtcnn_net, MtcnnDetector
import core.vision as vision




if __name__ == '__main__':

    # refer to your local model path 
    p_model = "./model_store/pnet_epoch.pt"
    r_model = "./model_store/rnet_epoch.pt"
    o_model = "./model_store/onet_epoch.pt"

    pnet, rnet, onet = create_mtcnn_net(p_model_path=p_model, r_model_path=r_model, o_model_path=o_model, use_cuda=True)
    mtcnn_detector = MtcnnDetector(pnet=pnet, rnet=rnet, onet=onet, min_face_size=24)

    img = cv2.imread("./test.jpg")
    b, g, r = cv2.split(img)
    img2 = cv2.merge([r, g, b])

    bboxs, landmarks = mtcnn_detector.detect_face(img)
    # print box_align

    vision.vis_face(img2,bboxs,landmarks)
Ejemplo n.º 3
0
import cv2
from core.detect import create_mtcnn_net, MtcnnDetector
import core.vision as vision

if __name__ == '__main__':

    pnet, rnet, onet = create_mtcnn_net(
        p_model_path="./model_store/pnet_epoch.pt",
        r_model_path="./model_store/rnet_epoch.pt",
        o_model_path="./model_store/onet_epoch.pt",
        use_cuda=True)
    mtcnn_detector = MtcnnDetector(pnet=pnet,
                                   rnet=rnet,
                                   onet=onet,
                                   min_face_size=24)

    img = cv2.imread(
        "/home/oljike/PycharmProjects/GlassDataset/Dataset/positive/with_glasses/17.jpg"
    )
    img_bg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    bboxs, landmarks = mtcnn_detector.detect_face(img)
    # print box_align

    vision.vis_face(img_bg, bboxs, landmarks)
Ejemplo n.º 4
0
import cv2
from setuptools.command.saveopts import saveopts

from core.detect import create_mtcnn_net, MtcnnDetector
from core.vision import vis_face

if __name__ == '__main__':
    pnet, rnet, onet = create_mtcnn_net(
        p_model_path="./original_model/pnet_epoch.pt",
        r_model_path="./original_model/rnet_epoch.pt",
        o_model_path="./original_model/onet_epoch.pt",
        use_cuda=False)
    mtcnn_detector = MtcnnDetector(pnet=pnet,
                                   rnet=rnet,
                                   onet=onet,
                                   min_face_size=24)
    img = cv2.imread("./test.jpg")
    img_bg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    #b, g, r = cv2.split(img)
    #img2 = cv2.merge([r, g, b])

    bboxs, landmarks = mtcnn_detector.detect_face(img)
    print(bboxs)
    # print box_align
    save_name = 'r_5.jpg'
    # for i in bboxs:
    #     print(bboxs.shape)
    vis_face(img_bg, bboxs, landmarks, save_name)
    img = cv2.imread(save_name)
    cv2.imshow("img", img)
    cv2.waitKey(0)
def gen_rnet_data(data_dir,
                  anno_file,
                  pnet_model_file,
                  prefix_path='',
                  use_cuda=True,
                  vis=False):
    """
    :param data_dir: train data
    :param anno_file:
    :param pnet_model_file:
    :param prefix_path:
    :param use_cuda:
    :param vis:
    :return:
    """

    # load trained pnet model
    pnet, _, _ = create_mtcnn_net(p_model_path=pnet_model_file,
                                  use_cuda=use_cuda)
    mtcnn_detector = MtcnnDetector(pnet=pnet, min_face_size=12)

    # load original_anno_file, length = 12880
    imagedb = ImageDB(anno_file, mode="test", prefix_path=prefix_path)
    imdb = imagedb.load_imdb()
    image_reader = TestImageLoader(imdb, 1, False)

    all_boxes = list()
    batch_idx = 0

    print('size:%d' % image_reader.size)
    for databatch in image_reader:
        if batch_idx % 100 == 0:
            print("%d images done" % batch_idx)
        im = databatch

        t = time.time()

        # obtain boxes and aligned boxes
        boxes, boxes_align = mtcnn_detector.detect_pnet(im=im)
        if boxes_align is None:
            all_boxes.append(np.array([]))
            batch_idx += 1
            continue
        if vis:
            rgb_im = cv2.cvtColor(np.asarray(im), cv2.COLOR_BGR2RGB)
            vision.vis_two(rgb_im, boxes, boxes_align)

        t1 = time.time() - t
        t = time.time()
        all_boxes.append(boxes_align)
        batch_idx += 1
        # if batch_idx == 100:
        # break
        # print("shape of all boxes {0}".format(all_boxes))
        # time.sleep(5)

    # save_path = model_store_path()
    # './model_store'
    save_path = '/home/quang/Downloads/mtcnn-pytorch/model_store'

    if not os.path.exists(save_path):
        os.mkdir(save_path)

    save_file = os.path.join(save_path, "detections_%d.pkl" % int(time.time()))
    with open(save_file, 'wb') as f:
        cPickle.dump(all_boxes, f, cPickle.HIGHEST_PROTOCOL)

    gen_rnet_sample_data(data_dir, anno_file, save_file, prefix_path)