Source code for DiatomTrack.utils.Transforms

import torch

import numpy as np


class ToTensor(object):
    """
    Creates tensor from array.
    """
    def __call__(self, sample):

        for key in sample.keys():
            if key != 'file_name':
                # transpose channels
                sample[key] = torch.from_numpy(
                    np.transpose(sample[key], (2, 0, 1)).astype(np.float32))

        return sample


class Normalize(object):
    """
    Normalizes an array.
    """
    def __call__(self, sample):

        for key in sample.keys():
            if key != 'file_name':
                sample[key] = (
                    sample[key] - np.min(sample[key])) / (
                        np.max(sample[key]) - np.min(sample[key]))
        return sample