| 287 | return scale_box, pha |
| 288 | |
| 289 | def rembg_predict_bbox(self, img, scale=1.0): |
| 290 | |
| 291 | height, width, _ = img.shape |
| 292 | |
| 293 | with torch.no_grad(): |
| 294 | img_rmbg = img[..., ::-1] # rgb2bgr |
| 295 | img_rmbg = remove(img_rmbg) |
| 296 | img_rmbg = img_rmbg[..., :3] |
| 297 | pha = copy.deepcopy(img_rmbg[..., -1:]) |
| 298 | |
| 299 | masks = copy.deepcopy(pha) |
| 300 | |
| 301 | masks[masks < 1.0] = 0.0 |
| 302 | masks[masks >= 1.0] = 1.0 |
| 303 | |
| 304 | # obtain bbox |
| 305 | _h, _w, _ = np.where(masks == 1) |
| 306 | |
| 307 | whwh = [ |
| 308 | _w.min().item(), |
| 309 | _h.min().item(), |
| 310 | _w.max().item(), |
| 311 | _h.max().item(), |
| 312 | ] |
| 313 | |
| 314 | box = Bbox(whwh) |
| 315 | |
| 316 | # scale box to 1.05 |
| 317 | scale_box = box.scale(scale=scale, width=width, height=height) |
| 318 | |
| 319 | return scale_box, pha[..., 0].astype(np.float32) / 255.0 |
| 320 | |
| 321 | def yolo_predict_bbox(self, img, scale=1.0, threshold=0.2): |
| 322 | if self.prior == None: |