from __future__ import print_function

import torch

import torch.nn.functional as F

from torchvision import datasets, transforms, models

import torch.utils.data.dataset

import numpy as np

from tqdm import tqdm

import pandas as pd

import utils

def predict(checkpoint, test_file, gpu_id=0):

    torch.cuda.set_device(gpu_id)

    model = models.densenet169(num_classes = 2).cuda()

    # Load the pretrained model

    model.load_state_dict(torch.load(checkpoint, map_location='cpu'))

    # Set the model in evaluation mode. In this case this is for the Dropout layers

    model = model.half()

    model.eval()

    normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])

    dataset_test = utils.CSVDataset(

            test_file,

            transforms.Compose([

                transforms.Resize(256),

                transforms.CenterCrop(224),

                transforms.ToTensor(),

                normalize,

            ]))

    test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=32, shuffle=False, num_workers=4)

    probs = []

    targets = []

    for data, target in tqdm(test_loader):

        # Send the data and label to the device

        data, target = data.half().cuda(), target.cuda()

        # Forward pass the data through the model

        output = model(data)

        output = F.softmax(output, dim=1)

        probs.extend(output.detach().cpu().numpy())

        targets.extend(target.detach().cpu().numpy())

    return probs, targets

if __name__ == '__main__':

    pass