From a8179ff53c4c2192167c7f0a6763d2ab5fa75f28 Mon Sep 17 00:00:00 2001 From: James Betker Date: Fri, 18 Dec 2020 08:53:18 -0700 Subject: [PATCH] Image label work --- codes/data/image_label_parser.py | 60 ++++++++++++++++++- .../ui/image_labeler/image_labeler_ui.py | 10 ++-- .../scripts/ui/image_labeler/label_editor.py | 33 ++++++++++ codes/train.py | 2 +- 4 files changed, 99 insertions(+), 6 deletions(-) create mode 100644 codes/scripts/ui/image_labeler/label_editor.py diff --git a/codes/data/image_label_parser.py b/codes/data/image_label_parser.py index bc9a7a28..89a79132 100644 --- a/codes/data/image_label_parser.py +++ b/codes/data/image_label_parser.py @@ -67,5 +67,63 @@ class VsNetImageLabeler: self.categories[binding]['labeledImages'].append(lbl) def save(self): - with open(self.label_file[0], "wb") as file: + with open(self.label_file[-1], "wb") as file: + file.write(json.dumps(self.categories)) + + +# A more compact format that is simpler to parse and understand. +class CompactJsonLabeler: + def __init__(self, lbl_files): + if not isinstance(lbl_files, list): + lbl_files = [lbl_files] + self.label_files = lbl_files + self.config, self.labels, self.label_map, self.images = None, None, None, None + for lfil in lbl_files: + with open(lfil, "r") as read_file: + # Format: + # { + # 'config': { 'dim' } + # 'labels': [{ 'label', 'key'}] <- ordered by label index. + # 'images': {'file': [{ 'lid', 'top', 'left' }} + # 'labelMap' {} + # } + parsed = json.loads(read_file.read()) + if self.config is None: + self.config = parsed['config'] + self.labels = parsed['labels'] + self.images = parsed['images'] + self.label_map = parsed['label_map'] + self.binding_map = {} + for i, lbl in enumerate(self.labels): + self.binding_map[lbl['key']] = i + else: + assert self.config == parsed['config'] + assert self.labels == parsed['labels'] + assert self.label_map == parsed['label_map'] + self.images.update(parsed['images']) # This will overwrite existing images, which is acceptable. + + def get_labeled_paths(self, base_path): + return [os.path.join(base_path, pth) for pth in self.images.keys()] + + def get_labels_as_tensor(self, hq, img_key, resize_factor): + _, h, w = hq.shape + labels = torch.zeros((1,h,w), dtype=torch.long) + mask = torch.zeros((1,h,w), dtype=torch.float) + lbl_list = self.images[img_key] + for patch_lbl in lbl_list: + t, l, h, w = patch_lbl['top'] // resize_factor, patch_lbl['left'] // resize_factor, \ + self.config['dim'] // resize_factor, self.config['dim'] // resize_factor + val = patch_lbl['labelValue'] + labels[:,t:t+h,l:l+w] = val + mask[:,t:t+h,l:l+w] = 1.0 + return labels, mask, self.str_labels + + def add_label(self, binding, img_name, top, left, dim): + lbl = {'lid': self.binding_map[binding], 'top': top, 'left': left} + if img_name not in self.images.keys(): + self.images[img_name] = [] + self.images[img_name].append(lbl) + + def save(self): + with open(self.label_file[-1], "wb") as file: file.write(json.dumps(self.categories)) diff --git a/codes/scripts/ui/image_labeler/image_labeler_ui.py b/codes/scripts/ui/image_labeler/image_labeler_ui.py index 63c3dfc9..1eb7d15d 100644 --- a/codes/scripts/ui/image_labeler/image_labeler_ui.py +++ b/codes/scripts/ui/image_labeler/image_labeler_ui.py @@ -37,7 +37,7 @@ def update_mode_label(): # Handles the "change mode" hotkey. Changes the classification label being targeted. def change_mode(event): global mode, pending_labels - mode += 1 + mode = (mode + 1) % len(labeler.str_labels) update_mode_label() @@ -88,7 +88,7 @@ def next_batch(): scale = hq.shape[-1] // res.shape[-1] # These are the confidence bounds. They apply to a post-softmax output. They are currently fixed. - conf_lower = .8 + conf_lower = .4 conf_upper = 1 valid_res = ((res > conf_lower) * (res < conf_upper)) * 1.0 # This results in a tensor of 0's where tensors are outside of the confidence bound, 1's where inside. @@ -96,8 +96,10 @@ def next_batch(): batch_sz = hq.shape[0] while cur_img < batch_sz: # Note: cur_img can (intentionally) be changed outside of this loop. # Build a random permutation for every image patch in the image. We will search for patches that fall within the confidence bound and yield them. - #permutation = torch.randperm(res.shape[-1] * res.shape[-2]) - for p in range(res.shape[-1]*res.shape[-2]): + permutation = torch.randperm(res.shape[-1] * res.shape[-2]) + for p in permutation: + p = p.item() + #for p in range(res.shape[-1]*res.shape[-2]): # Reconstruct a top & left coordinate. t = p // res.shape[-1] l = p % res.shape[-1] diff --git a/codes/scripts/ui/image_labeler/label_editor.py b/codes/scripts/ui/image_labeler/label_editor.py new file mode 100644 index 00000000..b081f40b --- /dev/null +++ b/codes/scripts/ui/image_labeler/label_editor.py @@ -0,0 +1,33 @@ +import orjson + +from data.image_label_parser import VsNetImageLabeler + + +# Translates from the label JSON output of the VS.NET UI to something more compact and usable. +def convert_from_vsnet_labels(): + labeler = VsNetImageLabeler(['F:\\4k6k\datasets\\ns_images\\512_unsupervised\\categories.json', + 'F:\\4k6k\datasets\\ns_images\\512_unsupervised\\categories_new.json', + 'F:\\4k6k\datasets\\ns_images\\512_unsupervised\\categories_new_new.json']) + # Proposed format: + # 'config': { 'dim' } + # 'labels': [{ 'label', 'key'}] <- ordered by label index. + # 'images': {'file': [{ 'lid', 'top', 'left' }} + # 'labelMap' {} + out_dict = { + 'config': { + 'dim': next(iter(labeler.labeled_images.values()))[0]['patch_width'] + }, + 'labels': [{'label': cat['label'], 'key': cat['keyBinding']} for cat in labeler.categories.values()], + } + out_dict['labelMap'] = {} + for i, lbl in enumerate(out_dict['labels']): + out_dict['labelMap'][lbl['label']] = i + out_dict['images'] = {} + for fname, ilbls in labeler.labeled_images.items(): + out_dict['images'][fname] = [{'lid': out_dict['labelMap'][il['label']], 'top': il['patch_top'], 'left': il['patch_left']} for il in ilbls] + with open("label_editor.json", 'wb') as fout: + fout.write(orjson.dumps(out_dict)) + + +if __name__ == '__main__': + convert_from_vsnet_labels() \ No newline at end of file diff --git a/codes/train.py b/codes/train.py index 53144678..19447aca 100644 --- a/codes/train.py +++ b/codes/train.py @@ -293,7 +293,7 @@ class Trainer: if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_imgset_structural_classifier.yml') + parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_exd_mi1_rrdb_bigboi_pretrain.yml') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0) args = parser.parse_args()