ZhengPeng7 / BiRefNet

[CAAI AIR'24] Bilateral Reference for High-Resolution Dichotomous Image Segmentation
https://www.birefnet.top
MIT License
322 stars 28 forks source link

How to remove black bacground of output image #21

Closed miladfa7 closed 2 months ago

miladfa7 commented 2 months ago

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 2 months ago

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.