26 lines
771 B
Python
26 lines
771 B
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
from PIL import Image
|
|
from torch.utils.data.dataloader import default_collate
|
|
|
|
|
|
def pil_collate_fn(self, batch):
|
|
batch_data = {}
|
|
for items in batch:
|
|
for key, item in items.items():
|
|
if isinstance(item, Image.Image):
|
|
if key not in batch_data:
|
|
batch_data[key] = []
|
|
batch_data[key].append(item)
|
|
else:
|
|
if key not in batch_data:
|
|
batch_data[key] = []
|
|
batch_data[key].append(item)
|
|
|
|
for key, item in batch_data.items():
|
|
if not all(isinstance(x, Image.Image) for x in item):
|
|
batch_data[key] = default_collate(item)
|
|
|
|
return batch_data
|