Spaces:
Sleeping
Sleeping
File size: 2,966 Bytes
8912244 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.models as models
from PIL import Image
import gradio as gr
# Device setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Image loader
def load_image(image, max_size=400):
transform = transforms.Compose([
transforms.Resize(max_size),
transforms.ToTensor(),
transforms.Lambda(lambda x: x[:3, :, :]), # remove alpha channel
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
image = transform(image).unsqueeze(0)
return image.to(device)
# Reconvert image
def tensor_to_pil(tensor):
image = tensor.cpu().clone().squeeze(0)
image = image * torch.tensor([0.229, 0.224, 0.225]).view(3,1,1) + torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)
image = image.clamp(0, 1)
image = transforms.ToPILImage()(image)
return image
# VGG19 Model
class VGGFeatures(nn.Module):
def __init__(self):
super(VGGFeatures, self).__init__()
vgg = models.vgg19(pretrained=True).features.eval()
self.slice = nn.Sequential(*list(vgg.children())[:21])
for param in self.slice.parameters():
param.requires_grad = False
def forward(self, x):
return self.slice(x)
# Gram Matrix
def gram_matrix(tensor):
b, c, h, w = tensor.size()
features = tensor.view(b * c, h * w)
G = torch.mm(features, features.t())
return G.div(b * c * h * w)
# Style Transfer Function
def style_transfer(content_img, style_img, steps=300, style_weight=1e6, content_weight=1):
content = load_image(content_img)
style = load_image(style_img)
generated = content.clone().requires_grad_(True)
model = VGGFeatures().to(device)
optimizer = torch.optim.LBFGS([generated])
content_features = model(content)
style_features = model(style)
style_gram = gram_matrix(style_features)
def closure():
optimizer.zero_grad()
generated_features = model(generated)
generated_gram = gram_matrix(generated_features)
content_loss = content_weight * nn.functional.mse_loss(generated_features, content_features)
style_loss = style_weight * nn.functional.mse_loss(generated_gram, style_gram)
total_loss = content_loss + style_loss
total_loss.backward()
return total_loss
for i in range(steps):
optimizer.step(closure)
return tensor_to_pil(generated)
# Gradio Interface
interface = gr.Interface(
fn=style_transfer,
inputs=[
gr.Image(type="pil", label="Content Image"),
gr.Image(type="pil", label="Style Image")
],
outputs=gr.Image(type="pil", label="Stylized Output"),
title="🎨 Neural Style Transfer",
description="Upload a content image and a style image. The model will apply the style to the content!"
)
if __name__ == "__main__":
interface.launch()
|