實現Dataset類：

• __init__：初始化信息

• __getitem__：傳入的圖片路徑和下標返回目錄下的圖片信息

• __len__：返回圖片的張數

from torch.utils.data import Dataset import os import cv2 as cvclass MyData(Dataset): # 繼承Datasetdef __init__(self, root_dir, label_dir):self.root_dir = root_dir # 定義根目錄self.label_dir = label_dir # 定義標簽目錄self.path = os.path.join(self.root_dir, self.label_dir) # 定義路徑self.img_path = os.listdir(self.path) # 定義圖片路徑def __getitem__(self, index):img_name = self.img_path[index] # 獲取每一張圖片的名稱img_item_path = os.path.join(self.root_dir, self.label_dir, img_name) # 將存放圖片的路徑進行拼接img = cv.imread(img_item_path) # 讀取路徑中的圖片label = self.label_dir # 圖片的標簽return img, labeldef __len__(self):return len(self.img_path) # 獲取有多少張圖片root_dir = "dataset/val" # 根目錄 ants_label_dir = "ants" # 存放螞蟻的目錄 bees_label_dir = "bees" # 存放蜜蜂的目錄 ants_dataset = MyData(root_dir, ants_label_dir) # 實例化 ants bees_label_dir = MyData(root_dir, bees_label_dir) # 實例化 bees print("ants_dataset：", len(ants_dataset)) print("bees_label_dir：", len(bees_label_dir))train_dataset = ants_dataset + bees_label_dir # 將兩個數據集存放到一起 print("train_dataset：", len(train_dataset))img1, label1 = train_dataset[69] # 查看第70張圖片和標簽信息img2, label2 = train_dataset[70] # 查看第71張圖片和標簽信息# 圖像顯示 cv.imshow("image1", img1) cv.imshow("image2", img2) cv.waitKey(0) cv.destroyAllWindows()

總結

以上是生活随笔為你收集整理的torch Dataset类的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。