How to remove black bacground of output image
Closed this issue · 1 comments
miladfa7 commented
After running the model and removal background, the output result has black background. My question is how to get output image without black, i mean save output image (image_preds[0]
) as transparent image.
The code:
def predict(self, image) :
images = [image]
image_shapes = [image.shape[:2] for image in images]
images = [array_to_pil_image(image, self.resolution) for image in images]
image_preprocessor = ImagePreprocessor(resolution=self.resolution)
images_proc = []
for image in images:
images_proc.append(image_preprocessor.proc(image))
images_proc = torch.cat([image_proc.unsqueeze(0) for image_proc in images_proc])
with torch.no_grad():
scaled_preds_tensor = self.model(images_proc.to(self.device))[-1].sigmoid()
preds = []
for image_shape, pred_tensor in zip(image_shapes, scaled_preds_tensor):
if self.device == 'cuda':
pred_tensor = pred_tensor.cpu()
preds.append(torch.nn.functional.interpolate(pred_tensor.unsqueeze(0), size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy())
image_preds = []
for image, pred in zip(images, preds):
image = image.resize(pred.shape[::-1])
pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1)
image_preds.append((pred * image).astype(np.uint8))
return image, image_preds[0]
ZhengPeng7 commented
Here, you get the returned value of both the predicted mask and the original image. With the mask (0 or 1), you can use mask * original image + (1 - mask) * new background
to replace the original background with a new image.