Add: model class

model.py

@@ -0,0 +1,111 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset
+from torchvision import datasets, transforms
+import matplotlib.pyplot as plt
+import numpy as np
+from PIL import Image
+import os
+
+class ColorNet(nn.Module):
+    DEFAULT_CHECKPOINT_PATH = "checkpoint/colornet.pt"
+
+    def __init__(self, checkpoint_path:str=DEFAULT_CHECKPOINT_PATH):
+        super(ColorNet, self).__init__()
+
+        self.encoder = nn.Sequential(
+            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
+            nn.ReLU()
+        )
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(64, 3, kernel_size=3, stride=1, padding=1),
+            nn.Sigmoid()  # to scale the output to [0, 1]
+        )
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.to(self.device)
+        
+        if os.path.exists(checkpoint_path):
+           self._load_model(checkpoint_path) 
+
+    def _load_model(self, path):
+        print("Loading ColorNet model...", end="")
+        self.load_state_dict(torch.load(path, map_location=self.device))
+        print("done.")
+
+    def forward(self, x):
+        x = x.to(self.device)
+        x = self.encoder(x)
+        x = self.decoder(x)
+        return x
+
+    def train_model(self, model, train_loader, criterion, optimizer, num_epochs=10):
+        for epoch in range(num_epochs):
+            model.train()
+            running_loss = 0.0
+            for inputs, _ in train_loader:
+                gray_images = transforms.Grayscale(num_output_channels=1)(inputs).to(self.device)
+                gray_images = gray_images.repeat(1,3,1,1)
+                color_images = inputs.to(self.device)
+
+                optimizer.zero_grad()
+
+                outputs = model(gray_images)
+                loss = criterion(outputs, color_images)
+                loss.backward()
+                optimizer.step()
+
+                running_loss += loss.item() * gray_images.size(0)
+
+            epoch_loss = running_loss / len(train_loader.dataset)
+            print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}')
+        
+        torch.save(model.state_dict(), self.DEFAULT_CHECKPOINT_PATH)
+
+
+    def colorize(self, input_path:str, output_path):
+        input_image = Image.open(input_path).convert("RGB")
+        input_image = transforms.ToTensor()(input_image).unsqueeze(0).to(self.device)
+
+        with torch.inference_mode():
+            output_image_tnsr = self(input_image)
+            output_image_tnsr = output_image_tnsr.squeeze(0).cpu()
+            output_image_tnsr = transforms.ToPILImage()(output_image_tnsr)
+
+            output_image_tnsr.save(output_path)
+
+    def visualize_results(model, test_loader, num_images=5):
+        model.eval()
+        with torch.no_grad():
+            data_iter = iter(test_loader)
+            images, _ = data_iter.next()
+            
+            # Get grayscale and colorized images
+            gray_images = images[:num_images]
+            colorized_images = model(gray_images)
+            
+            # Plotting the results
+            for i in range(num_images):
+                plt.subplot(3, num_images, i+1)
+                plt.imshow(gray_images[i].permute(1, 2, 0).squeeze(), cmap="gray")
+                plt.axis('off')
+
+                plt.subplot(3, num_images, num_images+i+1)
+                plt.imshow(colorized_images[i].permute(1, 2, 0))
+                plt.axis('off')
+
+                plt.subplot(3, num_images, 2*num_images+i+1)
+                plt.imshow(gray_images[i].permute(1, 2, 0).repeat(3, 1, 1).permute(1, 2, 0))
+                plt.axis('off')
+
+            plt.show()
+
+