Spaces:
Sleeping
Sleeping
| # import gradio as gr | |
| import os | |
| import glob | |
| import time | |
| import numpy as np | |
| from PIL import Image | |
| from skimage import io, color | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn, optim | |
| from torchvision import transforms | |
| import segmentation_models_pytorch as smp | |
| import gradio as gr | |
| ENCODER = 'resnet18' | |
| ENCODER_WEIGHTS = 'imagenet' | |
| ACTIVATION = 'tanh' | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| ResUNet = smp.Unet( | |
| encoder_name=ENCODER, | |
| encoder_weights=ENCODER_WEIGHTS, | |
| in_channels=1, | |
| classes=2, | |
| activation=ACTIVATION, | |
| ) | |
| class Discriminator(nn.Module): | |
| def __init__(self, in_channels): | |
| super(Discriminator, self).__init__() | |
| self.conv1 = nn.Sequential( | |
| nn.Conv2d(in_channels, 64, kernel_size=4, stride=2, padding=1), | |
| nn.LeakyReLU(0.2, inplace=True) | |
| ) | |
| self.conv2 = nn.Sequential( | |
| nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False), | |
| nn.BatchNorm2d(128), | |
| nn.LeakyReLU(0.2, inplace=True) | |
| ) | |
| self.conv3 = nn.Sequential( | |
| nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False), | |
| nn.BatchNorm2d(256), | |
| nn.LeakyReLU(0.2, inplace=True) | |
| ) | |
| self.conv4 = nn.Sequential( | |
| nn.Conv2d(256, 512, kernel_size=4, stride=1, padding=1, bias=False), | |
| nn.BatchNorm2d(512), | |
| nn.LeakyReLU(0.2, inplace=True) | |
| ) | |
| self.conv5 = nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=1) | |
| def forward(self, x): | |
| x = self.conv1(x) | |
| x = self.conv2(x) | |
| x = self.conv3(x) | |
| x = self.conv4(x) | |
| x = self.conv5(x) | |
| return x | |
| def init_weights(net, init_gain=0.02): | |
| def init(m): | |
| if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)): | |
| nn.init.normal_(m.weight.data, 0.0, init_gain) | |
| if m.bias is not None: | |
| nn.init.constant_(m.bias.data, 0.0) | |
| elif isinstance(m, (nn.BatchNorm2d)): | |
| nn.init.constant_(m.weight, 1) | |
| nn.init.constant_(m.bias, 0) | |
| net.apply(init) | |
| return net | |
| class GAN(nn.Module): | |
| def __init__(self, in_channels, out_channels, generator=None): | |
| super(GAN, self).__init__() | |
| self.generator = generator if generator else init_weights(ResUNet(in_channels, out_channels)) | |
| self.generator.to(device) | |
| self.discriminator = init_weights(Discriminator(in_channels=3).to(device)) | |
| self.GANLoss = nn.BCEWithLogitsLoss() | |
| self.L1Loss = nn.L1Loss() | |
| self.gen_optim = optim.Adam(self.generator.parameters(), lr=3e-5, betas=(0.5, 0.99)) | |
| self.disc_optim = optim.Adam(self.discriminator.parameters(), lr=3e-5, betas=(0.5, 0.99)) | |
| def forward(self, L, ab, train): | |
| self.L = L | |
| self.ab = ab | |
| # Generate fake images | |
| self.fake_ab = self.generator(L) | |
| fake_img = torch.cat((self.L, self.fake_ab), dim=1) | |
| disc_fake = self.discriminator(fake_img.detach()) | |
| disc_real = self.discriminator(torch.cat((self.L, self.ab), dim=1)) | |
| self.disc_loss = (self.GANLoss(disc_real, torch.ones_like(disc_real)) + self.GANLoss(disc_fake, torch.zeros_like(disc_fake)))/2 | |
| if train: | |
| self.disc_optim.zero_grad() | |
| self.disc_loss.backward() | |
| self.disc_optim.step() | |
| # Train generator | |
| disc_fake = self.discriminator(fake_img) | |
| gan_loss = self.GANLoss(disc_fake, torch.ones_like(disc_fake)) | |
| L1_loss = self.L1Loss(self.fake_ab, self.ab) * 100 | |
| self.gen_loss = gan_loss + L1_loss | |
| if train: | |
| self.gen_optim.zero_grad() | |
| self.gen_loss.backward() | |
| self.gen_optim.step() | |
| return self.gen_loss.item(), self.disc_loss.item() | |
| gan = GAN(in_channels=1, out_channels=2, generator=ResUNet).to(device) | |
| gan.load_state_dict(torch.load('gan.pth', map_location=torch.device('cpu'), weights_only=False)) | |
| def predict(img): | |
| L = Image.open(img).convert('L') | |
| L = transforms.Resize((256, 256), Image.BICUBIC)(L) | |
| L = transforms.ToTensor()(L) | |
| gan.eval() | |
| with torch.no_grad(): | |
| ab = gan.generator(L.unsqueeze(0).to(device)) | |
| ab = ab.squeeze(0).cpu() | |
| # Denormalize the L and ab channels | |
| L = L.numpy() * 100 # L is in range [0, 1], scale to [0, 100] | |
| ab = ab.numpy() * 128 # ab is in range [-1, 1], scale to [-128, 128] | |
| # Combine L and ab into a single LAB image | |
| lab_image = np.stack([L[0], ab[0], ab[1]], axis=-1) # Shape: [256, 256, 3] | |
| # Convert LAB to RGB using skimage's color.lab2rgb | |
| rgb_image = color.lab2rgb(lab_image) | |
| rgb_image = (rgb_image * 255).astype(np.uint8) | |
| rgb_image = Image.fromarray(rgb_image) | |
| return rgb_image | |
| image = gr.components.Image(type='filepath') | |
| output = gr.components.Image(type='pil') | |
| demo = gr.Interface(fn=predict, inputs=image, outputs=output, title='Colorize Black and White Images', examples=[['examples/1.jpg'], ['examples/2.jpg']]) | |
| demo.launch() |