(self, frames, th_step, kp2ds_all, use_valid_body_keypoints=True)
| 433 | return tpl_prompt, refer_prompt |
| 434 | |
| 435 | def get_mask(self, frames, th_step, kp2ds_all, use_valid_body_keypoints=True): |
| 436 | frame_num = len(frames) |
| 437 | masks = [None] * frame_num |
| 438 | if frame_num < th_step: |
| 439 | num_step = 1 |
| 440 | else: |
| 441 | num_step = (frame_num + th_step) // th_step |
| 442 | |
| 443 | for index in range(num_step): |
| 444 | chunk_start = index * th_step |
| 445 | each_frames = frames[chunk_start : chunk_start + th_step] |
| 446 | kp2ds = kp2ds_all[chunk_start : chunk_start + th_step] |
| 447 | if len(each_frames) == 0: |
| 448 | continue |
| 449 | |
| 450 | if len(each_frames) > 4: |
| 451 | key_frame_num = 4 |
| 452 | else: |
| 453 | key_frame_num = 1 |
| 454 | |
| 455 | key_frame_step = max(len(kp2ds) // key_frame_num, 1) |
| 456 | key_frame_index_list = list(range(0, len(kp2ds), key_frame_step))[:key_frame_num] |
| 457 | if use_valid_body_keypoints: |
| 458 | key_frame_index_list = [key_frame_index for key_frame_index in key_frame_index_list if self._is_valid_pose_meta(kp2ds[key_frame_index])] |
| 459 | if len(key_frame_index_list) == 0: |
| 460 | valid_frame_indices = [idx for idx, meta in enumerate(kp2ds) if self._is_valid_pose_meta(meta)] |
| 461 | if len(valid_frame_indices) > key_frame_num: |
| 462 | selected_indices = np.linspace(0, len(valid_frame_indices) - 1, key_frame_num, dtype=np.int32) |
| 463 | key_frame_index_list = [valid_frame_indices[idx] for idx in selected_indices] |
| 464 | else: |
| 465 | key_frame_index_list = valid_frame_indices |
| 466 | |
| 467 | key_points_index = [0, 1, 2, 5, 8, 11, 10, 13] |
| 468 | key_frame_body_points_list = [] |
| 469 | for key_frame_index in key_frame_index_list: |
| 470 | if use_valid_body_keypoints: |
| 471 | points = self._get_body_prompt_points(kp2ds[key_frame_index], kp2ds[0]["width"], kp2ds[0]["height"]) |
| 472 | else: |
| 473 | keypoints_body_list = [] |
| 474 | body_key_points = kp2ds[key_frame_index]["keypoints_body"] |
| 475 | for each_index in key_points_index: |
| 476 | each_keypoint = body_key_points[each_index] |
| 477 | if None is each_keypoint: |
| 478 | continue |
| 479 | keypoints_body_list.append(each_keypoint) |
| 480 | |
| 481 | if len(keypoints_body_list) == 0: |
| 482 | key_frame_body_points_list.append(np.zeros((0, 2), dtype=np.int32)) |
| 483 | continue |
| 484 | |
| 485 | keypoints_body = np.array(keypoints_body_list)[:, :2] |
| 486 | wh = np.array([[kp2ds[0]["width"], kp2ds[0]["height"]]]) |
| 487 | points = (keypoints_body * wh).astype(np.int32) |
| 488 | key_frame_body_points_list.append(points) |
| 489 | |
| 490 | chunk_masks = {} |
| 491 | sam_ran = False |
| 492 | if any(points.shape[0] > 0 for points in key_frame_body_points_list): |
no test coverage detected