File size: 22,228 Bytes
e0c7c25
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
import os, torch, numpy, base64, json, re, threading, random
from torch.utils.data import TensorDataset, DataLoader
from collections import defaultdict
from netdissect.easydict import EasyDict
from netdissect.modelconfig import create_instrumented_model
from netdissect.runningstats import RunningQuantile
from netdissect.dissection import safe_dir_name
from netdissect.zdataset import z_sample_for_model
from PIL import Image
from io import BytesIO

class DissectionProject:
    '''
    DissectionProject understand how to drive a GanTester within a
    dissection project directory structure: it caches data in files,
    creates image files, and translates data between plain python data
    types and the pytorch-specific tensors required by GanTester.
    '''
    def __init__(self, config, project_dir, path_url, public_host):
        print('config done', project_dir)
        self.use_cuda = torch.cuda.is_available()
        self.dissect = config
        self.project_dir = project_dir
        self.path_url = path_url
        self.public_host = public_host
        self.cachedir = os.path.join(self.project_dir, 'cache')
        self.tester = GanTester(
                config.settings, dissectdir=project_dir,
                device=torch.device('cuda') if self.use_cuda
                     else torch.device('cpu'))
        self.stdz = []

    def get_zs(self, size):
        if size <= len(self.stdz):
            return self.stdz[:size].tolist()
        z_tensor = self.tester.standard_z_sample(size)
        numpy_z = z_tensor.cpu().numpy()
        self.stdz = numpy_z
        return self.stdz.tolist()

    def get_z(self, id):
        if id < len(self.stdz):
            return self.stdz[id]
        return self.get_zs((id + 1) * 2)[id]

    def get_zs_for_ids(self, ids):
        max_id = max(ids)
        if max_id >= len(self.stdz):
            self.get_z(max_id)
        return self.stdz[ids]

    def get_layers(self):
        result = []
        layer_shapes = self.tester.layer_shapes()
        for layer in self.tester.layers:
            shape = layer_shapes[layer]
            result.append(dict(
                layer=layer,
                channels=shape[1],
                shape=[shape[2], shape[3]]))
        return result

    def get_units(self, layer):
        try:
            dlayer = [dl for dl in self.dissect['layers']
                    if dl['layer'] == layer][0]
        except:
            return None

        dunits = dlayer['units']
        result = [dict(unit=unit_num,
                       img='/%s/%s/s-image/%d-top.jpg' %
                        (self.path_url, layer, unit_num),
                       label=unit['iou_label'])
                  for unit_num, unit in enumerate(dunits)]
        return result

    def get_rankings(self, layer):
        try:
            dlayer = [dl for dl in self.dissect['layers']
                    if dl['layer'] == layer][0]
        except:
            return None
        result = [dict(name=ranking['name'],
                       metric=ranking.get('metric', None),
                       scores=ranking['score'])
                  for ranking in dlayer['rankings']]
        return result

    def get_levels(self, layer, quantiles):
        levels = self.tester.levels(
                layer, torch.from_numpy(numpy.array(quantiles)))
        return levels.cpu().numpy().tolist()

    def generate_images(self, zs, ids, interventions, return_urls=False):
        if ids is not None:
            assert zs is None
            zs = self.get_zs_for_ids(ids)
            if not interventions:
                # Do file caching when ids are given (and no ablations).
                imgdir = os.path.join(self.cachedir, 'img', 'id')
                os.makedirs(imgdir, exist_ok=True)
                exist = set(os.listdir(imgdir))
                unfinished = [('%d.jpg' % id) not in exist for id in ids]
                needed_z_tensor = torch.tensor(zs[unfinished]).float().to(
                        self.tester.device)
                needed_ids = numpy.array(ids)[unfinished]
                # Generate image files for just the needed images.
                if len(needed_z_tensor):
                    imgs = self.tester.generate_images(needed_z_tensor
                            ).cpu().numpy()
                    for i, img in zip(needed_ids, imgs):
                         Image.fromarray(img.transpose(1, 2, 0)).save(
                                 os.path.join(imgdir, '%d.jpg' % i), 'jpeg',
                                 quality=99, optimize=True, progressive=True)
                # Assemble a response.
                imgurls = ['/%s/cache/img/id/%d.jpg'
                      % (self.path_url, i) for i in ids]
                return [dict(id=i, d=d) for i, d in zip(ids, imgurls)]
        # No file caching when ids are not given (or ablations are applied)
        z_tensor = torch.tensor(zs).float().to(self.tester.device)
        imgs = self.tester.generate_images(z_tensor,
                intervention=decode_intervention_array(interventions,
                    self.tester.layer_shapes()),
                ).cpu().numpy()
        numpy_z = z_tensor.cpu().numpy()
        if return_urls:
            randdir = '%03d' % random.randrange(1000)
            imgdir = os.path.join(self.cachedir, 'img', 'uniq', randdir)
            os.makedirs(imgdir, exist_ok=True)
            startind = random.randrange(100000)
            imgurls = []
            for i, img in enumerate(imgs):
                filename = '%d.jpg' % (i + startind)
                Image.fromarray(img.transpose(1, 2, 0)).save(
                         os.path.join(imgdir, filename), 'jpeg',
                         quality=99, optimize=True, progressive=True)
                image_url_path = ('/%s/cache/img/uniq/%s/%s'
                      % (self.path_url, randdir, filename))
                imgurls.append(image_url_path)
                tweet_filename = 'tweet-%d.html' % (i + startind)
                tweet_url_path = ('/%s/cache/img/uniq/%s/%s'
                      % (self.path_url, randdir, tweet_filename))
                with open(os.path.join(imgdir, tweet_filename), 'w') as f:
                    f.write(twitter_card(image_url_path, tweet_url_path,
                        self.public_host))
            return [dict(d=d) for d in imgurls]
        imgurls = [img2base64(img.transpose(1, 2, 0)) for img in imgs]
        return [dict(d=d) for d in imgurls]

    def get_features(self, ids, masks, layers, interventions):
        zs = self.get_zs_for_ids(ids)
        z_tensor = torch.tensor(zs).float().to(self.tester.device)
        t_masks = torch.stack(
                [torch.from_numpy(mask_to_numpy(mask)) for mask in masks]
                )[:,None,:,:].to(self.tester.device)
        t_features = self.tester.feature_stats(z_tensor, t_masks,
                decode_intervention_array(interventions,
                    self.tester.layer_shapes()), layers)
        # Convert torch arrays to plain python lists before returning.
        return { layer: { key: value.cpu().numpy().tolist()
                          for key, value in feature.items() }
                 for layer, feature in t_features.items() }

    def get_featuremaps(self, ids, layers, interventions):
        zs = self.get_zs_for_ids(ids)
        z_tensor = torch.tensor(zs).float().to(self.tester.device)
        # Quantilized features are returned.
        q_features = self.tester.feature_maps(z_tensor,
                decode_intervention_array(interventions,
                    self.tester.layer_shapes()), layers)
        # Scale them 0-255 and return them.
        # TODO: turn them into pngs for returning.
        return { layer: [
            value.clamp(0, 1).mul(255).byte().cpu().numpy().tolist()
            for value in valuelist ]
            for layer, valuelist in q_features.items()
            if (not layers) or (layer in layers) }

    def get_recipes(self):
        recipedir = os.path.join(self.project_dir, 'recipe')
        if not os.path.isdir(recipedir):
            return []
        result = []
        for filename in os.listdir(recipedir):
            with open(os.path.join(recipedir, filename)) as f:
                result.append(json.load(f))
        return result




