Initial commit: DCNv4 custom op mirror setup
- Add enhanced README with project structure and quick start guide - Initialize repository with DCNv4 CUDA extension (PyTorch module) - Include classification, detection, and segmentation subdirectories - Reference upstream OpenGVLab DCNv4 implementation Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
This commit is contained in:
7
classification/dataset/__init__.py
Normal file
7
classification/dataset/__init__.py
Normal file
@@ -0,0 +1,7 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
|
||||
from .build import build_loader, build_loader2
|
||||
286
classification/dataset/build.py
Normal file
286
classification/dataset/build.py
Normal file
@@ -0,0 +1,286 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
|
||||
import os
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch.distributed as dist
|
||||
from torchvision import transforms
|
||||
from timm.data import Mixup
|
||||
from timm.data import create_transform
|
||||
from .cached_image_folder import ImageCephDataset
|
||||
from .samplers import SubsetRandomSampler, NodeDistributedSampler
|
||||
|
||||
try:
|
||||
from torchvision.transforms import InterpolationMode
|
||||
|
||||
def _pil_interp(method):
|
||||
if method == 'bicubic':
|
||||
return InterpolationMode.BICUBIC
|
||||
elif method == 'lanczos':
|
||||
return InterpolationMode.LANCZOS
|
||||
elif method == 'hamming':
|
||||
return InterpolationMode.HAMMING
|
||||
else:
|
||||
return InterpolationMode.BILINEAR
|
||||
except:
|
||||
from timm.data.transforms import _pil_interp
|
||||
|
||||
|
||||
class TTA(torch.nn.Module):
|
||||
|
||||
def __init__(self, size, scales=[1.0, 1.05, 1.1]):
|
||||
super().__init__()
|
||||
self.size = size
|
||||
self.scales = scales
|
||||
|
||||
def forward(self, img):
|
||||
out = []
|
||||
cc = transforms.CenterCrop(self.size)
|
||||
for scale in self.scales:
|
||||
size_ = int(scale * self.size)
|
||||
rs = transforms.Resize(size_, interpolation=_pil_interp('bicubic'))
|
||||
img_ = rs(img)
|
||||
img_ = cc(img_)
|
||||
out.append(img_)
|
||||
|
||||
return out
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}(size={self.size}, scale={self.scales})"
|
||||
|
||||
|
||||
def build_loader(config):
|
||||
config.defrost()
|
||||
dataset_train, config.MODEL.NUM_CLASSES = build_dataset('train',
|
||||
config=config)
|
||||
config.freeze()
|
||||
print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}"
|
||||
"successfully build train dataset")
|
||||
|
||||
dataset_val, _ = build_dataset('val', config=config)
|
||||
print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}"
|
||||
"successfully build val dataset")
|
||||
|
||||
dataset_test, _ = build_dataset('test', config=config)
|
||||
print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}"
|
||||
"successfully build test dataset")
|
||||
|
||||
num_tasks = dist.get_world_size()
|
||||
global_rank = dist.get_rank()
|
||||
|
||||
if dataset_train is not None:
|
||||
if config.DATA.IMG_ON_MEMORY:
|
||||
sampler_train = NodeDistributedSampler(dataset_train)
|
||||
else:
|
||||
if config.DATA.ZIP_MODE and config.DATA.CACHE_MODE == 'part':
|
||||
indices = np.arange(dist.get_rank(), len(dataset_train),
|
||||
dist.get_world_size())
|
||||
sampler_train = SubsetRandomSampler(indices)
|
||||
else:
|
||||
sampler_train = torch.utils.data.DistributedSampler(
|
||||
dataset_train,
|
||||
num_replicas=num_tasks,
|
||||
rank=global_rank,
|
||||
shuffle=True)
|
||||
|
||||
if dataset_val is not None:
|
||||
if config.TEST.SEQUENTIAL:
|
||||
sampler_val = torch.utils.data.SequentialSampler(dataset_val)
|
||||
else:
|
||||
sampler_val = torch.utils.data.distributed.DistributedSampler(
|
||||
dataset_val, shuffle=False)
|
||||
|
||||
if dataset_test is not None:
|
||||
if config.TEST.SEQUENTIAL:
|
||||
sampler_test = torch.utils.data.SequentialSampler(dataset_test)
|
||||
else:
|
||||
sampler_test = torch.utils.data.distributed.DistributedSampler(
|
||||
dataset_test, shuffle=False)
|
||||
|
||||
data_loader_train = torch.utils.data.DataLoader(
|
||||
dataset_train,
|
||||
sampler=sampler_train,
|
||||
batch_size=config.DATA.BATCH_SIZE,
|
||||
num_workers=config.DATA.NUM_WORKERS,
|
||||
pin_memory=config.DATA.PIN_MEMORY,
|
||||
drop_last=True,
|
||||
persistent_workers=True) if dataset_train is not None else None
|
||||
|
||||
data_loader_val = torch.utils.data.DataLoader(
|
||||
dataset_val,
|
||||
sampler=sampler_val,
|
||||
batch_size=config.DATA.BATCH_SIZE,
|
||||
shuffle=False,
|
||||
num_workers=config.DATA.NUM_WORKERS,
|
||||
pin_memory=config.DATA.PIN_MEMORY,
|
||||
drop_last=False,
|
||||
persistent_workers=True) if dataset_val is not None else None
|
||||
|
||||
data_loader_test = torch.utils.data.DataLoader(
|
||||
dataset_test,
|
||||
sampler=sampler_test,
|
||||
batch_size=config.DATA.BATCH_SIZE,
|
||||
shuffle=False,
|
||||
num_workers=config.DATA.NUM_WORKERS,
|
||||
pin_memory=config.DATA.PIN_MEMORY,
|
||||
drop_last=False,
|
||||
persistent_workers=True) if dataset_test is not None else None
|
||||
|
||||
# setup mixup / cutmix
|
||||
mixup_fn = None
|
||||
mixup_active = config.AUG.MIXUP > 0 or config.AUG.CUTMIX > 0. or config.AUG.CUTMIX_MINMAX is not None
|
||||
if mixup_active:
|
||||
mixup_fn = Mixup(mixup_alpha=config.AUG.MIXUP,
|
||||
cutmix_alpha=config.AUG.CUTMIX,
|
||||
cutmix_minmax=config.AUG.CUTMIX_MINMAX,
|
||||
prob=config.AUG.MIXUP_PROB,
|
||||
switch_prob=config.AUG.MIXUP_SWITCH_PROB,
|
||||
mode=config.AUG.MIXUP_MODE,
|
||||
label_smoothing=config.MODEL.LABEL_SMOOTHING,
|
||||
num_classes=config.MODEL.NUM_CLASSES)
|
||||
|
||||
return dataset_train, dataset_val, dataset_test, data_loader_train, \
|
||||
data_loader_val, data_loader_test, mixup_fn
|
||||
|
||||
|
||||
def build_loader2(config):
|
||||
config.defrost()
|
||||
dataset_train, config.MODEL.NUM_CLASSES = build_dataset('train',
|
||||
config=config)
|
||||
config.freeze()
|
||||
dataset_val, _ = build_dataset('val', config=config)
|
||||
dataset_test, _ = build_dataset('test', config=config)
|
||||
|
||||
data_loader_train = torch.utils.data.DataLoader(
|
||||
dataset_train,
|
||||
shuffle=True,
|
||||
batch_size=config.DATA.BATCH_SIZE,
|
||||
num_workers=config.DATA.NUM_WORKERS,
|
||||
pin_memory=config.DATA.PIN_MEMORY,
|
||||
drop_last=True,
|
||||
persistent_workers=True) if dataset_train is not None else None
|
||||
|
||||
data_loader_val = torch.utils.data.DataLoader(
|
||||
dataset_val,
|
||||
batch_size=config.DATA.BATCH_SIZE,
|
||||
shuffle=False,
|
||||
num_workers=config.DATA.NUM_WORKERS,
|
||||
pin_memory=config.DATA.PIN_MEMORY,
|
||||
drop_last=False,
|
||||
persistent_workers=True) if dataset_val is not None else None
|
||||
|
||||
data_loader_test = torch.utils.data.DataLoader(
|
||||
dataset_test,
|
||||
batch_size=config.DATA.BATCH_SIZE,
|
||||
shuffle=False,
|
||||
num_workers=config.DATA.NUM_WORKERS,
|
||||
pin_memory=config.DATA.PIN_MEMORY,
|
||||
drop_last=False,
|
||||
persistent_workers=True) if dataset_test is not None else None
|
||||
|
||||
# setup mixup / cutmix
|
||||
mixup_fn = None
|
||||
mixup_active = config.AUG.MIXUP > 0 or config.AUG.CUTMIX > 0. or config.AUG.CUTMIX_MINMAX is not None
|
||||
if mixup_active:
|
||||
mixup_fn = Mixup(mixup_alpha=config.AUG.MIXUP,
|
||||
cutmix_alpha=config.AUG.CUTMIX,
|
||||
cutmix_minmax=config.AUG.CUTMIX_MINMAX,
|
||||
prob=config.AUG.MIXUP_PROB,
|
||||
switch_prob=config.AUG.MIXUP_SWITCH_PROB,
|
||||
mode=config.AUG.MIXUP_MODE,
|
||||
label_smoothing=config.MODEL.LABEL_SMOOTHING,
|
||||
num_classes=config.MODEL.NUM_CLASSES)
|
||||
|
||||
return dataset_train, dataset_val, dataset_test, data_loader_train, \
|
||||
data_loader_val, data_loader_test, mixup_fn
|
||||
|
||||
|
||||
def build_dataset(split, config):
|
||||
transform = build_transform(split == 'train', config)
|
||||
dataset = None
|
||||
nb_classes = None
|
||||
prefix = split
|
||||
if config.DATA.DATASET == 'imagenet':
|
||||
if prefix == 'train' and not config.EVAL_MODE:
|
||||
root = os.path.join(config.DATA.DATA_PATH, 'train')
|
||||
dataset = ImageCephDataset(root,
|
||||
'train',
|
||||
transform=transform,
|
||||
on_memory=config.DATA.IMG_ON_MEMORY)
|
||||
elif prefix == 'val':
|
||||
root = os.path.join(config.DATA.DATA_PATH, 'val')
|
||||
dataset = ImageCephDataset(root, 'val', transform=transform)
|
||||
nb_classes = 1000
|
||||
elif config.DATA.DATASET == 'imagenet22K':
|
||||
if prefix == 'train':
|
||||
if not config.EVAL_MODE:
|
||||
root = config.DATA.DATA_PATH
|
||||
dataset = ImageCephDataset(root,
|
||||
'train',
|
||||
transform=transform,
|
||||
on_memory=config.DATA.IMG_ON_MEMORY)
|
||||
nb_classes = 21841
|
||||
elif prefix == 'val':
|
||||
root = os.path.join(config.DATA.DATA_PATH, 'val')
|
||||
dataset = ImageCephDataset(root, 'val', transform=transform)
|
||||
nb_classes = 1000
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f'build_dataset does support {config.DATA.DATASET}')
|
||||
|
||||
return dataset, nb_classes
|
||||
|
||||
|
||||
def build_transform(is_train, config):
|
||||
resize_im = config.DATA.IMG_SIZE > 32
|
||||
if is_train:
|
||||
# this should always dispatch to transforms_imagenet_train
|
||||
transform = create_transform(
|
||||
input_size=config.DATA.IMG_SIZE,
|
||||
is_training=True,
|
||||
color_jitter=config.AUG.COLOR_JITTER
|
||||
if config.AUG.COLOR_JITTER > 0 else None,
|
||||
auto_augment=config.AUG.AUTO_AUGMENT
|
||||
if config.AUG.AUTO_AUGMENT != 'none' else None,
|
||||
re_prob=config.AUG.REPROB,
|
||||
re_mode=config.AUG.REMODE,
|
||||
re_count=config.AUG.RECOUNT,
|
||||
interpolation=config.DATA.INTERPOLATION,
|
||||
)
|
||||
if not resize_im:
|
||||
# replace RandomResizedCropAndInterpolation with
|
||||
# RandomCrop
|
||||
transform.transforms[0] = transforms.RandomCrop(
|
||||
config.DATA.IMG_SIZE, padding=4)
|
||||
|
||||
return transform
|
||||
|
||||
t = []
|
||||
if resize_im:
|
||||
if config.TEST.CROP:
|
||||
size = int(1.0 * config.DATA.IMG_SIZE)
|
||||
t.append(
|
||||
transforms.Resize(size,
|
||||
interpolation=_pil_interp(
|
||||
config.DATA.INTERPOLATION)),
|
||||
# to maintain same ratio w.r.t. 224 images
|
||||
)
|
||||
t.append(transforms.CenterCrop(config.DATA.IMG_SIZE))
|
||||
elif config.AUG.RANDOM_RESIZED_CROP:
|
||||
t.append(
|
||||
transforms.RandomResizedCrop(
|
||||
(config.DATA.IMG_SIZE, config.DATA.IMG_SIZE),
|
||||
interpolation=_pil_interp(config.DATA.INTERPOLATION)))
|
||||
else:
|
||||
t.append(
|
||||
transforms.Resize(
|
||||
(config.DATA.IMG_SIZE, config.DATA.IMG_SIZE),
|
||||
interpolation=_pil_interp(config.DATA.INTERPOLATION)))
|
||||
t.append(transforms.ToTensor())
|
||||
t.append(transforms.Normalize(config.AUG.MEAN, config.AUG.STD))
|
||||
|
||||
return transforms.Compose(t)
|
||||
538
classification/dataset/cached_image_folder.py
Normal file
538
classification/dataset/cached_image_folder.py
Normal file
@@ -0,0 +1,538 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
|
||||
import io
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import json
|
||||
import math
|
||||
import mmcv
|
||||
import torch
|
||||
import logging
|
||||
import os.path as osp
|
||||
from PIL import Image
|
||||
from tqdm import tqdm, trange
|
||||
from abc import abstractmethod
|
||||
import torch.utils.data as data
|
||||
import torch.distributed as dist
|
||||
from mmcv.fileio import FileClient
|
||||
from .zipreader import is_zip_path, ZipReader
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
_ERROR_RETRY = 50
|
||||
|
||||
|
||||
def has_file_allowed_extension(filename, extensions):
|
||||
"""Checks if a file is an allowed extension.
|
||||
Args:
|
||||
filename (string): path to a file
|
||||
Returns:
|
||||
bool: True if the filename ends with a known image extension
|
||||
"""
|
||||
filename_lower = filename.lower()
|
||||
return any(filename_lower.endswith(ext) for ext in extensions)
|
||||
|
||||
|
||||
def find_classes(dir):
|
||||
classes = [
|
||||
d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))
|
||||
]
|
||||
classes.sort()
|
||||
class_to_idx = {classes[i]: i for i in range(len(classes))}
|
||||
return classes, class_to_idx
|
||||
|
||||
|
||||
def make_dataset(dir, class_to_idx, extensions):
|
||||
images = []
|
||||
dir = os.path.expanduser(dir)
|
||||
for target in sorted(os.listdir(dir)):
|
||||
d = os.path.join(dir, target)
|
||||
if not os.path.isdir(d):
|
||||
continue
|
||||
for root, _, fnames in sorted(os.walk(d)):
|
||||
for fname in sorted(fnames):
|
||||
if has_file_allowed_extension(fname, extensions):
|
||||
path = os.path.join(root, fname)
|
||||
item = (path, class_to_idx[target])
|
||||
images.append(item)
|
||||
|
||||
return images
|
||||
|
||||
|
||||
def make_dataset_with_ann(ann_file, img_prefix, extensions):
|
||||
images = []
|
||||
with open(ann_file, "r") as f:
|
||||
contents = f.readlines()
|
||||
for line_str in contents:
|
||||
path_contents = [c for c in line_str.split('\t')]
|
||||
im_file_name = path_contents[0]
|
||||
class_index = int(path_contents[1])
|
||||
assert str.lower(os.path.splitext(im_file_name)[-1]) in extensions
|
||||
item = (os.path.join(img_prefix, im_file_name), class_index)
|
||||
images.append(item)
|
||||
|
||||
return images
|
||||
|
||||
|
||||
class DatasetFolder(data.Dataset):
|
||||
"""A generic data loader where the samples are arranged in this way: ::
|
||||
root/class_x/xxx.ext
|
||||
root/class_x/xxy.ext
|
||||
root/class_x/xxz.ext
|
||||
root/class_y/123.ext
|
||||
root/class_y/nsdf3.ext
|
||||
root/class_y/asd932_.ext
|
||||
Args:
|
||||
root (string): Root directory path.
|
||||
loader (callable): A function to load a sample given its path.
|
||||
extensions (list[string]): A list of allowed extensions.
|
||||
transform (callable, optional): A function/transform that takes in
|
||||
a sample and returns a transformed version.
|
||||
E.g, ``transforms.RandomCrop`` for images.
|
||||
target_transform (callable, optional): A function/transform that takes
|
||||
in the target and transforms it.
|
||||
Attributes:
|
||||
samples (list): List of (sample path, class_index) tuples
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
root,
|
||||
loader,
|
||||
extensions,
|
||||
ann_file='',
|
||||
img_prefix='',
|
||||
transform=None,
|
||||
target_transform=None,
|
||||
cache_mode="no"):
|
||||
# image folder mode
|
||||
if ann_file == '':
|
||||
_, class_to_idx = find_classes(root)
|
||||
samples = make_dataset(root, class_to_idx, extensions)
|
||||
# zip mode
|
||||
else:
|
||||
samples = make_dataset_with_ann(os.path.join(root, ann_file),
|
||||
os.path.join(root, img_prefix),
|
||||
extensions)
|
||||
|
||||
if len(samples) == 0:
|
||||
raise (RuntimeError("Found 0 files in subfolders of: " + root +
|
||||
"\n" + "Supported extensions are: " +
|
||||
",".join(extensions)))
|
||||
|
||||
self.root = root
|
||||
self.loader = loader
|
||||
self.extensions = extensions
|
||||
|
||||
self.samples = samples
|
||||
self.labels = [y_1k for _, y_1k in samples]
|
||||
self.classes = list(set(self.labels))
|
||||
|
||||
self.transform = transform
|
||||
self.target_transform = target_transform
|
||||
|
||||
self.cache_mode = cache_mode
|
||||
if self.cache_mode != "no":
|
||||
self.init_cache()
|
||||
|
||||
def init_cache(self):
|
||||
assert self.cache_mode in ["part", "full"]
|
||||
n_sample = len(self.samples)
|
||||
global_rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
|
||||
samples_bytes = [None for _ in range(n_sample)]
|
||||
start_time = time.time()
|
||||
for index in range(n_sample):
|
||||
if index % (n_sample // 10) == 0:
|
||||
t = time.time() - start_time
|
||||
print(
|
||||
f'global_rank {dist.get_rank()} cached {index}/{n_sample} takes {t:.2f}s per block'
|
||||
)
|
||||
start_time = time.time()
|
||||
path, target = self.samples[index]
|
||||
if self.cache_mode == "full":
|
||||
samples_bytes[index] = (ZipReader.read(path), target)
|
||||
elif self.cache_mode == "part" and index % world_size == global_rank:
|
||||
samples_bytes[index] = (ZipReader.read(path), target)
|
||||
else:
|
||||
samples_bytes[index] = (path, target)
|
||||
self.samples = samples_bytes
|
||||
|
||||
def __getitem__(self, index):
|
||||
"""
|
||||
Args:
|
||||
index (int): Index
|
||||
Returns:
|
||||
tuple: (sample, target) where target is class_index of the target class.
|
||||
"""
|
||||
path, target = self.samples[index]
|
||||
sample = self.loader(path)
|
||||
if self.transform is not None:
|
||||
sample = self.transform(sample)
|
||||
if self.target_transform is not None:
|
||||
target = self.target_transform(target)
|
||||
|
||||
return sample, target
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def __repr__(self):
|
||||
fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
|
||||
fmt_str += ' Number of datapoints: {}\n'.format(self.__len__())
|
||||
fmt_str += ' Root Location: {}\n'.format(self.root)
|
||||
tmp = ' Transforms (if any): '
|
||||
fmt_str += '{0}{1}\n'.format(
|
||||
tmp,
|
||||
self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
|
||||
tmp = ' Target Transforms (if any): '
|
||||
fmt_str += '{0}{1}'.format(
|
||||
tmp,
|
||||
self.target_transform.__repr__().replace('\n',
|
||||
'\n' + ' ' * len(tmp)))
|
||||
|
||||
return fmt_str
|
||||
|
||||
|
||||
IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']
|
||||
|
||||
|
||||
def pil_loader(path):
|
||||
# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
|
||||
if isinstance(path, bytes):
|
||||
img = Image.open(io.BytesIO(path))
|
||||
elif is_zip_path(path):
|
||||
data = ZipReader.read(path)
|
||||
img = Image.open(io.BytesIO(data))
|
||||
else:
|
||||
with open(path, 'rb') as f:
|
||||
img = Image.open(f)
|
||||
return img.convert('RGB')
|
||||
|
||||
return img.convert('RGB')
|
||||
|
||||
|
||||
def accimage_loader(path):
|
||||
import accimage
|
||||
try:
|
||||
return accimage.Image(path)
|
||||
except IOError:
|
||||
# Potentially a decoding problem, fall back to PIL.Image
|
||||
return pil_loader(path)
|
||||
|
||||
|
||||
def default_img_loader(path):
|
||||
from torchvision import get_image_backend
|
||||
if get_image_backend() == 'accimage':
|
||||
return accimage_loader(path)
|
||||
else:
|
||||
return pil_loader(path)
|
||||
|
||||
|
||||
class CachedImageFolder(DatasetFolder):
|
||||
"""A generic data loader where the images are arranged in this way: ::
|
||||
root/dog/xxx.png
|
||||
root/dog/xxy.png
|
||||
root/dog/xxz.png
|
||||
root/cat/123.png
|
||||
root/cat/nsdf3.png
|
||||
root/cat/asd932_.png
|
||||
Args:
|
||||
root (string): Root directory path.
|
||||
transform (callable, optional): A function/transform that takes in an PIL image
|
||||
and returns a transformed version. E.g, ``transforms.RandomCrop``
|
||||
target_transform (callable, optional): A function/transform that takes in the
|
||||
target and transforms it.
|
||||
loader (callable, optional): A function to load an image given its path.
|
||||
Attributes:
|
||||
imgs (list): List of (image path, class_index) tuples
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
root,
|
||||
ann_file='',
|
||||
img_prefix='',
|
||||
transform=None,
|
||||
target_transform=None,
|
||||
loader=default_img_loader,
|
||||
cache_mode="no"):
|
||||
super(CachedImageFolder,
|
||||
self).__init__(root,
|
||||
loader,
|
||||
IMG_EXTENSIONS,
|
||||
ann_file=ann_file,
|
||||
img_prefix=img_prefix,
|
||||
transform=transform,
|
||||
target_transform=target_transform,
|
||||
cache_mode=cache_mode)
|
||||
self.imgs = self.samples
|
||||
|
||||
def __getitem__(self, index):
|
||||
"""
|
||||
Args:
|
||||
index (int): Index
|
||||
Returns:
|
||||
tuple: (image, target) where target is class_index of the target class.
|
||||
"""
|
||||
path, target = self.samples[index]
|
||||
image = self.loader(path)
|
||||
if self.transform is not None:
|
||||
img = self.transform(image)
|
||||
else:
|
||||
img = image
|
||||
if self.target_transform is not None:
|
||||
target = self.target_transform(target)
|
||||
|
||||
return img, target
|
||||
|
||||
|
||||
class ImageCephDataset(data.Dataset):
|
||||
|
||||
def __init__(self,
|
||||
root,
|
||||
split,
|
||||
parser=None,
|
||||
transform=None,
|
||||
target_transform=None,
|
||||
on_memory=False):
|
||||
if '22k' in root:
|
||||
# Imagenet 22k
|
||||
annotation_root = 'meta/'
|
||||
else:
|
||||
# Imagenet
|
||||
annotation_root = 'meta/'
|
||||
if parser is None or isinstance(parser, str):
|
||||
parser = ParserCephImage(root=root,
|
||||
split=split,
|
||||
annotation_root=annotation_root,
|
||||
on_memory=on_memory)
|
||||
self.parser = parser
|
||||
self.transform = transform
|
||||
self.target_transform = target_transform
|
||||
self._consecutive_errors = 0
|
||||
|
||||
def __getitem__(self, index):
|
||||
img, target = self.parser[index]
|
||||
self._consecutive_errors = 0
|
||||
if self.transform is not None:
|
||||
img = self.transform(img)
|
||||
if target is None:
|
||||
target = -1
|
||||
elif self.target_transform is not None:
|
||||
target = self.target_transform(target)
|
||||
return img, target
|
||||
|
||||
def __len__(self):
|
||||
return len(self.parser)
|
||||
|
||||
def filename(self, index, basename=False, absolute=False):
|
||||
return self.parser.filename(index, basename, absolute)
|
||||
|
||||
def filenames(self, basename=False, absolute=False):
|
||||
return self.parser.filenames(basename, absolute)
|
||||
|
||||
|
||||
class Parser:
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _filename(self, index, basename=False, absolute=False):
|
||||
pass
|
||||
|
||||
def filename(self, index, basename=False, absolute=False):
|
||||
return self._filename(index, basename=basename, absolute=absolute)
|
||||
|
||||
def filenames(self, basename=False, absolute=False):
|
||||
return [
|
||||
self._filename(index, basename=basename, absolute=absolute)
|
||||
for index in range(len(self))
|
||||
]
|
||||
|
||||
|
||||
class ParserCephImage(Parser):
|
||||
|
||||
def __init__(self,
|
||||
root,
|
||||
split,
|
||||
annotation_root,
|
||||
on_memory=False,
|
||||
**kwargs):
|
||||
super().__init__()
|
||||
|
||||
self.file_client = None
|
||||
self.kwargs = kwargs
|
||||
|
||||
self.root = root # dataset:s3://imagenet22k
|
||||
if '22k' in root:
|
||||
self.io_backend = 'petrel'
|
||||
with open(osp.join(annotation_root, '22k_class_to_idx.json'),
|
||||
'r') as f:
|
||||
self.class_to_idx = json.loads(f.read())
|
||||
with open(osp.join(annotation_root, '22k_label.txt'), 'r') as f:
|
||||
self.samples = f.read().splitlines()
|
||||
else:
|
||||
self.io_backend = 'disk'
|
||||
self.class_to_idx = None
|
||||
with open(osp.join(annotation_root, f'{split}.txt'), 'r') as f:
|
||||
self.samples = f.read().splitlines()
|
||||
local_rank = None
|
||||
local_size = None
|
||||
self._consecutive_errors = 0
|
||||
self.on_memory = on_memory
|
||||
if on_memory:
|
||||
self.holder = {}
|
||||
if local_rank is None:
|
||||
local_rank = int(os.environ.get('LOCAL_RANK', 0))
|
||||
if local_size is None:
|
||||
local_size = int(os.environ.get('LOCAL_SIZE', 1))
|
||||
self.local_rank = local_rank
|
||||
self.local_size = local_size
|
||||
self.rank = int(os.environ["RANK"])
|
||||
self.world_size = int(os.environ['WORLD_SIZE'])
|
||||
self.num_replicas = int(os.environ['WORLD_SIZE'])
|
||||
self.num_parts = local_size
|
||||
self.num_samples = int(
|
||||
math.ceil(len(self.samples) * 1.0 / self.num_replicas))
|
||||
self.total_size = self.num_samples * self.num_replicas
|
||||
self.total_size_parts = self.num_samples * self.num_replicas // self.num_parts
|
||||
self.load_onto_memory_v2()
|
||||
|
||||
def load_onto_memory(self):
|
||||
print("Loading images onto memory...", self.local_rank,
|
||||
self.local_size)
|
||||
if self.file_client is None:
|
||||
self.file_client = FileClient(self.io_backend, **self.kwargs)
|
||||
for index in trange(len(self.samples)):
|
||||
if index % self.local_size != self.local_rank:
|
||||
continue
|
||||
path, _ = self.samples[index].split(' ')
|
||||
path = osp.join(self.root, path)
|
||||
img_bytes = self.file_client.get(path)
|
||||
self.holder[path] = img_bytes
|
||||
|
||||
print("Loading complete!")
|
||||
|
||||
def load_onto_memory_v2(self):
|
||||
# print("Loading images onto memory...", self.local_rank, self.local_size)
|
||||
t = torch.Generator()
|
||||
t.manual_seed(0)
|
||||
indices = torch.randperm(len(self.samples), generator=t).tolist()
|
||||
# indices = range(len(self.samples))
|
||||
indices = [i for i in indices if i % self.num_parts == self.local_rank]
|
||||
# add extra samples to make it evenly divisible
|
||||
indices += indices[:(self.total_size_parts - len(indices))]
|
||||
assert len(indices) == self.total_size_parts
|
||||
|
||||
# subsample
|
||||
indices = indices[self.rank // self.num_parts:self.
|
||||
total_size_parts:self.num_replicas // self.num_parts]
|
||||
assert len(indices) == self.num_samples
|
||||
|
||||
if self.file_client is None:
|
||||
self.file_client = FileClient(self.io_backend, **self.kwargs)
|
||||
for index in tqdm(indices):
|
||||
if index % self.local_size != self.local_rank:
|
||||
continue
|
||||
path, _ = self.samples[index].split(' ')
|
||||
path = osp.join(self.root, path)
|
||||
img_bytes = self.file_client.get(path)
|
||||
|
||||
self.holder[path] = img_bytes
|
||||
|
||||
print("Loading complete!")
|
||||
|
||||
def __getitem__(self, index):
|
||||
if self.file_client is None:
|
||||
self.file_client = FileClient(self.io_backend, **self.kwargs)
|
||||
|
||||
filepath, target = self.samples[index].split(' ')
|
||||
filepath = osp.join(self.root, filepath)
|
||||
|
||||
try:
|
||||
if self.on_memory:
|
||||
img_bytes = self.holder[filepath]
|
||||
else:
|
||||
# pass
|
||||
img_bytes = self.file_client.get(filepath)
|
||||
img = mmcv.imfrombytes(img_bytes)[:, :, ::-1]
|
||||
except Exception as e:
|
||||
_logger.warning(
|
||||
f'Skipped sample (index {index}, file {filepath}). {str(e)}')
|
||||
self._consecutive_errors += 1
|
||||
if self._consecutive_errors < _ERROR_RETRY:
|
||||
return self.__getitem__((index + 1) % len(self))
|
||||
else:
|
||||
raise e
|
||||
self._consecutive_errors = 0
|
||||
|
||||
img = Image.fromarray(img)
|
||||
try:
|
||||
if self.class_to_idx is not None:
|
||||
target = self.class_to_idx[target]
|
||||
else:
|
||||
target = int(target)
|
||||
except:
|
||||
print('aaaaaaaaaaaa', filepath, target)
|
||||
exit()
|
||||
|
||||
return img, target
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def _filename(self, index, basename=False, absolute=False):
|
||||
filename, _ = self.samples[index].split(' ')
|
||||
filename = osp.join(self.root, filename)
|
||||
|
||||
return filename
|
||||
|
||||
|
||||
def get_temporal_info(date, miss_hour=False):
|
||||
try:
|
||||
if date:
|
||||
if miss_hour:
|
||||
pattern = re.compile(r'(\d*)-(\d*)-(\d*)', re.I)
|
||||
else:
|
||||
pattern = re.compile(r'(\d*)-(\d*)-(\d*) (\d*):(\d*):(\d*)',
|
||||
re.I)
|
||||
m = pattern.match(date.strip())
|
||||
|
||||
if m:
|
||||
year = int(m.group(1))
|
||||
month = int(m.group(2))
|
||||
day = int(m.group(3))
|
||||
x_month = math.sin(2 * math.pi * month / 12)
|
||||
y_month = math.cos(2 * math.pi * month / 12)
|
||||
if miss_hour:
|
||||
x_hour = 0
|
||||
y_hour = 0
|
||||
else:
|
||||
hour = int(m.group(4))
|
||||
x_hour = math.sin(2 * math.pi * hour / 24)
|
||||
y_hour = math.cos(2 * math.pi * hour / 24)
|
||||
return [x_month, y_month, x_hour, y_hour]
|
||||
else:
|
||||
return [0, 0, 0, 0]
|
||||
else:
|
||||
return [0, 0, 0, 0]
|
||||
except:
|
||||
return [0, 0, 0, 0]
|
||||
|
||||
|
||||
def get_spatial_info(latitude, longitude):
|
||||
if latitude and longitude:
|
||||
latitude = math.radians(latitude)
|
||||
longitude = math.radians(longitude)
|
||||
x = math.cos(latitude) * math.cos(longitude)
|
||||
y = math.cos(latitude) * math.sin(longitude)
|
||||
z = math.sin(latitude)
|
||||
return [x, y, z]
|
||||
else:
|
||||
return [0, 0, 0]
|
||||
114
classification/dataset/samplers.py
Normal file
114
classification/dataset/samplers.py
Normal file
@@ -0,0 +1,114 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
|
||||
import torch
|
||||
import os
|
||||
import math
|
||||
from torch.utils.data.sampler import Sampler
|
||||
import torch.distributed as dist
|
||||
import numpy as np
|
||||
|
||||
|
||||
class SubsetRandomSampler(torch.utils.data.Sampler):
|
||||
"""Samples elements randomly from a given list of indices, without replacement.
|
||||
|
||||
Arguments:
|
||||
indices (sequence): a sequence of indices
|
||||
"""
|
||||
|
||||
def __init__(self, indices):
|
||||
self.epoch = 0
|
||||
self.indices = indices
|
||||
|
||||
def __iter__(self):
|
||||
return (self.indices[i] for i in torch.randperm(len(self.indices)))
|
||||
|
||||
def __len__(self):
|
||||
return len(self.indices)
|
||||
|
||||
def set_epoch(self, epoch):
|
||||
self.epoch = epoch
|
||||
|
||||
|
||||
class NodeDistributedSampler(Sampler):
|
||||
"""Sampler that restricts data loading to a subset of the dataset.
|
||||
It is especially useful in conjunction with
|
||||
:class:`torch.nn.parallel.DistributedDataParallel`. In such case, each
|
||||
process can pass a DistributedSampler instance as a DataLoader sampler,
|
||||
and load a subset of the original dataset that is exclusive to it.
|
||||
.. note::
|
||||
Dataset is assumed to be of constant size.
|
||||
Arguments:
|
||||
dataset: Dataset used for sampling.
|
||||
num_replicas (optional): Number of processes participating in
|
||||
distributed training.
|
||||
rank (optional): Rank of the current process within num_replicas.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
dataset,
|
||||
num_replicas=None,
|
||||
rank=None,
|
||||
local_rank=None,
|
||||
local_size=None):
|
||||
if num_replicas is None:
|
||||
if not dist.is_available():
|
||||
raise RuntimeError(
|
||||
"Requires distributed package to be available")
|
||||
num_replicas = dist.get_world_size()
|
||||
if rank is None:
|
||||
if not dist.is_available():
|
||||
raise RuntimeError(
|
||||
"Requires distributed package to be available")
|
||||
rank = dist.get_rank()
|
||||
if local_rank is None:
|
||||
local_rank = int(os.environ.get('LOCAL_RANK', 0))
|
||||
if local_size is None:
|
||||
local_size = int(os.environ.get('LOCAL_SIZE', 1))
|
||||
self.dataset = dataset
|
||||
self.num_replicas = num_replicas
|
||||
self.num_parts = local_size
|
||||
self.rank = rank
|
||||
self.local_rank = local_rank
|
||||
self.epoch = 0
|
||||
self.num_samples = int(
|
||||
math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
|
||||
self.total_size = self.num_samples * self.num_replicas
|
||||
|
||||
self.total_size_parts = self.num_samples * self.num_replicas // self.num_parts
|
||||
|
||||
def __iter__(self):
|
||||
# deterministically shuffle based on epoch
|
||||
g = torch.Generator()
|
||||
g.manual_seed(self.epoch)
|
||||
|
||||
t = torch.Generator()
|
||||
t.manual_seed(0)
|
||||
|
||||
indices = torch.randperm(len(self.dataset), generator=t).tolist()
|
||||
# indices = range(len(self.dataset))
|
||||
indices = [i for i in indices if i % self.num_parts == self.local_rank]
|
||||
|
||||
# add extra samples to make it evenly divisible
|
||||
indices += indices[:(self.total_size_parts - len(indices))]
|
||||
assert len(indices) == self.total_size_parts
|
||||
|
||||
# subsample
|
||||
indices = indices[self.rank // self.num_parts:self.
|
||||
total_size_parts:self.num_replicas // self.num_parts]
|
||||
|
||||
index = torch.randperm(len(indices), generator=g).tolist()
|
||||
indices = list(np.array(indices)[index])
|
||||
|
||||
assert len(indices) == self.num_samples
|
||||
|
||||
return iter(indices)
|
||||
|
||||
def __len__(self):
|
||||
return self.num_samples
|
||||
|
||||
def set_epoch(self, epoch):
|
||||
self.epoch = epoch
|
||||
102
classification/dataset/zipreader.py
Normal file
102
classification/dataset/zipreader.py
Normal file
@@ -0,0 +1,102 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
|
||||
import os
|
||||
import zipfile
|
||||
import io
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from PIL import ImageFile
|
||||
|
||||
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||||
|
||||
|
||||
def is_zip_path(img_or_path):
|
||||
"""judge if this is a zip path"""
|
||||
return '.zip@' in img_or_path
|
||||
|
||||
|
||||
class ZipReader(object):
|
||||
"""A class to read zipped files"""
|
||||
zip_bank = dict()
|
||||
|
||||
def __init__(self):
|
||||
super(ZipReader, self).__init__()
|
||||
|
||||
@staticmethod
|
||||
def get_zipfile(path):
|
||||
zip_bank = ZipReader.zip_bank
|
||||
if path not in zip_bank:
|
||||
zfile = zipfile.ZipFile(path, 'r')
|
||||
zip_bank[path] = zfile
|
||||
return zip_bank[path]
|
||||
|
||||
@staticmethod
|
||||
def split_zip_style_path(path):
|
||||
pos_at = path.index('@')
|
||||
assert pos_at != -1, "character '@' is not found from the given path '%s'" % path
|
||||
|
||||
zip_path = path[0:pos_at]
|
||||
folder_path = path[pos_at + 1:]
|
||||
folder_path = str.strip(folder_path, '/')
|
||||
return zip_path, folder_path
|
||||
|
||||
@staticmethod
|
||||
def list_folder(path):
|
||||
zip_path, folder_path = ZipReader.split_zip_style_path(path)
|
||||
|
||||
zfile = ZipReader.get_zipfile(zip_path)
|
||||
folder_list = []
|
||||
for file_foler_name in zfile.namelist():
|
||||
file_foler_name = str.strip(file_foler_name, '/')
|
||||
if file_foler_name.startswith(folder_path) and \
|
||||
len(os.path.splitext(file_foler_name)[-1]) == 0 and \
|
||||
file_foler_name != folder_path:
|
||||
if len(folder_path) == 0:
|
||||
folder_list.append(file_foler_name)
|
||||
else:
|
||||
folder_list.append(file_foler_name[len(folder_path) + 1:])
|
||||
|
||||
return folder_list
|
||||
|
||||
@staticmethod
|
||||
def list_files(path, extension=None):
|
||||
if extension is None:
|
||||
extension = ['.*']
|
||||
zip_path, folder_path = ZipReader.split_zip_style_path(path)
|
||||
|
||||
zfile = ZipReader.get_zipfile(zip_path)
|
||||
file_lists = []
|
||||
for file_foler_name in zfile.namelist():
|
||||
file_foler_name = str.strip(file_foler_name, '/')
|
||||
if file_foler_name.startswith(folder_path) and \
|
||||
str.lower(os.path.splitext(file_foler_name)[-1]) in extension:
|
||||
if len(folder_path) == 0:
|
||||
file_lists.append(file_foler_name)
|
||||
else:
|
||||
file_lists.append(file_foler_name[len(folder_path) + 1:])
|
||||
|
||||
return file_lists
|
||||
|
||||
@staticmethod
|
||||
def read(path):
|
||||
zip_path, path_img = ZipReader.split_zip_style_path(path)
|
||||
zfile = ZipReader.get_zipfile(zip_path)
|
||||
data = zfile.read(path_img)
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def imread(path):
|
||||
zip_path, path_img = ZipReader.split_zip_style_path(path)
|
||||
zfile = ZipReader.get_zipfile(zip_path)
|
||||
data = zfile.read(path_img)
|
||||
try:
|
||||
im = Image.open(io.BytesIO(data))
|
||||
except:
|
||||
print("ERROR IMG LOADED: ", path_img)
|
||||
random_img = np.random.rand(224, 224, 3) * 255
|
||||
im = Image.fromarray(np.uint8(random_img))
|
||||
return im
|
||||
Reference in New Issue
Block a user