Fixes crop background color override (#2373)

fix crop color background
This commit is contained in:
Justin Kerr authored and GitHub committed 2023-08-27 11:50:57 -07:00
1 parent ae5d569376
commit d56338bce1
1 file changed
+7 -4
+7 -4
View File
@@ -101,8 +101,12 @@ class RGBRenderer(nn.Module):
else:
comp_rgb = torch.sum(weights * rgb, dim=-2)
accumulated_weight = torch.sum(weights, dim=-2)
if background_color == "random":
if BACKGROUND_COLOR_OVERRIDE is not None:
# This case must be before the others or the override is not properly applied
background_color = cls.get_background_color(
BACKGROUND_COLOR_OVERRIDE, shape=comp_rgb.shape, device=comp_rgb.device
)
elif background_color == "random":
# If background color is random, the predicted color is returned without blending,
# as if the background color was black.
return comp_rgb
@@ -110,8 +114,7 @@ class RGBRenderer(nn.Module):
elif background_color == "last_sample":
# Note, this is only supported for non-packed samples.
background_color = rgb[..., -1, :]
else:
background_color = cls.get_background_color(background_color, shape=comp_rgb.shape, device=comp_rgb.device)
assert isinstance(background_color, torch.Tensor)
comp_rgb = comp_rgb + background_color * (1.0 - accumulated_weight)
return comp_rgb