ClothingGAN / netdissect /serverstate.py
mfrashad's picture
Init code
97069e1
raw
history blame
22.2 kB
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)