mirror of https://github.com/JDAI-CV/fast-reid.git
51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
# encoding: utf-8
|
|
"""
|
|
@author: xingyu liao
|
|
@contact: sherlockliao01@gmail.com
|
|
"""
|
|
|
|
import glob
|
|
import os
|
|
|
|
from fastreid.data.datasets import DATASET_REGISTRY
|
|
from fastreid.data.datasets.bases import ImageDataset
|
|
|
|
|
|
__all__ = ["Hymenoptera"]
|
|
|
|
|
|
@DATASET_REGISTRY.register()
|
|
class Hymenoptera(ImageDataset):
|
|
"""This is a demo dataset for smoke test, you can refer to
|
|
https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html
|
|
"""
|
|
dataset_dir = 'hymenoptera_data'
|
|
dataset_name = "hyt"
|
|
|
|
def __init__(self, root='datasets', **kwargs):
|
|
self.root = root
|
|
self.dataset_dir = os.path.join(self.root, self.dataset_dir)
|
|
train_dir = os.path.join(self.dataset_dir, "train")
|
|
val_dir = os.path.join(self.dataset_dir, "val")
|
|
|
|
required_files = [
|
|
self.dataset_dir,
|
|
train_dir,
|
|
val_dir,
|
|
]
|
|
self.check_before_run(required_files)
|
|
|
|
train = self.process_dir(train_dir)
|
|
val = self.process_dir(val_dir)
|
|
|
|
super().__init__(train, val, [], **kwargs)
|
|
|
|
def process_dir(self, data_dir):
|
|
data = []
|
|
all_dirs = [d.name for d in os.scandir(data_dir) if d.is_dir()]
|
|
for dir_name in all_dirs:
|
|
all_imgs = glob.glob(os.path.join(data_dir, dir_name, "*.jpg"))
|
|
for img_name in all_imgs:
|
|
data.append([img_name, dir_name, '0'])
|
|
return data
|