add blur object button and multiple object support

app.py

@@ -7,50 +7,13 @@ import random
 import gradio as gr
 import numpy as np
 
+output_dict = {} # this dict is shared between segment and blur_background functions
+pred_label_unq = []
+
 
 def random_color_gen(n):
     return [tuple(random.randint(0,255) for i in range(3)) for i in range(n)]
 
-import math
-
-def get_gaussian_kernel(kernel_size=15, sigma=20, channels=3):
-    # Create a x, y coordinate grid of shape (kernel_size, kernel_size, 2)
-    x_coord = torch.arange(kernel_size)
-    x_grid = x_coord.repeat(kernel_size).view(kernel_size, kernel_size)
-    y_grid = x_grid.t()
-    xy_grid = torch.stack([x_grid, y_grid], dim=-1).float()
-
-    mean = (kernel_size - 1)/2.
-    variance = sigma**2.
-
-    # Calculate the 2-dimensional gaussian kernel which is
-    # the product of two gaussian distributions for two different
-    # variables (in this case called x and y)
-    gaussian_kernel = (1./(2.*math.pi*variance)) *\
-                      torch.exp(
-                          -torch.sum((xy_grid - mean)**2., dim=-1) /\
-                          (2*variance)
-                      )
-
-    # Make sure sum of values in gaussian kernel equals 1.
-    gaussian_kernel = gaussian_kernel / torch.sum(gaussian_kernel)
-
-    # Reshape to 2d depthwise convolutional weight
-    gaussian_kernel = gaussian_kernel.view(1, 1, kernel_size, kernel_size)
-    gaussian_kernel = gaussian_kernel.repeat(channels, 1, 1, 1)
-
-    gaussian_filter = nn.Conv2d(in_channels=channels, out_channels=channels,
-                                kernel_size=kernel_size, padding='same', groups=channels, bias=False)
-
-    gaussian_filter.weight.data = gaussian_kernel
-    gaussian_filter.weight.requires_grad = False
-    
-    return gaussian_filter
-
-
-output_dict = {} # this dict is shared between segment and blur_background functions
-pred_label_unq = []
-
 def segment(input_image):
     
     # prepare image for display
@@ -98,26 +61,54 @@ def segment(input_image):
     
     return bounding_box_img;
 
-def blur_background(input_image, label_name):
-    mask = output_dict[label_name]['mask']
-    mask = torch.tensor(mask).unsqueeze(0)
+
+def blur_object(input_image, label_name):
+
+    label_names = label_name.split(' ')
     
     input_tensor = T.ToTensor()(input_image).unsqueeze(0)
-    blur = get_gaussian_kernel()
+    blur = T.GaussianBlur(15, 20)
     blurred_tensor = blur(input_tensor)
+
+    final_img = input_tensor
+
+    for name in label_names:
+        mask = output_dict[name.strip()]['mask']
+        mask = torch.tensor(mask).unsqueeze(0)
     
-    final_img = blurred_tensor
-    final_img[:, :, mask.squeeze(0)] = input_tensor[:, :, mask.squeeze(0)];
+        final_img[:, :, mask.squeeze(0)] = blurred_tensor[:, :, mask.squeeze(0)];
     
     final_img = T.ToPILImage()(final_img.squeeze(0))
     
     return final_img;
+
+def blur_background(input_image, label_name):
+    label_names = label_name.split(' ')
+
+    input_tensor = T.ToTensor()(input_image).unsqueeze(0)
+    blur = T.GaussianBlur(15, 20)
+    blurred_tensor = blur(input_tensor)
+
+    final_img = blurred_tensor
+
+
+    for name in label_names:
+        mask = output_dict[name.strip()]['mask']
+        mask = torch.tensor(mask).unsqueeze(0)
     
+        final_img[:, :, mask.squeeze(0)] = input_tensor[:, :, mask.squeeze(0)];
     
+    final_img = T.ToPILImage()(final_img.squeeze(0))
     
-
+    return final_img;
     
     
+    
+
+############################
+""" User Interface """
+############################
+
 with gr.Blocks() as app:
     
     gr.Markdown("# Blur an objects background with AI")
@@ -126,26 +117,22 @@ with gr.Blocks() as app:
     with gr.Column():
         input_image = gr.Image(type='pil')
         b1 = gr.Button("Segment Image")
-        
-    
-    
+           
     with gr.Row():
-#         masked_image = gr.Image();
         bounding_box_image = gr.Image();
-    
-    
+     
     gr.Markdown("Now choose a label (eg: person1) from the above image of your desired object and input it below")
+    gr.Markdown("You can also input multiple labels separated by spaces (eg: person1 car1 handbag1)")
     with gr.Column():
         label_name = gr.Textbox()
-        b2 = gr.Button("Blur Backbround")
+        with gr.Row():
+            b2 = gr.Button("Blur Backbround")
+            b3 = gr.Button("Blur Object")
         result = gr.Image()
     
     b1.click(segment, inputs=input_image, outputs=bounding_box_image)
     b2.click(blur_background, inputs=[input_image, label_name], outputs=result)
+    b3.click(blur_object, inputs=[input_image, label_name], outputs=result)
     
-    
-    
-    
-#     instance_segmentation = gr.Interface(segment, inputs=input_image, outputs=['json', 'image'])
 
 app.launch(debug=True)