torchvision.datasets.ImageFolder
目录
继承关系
初始化方法：
一：find_classes
二：make_dataset
三： 写一个验证函数
四：loader
五：
六： __getitem__：
总结：
继承关系
class ImageFolder(DatasetFolder):
def __init__(
self,
root: str,
transform: Optional[Callable] = None,
target_transform: Optional[Callable] = None,
loader: Callable[[str], Any] = default_loader,
is_valid_file: Optional[Callable[[str], bool]] = None,
):
super().__init__(
root,
loader,
IMG_EXTENSIONS if is_valid_file is None else None,
transform=transform,
target_transform=target_transform,
is_valid_file=is_valid_file,
)
self.imgs = self.samples啥事没干
class DatasetFolder(VisionDataset):
def __init__(
self,
root: str,
loader: Callable[[str], Any],
extensions: Optional[Tuple[str, ...]] = None,
transform: Optional[Callable] = None,
target_transform: Optional[Callable] = None,
is_valid_file: Optional[Callable[[str], bool]] = None,
) -> None:
super().__init__(root, transform=transform, target_transform=target_transform)
classes, class_to_idx = self.find_classes(self.root)
samples = self.make_dataset(self.root, class_to_idx, extensions, is_valid_file)
self.loader = loader
self.extensions = extensions
self.classes = classes
self.class_to_idx = class_to_idx
self.samples = samples
self.targets = [s[1] for s in samples]
@staticmethod
def make_dataset(
directory: str,
class_to_idx: Dict[str, int],
extensions: Optional[Tuple[str, ...]] = None,
is_valid_file: Optional[Callable[[str], bool]] = None,
) -> List[Tuple[str, int]]:
"""Generates a list of samples of a form (path_to_sample, class).
return find_classes(directory)
def __getitem__(self, index: int) -> Tuple[Any, Any]:
"""
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) -> int:
return len(self.samples)初始化方法：
一：find_classes
classes, class_to_idx = self.find_classes(self.root)具体实现主要是：
def find_classes(directory: str) -> Tuple[List[str], Dict[str, int]]:
"""Finds the class folders in a dataset.
See :class:`DatasetFolder` for details.
"""
classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())
if not classes:
raise FileNotFoundError(f"Couldn't find any class folder in {directory}.")
class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
return classes, class_to_idx打断点进去看一下：
总的来说，就是根据路径：得到，文件名，数字索引。当然它将文件名表示为类别。
二：make_dataset
samples = self.make_dataset(self.root, class_to_idx, extensions, is_valid_file)参数说明: root: 文件地址
class_to_idx：类别索引
extensions：图片后缀 ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')
is_valid_file： 是一个可调用的函数： Optional[Callable[[str], bool]]
返回：图片路径，和类别（索引）
if is_valid_file(path):
item = path, class_index
instances.append(item)如果验证通过则会加到返回中，反之不会。
三： 写一个验证函数
class Check:
def __init__(self, key):
print('看看值')
print(key)
def __call__(self, *args, **kwargs):
return True使用：直接传进去
class TestDataset(torchvision.datasets.ImageFolder):
# 根路径,
def __init__(self, root, imgsz, cache, augment):
super().__init__(root=root, is_valid_file=Check)结果：
传的是一个图片地址，可以拿到图片做一些验证工作。
四：loader
self.loader = loader也是一个回调函数
loader: Callable[[str], Any],默认提供的是：
def default_loader(path: str) -> Any:
from torchvision import get_image_backend
if get_image_backend() == "accimage":
return accimage_loader(path)
else:
return pil_loader(path)def pil_loader(path: str) -> Image.Image:
# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
with open(path, "rb") as f:
img = Image.open(f)
return img.convert("RGB")一个读取图片的方法而已。
五：
self.classes = classes
self.class_to_idx = class_to_idx
self.samples = samples
self.targets = [s[1] for s in samples]我们其实只关心samples 图片和targers 标签
六： __getitem__：
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总结：
如果要重写，我觉得主要就是getitem方法。
首先：类给主要给我们提供了，文件读取的方法。我们可以直接拿到文件路径集合。
有了文件路径，就没必要用它的文件加载方法。yolo中有更高效的方法，如果是目标检测，我们可以重新指定标签，比如yolo中规定标签和图片名一样，便于找到。