class GanTester:
    '''
    GanTester holds on to a specific model to test.

    (1) loads and instantiates the GAN;
    (2) instruments it at every layer so that units can be ablated
    (3) precomputes z dimensionality, and output image dimensions.
    '''
    def __init__(self, args, dissectdir=None, device=None):
        self.cachedir = os.path.join(dissectdir, 'cache')
        self.device = device if device is not None else torch.device('cpu')
        self.dissectdir = dissectdir
        self.modellock = threading.Lock()

        # Load the generator from the pth file.
        args_copy = EasyDict(args)
        args_copy.edit = True
        model = create_instrumented_model(args_copy)
        model.eval()
        self.model = model

        # Get the set of layers of interest.
        # Default: all shallow children except last.
        self.layers = sorted(model.retained_features().keys())

        # Move it to CUDA if wanted.
        model.to(device)

        self.quantiles = {
            layer: load_quantile_if_present(os.path.join(self.dissectdir,
                safe_dir_name(layer)), 'quantiles.npz',
                device=torch.device('cpu'))
            for layer in self.layers }

    def layer_shapes(self):
        return self.model.feature_shape

    def standard_z_sample(self, size=100, seed=1, device=None):
        '''
        Generate a standard set of random Z as a (size, z_dimension) tensor.
        With the same random seed, it always returns the same z (e.g.,
        the first one is always the same regardless of the size.)
        '''
        result = z_sample_for_model(self.model, size)
        if device is not None:
            result = result.to(device)
        return result

    def reset_intervention(self):
        self.model.remove_edits()

    def apply_intervention(self, intervention):
        '''
        Applies an ablation recipe of the form [(layer, unit, alpha)...].
        '''
        self.reset_intervention()
        if not intervention:
            return
        for layer, (a, v) in intervention.items():
            self.model.edit_layer(layer, ablation=a, replacement=v)

    def generate_images(self, z_batch, intervention=None):
        '''
        Makes some images.
        '''
        with torch.no_grad(), self.modellock:
            batch_size = 10
            self.apply_intervention(intervention)
            test_loader = DataLoader(TensorDataset(z_batch[:,:,None,None]),
                batch_size=batch_size,
                pin_memory=('cuda' == self.device.type
                            and z_batch.device.type == 'cpu'))
            result_img = torch.zeros(
                    *((len(z_batch), 3) + self.model.output_shape[2:]),
                    dtype=torch.uint8, device=self.device)
            for batch_num, [batch_z,] in enumerate(test_loader):
                batch_z = batch_z.to(self.device)
                out = self.model(batch_z)
                result_img[batch_num*batch_size:
                        batch_num*batch_size+len(batch_z)] = (
                                (((out + 1) / 2) * 255).clamp(0, 255).byte())
            return result_img

    def get_layers(self):
        return self.layers

    def feature_stats(self, z_batch,
            masks=None, intervention=None, layers=None):
        feature_stat = defaultdict(dict)
        with torch.no_grad(), self.modellock:
            batch_size = 10
            self.apply_intervention(intervention)
            if masks is None:
                masks = torch.ones(z_batch.size(0), 1, 1, 1,
                        device=z_batch.device, dtype=z_batch.dtype)
            else:
                assert masks.shape[0] == z_batch.shape[0]
                assert masks.shape[1] == 1
            test_loader = DataLoader(
                TensorDataset(z_batch[:,:,None,None], masks),
                batch_size=batch_size,
                pin_memory=('cuda' == self.device.type
                    and z_batch.device.type == 'cpu'))
            processed = 0
            for batch_num, [batch_z, batch_m] in enumerate(test_loader):
                batch_z, batch_m = [
                        d.to(self.device) for d in [batch_z, batch_m]]
                # Run model but disregard output
                self.model(batch_z)
                processing = batch_z.shape[0]
                for layer, feature in self.model.retained_features().items():
                    if layers is not None:
                        if layer not in layers:
                            continue
                    # Compute max features touching mask
                    resized_max = torch.nn.functional.adaptive_max_pool2d(
                            batch_m,
                            (feature.shape[2], feature.shape[3]))
                    max_feature = (feature * resized_max).view(
                            feature.shape[0], feature.shape[1], -1
                            ).max(2)[0].max(0)[0]
                    if 'max' not in feature_stat[layer]:
                        feature_stat[layer]['max'] = max_feature
                    else:
                        torch.max(feature_stat[layer]['max'], max_feature,
                                    out=feature_stat[layer]['max'])
                    # Compute mean features weighted by overlap with mask
                    resized_mean = torch.nn.functional.adaptive_avg_pool2d(
                            batch_m,
                            (feature.shape[2], feature.shape[3]))
                    mean_feature = (feature * resized_mean).view(
                            feature.shape[0], feature.shape[1], -1
                            ).sum(2).sum(0) / (resized_mean.sum() + 1e-15)
                    if 'mean' not in feature_stat[layer]:
                        feature_stat[layer]['mean'] = mean_feature
                    else:
                        feature_stat[layer]['mean'] = (
                                processed * feature_mean[layer]['mean']
                                + processing * mean_feature) / (
                                        processed + processing)
                processed += processing
            # After summaries are done, also compute quantile stats
            for layer, stats in feature_stat.items():
                if self.quantiles.get(layer, None) is not None:
                    for statname in ['max', 'mean']:
                        stats['%s_quantile' % statname] = (
                            self.quantiles[layer].normalize(stats[statname]))
        return feature_stat

    def levels(self, layer, quantiles):
        return self.quantiles[layer].quantiles(quantiles)

    def feature_maps(self, z_batch, intervention=None, layers=None,
            quantiles=True):
        feature_map = defaultdict(list)
        with torch.no_grad(), self.modellock:
            batch_size = 10
            self.apply_intervention(intervention)
            test_loader = DataLoader(
                TensorDataset(z_batch[:,:,None,None]),
                batch_size=batch_size,
                pin_memory=('cuda' == self.device.type
                    and z_batch.device.type == 'cpu'))
            processed = 0
            for batch_num, [batch_z] in enumerate(test_loader):
                batch_z = batch_z.to(self.device)
                # Run model but disregard output
                self.model(batch_z)
                processing = batch_z.shape[0]
                for layer, feature in self.model.retained_features().items():
                    for single_featuremap in feature:
                        if quantiles:
                            feature_map[layer].append(self.quantiles[layer]
                                    .normalize(single_featuremap))
                        else:
                            feature_map[layer].append(single_featuremap)
        return feature_map

