# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import numbers import os import sys from collections import deque from collections.abc import Mapping import six try: from collections.abc import Sequence, Mapping except: from collections import Sequence, Mapping from ppdet.core.workspace import register, serializable from ppdet.utils.logger import setup_logger from ppdet.data.reader import BaseDataLoader, Compose import cv2 from imageio import imread, imwrite import numpy as np import paddle from paddle.framework import core logger = setup_logger(__name__) def default_collate_fn(batch): """ Default batch collating function for :code:`paddle.io.DataLoader`, get input data as a list of sample datas, each element in list if the data of a sample, and sample data should composed of list, dictionary, string, number, numpy array and paddle.Tensor, this function will parse input data recursively and stack number, numpy array and paddle.Tensor datas as batch datas. e.g. for following input data: [{'image': np.array(shape=[3, 224, 224]), 'label': 1}, {'image': np.array(shape=[3, 224, 224]), 'label': 3}, {'image': np.array(shape=[3, 224, 224]), 'label': 4}, {'image': np.array(shape=[3, 224, 224]), 'label': 5},] This default collate function zipped each number and numpy array field together and stack each field as the batch field as follows: {'image': np.array(shape=[4, 3, 224, 224]), 'label': np.array([1, 3, 4, 5])} Args: batch(list of sample data): batch should be a list of sample data. Returns: Batched data: batched each number, numpy array and paddle.Tensor in input data. """ sample = batch[0] if isinstance(sample, np.ndarray): batch = np.stack(batch, axis=0) return batch elif isinstance(sample, (paddle.Tensor, core.eager.Tensor)): return paddle.stack(batch, axis=0) elif isinstance(sample, numbers.Number): batch = np.array(batch) return batch elif isinstance(sample, (str, bytes)): return batch elif isinstance(sample, Mapping): return {key: default_collate_fn([d[key] for d in batch]) for key in sample} elif isinstance(sample, Sequence): sample_fields_num = len(sample) if not all(len(sample) == sample_fields_num for sample in iter(batch)): raise RuntimeError("fileds number not same among samples in a batch") return [default_collate_fn(fields) for fields in zip(*batch)] raise TypeError("batch data con only contains: tensor, numpy.ndarray, " "dict, list, number, but got {}".format(type(sample))) @register @serializable class MOTVideoStream: """ Load MOT dataset with MOT format from video stream. Args: video_stream (str): path or url of the video file, default ''. keep_ori_im (bool): whether to keep original image, default False. Set True when used during MOT model inference while saving images or video, or used in DeepSORT. """ def __init__(self, video_stream=None, keep_ori_im=False, **kwargs): self.video_stream = video_stream self.keep_ori_im = keep_ori_im self._curr_iter = 0 self.transform = None try: if video_stream == None: print('No video stream is specified, please check the --video_stream option.') raise FileNotFoundError("No video_stream is specified.") self.stream = cv2.VideoCapture(video_stream) if not self.stream.isOpened(): raise Exception("Open video stream Error!") except Exception as e: print('Failed to read {}.'.format(video_stream)) raise e self.videoframeraw_dir = os.path.splitext(os.path.basename(self.video_stream))[0] + '_raw' if not os.path.exists(self.videoframeraw_dir): os.makedirs(self.videoframeraw_dir) def set_kwargs(self, **kwargs): self.mixup_epoch = kwargs.get('mixup_epoch', -1) self.cutmix_epoch = kwargs.get('cutmix_epoch', -1) self.mosaic_epoch = kwargs.get('mosaic_epoch', -1) def set_transform(self, transform): self.transform = transform def set_epoch(self, epoch_id): self._epoch = epoch_id def parse_dataset(self): pass def __iter__(self): ct = 0 while True: ret, frame = self.stream.read() if ret: imgname = os.path.join(self.videoframeraw_dir, 'frame{}.png'.format(ct)) cv2.imwrite(imgname, frame) image = imread(imgname) rec = {'im_id': np.array([ct]), 'im_file': imgname} if self.keep_ori_im: rec.update({'keep_ori_im': 1}) rec['curr_iter'] = self._curr_iter self._curr_iter += 1 ct += 1 if self.transform: yield self.transform(rec) else: yield rec else: return @register @serializable class MOTImageStream: """ Load MOT dataset with MOT format from image stream. Args: keep_ori_im (bool): whether to keep original image, default False. Set True when used during MOT model inference while saving images or video, or used in DeepSORT. """ def __init__(self, sample_num=-1, keep_ori_im=False, **kwargs): self.keep_ori_im = keep_ori_im self._curr_iter = 0 self.transform = None self.imagequeue = deque() self.frameraw_dir = 'inputimages_raw' if not os.path.exists(self.frameraw_dir): os.makedirs(self.frameraw_dir) def add_image(self, image): self.imagequeue.append(image) def set_kwargs(self, **kwargs): self.mixup_epoch = kwargs.get('mixup_epoch', -1) self.cutmix_epoch = kwargs.get('cutmix_epoch', -1) self.mosaic_epoch = kwargs.get('mosaic_epoch', -1) def set_transform(self, transform): self.transform = transform def set_epoch(self, epoch_id): self._epoch = epoch_id def parse_dataset(self): pass def __iter__(self): ct = 0 while True: if self.imagequeue: frame = self.imagequeue.popleft() imgname = os.path.join(self.frameraw_dir, 'frame{}.png'.format(ct)) cv2.imwrite(imgname, frame) image = imread(imgname) rec = {'im_id': np.array([ct]), 'im_file': imgname} if self.keep_ori_im: rec.update({'keep_ori_im': 1}) rec['curr_iter'] = self._curr_iter self._curr_iter += 1 ct += 1 if self.transform: yield self.transform(rec) else: yield rec else: return @register class MOTVideoStreamReader: __shared__ = ['num_classes'] def __init__(self, sample_transforms=[], batch_size=1, drop_last=False, num_classes=1, **kwargs): self._sample_transforms = Compose(sample_transforms, num_classes=num_classes) self.batch_size = batch_size self.drop_last = drop_last self.num_classes = num_classes self.kwargs = kwargs def __call__( self, dataset, worker_num, ): self.dataset = dataset # get data self.dataset.set_transform(self._sample_transforms) # set kwargs self.dataset.set_kwargs(**self.kwargs) self.loader = iter(self.dataset) return self def __len__(self): return sys.maxint def __iter__(self): return self def to_tensor(self, batch): paddle.disable_static() if isinstance(batch, np.ndarray): batch = paddle.to_tensor(batch) elif isinstance(batch, Mapping): batch = {key: self.to_tensor(batch[key]) for key in batch} return batch def __next__(self): try: batch = [] for i in range(self.batch_size): batch.append(next(self.loader)) batch = default_collate_fn(batch) return self.to_tensor(batch) except StopIteration as e: raise e def next(self): # python2 compatibility return self.__next__()