|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torch.utils.data import Dataset, DataLoader
|
|
from torchvision import transforms
|
|
from PIL import Image
|
|
import os
|
|
from tqdm import tqdm
|
|
import matplotlib.pyplot as plt
|
|
from model import aeModel
|
|
|
|
class ImageDataset(Dataset):
|
|
def __init__(self, folder_path):
|
|
self.folder_path = folder_path
|
|
self.image_files = [f for f in os.listdir(folder_path) if f.endswith(('.jpg', '.jpeg', '.png'))]
|
|
self.transform = transforms.Compose([
|
|
transforms.Resize((64, 64)),
|
|
transforms.ToTensor(),
|
|
])
|
|
|
|
def __len__(self):
|
|
return len(self.image_files)
|
|
|
|
def __getitem__(self, idx):
|
|
img_path = os.path.join(self.folder_path, self.image_files[idx])
|
|
image = Image.open(img_path).convert('RGB')
|
|
image = self.transform(image)
|
|
return image
|
|
|
|
def train(model, dataloader, num_epochs, device):
|
|
criterion = nn.MSELoss()
|
|
optimizer = optim.Adam(model.parameters(), lr=1e-3)
|
|
|
|
for epoch in range(num_epochs):
|
|
model.train()
|
|
total_loss = 0
|
|
for batch in tqdm(dataloader, desc=f'Epoch {epoch+1}/{num_epochs}'):
|
|
batch = batch.to(device)
|
|
output = model(batch)
|
|
loss = criterion(output, batch)
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
total_loss += loss.item()
|
|
|
|
avg_loss = total_loss / len(dataloader)
|
|
print(f'Epoch [{epoch+1}/{num_epochs}], Average Loss: {avg_loss:.4f}')
|
|
|
|
|
|
def visualize_results(model, dataloader, device):
|
|
model.eval()
|
|
with torch.no_grad():
|
|
images = next(iter(dataloader))
|
|
images = images.to(device)
|
|
|
|
reconstructions = model(images)
|
|
fig, axes = plt.subplots(2, 5, figsize=(12, 6))
|
|
for i in range(5):
|
|
axes[0, i].imshow(images[i].cpu().permute(1, 2, 0))
|
|
axes[0, i].axis('off')
|
|
axes[1, i].imshow(reconstructions[i].cpu().permute(1, 2, 0))
|
|
axes[1, i].axis('off')
|
|
|
|
plt.tight_layout()
|
|
plt.show()
|
|
|
|
def main():
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Using device: {device}")
|
|
dataset = ImageDataset('dataset/images/')
|
|
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
|
|
model = aeModel().to(device)
|
|
|
|
num_epochs = 250
|
|
train(model, dataloader, num_epochs, device)
|
|
visualize_results(model, dataloader, device)
|
|
torch.save(model.state_dict(), 'autoencoder.pth')
|
|
|
|
if __name__ == "__main__":
|
|
main() |
