| 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: |
| 323 | from ultralytics import YOLO |
| 324 | |
| 325 | pdb.set_trace() |
| 326 | |
| 327 | height, width, _ = img.shape |
| 328 | |
| 329 | with torch.no_grad(): |
| 330 | results = yolo_seg(img[..., ::-1]) |
| 331 | for result in results: |
| 332 | masks = result.masks.data[result.boxes.cls == 0] |
| 333 | if masks.shape[0] >= 1: |
| 334 | masks[masks >= threshold] = 1 |
| 335 | masks[masks < threshold] = 0 |
| 336 | masks = masks.sum(dim=0) |
| 337 | |
| 338 | pha = masks.detach().cpu().numpy() |
| 339 | pha = cv2.resize(pha, (width, height), interpolation=cv2.INTER_AREA)[..., None] |
| 340 | |
| 341 | pha[pha >= 0.5] = 1 |
| 342 | pha[pha < 0.5] = 0 |
| 343 | |
| 344 | masks = copy.deepcopy(pha) |
| 345 | |
| 346 | pha = pha * 255.0 |
| 347 | # obtain bbox |
| 348 | _h, _w, _ = np.where(masks == 1) |
| 349 | |
| 350 | whwh = [ |
| 351 | _w.min().item(), |
| 352 | _h.min().item(), |
| 353 | _w.max().item(), |
| 354 | _h.max().item(), |
| 355 | ] |
| 356 | |
| 357 | box = Bbox(whwh) |
| 358 | |
| 359 | # scale box to 1.05 |
| 360 | scale_box = box.scale(scale=scale, width=width, height=height) |
| 361 | |
| 362 | return scale_box, pha[..., 0].astype(np.float32) / 255.0 |
| 363 | |
| 364 | def ratio_mapping(self, img): |
| 365 | |