(self, img_path, bbox, sup_res=True)
| 433 | return node_prompts, node_prompts_label |
| 434 | |
| 435 | def _forward(self, img_path, bbox, sup_res=True): |
| 436 | |
| 437 | img = self.get_img(img_path, sup_res) |
| 438 | |
| 439 | if bbox is None: |
| 440 | # bbox, pha = self.predict_bbox(img) |
| 441 | # bbox, pha = self.rembg_predict_bbox(img, 1.01) |
| 442 | # bbox, pha = self.yolo_predict_bbox(img) |
| 443 | bbox, pha = self.birefnet_predict_bbox(img, 1.01) |
| 444 | |
| 445 | box = bbox.to_whwh() |
| 446 | bbox = box.get_box() |
| 447 | |
| 448 | point_coords, point_coords_label = self.compute_coords(pha, bbox) |
| 449 | |
| 450 | self.image_predictor.set_image(img) |
| 451 | |
| 452 | masks, scores, logits = self.image_predictor.predict( |
| 453 | point_coords=point_coords, |
| 454 | point_labels=point_coords_label, |
| 455 | box=bbox, |
| 456 | multimask_output=False, |
| 457 | ) |
| 458 | |
| 459 | alpha = masks[0] |
| 460 | |
| 461 | # fill-mask NO USE |
| 462 | # alpha = fill_mask(alpha) |
| 463 | # alpha = erode_and_dialted( |
| 464 | # (alpha * 255).astype(np.uint8), kernel_size=3, iterations=3 |
| 465 | # ) |
| 466 | # alpha = alpha.astype(np.float32) / 255.0 |
| 467 | |
| 468 | img_float = img.astype(np.float32) / 255.0 |
| 469 | process_img = ( |
| 470 | img_float * alpha[..., None] + (1 - alpha[..., None]) * self.background |
| 471 | ) |
| 472 | process_img = (process_img * 255).astype(np.uint8) |
| 473 | |
| 474 | # using for draw box |
| 475 | # process_img = cv2.rectangle(process_img, bbox[:2], bbox[2:], (0, 0, 255), 2) |
| 476 | process_img = process_img.astype(np.float) / 255.0 |
| 477 | |
| 478 | process_pha_img = ( |
| 479 | img_float * pha[..., None] + (1 - pha[..., None]) * self.background |
| 480 | ) |
| 481 | |
| 482 | return SegmentOut( |
| 483 | masks=alpha, processed_img=process_img, alpha_img=process_pha_img[...] |
| 484 | ) |
| 485 | |
| 486 | @torch.no_grad() |
| 487 | def __call__(self, **inputs): |
no test coverage detected