Pytorch中使用ImageFolder讀取資料集時忽略特定檔案

2022-03-23 13:01:04

一、使用ImageFolder讀取資料集時忽略特定檔案

如果事先知道需要忽略哪些檔案，當然直接從資料集裡刪除就行了。但如果需要在程式執行時動態確認，或者篩選規則比較複雜，人工不好做，就需要讓ImageFolder在讀取時使用自定義的篩選規則。

ImageFolder有一個可選引數為is_valid_file，引數型別為可呼叫的函數，該函數傳入一個str引數，返回一個bool值。當返回值為True時保留該檔案，否則忽略。

例如，讀取時想要忽略所有檔名帶‘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'))

這裡定義了一個實現了__call__方法的Check類，相比於直接定義函數的好處在於可以在建構函式裡指定想要忽略的字元，並且能夠根據作業系統的不同把檔案目錄分隔符給確定了。

更加複雜的功能可以自行修改程式碼邏輯實現，但是要注意如果某個類別的所有檔案都被篩選掉了，ImageFolder會報FileNotFoundError錯誤。

如果想要忽略整個類別可以使用下面方法！！！

二、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

使用時，例如有mouse、cat、dog三個類別的資料集資料夾，只想讀取cat和dog，

程式碼如下：

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

到此這篇關於Pytorch中使用ImageFolder讀取資料集時忽略特定檔案的文章就介紹到這了,更多相關ImageFolder讀取資料集內容請搜尋it145.com以前的文章或繼續瀏覽下面的相關文章希望大家以後多多支援it145.com！

Pytorch中使用ImageFolder讀取資料集時忽略特定檔案

目錄

一、使用ImageFolder讀取資料集時忽略特定檔案

二、ImageFolder唯讀取部分類別資料夾

熱門文章