Method birefnet_predict_bbox

engine/SegmentAPI/SAM.py:227–287

(self, img, scale=1.0)

        return scale_box, pha[..., 0]

    def birefnet_predict_bbox(self, img, scale=1.0):

        # img: RGB-order

        if self.box_prior == None:
            from engine.BiRefNet.utils import check_state_dict

            birefnet = BiRefNet(bb_pretrained=False)
            state_dict = torch.load(BIREFNET_WEIGHT, map_location="cpu")
            state_dict = check_state_dict(state_dict)
            birefnet.load_state_dict(state_dict)
            device = avaliable_device()
            torch.set_float32_matmul_precision(["high", "highest"][0])

            birefnet.to(device)
            self.box_prior = birefnet
            self.box_prior.eval()
            self.box_transform = transforms.Compose(
                [
                    transforms.Resize((1024, 1024)),
                    transforms.ToTensor(),
                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
                ]
            )
            print("BiRefNet is ready to use.")
        else:
            device = avaliable_device()
            self.box_prior.to(device)

        height, width, _ = img.shape

        image = PIL.Image.fromarray(img)

        input_images = self.box_transform(image).unsqueeze(0).to("cuda")
        with torch.no_grad():
            preds = self.box_prior(input_images)[-1].sigmoid().cpu()
        pha = (preds[0]).squeeze(0).detach().numpy()

        pha = cv2.resize(pha, (width, height))

        masks = copy.deepcopy(pha[..., None])

        masks[masks < 0.3] = 0.0
        masks[masks >= 0.3] = 1.0

        # obtain bbox
        _h, _w, _ = np.where(masks == 1)

        whwh = [
            _w.min().item(),
            _h.min().item(),
            _w.max().item(),
            _h.max().item(),
        ]

        box = Bbox(whwh)

        # scale box to 1.05