Exemplo n.º 1
0
def test_image(args, model):
    if args.detector == 'dlib':
        import dlib
    elif args.detector == 'faceboxes':
        from utils.face_detector import FaceDetectorFaceboxes

    model.eval()

    device = torch.device("cuda" if args.gpu else "cpu")

    image = Image.open(args.image).convert('RGB')

    if args.resize > 0:
        image = resize(image, args.resize)

    detector = None
    if args.detector == 'dlib':
        detector = dlib.get_frontal_face_detector()
    elif args.detector == 'faceboxes':
        MODEL_PATH = 'model/faceboxes.pb'
        detector = FaceDetectorFaceboxes(MODEL_PATH, gpu_memory_fraction=0.25, visible_device_list='0')

    segmenter = Segmenter(model, device, detector, mode=args.detector)

    result = segmenter.segment(PIL2opencv(image), args.remove_small_area)
    result = opencv2PIL(result)

    if args.save:
        result.save(args.save)

    if not args.unshow:
        result.show()
        image.show()
Exemplo n.º 2
0
def test_video(args, model):
    if args.video == '0':
        cap = cv2.VideoCapture(0)
    else:
        cap = cv2.VideoCapture(args.video)

    w_win = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h_win = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    print(w_win, h_win)

    if args.resize > 0:
        short_size = args.resize
        if w_win > h_win:
            nw, nh = short_size, int(w_win * short_size / h_win)
        else:
            nw, nh = int(h_win * short_size / w_win), short_size
    else:
        nw, nh = w_win, h_win

    detector = None
    if args.detector == 'dlib':
        detector = dlib.get_frontal_face_detector()
    elif args.detector == 'faceboxes':
        MODEL_PATH = 'model/faceboxes.pb'
        detector = FaceDetectorFaceboxes(MODEL_PATH, gpu_memory_fraction=0.25, visible_device_list='0')

    device = torch.device("cuda" if args.gpu else "cpu")
    segmenter = Segmenter(model, device, detector, mode=args.detector)

    if args.save:
        fourcc = cv2.VideoWriter_fourcc(*'MJPG')
        out = cv2.VideoWriter(args.save, fourcc, 20, (nh, nw), True)

    while True:
        frame = cap.read()[1]

        if frame is None:
            break

        frame = cv2.resize(frame, (nh, nw))

        result = segmenter.segment(frame, args.remove_small_area)

        if args.save:
            out.write(result)


        if not args.unshow:
            cv2.imshow('image', result)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

    if args.save:
        out.release()