def transform(img, theta, outsize=None, oob=None):
"""
sampling from source coord by theta (xy_src = theta * xy_target !!!)
:param img: [BHWC] or [HWC]
:param theta: [2x3] or [B, 2x3] if img [BHWC] inverse theta!!!
:param outsize: [H,W] or None if none same HW to img
:param oob: outofbound color [C,]
:return:
"""
# tf.assert_type(img, tf_type=tf.float32)
if theta.ndim == 2:
# [2, 3]
# assert theta.shape == (2,3)
if img.ndim == 4:
# outdim : [BHWC]
return transform_4r(img, theta, outsize=outsize, oob=oob)
elif img.ndim == 3:
# outdim : [HWC]
return transform_3r(img, theta, outsize=outsize, oob=oob)
else:
raise ValueError('2d support in transform?')
elif theta.ndim == 3:
# theta [B,2,3]
# assert theta.shapes[1:3] == (2,3)
if img.ndim == 4:
if img.dims[0] is not None and theta.dims[0] is not None:
assert img.dims[0] == theta.dims[0]
else:
tf.assert_equal(img.shape[0], theta.shape[0])
# assert img.shapes[0] == theta.shapes[0], or tf.Assert()...
# outdim : [BHWC]
return tf.map_fn(
lambda x: transform_3r(x[0], x[1], outsize=outsize, oob=oob),
[img, theta],
dtype=tf.float32)
elif img.ndim == 3:
# one image to multi transform???
# [BHWC] !
return tf.map_fn(
lambda t: transform_3r(img, t, outsize=outsize, oob=oob),
[img, theta])
elif img.ndim == 2:
raise ValueError('2d support in transform? [HW]?')
def _rand_apply_fun(fun, imagez, p=0.5):
assert isinstance(imagez, (tuple, list))
shape = (imagez[0].dims[0],)
irand = tf.random_choice(shape, p=p)
applied = [tf.map_fn(fun, images) for images in imagez]
outs = tuple(tf.where(irand, images, ap) for images, ap in zip(imagez, applied))
return _standize_output_res(outs)
def getpixel(img, ind, name=None):
"""
no outof bound checking
:param img: [NHWC] or [HWC]
:param ind: [p, ij(2)] (height index, width index)
:param name:
:return: [N, p, C] or [p, C]
"""
if img.ndim == 3:
return _getpixel_3r(img, ind, name=name or 'getpixel')
elif img.ndim == 4:
return tf.map_fn(lambda x: _getpixel_3r(x, ind),
img,
name=name or 'getpixel')
else:
raise ValueError("need 2d support in getpixel?")
def _rand_crop_offsets(*imgs):
with tf.name_scope(None, 'rand_crop', list(imgs) + [sz]):
value = imgs[0]
size = tf.convert_to_tensor(sz, dtype=tf.int32, name="size")
shape = tf.shape(value)[1:3] # HW of BHWC
check = tf.Assert(tf.reduce_all(shape >= size), ["Need value.shape >= size, got", shape, size])
shape = control_flow_ops.with_dependencies([check], shape)
# assert same shape
for v in imgs:
vshape = tf.shape(v)[1:3] # assert v.ndim == 4
check = tf.Assert(tf.reduce_all(shape.equal(vshape)),
["Need same (H,W,?) image.shape[1:3] == otherimage.shape[1:3], got", shape, vshape])
shape = control_flow_ops.with_dependencies([check], shape)
limit = shape - size + 1
if value.dims[0] is None:
batchshape = tf.shape(value)[:1].append(2)
else:
batchshape = (value.dims[0], 2)
offsets = tf.random_uniform(batchshape, dtype=size.dtype, maxval=size.dtype.max) % limit # add seed
# offsets = tf.random_uniform(batchshape, maxval=limit, dtype=tf.int32)
# sz = size
size = size.append(-1)
def _3d_crop(args):
values, offset = args
offset = offset.append(0)
# outs = [tf.slice(img, offset, size) for img in values]
outs = []
for img in values:
out = tf.slice(img, offset, size)
out.set_shape(list(sz)+v.dims[-1:])
outs.append(out)
return outs
return tf.map_fn(_3d_crop, [imgs, offsets], dtype=[v.dtype for v in imgs]), offsets
def _flipud(img):
return tf.map_fn(tf.image.flip_up_down, img)
def _rot90(*imgs):
for im in imgs:
tf.assert_rank(im, 4)
return tuple(tf.map_fn(lambda x: _rot90_3r(x, k), im) for im in imgs)
def _transpose_img(img4d):
return tf.map_fn(tf.image.transpose_image, img4d)
def _fliplr(im):
return tf.map_fn(tf.image.flip_left_right, im)