Update app.py

app.py

@@ -110,35 +110,67 @@ def show_img(all_imgs, dropdown, bg, alpha_factor):
 ###########
 # Inference
 
-def inference(img: Image.Image, alpha: float = 0.4):
-    x = process_image(img)
-    with torch.no_grad():
-        mask_logits = model(x)[0].cpu().numpy()
+def inference(img, alpha_factor):
+    background = resize_pil(img) 
 
-    mask_probs = 1/(1+np.exp(-mask_logtis))
-    mask_preds = mask_probs > 0.5
+    img = process_pil(img)
 
-    mask_all = mask_preds.sum(axis=0)
-    mask_all = np.expand_dims(mask_all, axis=0)
-    mask_pred = np.concatenate((mask_all, mask_preds), axis=0)
-    labs = ["All"] + target_list 
+    mask = model(img.unsqueeze(0)) # we need a batch, hence we introduce an extra dimenation at position 0 (unsqueeze)
+    mask = mask[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
 
-    # saved all masks combined with unvisible xaxis und yaxis and without a white background
-    all_images =[]
-    for i in range(len(mask_preds)):
-        img_mask =(mask_preds[i]+255).astype(np.uint8)
-        pil_mask = Image.fromarray(img_mask)
-        all_images.append(pil_mask)
+    # All combined
+    mask_all =  mask_preds.sum(axis=0)
+    mask_all = np.expand_dims(mask_all, axis=0)
+    mask_all.shape
 
-    background = img.convert("RGBA")
-    overlay = (mask_preds[0]*255).astype(np.uint8)
-    overlay_img = Image.fromarray(overlay).convert("RGBA")
-    overlay_img.putalpha(int(alpha*255))
-    background.paste(overlay_img, (0,0), overlay_img)
+    # Concat all combined with normal preds
+    mask_preds = np.concatenate((mask_all, mask_preds),axis=0)
+    labs = ["ALL"] + target_list
+
+    fig, axes = plt.subplots(5, 4, figsize = (10,10))
+    
+    # save all mask_preds in all_mask
+    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()        
+
+    # plt to PIL
+    img_buf = io.BytesIO()
+    fig.savefig(img_buf, format='png')
+    im = Image.open(img_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])
+        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))
+
+    return im, all_images, background
 
-    mask_pil = Image.fromarray((mask_preds[0]*255).astype(np.uint8))
-    return mask_pil, background
     
 
 title = "Masterarbeit - Bauschadenerkennung"