Update app.py

app.py

@@ -55,7 +55,6 @@ print("Model ready!")
 ##################
 
 to_tensor = transforms.ToTensor()
-to_array = transforms.ToPILImage()
 resize = transforms.Resize((512,512))
 resize_small = transforms.Resize((369,369))
 normalize = transforms.Normalize(
@@ -91,19 +90,13 @@ def transparent(fg, bg, alpha_factor):
     return background
 
 def show_img(all_imgs, dropdown, bg, alpha_factor):
-    if all_imgs is None: 
-        return None 
         
-    idx = target_list_all.index(dropdown)
-    fg = all_imgs[idx]
+    idx = target_list_all.index(label)
+    fg = mask_images[idx].copy()
+    bg = bg.copy()
 
-    foreground = Image.open(fg)
-    background = np.array(bg)
-
-    background = Image.fromarray(bg)
-    new_alpha_factor = int(255*alpha_factor)
-    foreground.putalpha(new_alpha_factor)
-    background.paste(foreground, (0, 0), foreground)
+    fg = putalpha(int(255 * alpha))
+    bg.paste(fg, (0, 0), fg)
 
     return background
 
@@ -112,30 +105,19 @@ def show_img(all_imgs, dropdown, bg, alpha_factor):
 
 def inference(img):
     background = resize_pil(img) 
-
     img = process_pil(img).unsqueeze(0)
 
     with torch.no_grad():
-        mask = model(img)[0]
+        logits = model(img)[0]
 
 
     # Get probability values (logits to probs)
-    mask_probs = torch.sigmoid(mask)
-    mask_probs = mask_probs.detach().numpy()
-    mask_probs.shape
-
-    # Make binary mask
-    THRESHOLD = 0.5
-    mask_preds = mask_probs > THRESHOLD
-
-    # All combined
-    mask_all =  mask_preds.sum(axis=0)
-    mask_all = np.expand_dims(mask_all, axis=0)
-    mask_all.shape
+    probs = torch.sigmoid(logits).numpy()
+    mask = probs > 0.5
 
     # Concat all combined with normal preds
-    mask_preds = np.concatenate((mask_all, mask_preds),axis=0)
-    labs = ["ALL"] + target_list
+    mask_all = np.sum(masks, axis=0, keepdims=True)
+    masks = np.concatenate([mask_all, masks], axis=0)
 
     fig, axes = plt.subplots(5, 4, figsize = (10,10))
     
@@ -143,34 +125,31 @@ def inference(img):
     all_masks = []
 
     for i, ax in enumerate(axes.flat):
-        label = labs[i]
-        
-        all_masks.append(mask_preds[i])
-
-        ax.imshow(mask_preds[i])
-        ax.set_title(label)
-          
-    plt.tight_layout()        
+        ax.imshow(masks[i])
+        ax.set_title(target_list_all[i])
+        ax.axis("off")
+                 
 
     # plt to PIL
-    img_buf = io.BytesIO()
-    fig.savefig(img_buf, format='png')
-    im = Image.open(img_buf)
+    buf = io.BytesIO()
+    plt.tight_layout()
+    plt.savefig(buf, format='png')
+    plt.close()
+    preview = Image.open(buf)
 
     # Saved all masks combined with unvisible xaxis und yaxis and without a white 
     # background.
-    all_images = []
-    for i in range(len(all_masks)):
-        plt.figure()
-        fig = plt.imshow(all_masks[i])
+    mask_images = []
+    for m in masks:
+        fig = plt.figure()
+        plt.imshow(m)
         plt.axis('off')
-        fig.axes.get_xaxis().set_visible(False)
-        fig.axes.get_yaxis().set_visible(False)
-        img_buf = io.BytesIO()
-        plt.savefig(img_buf, bbox_inches='tight', pad_inches = 0, format='png')
-        all_images.append(Image.open(img_buf))
+        buf = io.BytesIO()
+        plt.savefig(buf, bbox_inches='tight', pad_inches = 0)
+        plt.close()
+        mask_images.append(Image.open(buf).convert("RGBA"))
 
-    return im, all_images, background
+    return preview, mask_images, background, mask_images
 
     
 
@@ -212,16 +191,16 @@ with gr.Blocks(title=title) as app:
         dropdown = gr.Dropdown(choices=target_list_all, label="Select Label", value="All")
         slider = gr.Slider(minimum=0, maximum=1, value=0.4, label="Alpha Factor")
 
-    all_masks = gr.Gallery(visible=False) 
-    background = gr.Image(visible=False)
+    mask_state = gr.Sate() 
+    bg_state = gr.State()
     
     gr.Button("1) Generate Masks").click(fn=inference, 
                                          inputs=[input_img],
-                                         outputs=[img, all_masks, background])
+                                         outputs=[preview, gr.Gallery(visible=False), bg_state, mask_state])
 
     gr.Button("2) Generate Transparent Mask (with Alpha Factor)").click(fn=show_img, 
-                                                                        inputs=[all_masks, dropdown, background, slider], 
-                                                                        outputs=[transparent_img])
+                                                                        inputs=[mask_state, dropdown, bg_state, slider], 
+                                                                        outputs=output)
    
      
 app.launch()