def load_quantile_if_present(outdir, filename, device):
    filepath = os.path.join(outdir, filename)
    if os.path.isfile(filepath):
        data = numpy.load(filepath)
        result = RunningQuantile(state=data)
        result.to_(device)
        return result
    return None

if __name__ == '__main__':
    test_main()

def mask_to_numpy(mask_record):
    # Detect a png image mask.
    bitstring = mask_record['bitstring']
    bitnumpy = None
    default_shape = (256, 256)
    if 'image/png;base64,' in bitstring:
        bitnumpy = base642img(bitstring)
        default_shape = bitnumpy.shape[:2]
    # Set up results
    shape = mask_record.get('shape', None)
    if not shape: # None or empty []
        shape = default_shape
    result = numpy.zeros(shape=shape, dtype=numpy.float32)
    bitbounds = mask_record.get('bitbounds', None)
    if not bitbounds: # None or empty []
        bitbounds = ([0] * len(result.shape)) + list(result.shape)
    start = bitbounds[:len(result.shape)]
    end = bitbounds[len(result.shape):]
    if bitnumpy is not None:
        if bitnumpy.shape[2] == 4:
            # Mask is any nontransparent bits in the alpha channel if present
            result[start[0]:end[0], start[1]:end[1]] = (bitnumpy[:,:,3] > 0)
        else:
            # Or any nonwhite pixels in the red channel if no alpha.
            result[start[0]:end[0], start[1]:end[1]] = (bitnumpy[:,:,0] < 255)
        return result
    else:
        # Or bitstring can be just ones and zeros.
        indexes = start.copy()
        bitindex = 0
        while True:
            result[tuple(indexes)] = (bitstring[bitindex] != '0')
            for ii in range(len(indexes) - 1, -1, -1):
                if indexes[ii] < end[ii] - 1:
                    break
                indexes[ii] = start[ii]
            else:
                assert (bitindex + 1) == len(bitstring)
                return result
            indexes[ii] += 1
            bitindex += 1

