Spaces:
Runtime error
Runtime error
File size: 6,231 Bytes
7734d5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 |
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.
import torch
from torch.utils.data.dataloader import DataLoader as torchDataLoader
from torch.utils.data.dataloader import default_collate
import os
import random
from .samplers import YoloBatchSampler
def get_yolox_datadir():
"""
get dataset dir of YOLOX. If environment variable named `YOLOX_DATADIR` is set,
this function will return value of the environment variable. Otherwise, use data
"""
yolox_datadir = os.getenv("YOLOX_DATADIR", None)
if yolox_datadir is None:
import yolox
yolox_path = os.path.dirname(os.path.dirname(yolox.__file__))
yolox_datadir = os.path.join(yolox_path, "datasets")
return yolox_datadir
class DataLoader(torchDataLoader):
"""
Lightnet dataloader that enables on the fly resizing of the images.
See :class:`torch.utils.data.DataLoader` for more information on the arguments.
Check more on the following website:
https://gitlab.com/EAVISE/lightnet/-/blob/master/lightnet/data/_dataloading.py
Note:
This dataloader only works with :class:`lightnet.data.Dataset` based datasets.
Example:
>>> class CustomSet(ln.data.Dataset):
... def __len__(self):
... return 4
... @ln.data.Dataset.resize_getitem
... def __getitem__(self, index):
... # Should return (image, anno) but here we return (input_dim,)
... return (self.input_dim,)
>>> dl = ln.data.DataLoader(
... CustomSet((200,200)),
... batch_size = 2,
... collate_fn = ln.data.list_collate # We want the data to be grouped as a list
... )
>>> dl.dataset.input_dim # Default input_dim
(200, 200)
>>> for d in dl:
... d
[[(200, 200), (200, 200)]]
[[(200, 200), (200, 200)]]
>>> dl.change_input_dim(320, random_range=None)
(320, 320)
>>> for d in dl:
... d
[[(320, 320), (320, 320)]]
[[(320, 320), (320, 320)]]
>>> dl.change_input_dim((480, 320), random_range=None)
(480, 320)
>>> for d in dl:
... d
[[(480, 320), (480, 320)]]
[[(480, 320), (480, 320)]]
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.__initialized = False
shuffle = False
batch_sampler = None
if len(args) > 5:
shuffle = args[2]
sampler = args[3]
batch_sampler = args[4]
elif len(args) > 4:
shuffle = args[2]
sampler = args[3]
if "batch_sampler" in kwargs:
batch_sampler = kwargs["batch_sampler"]
elif len(args) > 3:
shuffle = args[2]
if "sampler" in kwargs:
sampler = kwargs["sampler"]
if "batch_sampler" in kwargs:
batch_sampler = kwargs["batch_sampler"]
else:
if "shuffle" in kwargs:
shuffle = kwargs["shuffle"]
if "sampler" in kwargs:
sampler = kwargs["sampler"]
if "batch_sampler" in kwargs:
batch_sampler = kwargs["batch_sampler"]
# Use custom BatchSampler
if batch_sampler is None:
if sampler is None:
if shuffle:
sampler = torch.utils.data.sampler.RandomSampler(self.dataset)
# sampler = torch.utils.data.DistributedSampler(self.dataset)
else:
sampler = torch.utils.data.sampler.SequentialSampler(self.dataset)
batch_sampler = YoloBatchSampler(
sampler,
self.batch_size,
self.drop_last,
input_dimension=self.dataset.input_dim,
)
# batch_sampler = IterationBasedBatchSampler(batch_sampler, num_iterations =
self.batch_sampler = batch_sampler
self.__initialized = True
def close_mosaic(self):
self.batch_sampler.mosaic = False
def change_input_dim(self, multiple=32, random_range=(10, 19)):
"""This function will compute a new size and update it on the next mini_batch.
Args:
multiple (int or tuple, optional): values to multiply the randomly generated range by.
Default **32**
random_range (tuple, optional): This (min, max) tuple sets the range
for the randomisation; Default **(10, 19)**
Return:
tuple: width, height tuple with new dimension
Note:
The new size is generated as follows: |br|
First we compute a random integer inside ``[random_range]``.
We then multiply that number with the ``multiple`` argument,
which gives our final new input size. |br|
If ``multiple`` is an integer we generate a square size. If you give a tuple
of **(width, height)**, the size is computed
as :math:`rng * multiple[0], rng * multiple[1]`.
Note:
You can set the ``random_range`` argument to **None** to set
an exact size of multiply. |br|
See the example above for how this works.
"""
if random_range is None:
size = 1
else:
size = random.randint(*random_range)
if isinstance(multiple, int):
size = (size * multiple, size * multiple)
else:
size = (size * multiple[0], size * multiple[1])
self.batch_sampler.new_input_dim = size
return size
def list_collate(batch):
"""
Function that collates lists or tuples together into one list (of lists/tuples).
Use this as the collate function in a Dataloader, if you want to have a list of
items as an output, as opposed to tensors (eg. Brambox.boxes).
"""
items = list(zip(*batch))
for i in range(len(items)):
if isinstance(items[i][0], (list, tuple)):
items[i] = list(items[i])
else:
items[i] = default_collate(items[i])
return items
|