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