| 56 | |
| 57 | |
| 58 | class MultiClassificationProcessor(torch.utils.data.Dataset): |
| 59 | |
| 60 | def __init__(self, transform=None): |
| 61 | self.transformer_ = transform |
| 62 | self.extension_ = '.jpg .jpeg .png .bmp .webp .tif .eps' |
| 63 | # load category info |
| 64 | self.ctg_names_ = [] # ctg_idx to ctg_name |
| 65 | self.ctg_name2idx_ = OrderedDict() # ctg_name to ctg_idx |
| 66 | # load image infos |
| 67 | self.img_names_ = [] # img_idx to img_name |
| 68 | self.img_paths_ = [] # img_idx to img_path |
| 69 | self.img_labels_ = [] # img_idx to img_label |
| 70 | |
| 71 | self.srm = SRMConv2d_simple() |
| 72 | |
| 73 | def load_data_from_dir(self, dataset_list): |
| 74 | """Load image from folder. |
| 75 | |
| 76 | Args: |
| 77 | dataset_list: dataset list, each folder is a category, format is [file_root]. |
| 78 | """ |
| 79 | # load sample |
| 80 | for img_root in dataset_list: |
| 81 | ctg_name = os.path.basename(img_root) |
| 82 | self.ctg_name2idx_[ctg_name] = len(self.ctg_names_) |
| 83 | self.ctg_names_.append(ctg_name) |
| 84 | img_paths = [] |
| 85 | traverse_recursively(img_root, img_paths, self.extension_) |
| 86 | for img_path in img_paths: |
| 87 | img_name = os.path.basename(img_path) |
| 88 | self.img_names_.append(img_name) |
| 89 | self.img_paths_.append(img_path) |
| 90 | self.img_labels_.append(self.ctg_name2idx_[ctg_name]) |
| 91 | print('log: category is %d(%s), image num is %d' % (self.ctg_name2idx_[ctg_name], ctg_name, len(img_paths))) |
| 92 | |
| 93 | def load_data_from_txt(self, img_list_txt, ctg_list_txt): |
| 94 | """Load image from txt. |
| 95 | |
| 96 | Args: |
| 97 | img_list_txt: image txt, format is [file_path, ctg_idx]. |
| 98 | ctg_list_txt: category txt, format is [ctg_name, ctg_idx]. |
| 99 | """ |
| 100 | # check |
| 101 | assert os.path.exists(img_list_txt), 'log: does not exist: {}'.format(img_list_txt) |
| 102 | assert os.path.exists(ctg_list_txt), 'log: does not exist: {}'.format(ctg_list_txt) |
| 103 | |
| 104 | # load category |
| 105 | # : open category info file |
| 106 | with open(ctg_list_txt) as f: |
| 107 | ctg_infos = [line.strip() for line in f.readlines()] |
| 108 | # :load category name & category index |
| 109 | for ctg_info in ctg_infos: |
| 110 | tmp = ctg_info.split(' ') |
| 111 | ctg_name = tmp[0] |
| 112 | ctg_idx = int(tmp[-1]) |
| 113 | self.ctg_name2idx_[ctg_name] = ctg_idx |
| 114 | self.ctg_names_.append(ctg_name) |
| 115 |
no outgoing calls
no test coverage detected