def main():
    # define parameters
    num_class = 6
    batch_size = 1
    time_step = 32
    cnn_feat_size = 256     # AlexNet
    gaze_size = 3
    gaze_lstm_hidden_size = 64
    gaze_lstm_projected_size = 128
    # dataset_path = '../data/gaze_dataset'
    dataset_path = '../../gaze-net/gaze_dataset'
    img_size = (224, 224)
    time_skip = 2

    # define model
    arch = 'alexnet'
    extractor_model = FeatureExtractor(arch=arch)
    extractor_model.features = torch.nn.DataParallel(extractor_model.features)
    extractor_model.cuda()      # uncomment this line if using cpu
    extractor_model.eval()

    model = SpatialAttentionModel(num_class, cnn_feat_size,
                        gaze_size, gaze_lstm_hidden_size, gaze_lstm_projected_size)
    model.cuda()

    # load model from checkpoint
    model = load_checkpoint(model)


    trainGenerator = gaze_gen.GazeDataGenerator(validation_split=0.2)
    train_data = trainGenerator.flow_from_directory(dataset_path, subset='training', crop=False,
                    batch_size=batch_size, target_size= img_size, class_mode='sequence_pytorch',
                    time_skip=time_skip)
    # small dataset, error using validation split
    val_data = trainGenerator.flow_from_directory(dataset_path, subset='validation', crop=False,
                batch_size=batch_size, target_size= img_size, class_mode='sequence_pytorch',
                time_skip=time_skip)

    # start predict
    for i in range(10):
        print("start a new interaction")
        # img_seq: (ts,224,224,3), gaze_seq: (ts, 3), ouput: (ts, 6)
        # [img_seq, gaze_seq], target = next(val_data)
        [img_seq, gaze_seq], target = next(train_data)
        restart = True

        predict(img_seq, gaze_seq, extractor_model, model, restart=restart)
        print(target)
        for j in range(img_seq.shape[0]):
            # predict(img_seq[j], gaze_seq[j], None, model, restart=restart)
            # print(target[j])
            # restart = False
            img = img_seq[j,:,:,:]
            gazes = gaze_seq
            cv2.circle(img, (int(gazes[j,1]), int(gazes[j,2])), 10, (255,0,0),-1)
            cv2.imshow('ImageWindow', img)
            cv2.waitKey(33)
def main():
    # define parameters
    TRAIN = True
    num_class = 6
    batch_size = 1
    # time_step = 32
    epochs = 50
    cnn_feat_size = 256     # AlexNet
    gaze_size = 3
    gaze_lstm_hidden_size = 64
    gaze_lstm_projected_size = 128
    learning_rate = 0.0001
    momentum = 0.9
    weight_decay = 1e-4
    eval_freq = 1       # epoch
    print_freq = 1      # iteration
    # dataset_path = '../data/gaze_dataset'
    dataset_path = '../../gaze-net/gaze_dataset'
    img_size = (224, 224)
    log_path = '../log'
    logger = Logger(log_path, 'spatial')

    # define model
    arch = 'alexnet'
    extractor_model = FeatureExtractor(arch=arch)
    extractor_model.features = torch.nn.DataParallel(extractor_model.features)
    extractor_model.cuda()      # uncomment this line if using cpu
    extractor_model.eval()

    model = SpatialAttentionModel(num_class, cnn_feat_size,
                        gaze_size, gaze_lstm_hidden_size, gaze_lstm_projected_size)
    model.cuda()

    # define loss and optimizer
    # criterion = nn.CrossEntropyLoss()
    criterion = nn.CrossEntropyLoss().cuda()
    optimizer = torch.optim.SGD(model.parameters(), learning_rate,
                                momentum = momentum, weight_decay=weight_decay)

    # define generator
    trainGenerator = gaze_gen.GazeDataGenerator(validation_split=0.2)
    train_data = trainGenerator.flow_from_directory(dataset_path, subset='training', crop=False,
                    batch_size=batch_size, target_size= img_size, class_mode='sequence_pytorch')
    # small dataset, error using validation split
    val_data = trainGenerator.flow_from_directory(dataset_path, subset='validation', crop=False,
                batch_size=batch_size, target_size= img_size, class_mode='sequence_pytorch')
    # val_data = train_data

    def test(train_data):
        [img_seq, gaze_seq], target = next(train_data)
        img = img_seq[100,:,:,:]
        img_gamma = adjust_contrast(img)
        imsave('contrast.jpg', img_gamma)
        imsave('original.jpg', img)

    # test(train_data)
    print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
    # img_seq: (ts,224,224,3), gaze_seq: (ts, 3), ouput: (ts, 6)
    # [img_seq, gaze_seq], output = next(train_data)
    # print("gaze data shape")
    # print(img_seq.shape)
    # print(gaze_seq.shape)
    # print(output.shape)

    # start Training
    para = {'bs': batch_size, 'img_size': img_size, 'num_class': num_class,
            'print_freq': print_freq}
    if TRAIN:
        print("get into training mode")
        best_acc = 0

        for epoch in range(epochs):
            adjust_learning_rate(optimizer, epoch, learning_rate)
            print('Epoch: {}'.format(epoch))
            # train for one epoch
            train(train_data, extractor_model, model, criterion, optimizer, epoch, logger, para)

            # evaluate on validation set
            if epoch % eval_freq == 0 or epoch == epochs - 1:
                acc = validate(val_data, extractor_model, model, criterion, epoch, logger, para, False)
                is_best = acc > best_acc
                best_acc = max(acc, best_acc)
                save_checkpoint({
                    'epoch': epoch + 1,
                    'arch': arch,
                    'state_dict': model.state_dict(),
                    'best_acc': best_acc,
                    'optimizer': optimizer.state_dict(),
                }, is_best)
    else:
        model = load_checkpoint(model)
        print("get into testing and visualization mode")
        print("visualization for training data")
        vis_data_path = '../vis/train/'
        if not os.path.exists(vis_data_path):
            os.makedirs(vis_data_path)
        acc = validate(train_data, extractor_model, model, criterion, -1, \
                        logger, para, False, vis_data_path)
        print("visualization for validation data")
        vis_data_path = '../vis/val/'
        if not os.path.exists(vis_data_path):
            os.makedirs(vis_data_path)
        acc = validate(val_data, extractor_model, model, criterion, -1, \
                        logger, para, True, vis_data_path)