| 66 | |
| 67 | |
| 68 | class BBoxDataIterator: |
| 69 | def __init__(self, n, batch_size, ndim=2, produce_labels=False): |
| 70 | self.batch_size = batch_size |
| 71 | self.ndim = ndim |
| 72 | self.produce_labels = produce_labels |
| 73 | self.num_outputs = 2 if produce_labels else 1 |
| 74 | self.n = n |
| 75 | self.i = 0 |
| 76 | |
| 77 | def __len__(self): |
| 78 | return self.n |
| 79 | |
| 80 | def __iter__(self): |
| 81 | # return a copy, so that the iteration number doesn't collide |
| 82 | return BBoxDataIterator(self.n, self.batch_size, self.ndim, self.produce_labels) |
| 83 | |
| 84 | def __next__(self): |
| 85 | boxes = [] |
| 86 | labels = [] |
| 87 | bboxes = bboxes_data[self.ndim] |
| 88 | if self.i % 2 == 0: |
| 89 | boxes.append(np.array([bboxes[0], bboxes[1], bboxes[2]], dtype=np.float32)) |
| 90 | labels.append(np.array([1, 2, 3], dtype=np.int32)) |
| 91 | if self.batch_size > 1: |
| 92 | boxes.append(np.array([bboxes[2], bboxes[1]], dtype=np.float32)) |
| 93 | labels.append(np.array([2, 1], dtype=np.int32)) |
| 94 | for _ in range(self.batch_size - 2): |
| 95 | boxes.append(np.array([bboxes[2]], dtype=np.float32)) |
| 96 | labels.append(np.array([3], dtype=np.int32)) |
| 97 | else: |
| 98 | boxes.append(np.array([bboxes[2]], dtype=np.float32)) |
| 99 | labels.append(np.array([3], dtype=np.int32)) |
| 100 | if self.batch_size > 1: |
| 101 | boxes.append(np.array([bboxes[1], bboxes[2], bboxes[0]], dtype=np.float32)) |
| 102 | labels.append(np.array([2, 3, 1], dtype=np.int32)) |
| 103 | for _ in range(self.batch_size - 2): |
| 104 | boxes.append(np.array([bboxes[1]], dtype=np.float32)) |
| 105 | labels.append(np.array([2], dtype=np.int32)) |
| 106 | |
| 107 | if self.i < self.n: |
| 108 | self.i = self.i + 1 |
| 109 | outputs = [boxes] |
| 110 | if self.produce_labels: |
| 111 | outputs.append(labels) |
| 112 | return outputs |
| 113 | |
| 114 | self.i = 0 |
| 115 | raise StopIteration |
| 116 | |
| 117 | next = __next__ |
| 118 | |
| 119 | |
| 120 | class RandomBBoxCropSynthDataPipeline(Pipeline): |
no outgoing calls