Create app.py

app.py

@@ -0,0 +1,206 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+import kornia
+
+from PIL import Image
+import numpy as np
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+# Define model
+
+class BlurUpSample(nn.Module):
+    def __init__(self, c):
+        super(BlurUpSample, self).__init__()
+        self.blurpool =  kornia.filters.GaussianBlur2d((3, 3), (1.5, 1.5))
+        self.upsample = nn.Upsample(scale_factor=(2, 2), mode='bilinear', align_corners=False)
+
+    def forward(self, x):
+        x = self.blurpool(x)
+        x = self.upsample(x)
+
+        return x
+
+class DownLayer(nn.Module):
+    def __init__(self, c_in, c_out):
+        super(DownLayer, self).__init__()
+        self.maxblurpool = kornia.filters.MaxBlurPool2D(kernel_size=3)
+        self.conv1 = nn.Conv2d(c_in, c_out, kernel_size=3, stride=1, padding=1)
+        self.bn1 = nn.BatchNorm2d(c_out)
+        self.leakyrelu = nn.LeakyReLU(inplace=True)
+        self.conv2 = nn.Conv2d(c_out, c_out, kernel_size=3, stride=1, padding=1)
+        self.bn2 = nn.BatchNorm2d(c_out)
+
+    def forward(self, x):
+        x = self.maxblurpool(x)
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.leakyrelu(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.leakyrelu(x)
+        return x
+
+
+class UpLayer(nn.Module):
+    def __init__(self, c_in, c_out):
+        super(UpLayer, self).__init__()
+        self.upsample = BlurUpSample(c_in)
+        self.conv1 = nn.Conv2d(c_in+ c_out, c_out, kernel_size=3, stride=1, padding=1)
+        self.bn1 = nn.BatchNorm2d(c_out)
+        self.leakyrelu = nn.LeakyReLU(inplace=True)
+        self.conv2 = nn.Conv2d(c_out, c_out, kernel_size=3, stride=1, padding=1)
+        self.bn2 = nn.BatchNorm2d(c_out)
+
+    def forward(self, x, skip_x):
+        x = self.upsample(x)
+
+        dh = skip_x.size(2) - x.size(2)
+        dw = skip_x.size(3) - x.size(3)
+
+        x = F.pad(x, (dw // 2, dw - dw // 2, dh // 2, dh - dh // 2))
+
+        x = torch.cat([x, skip_x], dim=1)
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.leakyrelu(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.leakyrelu(x)
+        return x
+
+class Generator(nn.Module):
+  def __init__(self):
+    super(Generator, self).__init__()
+    self.conv1 = nn.Conv2d(5, 64, kernel_size=3, stride=1, padding=1)
+    self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
+    self.batchnorm1 = nn.BatchNorm2d(64)
+    self.leakyrelu = nn.LeakyReLU(inplace=True)
+    self.downlayer1 = DownLayer(64, 128)
+    self.downlayer2 = DownLayer(128, 256)
+    self.downlayer3 = DownLayer(256, 512)
+    self.downlayer4 = DownLayer(512, 1024)
+    self.uplayer1 = UpLayer(1024, 512)
+    self.uplayer2 = UpLayer(512, 256)
+    self.uplayer3 = UpLayer(256, 128)
+    self.uplayer4 = UpLayer(128, 64)
+    self.conv3 = nn.Conv2d(64, 3, kernel_size=1, stride=1, padding=0)
+
+  def forward(self, x):
+    #print(f'Input Shape: {x.shape}')
+    x1 = self.conv1(x)
+    x1 = self.batchnorm1(x1)
+    x1 = self.leakyrelu(x1)
+    x1 = self.conv2(x1)
+    x1 = self.batchnorm1(x1)
+    x1 = self.leakyrelu(x1)
+
+    #print(f'Processed Input Shape: {x.shape}')
+
+    x2 = self.downlayer1(x1)
+    x3 = self.downlayer2(x2)
+    x4 = self.downlayer3(x3)
+    x5 = self.downlayer4(x4)
+
+    #print(f'Done Downlayering... Shape: {x5.shape}')
+
+    x = self.uplayer1(x5, x4)
+    x = self.uplayer2(x, x3)
+    x = self.uplayer3(x, x2)
+    x = self.uplayer4(x, x1)
+    x = self.conv3(x)
+
+    #print(f'Output Shape: {x.shape}')
+    return x
+
+
+transform_resize = A.Compose([
+    A.Resize(512, 512),
+    ToTensorV2(),
+])
+
+
+
+
+# Load model
+generator_model = Generator()
+generator_model.load_state_dict(torch.load('age-transformation/large-aging-model.h5',map_location=torch.device(device)))
+generator_model.to(device)
+#generator_model.eval()
+print("")
+
+
+
+def age_filter(image, input_age, output_age):
+
+    resized_image = image.resize((512,512))
+    
+    input_image = transform_resize(image=np.array(image))['image']/255
+
+
+    #input_image=(dataset[0]['normalized_input_image'])
+    age_map1 = torch.full((1, 512, 512), input_age / 100)
+    age_map2 = torch.full((1, 512, 512),  output_age / 100)
+
+    input_tensor = torch.cat((input_image, age_map1,age_map2), dim=0)
+    
+    with torch.no_grad():
+        model_output = generator_model(input_tensor.unsqueeze(0).to(device))
+        
+    np_test = np.array(image)
+
+    new_image = (model_output.squeeze(0).cpu().permute(1,2,0).numpy()*255+np.array(resized_image)).astype('uint8')
+
+    sample_image = np.array(Image.fromarray(new_image).resize((np_test.shape[1],np_test.shape[0]))).astype('uint8')
+    return sample_image
+
+import gradio as gr
+from torchvision.transforms.functional import crop
+
+
+
+def crop_and_process_image(input_img,input_age,output_ag):
+    # Crop the image using the provided crop tool coordinates
+    processed_image = Image.fromarray(input_img)  # Modify this line to preprocess the cropped image
+
+    # Run the processed image through your model
+    output = age_filter(processed_image, input_age, output_ag)
+    # Return the output
+    return output
+
+# Define the input image component with the crop tool
+input_image = gr.Image(label="Input Image", interactive=True)
+
+# Define the output image component
+output_image = gr.Image(label="Output Image", type="pil")
+
+input_image.style(height=512, width=512)
+output_image.style(height=512, width=512)
+
+input_age = gr.Slider(label="Input Age")
+output_age = gr.Slider(label="Output Age")
+
+
+# Define the function to be called when the button is pressed
+def process_image(input_img,input_age,output_age):
+    # Convert the input image to a PyTorch tensor
+
+    # Call the crop_and_process_image function
+    output = crop_and_process_image(input_img,input_age,output_age)
+
+    # Convert the output tensor to a NumPy array and return it
+    output = Image.fromarray(output)
+    output.show()
+    return output
+
+
+
+# Create the Gradio interface
+gr.Interface(fn=process_image, inputs=[input_image,input_age,output_age], outputs=output_image, title="Image Crop and Process").launch(debug=True)