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