InvokeAI 2.2.4.5__py3-none-any.whl

vendored/clipseg/datasets/phrasecut.py

@@ -0,0 +1,380 @@
+
+import torch
+import numpy as np
+import os
+
+from os.path import join, isdir, isfile, expanduser
+from PIL import Image
+
+from torchvision import transforms
+from torchvision.transforms.transforms import Resize
+
+from torch.nn import functional as nnf
+from general_utils import get_from_repository
+
+from skimage.draw import polygon2mask
+
+PASCAL_SYNSETS = ['person.n.01', 'bird.n.01', 'cat.n.01', 'cattle.n.01', 'dog.n.01', 'horse.n.01', 'sheep.n.01', 'aeroplane.n.01', 'bicycle.n.01', 
+                    'boat.n.01', 'bus.n.01', 'car.n.01', 'motorcycle.n.01', 'train.n.01', 'bottle.n.01', 'chair.n.01', 'kitchen_table.n.01',  
+                    'breakfast_table.n.01', 'trestle_table.n.01','pot_plant.n.01', 'sofa.n.01', 'television.n.03']
+
+PASCAL_5I_SYNSETS_ORDERED = [
+    'aeroplane.n.01', 'bicycle.n.01',  'bird.n.01', 'vessel.n.02', 'bottle.n.01', 'bus.n.01', 'car.n.01', 
+     'cat.n.01',  'chair.n.01',  'cattle.n.01',   'table.n.02',  'dog.n.01', 'horse.n.01',  'motorcycle.n.01', 
+      'person.n.01', 'pot_plant.n.01', 'sheep.n.01',  'sofa.n.01', 'train.n.01', 'television.n.03']
+
+# Pascal-5i classes
+PASCAL_5I_CLASS_IDS = {
+    3: list(range(1, 16)),
+    2: list(range(1, 11)) + list(range(16, 21)),
+    1: list(range(1, 6)) + list(range(11, 21)),
+    0: list(range(6, 21))
+}
+
+
+def traverse_lemmas(synset):
+    """ list of all lemma names of hypernyms """
+
+    out = synset.lemma_names()
+    for h in synset.hypernyms():
+        out += traverse_lemmas(h)
+
+    return out
+
+
+def traverse_lemmas_hypo(synset, max_depth=None, depth=0):
+    """ list of all lemma names of hypernyms """
+
+    if synset.name() == 'person.n.01':
+        return ['person', 'human', 'man', 'woman', 'toddler', 'baby', 'body', 'child', 'infant'] 
+
+    out = [l.name() for l in synset.lemmas()]
+
+    # print(' '*depth, synset.name())
+
+    if max_depth is not None and depth >= max_depth:
+        return out
+
+    # print(' '*(3*depth), out)
+    for h in synset.hyponyms():
+        out += traverse_lemmas_hypo(h, max_depth, depth+1)
+
+    return out
+
+
+
+def random_crop_slices(origin_size, target_size):
+    """Gets slices of a random crop. """
+    assert origin_size[0] >= target_size[0] and origin_size[1] >= target_size[1], f'actual size: {origin_size}, target size: {target_size}'
+
+    offset_y = torch.randint(0, origin_size[0] - target_size[0] + 1, (1,)).item()  # range: 0 <= value < high
+    offset_x = torch.randint(0, origin_size[1] - target_size[1] + 1, (1,)).item()
+
+    return slice(offset_y, offset_y + target_size[0]), slice(offset_x, offset_x + target_size[1])
+
+
+def find_crop(seg, image_size, iterations=1000, min_frac=None, best_of=None):
+
+    best_crops = []
+    best_crop_not_ok = float('-inf'), None, None
+    min_sum = 0
+
+    seg = seg.astype('bool')
+    
+    if min_frac is not None:
+        #min_sum = seg.sum() * min_frac
+        min_sum = seg.shape[0] * seg.shape[1] * min_frac
+    
+    for iteration in range(iterations):
+        sl_y, sl_x = random_crop_slices(seg.shape, image_size)
+        seg_ = seg[sl_y, sl_x]
+        sum_seg_ = seg_.sum()
+
+        if sum_seg_ > min_sum:
+
+            if best_of is None:
+                return sl_y, sl_x, False
+            else:
+                best_crops += [(sum_seg_, sl_y, sl_x)]
+                if len(best_crops) >= best_of:
+                    best_crops.sort(key=lambda x:x[0], reverse=True)
+                    sl_y, sl_x = best_crops[0][1:]
+                    
+                    return sl_y, sl_x, False
+
+        else:
+            if sum_seg_ > best_crop_not_ok[0]:
+                best_crop_not_ok = sum_seg_, sl_y, sl_x
+        
+    else:
+        # return best segmentation found
+        return best_crop_not_ok[1:] + (best_crop_not_ok[0] <= min_sum,) 
+
+
+class PhraseCut(object):
+
+    def __init__(self, split, image_size=400, negative_prob=0, aug=None, aug_color=False, aug_crop=True,
+                 min_size=0, remove_classes=None, with_visual=False, only_visual=False, mask=None):
+        super().__init__()
+
+        self.negative_prob = negative_prob
+        self.image_size = image_size
+        self.with_visual = with_visual
+        self.only_visual = only_visual
+        self.phrase_form = '{}'
+        self.mask = mask
+        self.aug_crop = aug_crop
+        
+        if aug_color:
+            self.aug_color = transforms.Compose([
+                transforms.ColorJitter(0.5, 0.5, 0.2, 0.05),
+            ])
+        else:
+            self.aug_color = None
+
+        get_from_repository('PhraseCut', ['PhraseCut.tar'], integrity_check=lambda local_dir: all([
+            isdir(join(local_dir, 'VGPhraseCut_v0')),
+            isdir(join(local_dir, 'VGPhraseCut_v0', 'images')),
+            isfile(join(local_dir, 'VGPhraseCut_v0', 'refer_train.json')),
+            len(os.listdir(join(local_dir, 'VGPhraseCut_v0', 'images'))) in {108250, 108249}
+        ]))
+
+        from third_party.PhraseCutDataset.utils.refvg_loader import RefVGLoader
+        self.refvg_loader = RefVGLoader(split=split)
+
+        # img_ids where the size in the annotations does not match actual size
+        invalid_img_ids = set([150417, 285665, 498246, 61564, 285743, 498269, 498010, 150516, 150344, 286093, 61530, 
+                               150333, 286065, 285814, 498187, 285761, 498042])
+        
+        mean = [0.485, 0.456, 0.406]
+        std = [0.229, 0.224, 0.225]
+        self.normalize = transforms.Normalize(mean, std)
+
+        self.sample_ids = [(i, j) 
+                           for i in self.refvg_loader.img_ids 
+                           for j in range(len(self.refvg_loader.get_img_ref_data(i)['phrases']))
+                           if i not in invalid_img_ids]
+        
+
+        # self.all_phrases = list(set([p for i in self.refvg_loader.img_ids for p in self.refvg_loader.get_img_ref_data(i)['phrases']]))
+
+        from nltk.stem import WordNetLemmatizer
+        wnl = WordNetLemmatizer()        
+
+        # Filter by class (if remove_classes is set)
+        if remove_classes is None:
+            pass
+        else:
+            from nltk.corpus import wordnet
+
+            print('remove pascal classes...')
+
+            get_data = self.refvg_loader.get_img_ref_data  # shortcut
+            keep_sids = None
+
+            if remove_classes[0] == 'pas5i':
+                subset_id = remove_classes[1]
+
+                avoid = [PASCAL_5I_SYNSETS_ORDERED[i] for i in range(20) if i+1 not in PASCAL_5I_CLASS_IDS[subset_id]]
+      
+
+            elif remove_classes[0] == 'zs':
+                stop = remove_classes[1]
+                
+                from datasets.pascal_zeroshot import PASCAL_VOC_CLASSES_ZS
+
+                avoid = [c for class_set in PASCAL_VOC_CLASSES_ZS[:stop] for c in class_set]
+                print(avoid)
+
+            elif remove_classes[0] == 'aff':
+                # avoid = ['drink.v.01', 'sit.v.01', 'ride.v.02']
+                # all_lemmas = set(['drink', 'sit', 'ride'])
+                avoid = ['drink', 'drinks', 'drinking', 'sit', 'sits', 'sitting', 
+                         'ride', 'rides', 'riding',
+                         'fly', 'flies', 'flying', 'drive', 'drives', 'driving', 'driven', 
+                         'swim', 'swims', 'swimming',
+                         'wheels', 'wheel', 'legs', 'leg', 'ear', 'ears']
+                keep_sids = [(i, j) for i, j in self.sample_ids if 
+                             all(x not in avoid for x in get_data(i)['phrases'][j].split(' '))]
+
+            print('avoid classes:', avoid)
+
+
+            if keep_sids is None:
+                all_lemmas = [s for ps in avoid for s in traverse_lemmas_hypo(wordnet.synset(ps), max_depth=None)]
+                all_lemmas = list(set(all_lemmas))
+                all_lemmas = [h.replace('_', ' ').lower() for h in all_lemmas]
+                all_lemmas = set(all_lemmas)
+
+                # divide into multi word and single word
+                all_lemmas_s = set(l for l in all_lemmas if ' ' not in l)
+                all_lemmas_m = set(l for l in all_lemmas if l not in all_lemmas_s)
+
+                # new3
+                phrases = [get_data(i)['phrases'][j] for i, j in self.sample_ids]
+                remove_sids = set((i,j) for (i,j), phrase in zip(self.sample_ids, phrases)
+                                  if any(l in phrase for l in all_lemmas_m) or 
+                                  len(set(wnl.lemmatize(w) for w in phrase.split(' ')).intersection(all_lemmas_s)) > 0
+                )
+                keep_sids = [(i, j) for i, j in self.sample_ids if (i,j) not in remove_sids]
+
+            print(f'Reduced to {len(keep_sids) / len(self.sample_ids):.3f}')
+            removed_ids = set(self.sample_ids) - set(keep_sids)
+
+            print('Examples of removed', len(removed_ids))
+            for i, j in list(removed_ids)[:20]:
+                print(i, get_data(i)['phrases'][j])
+
+            self.sample_ids = keep_sids
+
+        from itertools import groupby
+        samples_by_phrase = [(self.refvg_loader.get_img_ref_data(i)['phrases'][j], (i, j)) 
+                             for i, j in self.sample_ids]
+        samples_by_phrase = sorted(samples_by_phrase)
+        samples_by_phrase = groupby(samples_by_phrase, key=lambda x: x[0])
+        
+        self.samples_by_phrase = {prompt: [s[1] for s in prompt_sample_ids] for prompt, prompt_sample_ids in samples_by_phrase}
+
+        self.all_phrases = list(set(self.samples_by_phrase.keys()))
+
+
+        if self.only_visual:
+            assert self.with_visual
+            self.sample_ids = [(i, j) for i, j in self.sample_ids
+                               if len(self.samples_by_phrase[self.refvg_loader.get_img_ref_data(i)['phrases'][j]]) > 1]
+
+        # Filter by size (if min_size is set)
+        sizes = [self.refvg_loader.get_img_ref_data(i)['gt_boxes'][j] for i, j in self.sample_ids]
+        image_sizes = [self.refvg_loader.get_img_ref_data(i)['width'] * self.refvg_loader.get_img_ref_data(i)['height'] for i, j in self.sample_ids]
+        #self.sizes = [sum([(s[2] - s[0]) * (s[3] - s[1]) for s in size]) for size in sizes]
+        self.sizes = [sum([s[2] * s[3] for s in size]) / img_size for size, img_size in zip(sizes, image_sizes)]
+
+        if min_size:
+            print('filter by size')
+
+        self.sample_ids = [self.sample_ids[i] for i in range(len(self.sample_ids)) if self.sizes[i] > min_size]
+
+        self.base_path = join(expanduser('~/datasets/PhraseCut/VGPhraseCut_v0/images/'))
+
+    def __len__(self):
+        return len(self.sample_ids)
+
+
+    def load_sample(self, sample_i, j):
+
+        img_ref_data = self.refvg_loader.get_img_ref_data(sample_i)
+
+        polys_phrase0 = img_ref_data['gt_Polygons'][j]
+        phrase = img_ref_data['phrases'][j]
+        phrase = self.phrase_form.format(phrase)
+
+        masks = []
+        for polys in polys_phrase0:
+            for poly in polys:
+                poly = [p[::-1] for p in poly]  # swap x,y
+                masks += [polygon2mask((img_ref_data['height'], img_ref_data['width']), poly)]
+
+        seg = np.stack(masks).max(0)
+        img = np.array(Image.open(join(self.base_path, str(img_ref_data['image_id']) + '.jpg')))
+
+        min_shape = min(img.shape[:2])
+
+        if self.aug_crop:
+            sly, slx, exceed = find_crop(seg, (min_shape, min_shape), iterations=50, min_frac=0.05)
+        else:
+            sly, slx = slice(0, None), slice(0, None)
+    
+        seg = seg[sly, slx]
+        img = img[sly, slx]
+
+        seg = seg.astype('uint8')
+        seg = torch.from_numpy(seg).view(1, 1, *seg.shape)
+
+        if img.ndim == 2:
+            img = np.dstack([img] * 3)
+
+        img = torch.from_numpy(img).permute(2,0,1).unsqueeze(0).float()
+
+        seg = nnf.interpolate(seg, (self.image_size, self.image_size), mode='nearest')[0,0]
+        img = nnf.interpolate(img, (self.image_size, self.image_size), mode='bilinear', align_corners=True)[0]
+
+        # img = img.permute([2,0, 1])
+        img = img / 255.0
+
+        if self.aug_color is not None:
+            img = self.aug_color(img)
+
+        img = self.normalize(img)
+
+
+
+        return img, seg, phrase
+
+    def __getitem__(self, i):
+ 
+        sample_i, j = self.sample_ids[i]
+
+        img, seg, phrase = self.load_sample(sample_i, j)
+
+        if self.negative_prob > 0:
+            if torch.rand((1,)).item() < self.negative_prob:
+
+                new_phrase = None
+                while new_phrase is None or new_phrase == phrase:
+                    idx = torch.randint(0, len(self.all_phrases), (1,)).item()
+                    new_phrase = self.all_phrases[idx]
+                phrase = new_phrase
+                seg = torch.zeros_like(seg)
+
+        if self.with_visual:
+            # find a corresponding visual image
+            if phrase in self.samples_by_phrase and len(self.samples_by_phrase[phrase]) > 1:
+                idx = torch.randint(0, len(self.samples_by_phrase[phrase]), (1,)).item()
+                other_sample = self.samples_by_phrase[phrase][idx]
+                #print(other_sample)
+                img_s, seg_s, _ = self.load_sample(*other_sample)
+
+                from datasets.utils import blend_image_segmentation
+
+                if self.mask in {'separate', 'text_and_separate'}:
+                    # assert img.shape[1:] == img_s.shape[1:] == seg_s.shape == seg.shape[1:]
+                    add_phrase = [phrase] if self.mask == 'text_and_separate' else []
+                    vis_s = add_phrase + [img_s, seg_s, True]
+                else:
+                    if self.mask.startswith('text_and_'):
+                        mask_mode = self.mask[9:]
+                        label_add = [phrase]
+                    else:
+                        mask_mode = self.mask
+                        label_add = []
+
+                    masked_img_s = torch.from_numpy(blend_image_segmentation(img_s, seg_s, mode=mask_mode, image_size=self.image_size)[0])
+                    vis_s = label_add + [masked_img_s, True]
+                
+            else:
+                # phrase is unique
+                vis_s = torch.zeros_like(img)
+
+                if self.mask in {'separate', 'text_and_separate'}:
+                    add_phrase = [phrase] if self.mask == 'text_and_separate' else []
+                    vis_s = add_phrase + [vis_s, torch.zeros(*vis_s.shape[1:], dtype=torch.uint8), False]
+                elif self.mask.startswith('text_and_'):
+                    vis_s = [phrase, vis_s, False]
+                else:
+                    vis_s = [vis_s, False]
+        else:
+            assert self.mask == 'text'
+            vis_s = [phrase]
+        
+        seg = seg.unsqueeze(0).float()
+
+        data_x = (img,) + tuple(vis_s)
+
+        return data_x, (seg, torch.zeros(0), i)
+
+
+class PhraseCutPlus(PhraseCut):
+
+    def __init__(self, split, image_size=400, aug=None, aug_color=False, aug_crop=True, min_size=0, remove_classes=None, only_visual=False, mask=None):
+        super().__init__(split, image_size=image_size, negative_prob=0.2, aug=aug, aug_color=aug_color, aug_crop=aug_crop, min_size=min_size, 
+                         remove_classes=remove_classes, with_visual=True, only_visual=only_visual, mask=mask)

