| 69 | |
| 70 | class TrainDataset(BaseDataset): |
| 71 | def __init__(self, root_dataset, odgt, opt, batch_per_gpu=1, **kwargs): |
| 72 | super(TrainDataset, self).__init__(odgt, opt, **kwargs) |
| 73 | self.root_dataset = root_dataset |
| 74 | # down sampling rate of segm labe |
| 75 | self.segm_downsampling_rate = opt.segm_downsampling_rate |
| 76 | self.batch_per_gpu = batch_per_gpu |
| 77 | |
| 78 | # classify images into two classes: 1. h > w and 2. h <= w |
| 79 | self.batch_record_list = [[], []] |
| 80 | |
| 81 | # override dataset length when trainig with batch_per_gpu > 1 |
| 82 | self.cur_idx = 0 |
| 83 | self.if_shuffled = False |
| 84 | |
| 85 | def _get_sub_batch(self): |
| 86 | while True: |