def decode_intervention_array(interventions, layer_shapes):
    result = {}
    for channels in [decode_intervention(intervention, layer_shapes)
            for intervention in (interventions or [])]:
        for layer, channel in channels.items():
            if layer not in result:
                result[layer] = channel
                continue
            accum = result[layer]
            newalpha = 1 - (1 - channel[:1]) * (1 - accum[:1])
            newvalue = (accum[1:] * accum[:1] * (1 - channel[:1]) +
                    channel[1:] * channel[:1]) / (newalpha + 1e-40)
            accum[:1] = newalpha
            accum[1:] = newvalue
    return result

def decode_intervention(intervention, layer_shapes):
    # Every plane of an intervention is a solid choice of activation
    # over a set of channels, with a mask applied to alpha-blended channels
    # (when the mask resolution is different from the feature map, it can
    # be either a max-pooled or average-pooled to the proper resolution).
    # This can be reduced to a single alpha-blended featuremap.
    if intervention is None:
        return None
    mask = intervention.get('mask', None)
    if mask:
        mask = torch.from_numpy(mask_to_numpy(mask))
    maskpooling = intervention.get('maskpooling', 'max')
    channels = {}  # layer -> ([alpha, val], c)
    for arec in intervention.get('ablations', []):
        unit = arec['unit']
        layer = arec['layer']
        alpha = arec.get('alpha', 1.0)
        if alpha is None:
            alpha = 1.0
        value = arec.get('value', 0.0)
        if value is None:
            value = 0.0
        if alpha != 0.0 or value != 0.0:
            if layer not in channels:
                channels[layer] = torch.zeros(2, *layer_shapes[layer][1:])
            channels[layer][0, unit] = alpha
            channels[layer][1, unit] = value
    if mask is not None:
        for layer in channels:
            layer_shape = layer_shapes[layer][2:]
            if maskpooling == 'mean':
                layer_mask = torch.nn.functional.adaptive_avg_pool2d(
                    mask[None,None,...], layer_shape)[0]
            else:
                layer_mask = torch.nn.functional.adaptive_max_pool2d(
                    mask[None,None,...], layer_shape)[0]
            channels[layer][0] *= layer_mask
    return channels

