Autoencoder fixed

autoencoder.py

 import torch
-import torchvision
 from torch import nn
-from torch.autograd import Variable
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, random_split, Dataset
 from torchvision import transforms
-from torchvision.utils import save_image
 from torchvision.datasets import ImageFolder
 import os
+from sklearn import svm
+from sklearn.metrics import accuracy_score
+from sklearn.utils.class_weight import compute_class_weight
+import numpy as np
 
 if not os.path.exists('./dc_img'):
     os.mkdir('./dc_img')
@@ -17,6 +18,35 @@ def to_img(x):
     x = x.view(x.size(0), 3, 256, 256)  # Anpassung der Form für RGB-Bilder
     return x
 
+class LimitedImageFolder(Dataset):
+    def __init__(self, root, transform=None, limit_per_class=10000):
+        self.root = root
+        self.transform = transform
+        self.limit_per_class = limit_per_class
+
+        self.image_folder = ImageFolder(root=self.root, transform=self.transform)
+        self.class_indices = self._limit_per_class()
+
+    def _limit_per_class(self):
+        class_indices = {}
+        for i, (image_path, class_label) in enumerate(self.image_folder.imgs):
+            if class_label not in class_indices:
+                class_indices[class_label] = []
+
+            if len(class_indices[class_label]) < self.limit_per_class:
+                class_indices[class_label].append(i)
+
+        return class_indices
+
+    def __getitem__(self, index):
+        original_index = self.class_indices[index // self.limit_per_class][index % self.limit_per_class]
+        return self.image_folder[original_index]
+
+
+    def __len__(self):
+        return len(self.class_indices) * self.limit_per_class
+
+
 num_epochs = 100
 batch_size = 128
 learning_rate = 1e-3
@@ -29,8 +59,16 @@ img_transform = transforms.Compose([
 
 # Anpassung des Datasets auf ImageFolder
 data_dir = 'Plant_leave_diseases_dataset_without_augmentation'  # Setze das Verzeichnis deines Bild-Datasets hier ein
-dataset = ImageFolder(root=data_dir, transform=img_transform)
-dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+dataset = LimitedImageFolder(root=data_dir, transform=img_transform, limit_per_class=10)
+
+# Aufteilung in Trainings- und Testset
+train_size = int(0.8 * len(dataset))
+test_size = len(dataset) - train_size
+train_dataset, test_dataset = random_split(dataset, [train_size, test_size])
+
+# DataLoader für Trainings- und Testset
+train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
 
 class autoencoder(nn.Module):
     def __init__(self):
@@ -53,6 +91,7 @@ class autoencoder(nn.Module):
         )
 
     def forward(self, x):
+        #print("Hallo")
         x = self.encoder(x)
         x = self.decoder(x)
         return x
@@ -65,9 +104,9 @@ optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=
 
 for epoch in range(num_epochs):
     total_loss = 0
-    for data in dataloader:
-        img, _ = data
-        img = Variable(img).to(device)
+    for data in train_dataloader:
+        img, labels = data  # Achtung: Hier gehe ich davon aus, dass deine DataLoader die Labels zurückgeben
+        img = img.to(device)
         output = model(img)
         loss = criterion(output, img)
         optimizer.zero_grad()
@@ -75,9 +114,50 @@ for epoch in range(num_epochs):
         optimizer.step()
         total_loss += loss.data
 
-    print('epoch [{}/{}], loss:{:.4f}'.format(epoch+1, num_epochs, total_loss))
-    if epoch % 10 == 0:
-        pic = to_img(output.cpu().data)
-        save_image(pic, './dc_img/image_{}.png'.format(epoch))
+    print('Autoencoder: epoch [{}/{}], loss:{:.4f}'.format(epoch+1, num_epochs, total_loss))
+
+# Wende den Autoencoder auf den Trainings- und Testdatensatz an und extrahiere den Latent-Space
+model.eval()
+with torch.no_grad():
+    train_latent = []
+    train_labels = []
+    for data in train_dataloader:
+        img, labels = data
+        img = img.to(device)
+        latent = model.encoder(img)
+        train_latent.append(latent.cpu().numpy())
+        train_labels.extend(labels.numpy())
+
+    test_latent = []
+    test_labels = []
+    for data in test_dataloader:
+        img, labels = data
+        img = img.to(device)
+        latent = model.encoder(img)
+        test_latent.append(latent.cpu().numpy())
+        test_labels.extend(labels.numpy())
+unique_classes = torch.unique(torch.tensor(train_labels))
+
+# Konvertiere Latent-Space-Daten in Tensoren
+train_latent = torch.cat([torch.from_numpy(latent) for latent in train_latent], dim=0)
+test_latent = torch.cat([torch.from_numpy(latent) for latent in test_latent], dim=0)
+
+# Flatten Sie die Latent-Space-Daten (optional)
+train_latent = train_latent.view(train_latent.size(0), -1)
+test_latent = test_latent.view(test_latent.size(0), -1)
+
+# Berechne die Klassen-Gewichte basierend auf der Anzahl der Bilder pro Klasse
+class_weights = compute_class_weight('balanced', classes=np.unique(train_labels), y=train_labels)
+# Konvertiere die Gewichte in ein Dictionary
+class_weight_dict = {class_idx: weight for class_idx, weight in enumerate(class_weights)}
+
+# Trainiere eine SVM auf den Latent-Space-Daten
+svm_classifier = svm.SVC(class_weight=class_weight_dict)
+svm_classifier.fit(train_latent, train_labels)
+
+# Klassifiziere den Testdatensatz mit der trainierten SVM
+predicted_labels = svm_classifier.predict(test_latent)
 
-torch.save(model.state_dict(), './conv_autoencoder.pth')
+# Berechne die Genauigkeit
+accuracy = accuracy_score(test_labels, predicted_labels)
+print(f'SVM Accuracy: {accuracy}')
\ No newline at end of file