|
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() |
|
) |
|
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() |
|
|
|
|
|
gray_images = images[:num_images] |
|
colorized_images = model(gray_images) |
|
|
|
|
|
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() |
|
|
|
|
|
|
