File size: 13,571 Bytes
e0b74e6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335

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



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 datasets.generate_lvis_oneshot import PASCAL_SYNSETS, traverse_lemmas, traverse_lemmas_hypo
            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]
                from datasets.generate_lvis_oneshot import PASCAL_5I_SYNSETS_ORDERED, PASCAL_5I_CLASS_IDS
                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)