| 27 | VIDEO_READER_TIMEOUT = 20 |
| 28 | |
| 29 | def get_random_mask(shape, image_start_only=False): |
| 30 | f, c, h, w = shape |
| 31 | mask = torch.zeros((f, 1, h, w), dtype=torch.uint8) |
| 32 | |
| 33 | if not image_start_only: |
| 34 | if f != 1: |
| 35 | mask_index = np.random.choice([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], p=[0.05, 0.2, 0.2, 0.2, 0.05, 0.05, 0.05, 0.1, 0.05, 0.05]) |
| 36 | else: |
| 37 | mask_index = np.random.choice([0, 1, 7, 8], p = [0.2, 0.7, 0.05, 0.05]) |
| 38 | if mask_index == 0: |
| 39 | center_x = torch.randint(0, w, (1,)).item() |
| 40 | center_y = torch.randint(0, h, (1,)).item() |
| 41 | block_size_x = torch.randint(w // 4, w // 4 * 3, (1,)).item() # 方块的宽度范围 |
| 42 | block_size_y = torch.randint(h // 4, h // 4 * 3, (1,)).item() # 方块的高度范围 |
| 43 | |
| 44 | start_x = max(center_x - block_size_x // 2, 0) |
| 45 | end_x = min(center_x + block_size_x // 2, w) |
| 46 | start_y = max(center_y - block_size_y // 2, 0) |
| 47 | end_y = min(center_y + block_size_y // 2, h) |
| 48 | mask[:, :, start_y:end_y, start_x:end_x] = 1 |
| 49 | elif mask_index == 1: |
| 50 | mask[:, :, :, :] = 1 |
| 51 | elif mask_index == 2: |
| 52 | mask_frame_index = np.random.randint(1, 5) |
| 53 | mask[mask_frame_index:, :, :, :] = 1 |
| 54 | elif mask_index == 3: |
| 55 | mask_frame_index = np.random.randint(1, 5) |
| 56 | mask[mask_frame_index:-mask_frame_index, :, :, :] = 1 |
| 57 | elif mask_index == 4: |
| 58 | center_x = torch.randint(0, w, (1,)).item() |
| 59 | center_y = torch.randint(0, h, (1,)).item() |
| 60 | block_size_x = torch.randint(w // 4, w // 4 * 3, (1,)).item() # 方块的宽度范围 |
| 61 | block_size_y = torch.randint(h // 4, h // 4 * 3, (1,)).item() # 方块的高度范围 |
| 62 | |
| 63 | start_x = max(center_x - block_size_x // 2, 0) |
| 64 | end_x = min(center_x + block_size_x // 2, w) |
| 65 | start_y = max(center_y - block_size_y // 2, 0) |
| 66 | end_y = min(center_y + block_size_y // 2, h) |
| 67 | |
| 68 | mask_frame_before = np.random.randint(0, f // 2) |
| 69 | mask_frame_after = np.random.randint(f // 2, f) |
| 70 | mask[mask_frame_before:mask_frame_after, :, start_y:end_y, start_x:end_x] = 1 |
| 71 | elif mask_index == 5: |
| 72 | mask = torch.randint(0, 2, (f, 1, h, w), dtype=torch.uint8) |
| 73 | elif mask_index == 6: |
| 74 | num_frames_to_mask = random.randint(1, max(f // 2, 1)) |
| 75 | frames_to_mask = random.sample(range(f), num_frames_to_mask) |
| 76 | |
| 77 | for i in frames_to_mask: |
| 78 | block_height = random.randint(1, h // 4) |
| 79 | block_width = random.randint(1, w // 4) |
| 80 | top_left_y = random.randint(0, h - block_height) |
| 81 | top_left_x = random.randint(0, w - block_width) |
| 82 | mask[i, 0, top_left_y:top_left_y + block_height, top_left_x:top_left_x + block_width] = 1 |
| 83 | elif mask_index == 7: |
| 84 | center_x = torch.randint(0, w, (1,)).item() |
| 85 | center_y = torch.randint(0, h, (1,)).item() |
| 86 | a = torch.randint(min(w, h) // 8, min(w, h) // 4, (1,)).item() # 长半轴 |