Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 2,381 Bytes
a29a059 42acc98 a29a059 c2297fd a29a059 |
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 |
import gradio as gr
import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision
import kornia as K
def inference(file1,num_iters):
img: np.ndarray = cv2.imread(file1.name)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
img = img + np.random.normal(loc=0.0, scale=0.1, size=img.shape)
img = np.clip(img, 0.0, 1.0)
# convert to torch tensor
noisy_image = K.utils.image_to_tensor(img).squeeze()
class TVDenoise(torch.nn.Module):
def __init__(self, noisy_image):
super(TVDenoise, self).__init__()
self.l2_term = torch.nn.MSELoss(reduction='mean')
self.regularization_term = K.losses.TotalVariation()
# create the variable which will be optimized to produce the noise free image
self.clean_image = torch.nn.Parameter(data=noisy_image.clone(), requires_grad=True)
self.noisy_image = noisy_image
def forward(self):
# print(self.l2_term(self.clean_image, self.noisy_image))
# print(self.regularization_term(self.clean_image))
return self.l2_term(self.clean_image, self.noisy_image) + 0.0001 * self.regularization_term(self.clean_image)
def get_clean_image(self):
return self.clean_image
tv_denoiser = TVDenoise(noisy_image)
# define the optimizer to optimize the 1 parameter of tv_denoiser
optimizer = torch.optim.SGD(tv_denoiser.parameters(), lr=0.1, momentum=0.9)
for i in range(int(num_iters)):
optimizer.zero_grad()
loss = torch.mean(tv_denoiser())
if i % 50 == 0:
print("Loss in iteration {} of {}: {:.3f}".format(i, num_iters, loss.item()))
loss.backward()
optimizer.step()
img_clean: np.ndarray = K.utils.tensor_to_image(tv_denoiser.get_clean_image())
return img, img_clean
examples = [ ["doraemon.png",2000]
]
inputs = [
gr.Image(type='file', label='Input Image'),
gr.Slider(minimum=50, maximum=10000, step=50, default=500, label="num_iters")
]
outputs = [
gr.Image(type='file', label='Noised Image'),
gr.Image(type='file', label='Denoised Image'),
]
title = "Denoise image using total variation"
demo_app = gr.Interface(
fn=inference,
inputs=inputs,
outputs=outputs,
title=title,
examples=examples,
theme='huggingface',
)
demo_app.launch(debug=True) |