#!/usr/bin/env python
# coding: utf-8


# get_ipython().run_line_magic('matplotlib', 'inline')
# get_ipython().run_line_magic('config', "InlineBackend.figure_format = 'retina'")

import matplotlib.pyplot as plt

import numpy as np
import torch
import os
from torch import nn
from torch import optim
import torch.nn.functional as F
from torchvision import datasets, transforms, models
from PIL import Image
from torch.autograd import Variable
from binaryClassifier import Net


'''Update the data_dir to input data location and the model location'''

classes = ['typing','noTyping']

model = Net()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

checkpoint = torch.load('model_cifar.pt') #Update the saved model location
model.load_state_dict(checkpoint)
model.eval()



data_dir = '/Users/kesina/Desktop/UsageTesting-Repo/video_data/6pm-video-signin-2/test' # Update the data path. The images needs
# to be in a subfolder within the location given here.

test_transforms = transforms.Compose([transforms.Resize((224, 224)),
            transforms.ToTensor()]) # need to transform the image with same parameters as training image


'''Predicts the image label'''
def predict_image(image):
    image_tensor = test_transforms(image).float()
    image_tensor = image_tensor.unsqueeze_(0)
    input = Variable(image_tensor)
    input = input.to(device)
    output = model(input)
    index = output.data.cpu().numpy().argmax()
    return index



class ImageFolderWithPaths(datasets.ImageFolder):
    """Custom dataset that includes image file paths. Extends
    torchvision.datasets.ImageFolder
    """


    # override the __getitem__ method. this is the method that dataloader calls
    def __getitem__(self, index):
        # this is what ImageFolder normally returns
        original_tuple = super(ImageFolderWithPaths, self).__getitem__(index)
        # the image file path
        path = self.imgs[index][0]
        # make a new tuple that includes original and the path
        tuple_with_path = (original_tuple + (path,))
        return tuple_with_path

# instantiate the dataset and dataloader
dataset = ImageFolderWithPaths(data_dir,transform=test_transforms) # our custom dataset
dataloader = torch.utils.data.DataLoader(dataset)



# iterate over data

for images, labels, paths in dataloader:

    to_pil = transforms.ToPILImage()
    for i in range(len(images)):
        image = to_pil(images[i])
        index = predict_image(image)
        print(paths , ":",str(classes[index]))
