Upload 4 files

model weights, inference code

autoencoder.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c0e83451feb425b0c1e6d795c9e94ec2f10ef0444bb979c02b25de0ae76bfd71
+size 11511830

autoencoder.py

@@ -0,0 +1,81 @@
+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") # if ur not using nvidia for inference, are you a freak who uses directml :eww:
+    print(f"Using device: {device}")
+    dataset = ImageDataset('dataset/images/')
+    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
+    model = aeModel().to(device)
+    #model.load_state_dict(torch.load('autoencoder_250.pth'))
+    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()

autoencoder_inf.py

@@ -0,0 +1,60 @@
+import torch
+from torchvision import transforms
+from PIL import Image
+import matplotlib.pyplot as plt
+from model import aeModel
+
+def load_model(model_path, device):
+    model = aeModel().to(device)
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.eval()
+    return model
+
+def process_single_image(image_path, model, device):
+    transform = transforms.Compose([
+        transforms.Resize((64, 64)),
+        transforms.ToTensor(),
+    ])
+    
+    image = Image.open(image_path).convert('RGB')
+    image_tensor = transform(image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        encoded = model.encode(image_tensor)
+        reconstruction = model.decode(encoded)
+        
+        print(f'Original shape: {image_tensor.shape}')
+        print(f'Encoded shape: {encoded.shape}')
+        print(f'Decoded shape: {reconstruction.shape}')
+
+    return image_tensor.squeeze(0).cpu(), reconstruction.squeeze(0).cpu()
+
+def visualize_original_and_reconstruction(original, reconstruction):
+    original = torch.clamp(original, 0, 1)
+    reconstruction = torch.clamp(reconstruction, 0, 1)
+    
+    fig, axes = plt.subplots(1, 2, figsize=(8, 4))
+    
+    axes[0].imshow(original.permute(1, 2, 0))
+    axes[0].set_title("Original")
+    axes[0].axis("off")
+    
+    axes[1].imshow(reconstruction.permute(1, 2, 0))
+    axes[1].set_title("Decoded")
+    axes[1].axis("off")
+    
+    plt.tight_layout()
+    plt.show()
+
+
+if __name__ == "__main__":
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    model_path = 'autoencoder.pth'
+    model = load_model(model_path, device)
+
+    image_path = r"dataset\images\proof_2.png"
+ 
+    original, reconstruction = process_single_image(image_path, model, device)
+    visualize_original_and_reconstruction(original, reconstruction)

model.py

@@ -0,0 +1,120 @@
+import torch
+from torch import nn
+
+class SelfAttention(nn.Module):
+    def __init__(self, in_channels):
+        super(SelfAttention, self).__init__()
+        self.query = nn.Conv2d(in_channels, in_channels//8, 1)
+        self.key = nn.Conv2d(in_channels, in_channels//8, 1)
+        self.value = nn.Conv2d(in_channels, in_channels, 1)
+        self.gamma = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x):
+        batch_size, C, H, W = x.size()
+        
+        q = self.query(x).view(batch_size, -1, H*W).permute(0, 2, 1)
+        k = self.key(x).view(batch_size, -1, H*W)
+        v = self.value(x).view(batch_size, -1, H*W)
+        
+        attention = torch.bmm(q, k)
+        attention = torch.softmax(attention, dim=-1)
+        
+        out = torch.bmm(v, attention.permute(0, 2, 1))
+        out = out.view(batch_size, C, H, W)
+        
+        return self.gamma * out + x
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
+        self.bn1 = nn.BatchNorm2d(channels)
+        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
+        self.bn2 = nn.BatchNorm2d(channels)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        residual = x
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += residual
+        out = self.relu(out)
+        return out
+
+class aeModel(nn.Module):
+    def __init__(self):
+        super(aeModel, self).__init__()
+
+        self.encoder = nn.ModuleList([
+            nn.Sequential(
+                nn.Conv2d(3, 32, 3, stride=2, padding=1),
+                nn.BatchNorm2d(32),
+                nn.ReLU(),
+                ResidualBlock(32)
+            ),
+            nn.Sequential(
+                nn.Conv2d(32, 64, 3, stride=2, padding=1),
+                nn.BatchNorm2d(64),
+                nn.ReLU(),
+                ResidualBlock(64)
+            ),
+            nn.Sequential(
+                nn.Conv2d(64, 128, 3, stride=2, padding=1),
+                nn.BatchNorm2d(128),
+                nn.ReLU(),
+                ResidualBlock(128),
+                SelfAttention(128)
+            ),
+            nn.Sequential(
+                nn.Conv2d(128, 256, 3, stride=2, padding=1),
+                nn.BatchNorm2d(256),
+                nn.ReLU(),
+                ResidualBlock(256),
+                SelfAttention(256)
+            )
+        ])
+
+        self.decoder = nn.ModuleList([
+            nn.Sequential(
+                nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1),
+                nn.BatchNorm2d(128),
+                nn.ReLU(),
+                ResidualBlock(128),
+                SelfAttention(128)
+            ),
+            nn.Sequential(
+                nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
+                nn.BatchNorm2d(64),
+                nn.ReLU(),
+                ResidualBlock(64)
+            ),
+            nn.Sequential(
+                nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1),
+                nn.BatchNorm2d(32),
+                nn.ReLU(),
+                ResidualBlock(32)
+            ),
+            nn.Sequential(
+                nn.ConvTranspose2d(32, 3, 3, stride=2, padding=1, output_padding=1),
+                nn.Sigmoid()
+            )
+        ])
+
+    def forward(self, x):
+        for encoder_block in self.encoder:
+            x = encoder_block(x)
+        
+        for decoder_block in self.decoder:
+            x = decoder_block(x)
+        
+        return x
+    
+    def encode(self, x):
+        for encoder_block in self.encoder:
+            x = encoder_block(x)
+        return x
+    
+    def decode(self, x):
+        for decoder_block in self.decoder:
+            x = decoder_block(x)
+        return x