Spaces:
Build error
Build error
| 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) | |