Beispiel #1
0
def main():
    learning_rate = 4 * 1e-4

    if args.start_from:
        start_from = torch.load(args.start_from, map_location=device)
        nstack = start_from['nstack']
        nfeatures = start_from['nfeatures']
        nlandmarks = start_from['nlandmarks']
        best_val_loss = start_from['best_val_loss']
        eyenet = EyeNet(nstack=nstack,
                        nfeatures=nfeatures,
                        nlandmarks=nlandmarks).to(device)
        optimizer = torch.optim.Adam(eyenet.parameters(), lr=learning_rate)
        eyenet.load_state_dict(start_from['model_state_dict'])
        optimizer.load_state_dict(start_from['optimizer_state_dict'])
    elif os.path.exists(args.out):
        raise Exception(f'Out file {args.out} already exists.')
    else:
        nstack = args.nstack
        nfeatures = args.nfeatures
        nlandmarks = args.nlandmarks
        best_val_loss = float('inf')
        eyenet = EyeNet(nstack=nstack,
                        nfeatures=nfeatures,
                        nlandmarks=nlandmarks).to(device)
        optimizer = torch.optim.Adam(eyenet.parameters(), lr=learning_rate)

    train(eyenet=eyenet,
          optimizer=optimizer,
          nepochs=args.nepochs,
          best_val_loss=best_val_loss,
          checkpoint_fn=args.out)
def init_model(transform):
    global face_cascade, landmarks_detector, checkpoint, nstack, nfeatures, nlandmarks, eyenet
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(device)
    # dirname = os.path.dirname(__file__)
    dirname = pathToProject
    face_cascade = cv2.CascadeClassifier(
        os.path.join(dirname, 'lbpcascade_frontalface_improved.xml'))
    landmarks_detector = dlib.shape_predictor(
        os.path.join(dirname, 'shape_predictor_5_face_landmarks.dat'))
    # face_cascade = cv2.CascadeClassifier(dirname + 'lbpcascade_frontalface_improved.xml')
    # landmarks_detector = dlib.shape_predictor(dirname +'shape_predictor_5_face_landmarks.dat')

    checkpoint = torch.load('checkpoint.pt', map_location=device)
    # checkpoint = torch.load(dirname + 'checkpoint.pt', map_location=device)
    nstack = checkpoint['nstack']
    nfeatures = checkpoint['nfeatures']
    nlandmarks = checkpoint['nlandmarks']
    eyenet = EyeNet(nstack=nstack, nfeatures=nfeatures,
                    nlandmarks=nlandmarks).to(device)
    eyenet.load_state_dict(checkpoint['model_state_dict'])
    return None, None
import os
import numpy as np
import cv2
from util.preprocess import gaussian_2d
from matplotlib import pyplot as plt
import util.gaze

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
dataset = MPIIGaze()
checkpoint = torch.load('checkpoint.pt', map_location=device)
nstack = checkpoint['nstack']
nfeatures = checkpoint['nfeatures']
nlandmarks = checkpoint['nlandmarks']
eyenet = EyeNet(nstack=nstack, nfeatures=nfeatures,
                nlandmarks=nlandmarks).to(device)
eyenet.load_state_dict(checkpoint['model_state_dict'])

with torch.no_grad():
    errors = []

    print('N', len(dataset))
    for i, sample in enumerate(dataset):
        print(i)
        x = torch.tensor([sample['img']]).float().to(device)

        heatmaps_pred, landmarks_pred, gaze_pred = eyenet.forward(x)

        gaze = sample['gaze'].reshape((1, 2))
        gaze_pred = np.asarray(gaze_pred.cpu().numpy())

        if sample['side'] == 'right':