utils

import torchvision
import matplotlib.pyplot as plt
import numpy as np
from torchsummary import summary
import torch


def get_num_correct(preds, labels):
    return preds.argmax(dim=1).eq(labels).sum().item()

def display_image(image, label):
    plt.imshow(image.squeeze(), cmap='gray')
    print('label:', label)
    plt.waitforbuttonpress()


def display_image_grid(images, labels):
    grid = torchvision.utils.make_grid(images, nrow=10)
    plt.figure(figsize=(15, 15))
    plt.imshow(np.transpose(grid, (1, 2, 0)))
    print('labels:', labels)
    plt.waitforbuttonpress()


def get_device():
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    return device


def get_summary(networ):
    model = networ().to(get_device())
    summary(model, input_size=(1, 28, 28))