# -*- coding: utf-8 -*-

"""
@date: 2022/4/3 下午4:42
@file: collate.py
@author: zj
@description:
"""

import torch

import numpy as np
from typing import Tuple, Optional, Any


def fast_collate(batch, memory_format: Optional[torch.memory_format] = torch.contiguous_format) \
        -> Tuple[Any, Any]:
    imgs = [img[0] for img in batch]
    targets = torch.tensor([target[1] for target in batch], dtype=torch.int64)
    w = imgs[0].size[0]
    h = imgs[0].size[1]
    tensor = torch.zeros((len(imgs), 3, h, w), dtype=torch.uint8).contiguous(memory_format=memory_format)
    for i, img in enumerate(imgs):
        nump_array = np.asarray(img, dtype=np.uint8)
        if (nump_array.ndim < 3):
            nump_array = np.expand_dims(nump_array, axis=-1)
        nump_array = np.rollaxis(nump_array, 2)
        tensor[i] += torch.from_numpy(nump_array)
    return tensor, targets
