如果事先知道需要忽略哪些文件，當(dāng)然直接從數(shù)據(jù)集里刪除就行了。但如果需要在程序運(yùn)行時(shí)動(dòng)態(tài)確認(rèn)，或者篩選規(guī)則比較復(fù)雜，人工不好做，就需要讓ImageFolder在讀取時(shí)使用自定義的篩選規(guī)則。

ImageFolder有一個(gè)可選參數(shù)為is_valid_file，參數(shù)類型為可調(diào)用的函數(shù)，該函數(shù)傳入一個(gè)str參數(shù)，返回一個(gè)bool值。當(dāng)返回值為True時(shí)保留該文件，否則忽略。

例如，讀取時(shí)想要忽略所有文件名帶‘invalid’的文件，

代碼如下：

import platform
from torchvision.datasets import ImageFolder


class Check(object):
? ? def __init__(self,
? ? ? ? ? ? ? ? ?key_word: str):
? ? ? ? self.key_word = key_word
? ? ? ? self.separator = '\\' if platform.system() == 'Windows' else '/'

? ? def __call__(self,?
? ? ?? ??? ??? ? file_name: str) -> bool:
? ? ? ? folders = file_name.split(self.separator)
? ? ? ? return folders[-1].find(self.key_word) < 0

dataset = ImageFolder('./data', is_valid_file=Check('invalid'))

這里定義了一個(gè)實(shí)現(xiàn)了__call__方法的Check類，相比于直接定義函數(shù)的好處在于可以在構(gòu)造函數(shù)里指定想要忽略的字符，并且能夠根據(jù)操作系統(tǒng)的不同把文件目錄分隔符給確定了。

更加復(fù)雜的功能可以自行修改代碼邏輯實(shí)現(xiàn)，但是要注意如果某個(gè)類別的所有文件都被篩選掉了，ImageFolder會(huì)報(bào)FileNotFoundError錯(cuò)誤。

如果想要忽略整個(gè)類別可以使用下面方法?。?！

二、ImageFolder只讀取部分類別文件夾

直接繼承并且重寫ImageFolder類的find_classes方法即可

from torchvision.datasets.folder import *
from typing import *


class FilterableImageFolder(ImageFolder):
? ? 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,
? ? ? ? ? ? valid_classes: List = None
? ? ):
? ? ? ? self.valid_classes = valid_classes
? ? ? ? super(FilterableImageFolder, self).__init__(root, transform, target_transform, loader, is_valid_file)

? ? def find_classes(self, directory: str) -> Tuple[List[str], Dict[str, int]]:
? ? ? ? classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())
? ? ? ? #增加了這下面這句
? ? ? ? classes = [valid_class for valid_class in classes if valid_class in self.valid_classes]
? ? ? ? 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

使用時(shí)，例如有mouse、cat、dog三個(gè)類別的數(shù)據(jù)集文件夾，只想讀取cat和dog，

代碼如下：

dataset = FilterableImageFolder('./data', valid_classes=['cat', 'dog'])

到此這篇關(guān)于Pytorch中使用ImageFolder讀取數(shù)據(jù)集時(shí)忽略特定文件的文章就介紹到這了,更多相關(guān)ImageFolder讀取數(shù)據(jù)集內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: