def _render(self,meta=None):
if meta=="OR":
self.val=None
else: # Assuming if torch.Tensor and is not Union type it is an image
if meta:
st.write(meta)
input_=F.pil_to_tensor(self.image).float()/255
input_=torch.stack(1*[input_])
self.val=input_
def image64_to_tensor(data, size):
resize = Resize([size, size])
res = resize(PILImage.create(BytesIO(data)))
tn = pil_to_tensor(res).to(torch.float)
tn = normalize_tensor(tn)
# print("post norm: mean: ", mean, " std: ", std)
tn.unsqueeze_(0)
return tn
def read_image1():
image_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
"assets", "encode_jpeg",
"grace_hopper_517x606.jpg")
image = Image.open(image_path)
image = image.resize((224, 224))
x = F.pil_to_tensor(image)
x = F.convert_image_dtype(x)
return x.view(1, 3, 224, 224)
def get_image(self, rel_path: str, size: Tuple[int, int]) -> torch.Tensor:
import os
from PIL import Image
from torchvision.transforms import functional as F
data_dir = os.path.join(os.path.dirname(__file__), "assets")
path = os.path.join(data_dir, *rel_path.split("/"))
image = Image.open(path).convert("RGB").resize(size, Image.BILINEAR)
return F.convert_image_dtype(F.pil_to_tensor(image))
def flip_only_bboxes(img, bboxs, p):
img = F.pil_to_tensor(img)
flip_img = torch.zeros_like(img)
for bbox in bboxs:
if torch.rand(1) < p:
min_x, min_y, max_x, max_y = bbox
min_x, min_y, max_x, max_y = int(min_x.item()), int(min_y.item()), int(max_x.item()), int(max_y.item())
flip_img[:, min_y:max_y+1, min_x:max_x+1] = F.hflip(img[:, min_y:max_y+1, min_x:max_x+1])
return F.to_pil_image(torch.where(flip_img != 0, flip_img, img))
def random_occlusion(self, img, img_nogt, new_gt):
"""
erase(img: torch.Tensor,
i: int,
j: int,
h: int,
w: int,
v: torch.Tensor,
inplace: bool = False)
"""
occ_i, occ_j, MaxRow, MaxCol = self.affine_img_mask()
tensor_img = TF.pil_to_tensor(img)
tensor_img_nogt = TF.pil_to_tensor(img_nogt)
img_crop_resize = self.crop_resize(tensor_img_nogt, occ_i, occ_j,
MaxRow, MaxCol, new_gt)
img_erase = TF.erase(tensor_img, occ_i, occ_j, MaxRow - occ_i,
MaxCol - occ_j, 0)
self.all_param.append((occ_i, occ_j, MaxRow, MaxCol))
return img_erase, img_crop_resize
def cutout_only_bboxes(img, bboxs, p, pad_size, replace):
img = F.pil_to_tensor(img)
cutout_img = img.clone()
for bbox in bboxs:
if torch.rand(1) < p:
min_x, min_y, max_x, max_y = bbox
min_x, min_y, max_x, max_y = int(min_x.item()), int(min_y.item()), int(max_x.item()), int(max_y.item())
cutout_x, cutout_y = torch.randint(low=min_x, high=max_x, size=(1,)), torch.randint(low=min_y, high=max_y, size=(1,))
y_min, y_max = int(torch.clamp(cutout_y-pad_size, min_y, max_y).item()), int(torch.clamp(cutout_y+pad_size, min_y, max_y).item())
x_min, x_max = int(torch.clamp(cutout_x-pad_size, min_x, max_x).item()), int(torch.clamp(cutout_x+pad_size, min_x, max_x).item())
cutout_img[:, y_min:y_max, x_min:x_max] = replace
return F.to_pil_image(cutout_img)
def forward(
self,
image: Tensor,
target: Optional[Dict[str, Tensor]] = None
) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
if isinstance(image, torch.Tensor):
if image.ndimension() not in {2, 3}:
raise ValueError(
f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions."
)
elif image.ndimension() == 2:
image = image.unsqueeze(0)
r = torch.rand(7)
if r[0] < self.p:
image = self._brightness(image)
contrast_before = r[1] < 0.5
if contrast_before:
if r[2] < self.p:
image = self._contrast(image)
if r[3] < self.p:
image = self._saturation(image)
if r[4] < self.p:
image = self._hue(image)
if not contrast_before:
if r[5] < self.p:
image = self._contrast(image)
if r[6] < self.p:
channels, _, _ = F.get_dimensions(image)
permutation = torch.randperm(channels)
is_pil = F._is_pil_image(image)
if is_pil:
image = F.pil_to_tensor(image)
image = F.convert_image_dtype(image)
image = image[..., permutation, :, :]
if is_pil:
image = F.to_pil_image(image)
return image, target
def translate_bbox(img, bboxs, pixels, replace, shift_horizontal):
img = F.pil_to_tensor(img)
_, h, w = img.shape
translate_bboxs = []
if shift_horizontal:
for bbox in bboxs:
min_x, min_y, max_x, max_y = bbox
translate_min_x, translate_max_x = torch.clamp(min_x+pixels, 0, w), torch.clamp(max_x+pixels, 0, w)
translate_min_x, translate_max_x = int(translate_min_x.item()), int(translate_max_x.item())
translate_bboxs.append(torch.FloatTensor([translate_min_x, min_y, translate_max_x, max_y]))
else:
for bbox in bboxs:
min_x, min_y, max_x, max_y = bbox
translate_min_y, translate_max_y = torch.clamp(min_y+pixels, 0, h), torch.clamp(max_y+pixels, 0, h)
translate_min_y, translate_max_y = int(translate_min_y.item()), int(translate_max_y.item())
translate_bboxs.append(torch.FloatTensor([min_x, translate_min_y, max_x, translate_max_y]))
return torch.stack(translate_bboxs)
def _transform(self, input: Any, params: Dict[str, Any]) -> Any:
if not (isinstance(input, (features.Image, PIL.Image.Image))
or is_simple_tensor(input)):
return input
image = input
if isinstance(input, PIL.Image.Image):
image = _F.pil_to_tensor(image)
output = image[..., params["permutation"], :, :]
if isinstance(input, features.Image):
output = features.Image.new_like(
input, output, color_space=features.ColorSpace.OTHER)
elif isinstance(input, PIL.Image.Image):
output = _F.to_pil_image(output)
return output
def random_morph(trainData, op_name=None):
# if op_name not in known_patterns:
# raise Exception("Unknown pattern " + op_name + "!")
dil = ImageMorph.MorphOp(op_name="dilation8")
ero = ImageMorph.MorphOp(op_name="erosion4")
ops = [dil, ero]
for i in range(len(trainData)):
r = torch.randint(2, (1, 1)).item()
if r == 1:
continue
data = F.rgb_to_grayscale(trainData[i].clone(), 1)
_, data = ero.apply(F.to_pil_image(data, mode="L"))
data = ImageOps.colorize(data, "black", "white")
trainData[i] = F.pil_to_tensor(data)
def apply_recursively(obj: Any) -> Any:
if isinstance(obj, MultiCropResult):
crops = obj
if isinstance(obj[0], PIL.Image.Image):
crops = [pil_to_tensor(crop) for crop in crops] # type: ignore[assignment]
batch = torch.stack(crops)
if isinstance(obj[0], features.Image):
batch = features.Image.new_like(obj[0], batch)
return batch
elif isinstance(obj, collections.abc.Sequence) and not isinstance(obj, str):
return [apply_recursively(item) for item in obj]
elif isinstance(obj, collections.abc.Mapping):
return {key: apply_recursively(item) for key, item in obj.items()}
else:
return obj
def __init__(self, model, img_path, input_resolution=(384,384), topk=3, tau=0.001, batch_size=16, mask_size=4):
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.model = model.eval().to(self.device)
self.topk = topk
self.tau = tau
self.batch_size = batch_size
self.mask_size = mask_size
# Getting our input image and storing initial outputs
self.image = Image.open(img_path)
image_tensor = TF.pil_to_tensor(self.image.resize(input_resolution)).float()
self.image_tensor = TF.normalize(image_tensor, mean=image_tensor.mean([1,2]), std=image_tensor.std([1,2]))
self.image_tensor.requires_grad = True
self.original_output = torch.softmax(self.get_model_preds(), dim=1)
self.masked_img = None
if self.device == 'cuda': # empty cache to free up memory
torch.cuda.empty_cache()
def __getitem__(self, index):
image_path = self.image_paths[index]
image = Image.open(image_path)
image = self.initial_transforms(image.convert("RGB")) ## to PIL.rgb
if self.do_augmentations:
image = self.augmentations(image)
image = TF.pil_to_tensor(
image) ## save the image tensor for visualization
data = self.normalize(self.to_tensor(TF.to_pil_image(image)))
label, path = image_path.split('/')[-2:]
image_id = path.split('.')[0] ## str
if self.return_gt_label:
assert label in self.classes
n = len(self.classes[label])
rand_idx = torch.randperm(n)[0]
target = self.classes[label][rand_idx]
return data, target, self.classes_str.index(label), image, image_id
else:
return data, image, image_id
def _transform_1(self, *, url: str = None, data: bytes = None):
assert url is not None or data is not None
fp = data if data is not None else url
if url is not None and url.startswith('data:,'):
raise RequestIgnored('Not an image')
try:
img = getSmallImage(fp, self.__image_size, self.__image_size)
except ImageError as ex:
raise RequestIgnored(ex) from ex
_input = pil_to_tensor(img.convert('RGB'))
assert 'uint8' in str(_input.dtype)
_input = _input / 255.
assert 'float32' in str(_input.dtype)
_input = normalize(_input, self.__mean, self.__std, inplace=True)
assert len(_input.shape) == 3, _input.shape
_input = _input.unsqueeze(0)
assert len(_input.shape) == 4, _input.shape
return _input
def main(model_path, input_dir, output_dir, gpu):
ptu.set_gpu_mode(gpu)
model_dir = Path(model_path).parent
model, variant = load_model(model_path)
model.to(ptu.device)
normalization_name = variant["dataset_kwargs"]["normalization"]
normalization = STATS[normalization_name]
cat_names, cat_colors = dataset_cat_description(ADE20K_CATS_PATH)
input_dir = Path(input_dir)
output_dir = Path(output_dir)
output_dir.mkdir(exist_ok=True)
list_dir = list(input_dir.iterdir())
for filename in tqdm(list_dir, ncols=80):
pil_im = Image.open(filename).copy()
im = F.pil_to_tensor(pil_im).float() / 255
im = F.normalize(im, normalization["mean"], normalization["std"])
im = im.to(ptu.device).unsqueeze(0)
im_meta = dict(flip=False)
logits = inference(
model,
[im],
[im_meta],
ori_shape=im.shape[2:4],
window_size=variant["inference_kwargs"]["window_size"],
window_stride=variant["inference_kwargs"]["window_stride"],
batch_size=2,
)
seg_map = logits.argmax(0, keepdim=True)
seg_rgb = seg_to_rgb(seg_map, cat_colors)
seg_rgb = (255 * seg_rgb.cpu().numpy()).astype(np.uint8)
pil_seg = Image.fromarray(seg_rgb[0])
pil_blend = Image.blend(pil_im, pil_seg, 0.5).convert("RGB")
pil_blend.save(output_dir / filename.name)
def _rotate_bbox(img, bboxs, degrees):
img = F.pil_to_tensor(img)
_, h, w = img.shape
rotate_bboxs = []
rotate_matrix = torch.FloatTensor([[math.cos(degrees*math.pi/180), math.sin(degrees*math.pi/180)],
[-math.sin(degrees*math.pi/180), math.cos(degrees*math.pi/180)]])
for bbox in bboxs:
min_x, min_y, max_x, max_y = bbox
rel_min_x, rel_max_x, rel_min_y, rel_max_y = min_x-w/2, max_x-w/2, min_y-h/2, max_y-h/2
coords = torch.FloatTensor([[rel_min_x, rel_min_y],
[rel_min_x, rel_max_y],
[rel_max_x, rel_max_y],
[rel_max_x, rel_min_y]])
rotate_coords = torch.matmul(rotate_matrix, coords.t()).t()
x_min, y_min = torch.min(rotate_coords, dim=0)[0]
x_max, y_max = torch.max(rotate_coords, dim=0)[0]
rotate_min_x, rotate_max_x = torch.clamp(x_min+w/2, 0, w),torch.clamp(x_max+w/2, 0, w)
rotate_min_y, rotate_max_y = torch.clamp(y_min+h/2, 0, h),torch.clamp(y_max+h/2, 0, h)
rotate_bboxs.append(torch.FloatTensor([rotate_min_x, rotate_min_y, rotate_max_x, rotate_max_y]))
return torch.stack(rotate_bboxs)
def bbox_cutout(img, bboxs, pad_fraction, replace_with_mean):
img = F.pil_to_tensor(img)
_, h, w = img.shape
random_index = torch.randint(bboxs.size(0), size=(1,)).item()
chosen_bbox = bboxs[random_index]
min_x, min_y, max_x, max_y = chosen_bbox
min_x, min_y, max_x, max_y = int(min_x.item()), int(min_y.item()), int(max_x.item()), int(max_y.item())
if (min_x == max_x) or (min_y == max_y): return F.to_pil_image(img)
mask_x, mask_y = torch.randint(low=min_x, high=max_x, size=(1,)), torch.randint(low=min_y, high=max_y, size=(1,))
mask_w, mask_h = pad_fraction * w / 2, pad_fraction * h / 2
x_min, x_max = int(torch.clamp(mask_x-mask_w, 0, w).item()), int(torch.clamp(mask_x+mask_w, 0, w).item())
y_min, y_max = int(torch.clamp(mask_y-mask_h, 0, h).item()), int(torch.clamp(mask_y+mask_h, 0, h).item())
if replace_with_mean == True: replace = torch.mean(img[:, min_y:max_y, min_x:max_x]).item()
else: replace = 128
cutout_img = img.clone()
cutout_img[:, y_min:y_max, x_min:x_max] = replace
return F.to_pil_image(cutout_img)
def shear_with_bboxes(img, bboxs, level, replace, shift_horizontal):
img = F.pil_to_tensor(img)
_, h, w = img.shape
shear_bboxs = []
if shift_horizontal:
shear_matrix = torch.FloatTensor([[1, -level],
[0, 1]])
for bbox in bboxs:
min_x, min_y, max_x, max_y = bbox
coords = torch.FloatTensor([[min_x, min_y],
[min_x, max_y],
[max_x, max_y],
[max_x, min_y]])
shear_coords = torch.matmul(shear_matrix, coords.t()).t()
x_min, y_min = torch.min(shear_coords, dim=0)[0]
x_max, y_max = torch.max(shear_coords, dim=0)[0]
shear_min_x, shear_max_x = torch.clamp(x_min, 0, w), torch.clamp(x_max, 0, w)
shear_min_y, shear_max_y = torch.clamp(y_min, 0, h), torch.clamp(y_max, 0, h)
shear_bboxs.append(torch.FloatTensor([shear_min_x, shear_min_y, shear_max_x, shear_max_y]))
else:
shear_matrix = torch.FloatTensor([[1, 0],
[-level, 1]])
for bbox in bboxs:
min_x, min_y, max_x, max_y = bbox
coords = torch.FloatTensor([[min_x, min_y],
[min_x, max_y],
[max_x, max_y],
[max_x, min_y]])
shear_coords = torch.matmul(shear_matrix, coords.t()).t()
x_min, y_min = torch.min(shear_coords, dim=0)[0]
x_max, y_max = torch.max(shear_coords, dim=0)[0]
shear_min_x, shear_max_x = torch.clamp(x_min, 0, w), torch.clamp(x_max, 0, w)
shear_min_y, shear_max_y = torch.clamp(y_min, 0, h), torch.clamp(y_max, 0, h)
shear_bboxs.append(torch.FloatTensor([shear_min_x, shear_min_y, shear_max_x, shear_max_y]))
return torch.stack(shear_bboxs)
def __call__(self, image, mask):
return image, F.pil_to_tensor(mask)
def solarize_add(img, addition, threshold):
img = F.pil_to_tensor(img)
added_img = img + addition
added_img = torch.clamp(added_img, 0, 255)
return F.to_pil_image(torch.where(img < threshold, added_img, img))
def forward(self, *inputs: Any) -> Any:
sample = inputs if len(inputs) > 1 else inputs[0]
id, orig_image = self._extract_image(sample)
num_channels, height, width = get_image_dimensions(orig_image)
fill = self._parse_fill(orig_image, num_channels)
if isinstance(orig_image, torch.Tensor):
image = orig_image
else: # isinstance(input, PIL.Image.Image):
image = pil_to_tensor(orig_image)
augmentation_space = self._AUGMENTATION_SPACE if self.all_ops else self._PARTIAL_AUGMENTATION_SPACE
orig_dims = list(image.shape)
batch = image.view([1] * max(4 - image.ndim, 0) + orig_dims)
batch_dims = [batch.size(0)] + [1] * (batch.ndim - 1)
# Sample the beta weights for combining the original and augmented image. To get Beta, we use a Dirichlet
# with 2 parameters. The 1st column stores the weights of the original and the 2nd the ones of augmented image.
m = self._sample_dirichlet(
torch.tensor([self.alpha, self.alpha],
device=batch.device).expand(batch_dims[0], -1))
# Sample the mixing weights and combine them with the ones sampled from Beta for the augmented images.
combined_weights = self._sample_dirichlet(
torch.tensor(
[self.alpha] * self.mixture_width, device=batch.device).expand(
batch_dims[0], -1)) * m[:, 1].view([batch_dims[0], -1])
mix = m[:, 0].view(batch_dims) * batch
for i in range(self.mixture_width):
aug = batch
depth = self.chain_depth if self.chain_depth > 0 else int(
torch.randint(low=1, high=4, size=(1, )).item())
for _ in range(depth):
transform_id, (
magnitudes_fn,
signed) = self._get_random_item(augmentation_space)
magnitudes = magnitudes_fn(self._PARAMETER_MAX, height, width)
if magnitudes is not None:
magnitude = float(magnitudes[int(
torch.randint(self.severity, ()))])
if signed and torch.rand(()) <= 0.5:
magnitude *= -1
else:
magnitude = 0.0
aug = self._apply_image_transform(
aug,
transform_id,
magnitude,
interpolation=self.interpolation,
fill=fill)
mix.add_(combined_weights[:, i].view(batch_dims) * aug)
mix = mix.view(orig_dims).to(dtype=image.dtype)
if isinstance(orig_image, features.Image):
mix = features.Image.new_like(orig_image, mix)
elif isinstance(orig_image, PIL.Image.Image):
mix = to_pil_image(mix)
return _put_into_sample(sample, id, mix)
class Watermark(BasicObject):
r"""Watermark class that is used for backdoor attacks.
Note:
Images with alpha channel are supported.
In this case, :attr:`mark_alpha` will be multiplied.
Warning:
:attr:`mark_random_init` and :attr:`mark_scattered` can't be used together.
Args:
mark_path (str):
| Path to watermark image or npy file.
There are some preset marks in the package.
| Defaults to ``'square_white.png'``.
.. table::
:widths: auto
+-----------------------------+---------------------+
| mark_path | mark image |
+=============================+=====================+
| ``'apple_black.png'`` | |apple_black| |
+-----------------------------+---------------------+
| ``'apple_white.png'`` | |apple_white| |
+-----------------------------+---------------------+
| ``'square_black.png'`` | |square_black| |
+-----------------------------+---------------------+
| ``'square_white.png'`` | |square_white| |
+-----------------------------+---------------------+
| ``'watermark_black.png'`` | |watermark_black| |
+-----------------------------+---------------------+
| ``'watermark_white.png'`` | |watermark_white| |
+-----------------------------+---------------------+
data_shape (list[int]): The shape of image data ``[C, H, W]``.
See Also:
Usually passed by ``dataset.data_shape``.
See :attr:`data_shape` from
:class:`trojanvision.datasets.ImageSet`.
mark_background_color (str | torch.Tensor): Mark background color.
If :class:`str`, choose from ``['auto', 'black', 'white']``;
else, it shall be 1-dim tensor ranging in ``[0, 1]``.
It's ignored when alpha channel in watermark image.
Defaults to ``'auto'``.
mark_alpha (float): Mark opacity. Defaults to ``1.0``.
mark_height (int): Mark resize height. Defaults to ``3``.
mark_width (int): Mark resize width. Defaults to ``3``.
Note:
:attr:`self.mark_height` and :attr:`self.mark_width` will be different
from the passed argument values
when :attr:`mark_scattered` is ``True``.
mark_height_offset (int): Mark height offset. Defaults to ``0``.
mark_width_offset (int): Mark width offset. Defaults to ``0``.
Note:
:attr:`mark_height_offset` and
:attr:`mark_width_offset` will be ignored
when :attr:`mark_random_pos` is ``True``.
mark_random_init (bool): Whether to randomly set pixel values of watermark,
which means only using the mark shape from the watermark image.
Defaults to ``False``.
mark_random_pos (bool): Whether to add mark at random location when calling :meth:`add_mark()`.
If ``True``, :attr:`mark_height_offset` and :attr:`mark_height_offset` will be ignored.
Defaults to ``False``.
mark_scattered (bool): Random scatter mark pixels
in the entire image to get the watermark. Defaults to ``False``.
mark_scattered_height (int | None): Scattered mark height. Defaults to data_shape[1].
mark_scattered_width (int | None): Scattered mark width. Defaults to data_shape[2].
Note:
- The random scatter process only occurs once at watermark initialization.
:meth:`add_mark()` will still add the same scattered mark to images.
- Mark image will first resize to ``(mark_height, mark_width)`` and then
scattered to ``(mark_scattered_height, mark_scattered_width)``.
If they are the same, it's actually pixel shuffling.
- :attr:`self.mark_height` and :attr:`self.mark_width` will be set to scattered version.
add_mark_fn (~collections.abc.Callable | None):
Customized function to add mark to images for :meth:`add_mark()` to call.
``add_mark_fn(_input, mark_random_pos=mark_random_pos, mark_alpha=mark_alpha, **kwargs)``
Defaults to ``None``.
Attributes:
mark (torch.Tensor): Mark float tensor with shape
``(data_shape[0] + 1, mark_height, mark_width)``
(last dimension is alpha channel).
mark_alpha (float): Mark opacity. Defaults to ``1.0``.
mark_height (int): Mark resize height. Defaults to ``3``.
mark_width (int): Mark resize width. Defaults to ``3``.
Note:
:attr:`self.mark_height` and :attr:`self.mark_width` will be different
from the passed argument values
when :attr:`mark_scattered` is ``True``.
mark_height_offset (int): Mark height offset. Defaults to ``0``.
mark_width_offset (int): Mark width offset. Defaults to ``0``.
Note:
:attr:`mark_height_offset` and
:attr:`mark_width_offset` will be ignored
when :attr:`mark_random_pos` is ``True``.
mark_random_init (bool): Whether to randomly set pixel values of watermark,
which means only using the mark shape from the watermark image.
Defaults to ``False``.
mark_random_pos (bool): Whether to add mark at random location when calling :meth:`add_mark()`.
If ``True``, :attr:`mark_height_offset` and :attr:`mark_height_offset` will be ignored.
Defaults to ``False``.
mark_scattered (bool): Random scatter mark pixels
in the entire image to get the watermark. Defaults to ``False``.
mark_scattered_height (int): Scattered mark height. Defaults to data_shape[1].
mark_scattered_width (int): Scattered mark width. Defaults to data_shape[2].
add_mark_fn (~collections.abc.Callable | None):
Customized function to add mark to images for :meth:`add_mark()` to call.
``add_mark_fn(_input, mark_random_pos=mark_random_pos, mark_alpha=mark_alpha, **kwargs)``
Defaults to ``None``.
.. |apple_black| image:: ../../../trojanvision/marks/apple_black.png
:height: 50px
:width: 50px
.. |apple_white| image:: ../../../trojanvision/marks/apple_white.png
:height: 50px
:width: 50px
.. |square_black| image:: ../../../trojanvision/marks/square_black.png
:height: 50px
:width: 50px
.. |square_white| image:: ../../../trojanvision/marks/square_white.png
:height: 50px
:width: 50px
.. |watermark_black| image:: ../../../trojanvision/marks/watermark_black.png
:height: 50px
:width: 50px
.. |watermark_white| image:: ../../../trojanvision/marks/watermark_white.png
:height: 50px
:width: 50px
"""
name: str = 'mark'
@staticmethod
def add_argument(group: argparse._ArgumentGroup):
r"""Add watermark arguments to argument parser group.
View source to see specific arguments.
Note:
This is the implementation of adding arguments.
For users, please use :func:`add_argument()` instead, which is more user-friendly.
"""
group.add_argument('--mark_background_color',
choices=['auto', 'black', 'white'],
help='background color in watermark image. '
'It\'s ignored when alpha channel is in watermark image. '
'(default: "auto")')
group.add_argument('--mark_path', help='watermark path (image or npy file), '
'default: "square_white.png")')
group.add_argument('--mark_alpha', type=float,
help='mark opacity (default: 1.0)')
group.add_argument('--mark_height', type=int,
help='mark height (default: 3)')
group.add_argument('--mark_width', type=int,
help='mark width (default: 3)')
group.add_argument('--mark_height_offset', type=int,
help='mark height offset (default: 0)')
group.add_argument('--mark_width_offset', type=int,
help='mark width offset (default: 0)')
group.add_argument('--mark_random_pos', action='store_true',
help='Random offset Location for add_mark.')
group.add_argument('--mark_random_init', action='store_true',
help='random values for mark pixel.')
group.add_argument('--mark_scattered',
action='store_true', help='Random scatter mark pixels.')
group.add_argument('--mark_scattered_height', type=int,
help='Scattered mark height (default: same as input image)')
group.add_argument('--mark_scattered_width', type=int,
help='Scattered mark width (default: same as input image)')
return group
def __init__(self, mark_path: str = 'square_white.png',
data_shape: list[int] = None, mark_background_color: str | torch.Tensor = 'auto',
mark_alpha: float = 1.0, mark_height: int = 3, mark_width: int = 3,
mark_height_offset: int = 0, mark_width_offset: int = 0,
mark_random_init: bool = False, mark_random_pos: bool = False,
mark_scattered: bool = False,
mark_scattered_height: int = None,
mark_scattered_width: int = None,
add_mark_fn: Callable[..., torch.Tensor] = None, **kwargs):
super().__init__(**kwargs)
self.param_list: dict[str, list[str]] = {}
self.param_list['mark'] = ['mark_path',
'mark_alpha', 'mark_height', 'mark_width',
'mark_random_init', 'mark_random_pos',
'mark_scattered']
if not mark_random_pos:
self.param_list['mark'].extend(['mark_height_offset', 'mark_width_offset'])
assert mark_height > 0 and mark_width > 0
# --------------------------------------------------- #
self.mark_alpha = mark_alpha
self.mark_path = mark_path
self.mark_height = mark_height
self.mark_width = mark_width
self.mark_height_offset = mark_height_offset
self.mark_width_offset = mark_width_offset
self.mark_random_init = mark_random_init
self.mark_random_pos = mark_random_pos
self.mark_scattered = mark_scattered
self.mark_scattered_height = mark_scattered_height or data_shape[1]
self.mark_scattered_width = mark_scattered_width or data_shape[2]
self.add_mark_fn = add_mark_fn
self.data_shape = data_shape
# --------------------------------------------------- #
self.mark = self.load_mark(mark_img=mark_path,
mark_background_color=mark_background_color)
def add_mark(self, _input: torch.Tensor, mark_random_pos: bool = None,
mark_alpha: float = None, mark: torch.Tensor = None,
**kwargs) -> torch.Tensor:
r"""Main method to add watermark to a batched input image tensor ranging in ``[0, 1]``.
Call :attr:`self.add_mark_fn()` instead if it's not ``None``.
Args:
_input (torch.Tensor): Batched input tensor
ranging in ``[0, 1]`` with shape ``(N, C, H, W)``.
mark_random_pos (bool | None): Whether to add mark at random location.
Defaults to :attr:`self.mark_random_pos`.
mark_alpha (float | None): Mark opacity. Defaults to :attr:`self.mark_alpha`.
mark (torch.Tensor | None): Mark tensor. Defaults to :attr:`self.mark`.
**kwargs: Keyword arguments passed to `self.add_mark_fn()`.
"""
mark_alpha = mark_alpha if mark_alpha is not None else self.mark_alpha
mark = mark if mark is not None else self.mark
mark_random_pos = mark_random_pos if mark_random_pos is not None else self.mark_random_pos
if callable(self.add_mark_fn):
return self.add_mark_fn(_input, mark_random_pos=mark_random_pos,
mark_alpha=mark_alpha, **kwargs)
trigger_input = _input.clone()
mark = mark.clone().to(device=_input.device)
mark_rgb_channel = mark[..., :-1, :, :]
mark_alpha_channel = mark[..., -1, :, :].unsqueeze(-3)
mark_alpha_channel *= mark_alpha
if mark_random_pos:
batch_size = _input.size(0)
h_start = torch.randint(high=_input.size(-2) - self.mark_height, size=[batch_size])
w_start = torch.randint(high=_input.size(-1) - self.mark_width, size=[batch_size])
h_end, w_end = h_start + self.mark_height, w_start + self.mark_width
for i in range(len(_input)): # TODO: any parallel approach?
org_patch = _input[i, :, h_start[i]:h_end[i], w_start[i]:w_end[i]]
trigger_patch = org_patch + mark_alpha_channel * (mark_rgb_channel - org_patch)
trigger_input[i, :, h_start[i]:h_end[i], w_start[i]:w_end[i]] = trigger_patch
return trigger_input
h_start, w_start = self.mark_height_offset, self.mark_width_offset
h_end, w_end = h_start + self.mark_height, w_start + self.mark_width
org_patch = _input[..., h_start:h_end, w_start:w_end]
trigger_patch = org_patch + mark_alpha_channel * (mark_rgb_channel - org_patch)
trigger_input[..., h_start:h_end, w_start:w_end] = trigger_patch
return trigger_input
def get_mask(self) -> torch.Tensor:
mask = torch.zeros(self.data_shape[-2:], device=self.mark.device)
h_start, w_start = self.mark_height_offset, self.mark_width_offset
h_end, w_end = h_start + self.mark_height, w_start + self.mark_width
mask[h_start:h_end, w_start:w_end].copy_(self.mark[-1])
return mask
@staticmethod
def scatter_mark(mark_unscattered: torch.Tensor,
mark_scattered_shape: list[int]) -> torch.Tensor:
r"""Scatter the original mark tensor to a provided shape.
If the shape are the same, it becomes a pixel shuffling process.
Args:
mark_unscattered (torch.Tensor): The unscattered mark tensor
with shape ``(data_shape[0] + 1, mark_height, mark_width)``
mark_scattered_shape (list[int]): The scattered mark shape
``(data_shape[0] + 1, mark_scattered_height, mark_scattered_width)``
Returns:
torch.Tensor: The scattered mark with shape :attr:`mark_scattered_shape`.
"""
assert mark_scattered_shape[1] >= mark_unscattered.size(1), \
f'mark_scattered_height={mark_scattered_shape[1]:d} >= mark_height={mark_unscattered.size(1):d}'
assert mark_scattered_shape[2] >= mark_unscattered.size(2), \
f'mark_scattered_width={mark_scattered_shape[2]:d} >= mark_width={mark_unscattered.size(2):d}'
pixel_num = mark_unscattered[0].numel()
mark = torch.zeros(mark_scattered_shape, device=env['device'])
idx = torch.randperm(mark[0].numel())[:pixel_num]
mark.flatten(1)[:, idx].copy_(mark_unscattered.flatten(1))
return mark
# ------------------------------ I/O --------------------------- #
def load_mark(
self,
mark_img: str | Image.Image | np.ndarray | torch.Tensor,
mark_background_color: None | str | torch.Tensor = 'auto',
already_processed: bool = False
) -> torch.Tensor:
r"""Load watermark tensor from image :attr:`mark_img`,
scale by calling :any:`PIL.Image.Image.resize`
and transform to ``(channel + 1, height, width)`` with alpha channel.
Args:
mark_img (PIL.Image.Image | str): Pillow image instance or file path.
mark_background_color (str | torch.Tensor | None): Mark background color.
If :class:`str`, choose from ``['auto', 'black', 'white']``;
else, it shall be 1-dim tensor ranging in ``[0, 1]``.
It's ignored when alpha channel in watermark image.
Defaults to ``'auto'``.
already_processed (bool):
If ``True``, will just load :attr:`mark_img` as :attr:`self.mark`.
Defaults to ``False``.
Returns:
torch.Tensor: Watermark tensor ranging in ``[0, 1]``
with shape ``(channel + 1, height, width)`` with alpha channel.
"""
if isinstance(mark_img, str):
if mark_img.endswith('.npy'):
mark_img = np.load(mark_img)
else:
if not os.path.isfile(mark_img) and \
not os.path.isfile(mark_img := os.path.join(dir_path, mark_img)):
raise FileNotFoundError(mark_img.removeprefix(dir_path))
mark_img = F.convert_image_dtype(F.pil_to_tensor(Image.open(mark_img)))
def forward(
self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
image = F.pil_to_tensor(image)
return image, target
def __call__(self, image_name, country='ng'):
path = '/'.join([self.base_path, country, image_name])
img = Image.open(path)
img = img.resize((self.width, self.height))
return TF.pil_to_tensor(img.convert('RGB')).type(torch.FloatTensor)/255
def pil(buffer: io.IOBase) -> features.Image:
return features.Image(pil_to_tensor(PIL.Image.open(buffer)))
def __call__(self, image, target):
image = F.pil_to_tensor(image)
image = F.convert_image_dtype(image)
target = torch.as_tensor(np.array(target), dtype=torch.int64)
return image, target
transform = nn.Sequential()
return transform
st.markdown("# Kornia Augmentations Demo")
st.sidebar.markdown(
"[Kornia](https://github.com/kornia/kornia) is a *differentiable* computer vision library for PyTorch."
)
uploaded_file = st.sidebar.file_uploader("Choose a file")
if uploaded_file is not None:
im = Image.open(uploaded_file)
else:
im = Image.open("./images/pretty_bird.jpg")
scaler = int(im.height / 2)
st.sidebar.image(im, caption="Input Image", width=256)
image = F.pil_to_tensor(im).float() / 255
# batch size is just for show
batch_size = st.sidebar.slider("batch_size",
min_value=4,
max_value=16,
value=8)
gpu = st.sidebar.checkbox("Use GPU!", value=True)
if not gpu:
st.sidebar.markdown("With Kornia you do ops on the GPU!")
device = torch.device("cpu")
else:
if not IS_LOCAL:
st.sidebar.markdown(
"(GPU Not available on hosted demo, try on your local!)")
st.sidebar.markdown(
def __call__(self, image, target):
image = functional.pil_to_tensor(image)
target = torch.as_tensor(np.array(target), dtype=torch.int64)
return image, target
def read_tensor(fp: str) -> torch.Tensor:
tensor = F.convert_image_dtype(F.pil_to_tensor(Image.open(fp)))
return tensor.unsqueeze(0) if tensor.dim() == 2 else tensor
