fix overlay alpha computation

inference/interact/interactive_utils.py

@@ -107,12 +107,12 @@ def overlay_layer(image, mask, layer, target_object):
     # insert a layer between foreground and background
     # The CPU version is less accurate because we are using the hard mask
     # The GPU version has softer edges as it uses soft probabilities
-    obj_mask = (np.isin(mask, target_object)).astype(np.float32)
-    layer_alpha = layer[:, :, 3].astype(np.float32) / 255
+    obj_mask = (np.isin(mask, target_object)).astype(np.float32)[:, :, np.newaxis]
+    layer_alpha = layer[:, :, 3].astype(np.float32)[:, :, np.newaxis] / 255
     layer_rgb = layer[:, :, :3]
-    background_alpha = np.maximum(obj_mask, layer_alpha)[:,:,np.newaxis]
-    obj_mask = obj_mask[:,:,np.newaxis]
-    im_overlay = (image*(1-background_alpha) + layer_rgb*(1-obj_mask) + image*obj_mask).clip(0, 255)
+    background_alpha = np.maximum(obj_mask, layer_alpha)
+    im_overlay = (image * (1 - background_alpha) + layer_rgb * (1 - obj_mask) * layer_alpha +
+                  image * obj_mask).clip(0, 255)
     return im_overlay.astype(image.dtype)
 
 def overlay_davis_torch(image, mask, alpha=0.5, fade=False):
@@ -156,28 +156,25 @@ def overlay_popup_torch(image, mask, target_object):
 
     return im_overlay
 
-def overlay_layer_torch(image, mask, layer, target_object):
+def overlay_layer_torch(image, prob, layer, target_object):
     # insert a layer between foreground and background
     # The CPU version is less accurate because we are using the hard mask
     # The GPU version has softer edges as it uses soft probabilities
     image = image.permute(1, 2, 0)
 
     if len(target_object) == 0:
-        obj_mask = torch.zeros_like(mask[0])
+        obj_mask = torch.zeros_like(prob[0]).unsqueeze(2)
     else:
-        # I should not need to convert this to numpy.
-        # uUsing list works most of the time but consistently fails
-        # if I include first object -> exclude it -> include it again.
-        # I check everywhere and it makes absolutely no sense.
-        # I am blaming this on PyTorch and calling it a day
-        obj_mask = mask[np.array(target_object,dtype=np.int32)].sum(0)
-    layer_alpha = layer[:, :, 3]
+        # TODO: figure out why we need to convert this to numpy array
+        obj_mask = prob[np.array(target_object, dtype=np.int32)].sum(0).unsqueeze(2)
+    layer_alpha = layer[:, :, 3].unsqueeze(2)
     layer_rgb = layer[:, :, :3]
-    background_alpha = torch.maximum(obj_mask, layer_alpha).unsqueeze(2)
-    obj_mask = obj_mask.unsqueeze(2)
-    im_overlay = (image*(1-background_alpha) + layer_rgb*(1-obj_mask) + image*obj_mask).clip(0, 1)
+    background_alpha = torch.maximum(obj_mask, layer_alpha)
+    im_overlay = (image * (1 - background_alpha) + layer_rgb * (1 - obj_mask) * layer_alpha +
+                  image * obj_mask).clip(0, 1)
 
-    im_overlay = (im_overlay*255).cpu().numpy()
+    im_overlay = (im_overlay * 255).cpu().numpy()
     im_overlay = im_overlay.astype(np.uint8)
 
     return im_overlay
+