vendored/clipseg/datasets/utils.py

@@ -0,0 +1,68 @@
+
+import numpy as np
+import torch
+
+
+def blend_image_segmentation(img, seg, mode, image_size=224):
+
+
+    if mode in {'blur_highlight', 'blur3_highlight', 'blur3_highlight01', 'blur_highlight_random', 'crop'}:
+        if isinstance(img, np.ndarray):
+            img = torch.from_numpy(img)
+
+        if isinstance(seg, np.ndarray):
+            seg = torch.from_numpy(seg)            
+
+    if mode == 'overlay':
+        out = img * seg
+        out = [out.astype('float32')]
+    elif mode == 'highlight':
+        out = img * seg[None, :, :] * 0.85 + 0.15 * img
+        out = [out.astype('float32')]
+    elif mode == 'highlight2':
+        img = img / 2
+        out = (img+0.1) * seg[None, :, :] + 0.3 * img
+        out = [out.astype('float32')]
+    elif mode == 'blur_highlight':
+        from evaluation_utils import img_preprocess
+        out  = [img_preprocess((None, [img], [seg]), blur=1, bg_fac=0.5).numpy()[0] - 0.01]
+    elif mode == 'blur3_highlight':
+        from evaluation_utils import img_preprocess
+        out  = [img_preprocess((None, [img], [seg]), blur=3, bg_fac=0.5).numpy()[0] - 0.01]
+    elif mode == 'blur3_highlight01':
+        from evaluation_utils import img_preprocess
+        out  = [img_preprocess((None, [img], [seg]), blur=3, bg_fac=0.1).numpy()[0] - 0.01]                
+    elif mode == 'blur_highlight_random':
+        from evaluation_utils import img_preprocess
+        out  = [img_preprocess((None, [img], [seg]), blur=0 + torch.randint(0, 3, (1,)).item(), bg_fac=0.1 + 0.8*torch.rand(1).item()).numpy()[0] - 0.01]               
+    elif mode == 'crop':
+        from evaluation_utils import img_preprocess
+        out  = [img_preprocess((None, [img], [seg]), blur=1, center_context=0.1, image_size=image_size)[0].numpy()]  
+    elif mode == 'crop_blur_highlight':
+        from evaluation_utils import img_preprocess
+        out  = [img_preprocess((None, [img], [seg]), blur=3, center_context=0.1, bg_fac=0.1, image_size=image_size)[0].numpy()]  
+    elif mode == 'crop_blur_highlight352':
+        from evaluation_utils import img_preprocess
+        out  = [img_preprocess((None, [img], [seg]), blur=3, center_context=0.1, bg_fac=0.1, image_size=352)[0].numpy()]          
+    elif mode == 'shape':
+        out = [np.stack([seg[:, :]]*3).astype('float32')]
+    elif mode == 'concat':
+        out = [np.concatenate([img, seg[None, :, :]]).astype('float32')]
+    elif mode == 'image_only':
+        out = [img.astype('float32')]
+    elif mode == 'image_black':
+        out = [img.astype('float32')*0]        
+    elif mode is None:
+        out = [img.astype('float32')]
+    elif mode == 'separate':
+        out = [img.astype('float32'), seg.astype('int64')]
+    elif mode == 'separate_img_black':
+        out = [img.astype('float32')*0, seg.astype('int64')]        
+    elif mode == 'separate_seg_ones':
+        out = [img.astype('float32'), np.ones_like(seg).astype('int64')]                
+    elif mode == 'separate_both_black':
+        out = [img.astype('float32')*0, seg.astype('int64')*0]        
+    else:
+        raise ValueError(f'invalid mode: {mode}')
+
+    return out