update code and readme

hwdb.py

@@ -1,4 +1,5 @@
 import os
+import random
 from torch.utils.data import DataLoader
 import torchvision.transforms as transforms
 import torchvision.datasets as datasets
@@ -6,7 +7,7 @@ import matplotlib.pyplot as plt
 
 
 class HWDB(object):
-    def __init__(self, transform, path='./data'):
+    def __init__(self,path, transform):
         # 预处理过程
 
         traindir = os.path.join(path, 'train')
@@ -16,6 +17,8 @@ class HWDB(object):
         self.testset = datasets.ImageFolder(testdir, transform)
         self.train_size = len(self.trainset)
         self.test_size = len(self.testset)
+        self.num_classes = len(self.trainset.classes)
+        self.class_to_idx = self.trainset.class_to_idx
 
     def get_sample(self, index=0):
         sample = self.trainset[index]
@@ -30,21 +33,14 @@ class HWDB(object):
 
 if __name__ == '__main__':
     transform = transforms.Compose([
-        # transforms.Grayscale(),
         transforms.Resize((64, 64)),
         transforms.ToTensor(),
     ])
-    dataset = HWDB(transform=transform, path=r'data')
-    print(dataset.train_size)
-    print(dataset.test_size)
-    for i in [1020, 120, 2000, 6000, 1000]:
-        img, label = dataset.get_sample(i)
-        img = img[0]
-        print(label)
-        plt.imshow(img, cmap='gray')
-        plt.show()
-
-    train_loader, test_loader = dataset.get_loader()
-    # for (img, label) in train_loader:
-    #     print(img)
-    #     print(label)
+    dataset = HWDB(path=r'data', transform=transform)
+    print("训练集数量：", dataset.train_size)
+    print("测试集数量：", dataset.test_size)
+    print("类别数量：", dataset.num_classes)
+    index = random.randint(0, dataset.train_size)
+    img = dataset.get_sample(index)[0][0]
+    plt.imshow(img, cmap='gray')
+    plt.show()