def test_output(image, jpeg_quality=50, rounding_approximation=None):
jpg = DJPG(rounding_approximation=rounding_approximation, rounding_approximation_steps=5)
print(jpg)
batch_x = np.expand_dims(image, 0)
batch_y = jpg.process(batch_x / 255, jpeg_quality)
n_images = batch_x.shape[0]
batch_j = np.zeros_like(batch_x)
for n in range(n_images):
io.imwrite('/tmp/patch_{}.jpg'.format(n), (batch_x[n].squeeze()).astype(np.uint8), quality=jpeg_quality)
batch_j[n] = io.imread('/tmp/patch_{}.jpg'.format(n))
for n in range(n_images):
plt.subplot(n_images, 3, 1 + n*3)
plotting.quickshow(batch_x[n].squeeze() / np.max(np.abs(batch_x)), 'Input')
plt.subplot(n_images, 3, 2 + n*3)
plotting.quickshow(batch_y[n].squeeze() / np.max(np.abs(batch_y)), 'dJPEG Model')
plt.subplot(n_images, 3, 3 + n*3)
plotting.quickshow(batch_j[n].squeeze() / np.max(np.abs(batch_j)), 'libJPG Codec')
plt.show()
def test_quality(image, rounding_approximation=None, n_quality_levels=91):
jpg = DJPG(rounding_approximation=rounding_approximation,
rounding_approximation_steps=5)
print(jpg)
batch_x = np.expand_dims(image[0:1024, 0:1024, :], 0)
psnrs_y, psnrs_j = [], []
quality_levels = np.unique(
np.round(np.linspace(10, 100,
n_quality_levels)).astype(np.int32)).tolist()
print('Using quality levels: {}'.format(quality_levels))
for jpeg_quality in quality_levels:
batch_y = jpg.process(batch_x / 255, jpeg_quality)
batch_y = np.round(255 * batch_y) / 255
io.imwrite('/tmp/patch.jpg', (batch_x.squeeze()).astype(np.uint8),
quality=jpeg_quality,
subsampling='4:4:4')
batch_j = io.imread('/tmp/patch.jpg')
psnrs_y.append(
compare_psnr(batch_x.squeeze(), 255 * batch_y.squeeze(), 255))
psnrs_j.append(compare_psnr(batch_x.squeeze(), batch_j.squeeze(), 255))
# Plot
plt.figure(figsize=(6, 6))
plt.plot(psnrs_y, psnrs_j, 'bo', alpha=0.25)
plt.plot([30, 50], [30, 50], 'k:')
plt.xlabel('PSNR for dJPEG')
plt.ylabel('PSNR for libJPEG')
plt.xlim([30, 60])
plt.ylim([30, 50])
if rounding_approximation is None:
plt.title('dJPEG vs libJPEG quality (with standard rounding)'.format(
rounding_approximation))
else:
plt.title('dJPEG vs libJPEG quality (with {} rounding approx.)'.format(
rounding_approximation))
for i, q in enumerate(quality_levels):
if q % 10 == 0:
plt.plot(psnrs_y[i], psnrs_j[i], 'ko')
plt.text(psnrs_y[i] + 1, psnrs_j[i] - 0.25, 'Q{:02d}'.format(q))
plt.show()
def construct_models(nip_model, patch_size=128, distribution_jpeg=50, distribution_down='pool', loss_metric='L2', jpeg_approx='sin'):
"""
Setup the TF model of the entire acquisition and distribution workflow.
:param nip_model: name of the NIP class
:param patch_size: patch size for manipulation training (raw patch - rgb patches will be 4 times as big)
:param distribution_jpeg: JPEG quality level in the distribution channel
:param distribution_down: Sub-sampling method in the distribution channel ('pool' or 'bilin')
:param loss_metric: NIP loss metric: L2, L1 or SSIM
"""
# Sanitize inputs
if patch_size < 32 or patch_size > 512:
raise ValueError('The patch size ({}) looks incorrect, typical values should be >= 32 and <= 512'.format(patch_size))
if distribution_jpeg < 1 or distribution_jpeg > 100:
raise ValueError('Invalid JPEG quality level ({})'.format(distribution_jpeg))
if not issubclass(getattr(pipelines, nip_model), pipelines.NIPModel):
supported_nips = [x for x in dir(pipelines) if x != 'NIPModel' and type(getattr(pipelines, x)) is type and issubclass(getattr(pipelines, x), pipelines.NIPModel)]
raise ValueError('Invalid NIP model ({})! Available NIPs: ({})'.format(nip_model, supported_nips))
if loss_metric not in ['L2', 'L1', 'SSIM']:
raise ValueError('Invalid loss metric ({})!'.format(loss_metric))
tf.reset_default_graph()
sess = tf.Session()
# The pipeline -----------------------------------------------------------------------------------------------------
model = getattr(pipelines, nip_model)(sess, tf.get_default_graph(), patch_size=patch_size, loss_metric=loss_metric)
print('NIP network: {}'.format(model.summary()))
# Several paths for post-processing --------------------------------------------------------------------------------
with tf.name_scope('distribution'):
# Sharpen
im_shr = tf_helpers.tf_sharpen(model.y, 0, hsv=True)
# Bilinear resampling
im_res = tf.image.resize_images(model.y, [tf.shape(model.y)[1] // 2, tf.shape(model.y)[1] // 2])
im_res = tf.image.resize_images(im_res, [tf.shape(model.y)[1], tf.shape(model.y)[1]])
# Gaussian filter
im_gauss = tf_helpers.tf_gaussian(model.y, 5, 4)
# Mild JPEG
tf_jpg = DJPG(sess, tf.get_default_graph(), model.y, None, quality=80, rounding_approximation=jpeg_approx)
im_jpg = tf_jpg.y
# Setup operations for detection
operations = (model.y, im_shr, im_gauss, im_jpg, im_res)
forensics_classes = ['native', 'sharpen', 'gaussian', 'jpg', 'resample']
n_classes = len(operations)
# Concatenate outputs from multiple post-processing paths ------------------------------------------------------
y_concat = tf.concat(operations, axis=0)
# Add sub-sampling and JPEG compression in the channel ---------------------------------------------------------
if distribution_down == 'pool':
imb_down = tf.nn.avg_pool(y_concat, [1, 2, 2, 1], [1, 2, 2, 1], 'SAME', name='post_downsample')
elif distribution_down == 'bilin':
imb_down = tf.image.resize_images(y_concat, [tf.shape(y_concat)[1] // 2, tf.shape(y_concat)[1] // 2])
else:
raise ValueError('Unsupported channel down-sampling {}'.format(distribution_down))
jpg = DJPG(sess, tf.get_default_graph(), imb_down, model.x, quality=distribution_jpeg, rounding_approximation=jpeg_approx)
imb_out = jpg.y
# Add manipulation detection
fan = FAN(sess, tf.get_default_graph(), n_classes=n_classes, x=imb_out, nip_input=model.x, n_convolutions=4)
print('Forensics network parameters: {:,}'.format(fan.count_parameters()))
# Setup a combined loss and training op
with tf.name_scope('combined_optimization') as scope:
nip_fw = tf.placeholder(tf.float32, name='nip_weight')
lr = tf.placeholder(tf.float32, name='learning_rate')
loss = fan.loss + nip_fw * model.loss
adam = tf.train.AdamOptimizer(learning_rate=lr, name='adam')
opt = adam.minimize(loss, name='opt_combined')
# Initialize all variables
sess.run(tf.global_variables_initializer())
tf_ops = {
'sess': sess,
'nip': model,
'fan': fan,
'loss': loss,
'opt': opt,
'lr': lr,
'lambda': nip_fw,
'operations': operations,
}
distribution = {
'forensics_classes': forensics_classes,
'channel_jpeg_quality': distribution_jpeg,
'channel_downsampling': distribution_down,
'jpeg_approximation': jpeg_approx
}
return tf_ops, distribution
def construct_models(nip_model,
patch_size=128,
trainable=None,
distribution=None,
manipulations=None,
loss_metric='L2'):
"""
Setup the TF model of the entire acquisition and distribution workflow.
:param nip_model: name of the NIP class
:param patch_size: patch size for manipulation training (raw patch - rgb patches will be 4 times as big)
:param distribution: definition of the dissemination channel (set to None for the default down+jpeg(50))
:param loss_metric: NIP loss metric: L2, L1 or SSIM
"""
# Sanitize inputs
if patch_size < 16 or patch_size > 512:
raise ValueError(
'The patch size ({}) looks incorrect, typical values should be >= 16 and <= 512'
.format(patch_size))
trainable = trainable or {}
# Setup a default distribution channel
if distribution is None:
distribution = {
'downsampling': 'pool',
'compression': 'jpeg',
'compression_params': {
'quality': 50,
'rounding_approximation': 'sin'
}
}
if 'dcn' in trainable and distribution['compression'] != 'dcn':
raise ValueError(
'Cannot make DCN trainable given current compression model: {}'.
format(distribution['compression']))
if distribution['compression'] == 'jpeg' and (
distribution['compression_params']['quality'] < 1
or distribution['compression_params']['quality'] > 100):
raise ValueError('Invalid JPEG quality level ({})'.format(
distribution['compression_params']['quality']))
if not issubclass(getattr(pipelines, nip_model), pipelines.NIPModel):
supported_nips = [
x for x in dir(pipelines)
if x != 'NIPModel' and type(getattr(pipelines, x)) is type
and issubclass(getattr(pipelines, x), pipelines.NIPModel)
]
raise ValueError('Invalid NIP model ({})! Available NIPs: ({})'.format(
nip_model, supported_nips))
if loss_metric not in ['L2', 'L1', 'SSIM']:
raise ValueError('Invalid loss metric ({})!'.format(loss_metric))
tf.reset_default_graph()
sess = tf.Session()
# The pipeline -----------------------------------------------------------------------------------------------------
model = getattr(pipelines, nip_model)(sess,
tf.get_default_graph(),
patch_size=patch_size,
loss_metric=loss_metric)
print('NIP network: {}'.format(model.summary()))
# Several paths for post-processing --------------------------------------------------------------------------------
with tf.name_scope('distribution'):
# Parse manipulation specs
manipulations = manipulations or [
'sharpen', 'resample', 'gaussian', 'jpeg'
]
strengths = {
'sharpen': 1,
'resample': 50,
'gaussian': 0.83,
'jpeg': 80,
'awgn': 5.1,
'gamma': 3,
'median': 3
}
manipulations_set = set()
for m in manipulations:
spec = m.split(':')
manipulations_set.add(spec[0])
if len(spec) > 1:
strengths[spec[0]] = float(spec[-1])
if any(x not in SUPPORTED_MANIPULATIONS for x in manipulations_set):
raise ValueError(
'Unsupported manipulation requested! Available: {}'.format(
SUPPORTED_MANIPULATIONS))
operations = [model.y]
forensics_classes = ['native']
# Sharpen
if 'sharpen' in manipulations_set:
im_shr = tf_helpers.manipulation_sharpen(model.y,
strengths['sharpen'],
hsv=True)
operations.append(im_shr)
forensics_classes.append('sharpen:{}'.format(strengths['sharpen']))
# Bilinear resampling
if 'resample' in manipulations_set:
im_res = tf_helpers.manipulation_resample(model.y,
strengths['resample'])
operations.append(im_res)
forensics_classes.append('resample:{}'.format(
strengths['resample']))
# Gaussian filter
if 'gaussian' in manipulations_set:
im_gauss = tf_helpers.manipulation_gaussian(
model.y, 5, strengths['gaussian'])
operations.append(im_gauss)
forensics_classes.append('gaussian:{}'.format(
strengths['gaussian']))
# Mild JPEG
if 'jpeg' in manipulations_set:
tf_jpg = DJPG(sess,
tf.get_default_graph(),
model.y,
None,
quality=strengths['jpeg'],
rounding_approximation='soft')
operations.append(tf_jpg.y)
forensics_classes.append('jpeg:{}'.format(strengths['jpeg']))
# AWGN
if 'awgn' in manipulations_set:
im_awgn = tf_helpers.manipulation_awgn(model.y,
strengths['awgn'] / 255)
operations.append(im_awgn)
forensics_classes.append('awgn:{}'.format(strengths['awgn']))
# Gamma + inverse
if 'gamma' in manipulations_set:
im_gamma = tf_helpers.manipulation_gamma(model.y,
strengths['gamma'])
operations.append(im_gamma)
forensics_classes.append('gamma:{}'.format(strengths['gamma']))
# Median
if 'median' in manipulations_set:
im_median = tf_helpers.manipulation_median(model.y,
strengths['median'])
operations.append(im_median)
forensics_classes.append('median:{}'.format(strengths['median']))
n_classes = len(operations)
assert len(forensics_classes) == n_classes
# Concatenate outputs from multiple post-processing paths ------------------------------------------------------
y_concat = tf.concat(operations, axis=0)
# Add sub-sampling and lossy compression in the channel --------------------------------------------------------
down_patch_size = 2 * patch_size if distribution[
'downsampling'] == 'none' else patch_size
if distribution['downsampling'] == 'pool':
imb_down = tf.nn.avg_pool(y_concat, [1, 2, 2, 1], [1, 2, 2, 1],
'SAME',
name='post_downsample')
elif distribution['downsampling'] == 'bilin':
imb_down = tf.image.resize_images(
y_concat,
[tf.shape(y_concat)[1] // 2,
tf.shape(y_concat)[1] // 2])
elif distribution['downsampling'] == 'none':
imb_down = y_concat
else:
raise ValueError('Unsupported channel down-sampling {}'.format(
distribution['downsampling']))
if distribution['compression'] == 'jpeg':
print(
'Channel compression: JPEG({quality}, {rounding_approximation})'
.format(**distribution['compression_params']))
dist_compression = DJPG(sess, tf.get_default_graph(), imb_down,
model.x,
**distribution['compression_params'])
imb_out = dist_compression.y
elif distribution['compression'] == 'dcn':
print('Channel compression: DCN from {dirname}'.format(
**distribution['compression_params']))
if 'dirname' in distribution['compression_params']:
model_directory = distribution['compression_params']['dirname']
dist_compression = codec.restore_model(
model_directory,
down_patch_size,
sess=sess,
graph=tf.get_default_graph(),
x=imb_down,
nip_input=model.x)
else:
# TODO Not tested yet
raise NotImplementedError(
'DCN models should be restored from a pre-training session!'
)
# dist_compression = compression.TwitterDCN(sess, tf.get_default_graph(), x=imb_down, nip_input=model.x, patch_size=down_patch_size, **distribution['compression_params'])
imb_out = dist_compression.y
elif distribution['compression'] == 'none':
dist_compression = None
imb_out = imb_down
else:
raise ValueError('Unsupported channel compression {}'.format(
distribution['compression']))
# Add manipulation detection
fan = FAN(sess,
tf.get_default_graph(),
n_classes=n_classes,
x=imb_out,
nip_input=model.x,
n_convolutions=4)
print('Forensics network parameters: {:,}'.format(fan.count_parameters()))
# Setup a combined loss and training op
with tf.name_scope('combined_optimization') as scope:
lambda_nip = tf.placeholder(tf.float32, name='lambda_nip')
lambda_dcn = tf.placeholder(tf.float32, name='lambda_dcn')
lr = tf.placeholder(tf.float32, name='learning_rate')
loss = fan.loss
if 'nip' in trainable:
loss += lambda_nip * model.loss
if 'dcn' in trainable:
loss += lambda_dcn * dist_compression.loss
adam = tf.train.AdamOptimizer(learning_rate=lr, name='adam')
# List parameters that need to be optimized
parameters = []
parameters.extend(fan.parameters)
if 'nip' in trainable:
parameters.extend(model.parameters)
if 'dcn' in trainable:
parameters.extend(dist_compression.parameters)
opt = adam.minimize(loss, name='opt_combined', var_list=parameters)
tf_ops = {
'sess': sess,
'nip': model,
'fan': fan,
'loss': loss,
'opt': opt,
'lr': lr,
'lambda_nip': lambda_nip,
'lambda_dcn': lambda_dcn,
'operations': operations,
'dcn': dist_compression
}
dist = {'forensics_classes': forensics_classes}
dist.update(distribution)
return tf_ops, dist