| 20 | from detectors.detector_factory import detector_factory |
| 21 | |
| 22 | class PrefetchDataset(torch.utils.data.Dataset): |
| 23 | def __init__(self, opt, dataset, pre_process_func): |
| 24 | self.images = dataset.images |
| 25 | self.load_image_func = dataset.coco.loadImgs |
| 26 | self.img_dir = dataset.img_dir |
| 27 | self.pre_process_func = pre_process_func |
| 28 | self.opt = opt |
| 29 | |
| 30 | def __getitem__(self, index): |
| 31 | img_id = self.images[index] |
| 32 | img_info = self.load_image_func(ids=[img_id])[0] |
| 33 | img_path = os.path.join(self.img_dir, img_info['file_name']) |
| 34 | image = cv2.imread(img_path) |
| 35 | images, meta = {}, {} |
| 36 | for scale in opt.test_scales: |
| 37 | if opt.task == 'ddd': |
| 38 | images[scale], meta[scale] = self.pre_process_func( |
| 39 | image, scale, img_info['calib']) |
| 40 | else: |
| 41 | images[scale], meta[scale] = self.pre_process_func(image, scale) |
| 42 | return img_id, {'images': images, 'image': image, 'meta': meta} |
| 43 | |
| 44 | def __len__(self): |
| 45 | return len(self.images) |
| 46 | |
| 47 | def prefetch_test(opt): |
| 48 | os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str |
no outgoing calls