(self, img_dir=None,out_path=None)
| 35 | |
| 36 | class INFERENCE(torch.utils.data.Dataset): |
| 37 | def __init__(self, img_dir=None,out_path=None): |
| 38 | |
| 39 | self.output_path = out_path |
| 40 | |
| 41 | self.img_dir = img_dir |
| 42 | |
| 43 | self.is_vid = False |
| 44 | |
| 45 | rank, _ = get_dist_info() |
| 46 | if self.img_dir.endswith('.mp4'): |
| 47 | self.is_vid = True |
| 48 | img_name = self.img_dir.split('/')[-1][:-4] |
| 49 | # self.img_dir = self.img_dir[:-4] |
| 50 | else: |
| 51 | img_name = self.img_dir.split('/')[-1] |
| 52 | self.img_name = img_name+'_out' |
| 53 | self.output_path = os.path.join(self.output_path,self.img_name) |
| 54 | os.makedirs(self.output_path, exist_ok=True) |
| 55 | self.tmp_dir = os.path.join(self.output_path, 'temp_img') |
| 56 | os.makedirs(self.tmp_dir, exist_ok=True) |
| 57 | self.result_img_dir = os.path.join(self.output_path, 'res_img') |
| 58 | |
| 59 | |
| 60 | if not self.is_vid: |
| 61 | if rank == 0: |
| 62 | image_files = sorted(glob(self.img_dir + '/*.jpg') + glob(self.img_dir + '/*.png')) |
| 63 | for i, image_file in enumerate(image_files): |
| 64 | new_name = os.path.join(self.tmp_dir, '%06d.png'%i) |
| 65 | shutil.copy(image_file, new_name) |
| 66 | dist.barrier() |
| 67 | else: |
| 68 | if rank == 0: |
| 69 | video_to_images(self.img_dir, self.tmp_dir) |
| 70 | dist.barrier() |
| 71 | self.img_paths = sorted(glob(self.tmp_dir+'/*',recursive=True)) |
| 72 | self.score_threshold = 0.2 |
| 73 | self.resolution = [720 ,1280] # AGORA test |
| 74 | self.format = DefaultFormatBundle() |
| 75 | self.normalize = Normalize(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]) |
| 76 | |
| 77 | def __len__(self): |
| 78 | return len(self.img_paths) |
nothing calls this directly
no test coverage detected