def init(args):
'''Returns initialized Lsb (writer) object from arguments passed from
command line.
'''
wm = create_watermark(get_correct_wm(args, __name__.split('.')[-1]))
return Lsb(args.bands, args.dest_dir, args.format, wm, args.suffix,
args.position)
def _create_watermarked(self, src_img):
src_img.load()
bands = src_img.split()
names = src_img.getbands()
bands_wm = []
src_width, src_height = src_img.size
wm_width, wm_height = self.wm.img.size
if src_width == wm_width and src_height == wm_height:
wm = self.wm
else:
wm = create_watermark(self.wm.img,
width=src_width,
height=src_height,
position=self.position)
for name, band in zip(names, bands):
if self._band_is_used(name):
band_wm = Image.new('L', src_img.size)
band_wm.putdata([
convert(orig_px, wm_px,
wm.threshold) for orig_px, wm_px in zip(
band.getdata(), wm.band.getdata())
])
else:
band_wm = band
bands_wm.append(band_wm)
dst_img = Image.merge(src_img.mode, bands_wm)
return dst_img
def run_writer_lsb():
filepath = os.path.join(os.path.dirname(__file__), '..', 'test', 'data', 'shape1-rgb-l9.png')
wm_filepath = os.path.join(os.path.dirname(__file__), '..', 'test', 'data', 'shape1-rgb-l9.png')
wm = create_watermark(wm_filepath)
dest_dir = os.path.join(os.path.dirname(__file__), '..', 'tmp')
writer = watermarks.core.writers.lsb.Lsb(dest_dir, 'png', wm)
writer.run([filepath])
def run_writer_and_assert(dst_dir, writer_class, filename, wm_filename,
wm_data=None, ext=None, width=None, height=None,
position='', bands_are_filtered=False, **kwargs):
base, f_ext = os.path.splitext(filename)
ext = ext or f_ext
filepath = os.path.join(DATA_DIR, filename)
wm_filepath = os.path.join(DATA_DIR, wm_filename)
suffix = '_watermarked_test'
if not width or not height:
img = Image.open(filepath)
width, height = img.size
wm = create_watermark(wm_filepath, width=width, height=height, position=position)
writer = writer_class(destination=dst_dir, format_=ext.lstrip('.'),
wm=wm, suffix=suffix, position=position, **kwargs)
results = list(writer.run([filepath]))
if wm_data is None:
assert_equal(len(results), 0)
return
assert_equal(len(results), 1)
res_filepath = results[0]
res_img = Image.open(res_filepath)
res_img.load()
for i, band in enumerate(res_img.split()):
expected_data = wm_data[i] if bands_are_filtered else wm_data
assert_equal(list(band.getdata()), expected_data)
def test_dir(dst_dir):
filenames = ['gen-{0}-rgb.png'.format(IM_PREFIX), 'gen-{0}-g.png'.format(IM_PREFIX)]
data_dir_path = create_data_dir(os.path.join(dst_dir, 'writer_dir'), filenames)
filepath = os.path.join(DATA_DIR, 'shape1-g-l0.png')
suffix = '_watermarked_test'
wm = create_watermark(os.path.join(DATA_DIR, 'shape1-g-l0.png'))
lsb = Lsb(destination=dst_dir, format_='png', wm=wm, suffix=suffix)
generated_filepaths = lsb.run([data_dir_path, filepath])
generated_filenames = set([os.path.basename(f) for f in generated_filepaths])
assert_true('gen-{0}-rgb{1}.png'.format(IM_PREFIX, suffix) in generated_filenames)
assert_true('gen-{0}-g{1}.png'.format(IM_PREFIX, suffix) in generated_filenames)
assert_true('shape1-g-l0{0}.png'.format(suffix) in generated_filenames)
assert_equal(len(generated_filenames), 3)
def test_dir(dst_dir):
filenames = [
'gen-{0}-rgb.png'.format(IM_PREFIX), 'gen-{0}-g.png'.format(IM_PREFIX)
]
data_dir_path = create_data_dir(os.path.join(dst_dir, 'writer_dir'),
filenames)
filepath = os.path.join(DATA_DIR, 'shape1-g-l0.png')
suffix = '_watermarked_test'
wm = create_watermark(os.path.join(DATA_DIR, 'shape1-g-l0.png'))
lsb = Lsb(destination=dst_dir, format_='png', wm=wm, suffix=suffix)
generated_filepaths = lsb.run([data_dir_path, filepath])
generated_filenames = set(
[os.path.basename(f) for f in generated_filepaths])
assert_true(
'gen-{0}-rgb{1}.png'.format(IM_PREFIX, suffix) in generated_filenames)
assert_true(
'gen-{0}-g{1}.png'.format(IM_PREFIX, suffix) in generated_filenames)
assert_true('shape1-g-l0{0}.png'.format(suffix) in generated_filenames)
assert_equal(len(generated_filenames), 3)
def _create_watermarked(self, src_img):
src_img.load()
bands = src_img.split()
names = src_img.getbands()
bands_wm = []
src_width, src_height = src_img.size
wm_width, wm_height = self.wm.img.size
if src_width == wm_width and src_height == wm_height:
wm = self.wm
else:
wm = create_watermark(self.wm.img, width=src_width,
height=src_height, position=self.position)
for name, band in zip(names, bands):
if self._band_is_used(name):
band_wm = Image.new('L', src_img.size)
band_wm.putdata([convert(orig_px, wm_px, wm.threshold)
for orig_px, wm_px
in zip(band.getdata(), wm.band.getdata())
])
else:
band_wm = band
bands_wm.append(band_wm)
dst_img = Image.merge(src_img.mode, bands_wm)
return dst_img
def test_unsupported_format():
create_watermark(os.path.join(DATA_DIR, 'unsupported_format.tiff'))
def init(args):
'''Returns initialized Lsb (writer) object from arguments passed from
command line.
'''
wm = create_watermark(get_correct_wm(args, __name__.split('.')[-1]))
return Lsb(args.bands, args.dest_dir, args.format, wm, args.suffix, args.position)