def img2base64(imgarray, for_html=True, image_format='jpeg'):
    '''
    Converts a numpy array to a jpeg base64 url
    '''
    input_image_buff = BytesIO()
    Image.fromarray(imgarray).save(input_image_buff, image_format,
            quality=99, optimize=True, progressive=True)
    res = base64.b64encode(input_image_buff.getvalue()).decode('ascii')
    if for_html:
        return 'data:image/' + image_format + ';base64,' + res
    else:
        return res

def base642img(stringdata):
    stringdata = re.sub('^(?:data:)?image/\w+;base64,', '', stringdata)
    im = Image.open(BytesIO(base64.b64decode(stringdata)))
    return numpy.array(im)

def twitter_card(image_path, tweet_path, public_host):
    return '''\
<!doctype html>
<html>
<head>
<meta name="twitter:card" content="summary_large_image" />
<meta name="twitter:title" content="Painting with GANs from MIT-IBM Watson AI Lab" />
<meta name="twitter:description" content="This demo lets you modify a selection of meaningful GAN units for a generated image by simply painting." />
<meta name="twitter:image" content="http://{public_host}{image_path}" />
<meta name="twitter:url" content="http://{public_host}{tweet_path}" />
<meta http-equiv="refresh" content="10; url=http://bit.ly/ganpaint">
</head>
<style>
body {{ font: 12px Arial, sans-serif; }}
</style>
<body>
<center>
<h1>Painting with GANs from MIT-IBM Watson AI Lab</h1>
<p>This demo lets you modify a selection of meatningful GAN units for a generated image by simply painting.</p>
<img src="{image_path}">
<p>Redirecting to
<a href="http://bit.ly/ganpaint">GANPaint</a>
</p>
</center>
</body>
'''.format(
        image_path=image_path,
        tweet_path=tweet_path,
        public_host=public_host)