def _test_adjust_fn(self, fn, fn_pil, fn_t, configs, tol=2.0 + 1e-10, agg_method="max", dts=(None, torch.float32, torch.float64)): script_fn = torch.jit.script(fn) torch.manual_seed(15) tensor, pil_img = self._create_data(26, 34, device=self.device) batch_tensors = self._create_data_batch(16, 18, num_samples=4, device=self.device) for dt in dts: if dt is not None: tensor = F.convert_image_dtype(tensor, dt) batch_tensors = F.convert_image_dtype(batch_tensors, dt) for config in configs: adjusted_tensor = fn_t(tensor, **config) adjusted_pil = fn_pil(pil_img, **config) scripted_result = script_fn(tensor, **config) msg = "{}, {}".format(dt, config) self.assertEqual(adjusted_tensor.dtype, scripted_result.dtype, msg=msg) self.assertEqual(adjusted_tensor.size()[1:], adjusted_pil.size[::-1], msg=msg) rbg_tensor = adjusted_tensor if adjusted_tensor.dtype != torch.uint8: rbg_tensor = F.convert_image_dtype(adjusted_tensor, torch.uint8) # Check that max difference does not exceed 2 in [0, 255] range # Exact matching is not possible due to incompatibility convert_image_dtype and PIL results self.approxEqualTensorToPIL(rbg_tensor.float(), adjusted_pil, tol=tol, msg=msg, agg_method=agg_method) atol = 1e-6 if adjusted_tensor.dtype == torch.uint8 and "cuda" in torch.device(self.device).type: atol = 1.0 self.assertTrue(adjusted_tensor.allclose(scripted_result, atol=atol), msg=msg) self._test_fn_on_batch(batch_tensors, fn, **config)
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 __getitem__(self, idx): img_path = self.img_path + self.img_dir + self.input_filenames.iloc[idx, 0] image = torchvision.io.read_image(img_path) image = TF.convert_image_dtype(image, torch.float) # Read in the attribute labels for the current input image attributes = self.img_labels.iloc[idx,] attributes = torch.tensor(attributes) attributes = torch.gt(attributes, 0).float() if self.transform: if self.train: image, image2 = self.transform(image) return (image, image2), attributes elif self.test: image = self.transform(image) return (image, image), attributes
def __call__(self, image, target): image = functional.convert_image_dtype(image, self.dtype) return image, target
colors = ["blue", "yellow"] result = draw_bounding_boxes(dog1_int, boxes, colors=colors, width=5) show(result) ##################################### # Naturally, we can also plot bounding boxes produced by torchvision detection # models. Here is demo with a Faster R-CNN model loaded from # :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` # model. You can also try using a RetinaNet with # :func:`~torchvision.models.detection.retinanet_resnet50_fpn`. from torchvision.models.detection import fasterrcnn_resnet50_fpn from torchvision.transforms.functional import convert_image_dtype batch_int = torch.stack([dog1_int, dog2_int]) batch = convert_image_dtype(batch_int, dtype=torch.float) model = fasterrcnn_resnet50_fpn(pretrained=True, progress=False) model = model.eval() outputs = model(batch) print(outputs) ##################################### # Let's plot the boxes detected by our model. We will only plot the boxes with a # score greater than a given threshold. threshold = .8 dogs_with_boxes = [ draw_bounding_boxes(dog_int, boxes=output['boxes'][output['scores'] > threshold],
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
def forward( self, image: Tensor, target: Optional[Dict[str, Tensor]] = None ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]: image = F.convert_image_dtype(image, self.dtype) return image, target
##################################### # Naturally, we can also plot bounding boxes produced by torchvision detection # models. Here is demo with a Faster R-CNN model loaded from # :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` # model. You can also try using a RetinaNet with # :func:`~torchvision.models.detection.retinanet_resnet50_fpn`, an SSDlite with # :func:`~torchvision.models.detection.ssdlite320_mobilenet_v3_large` or an SSD with # :func:`~torchvision.models.detection.ssd300_vgg16`. For more details # on the output of such models, you may refer to :ref:`instance_seg_output`. from torchvision.models.detection import fasterrcnn_resnet50_fpn from torchvision.transforms.functional import convert_image_dtype batch_int = torch.stack([dog1_int, dog2_int]) batch = convert_image_dtype(batch_int, dtype=torch.float) model = fasterrcnn_resnet50_fpn(pretrained=True, progress=False) model = model.eval() outputs = model(batch) print(outputs) ##################################### # Let's plot the boxes detected by our model. We will only plot the boxes with a # score greater than a given threshold. score_threshold = .8 dogs_with_boxes = [ draw_bounding_boxes(dog_int, boxes=output['boxes'][output['scores'] > score_threshold], width=4) for dog_int, output in zip(batch_int, outputs)
# to (1600,1600,4) im_perm_lis.append(im_start_lis[i].permute(1, 2, 0)) batch_int = list() for i, j in zip(range(0, len(im_start_lis), 2), range(1, len(im_start_lis), 2)): # Gather inputs for batches. The model is run on each batch for the # comparison. batch_int.append(torch.stack([im_start_lis[i], im_start_lis[j]])) #batch_lis.append(batch(convert_image_dtype(batch_int[i], dtype=torch.float))) batch = list() for i in range(len(batch_int)): batch.append(convert_image_dtype(batch_int[i], dtype=torch.float)) # Creates outputs from the model that can be used with other functions and # classes created later on. # This will grab the first of the pair of images to be compared. def get_first_outputs(image, model, threshold): with torch.no_grad(): # forward pass of the image through the modle outputs = model(image) # get all the scores scores = list(outputs[0]['scores'].detach().cpu().numpy()) # index of those scores which are above a certain threshold thresholded_preds_indices = [ scores.index(i) for i in scores if i > threshold
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
dtype=torch.float) colors = ["blue", "yellow"] result = draw_bounding_boxes(dog1_int, boxes, colors=colors, width=5) show(result) ##################################### # Naturally, we can also plot bounding boxes produced by torchvision detection # models. Here is demo with a Faster R-CNN model loaded from # :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` # model. You can also try using a RetinaNet with # :func:`~torchvision.models.detection.retinanet_resnet50_fpn`. from torchvision.models.detection import fasterrcnn_resnet50_fpn from torchvision.transforms.functional import convert_image_dtype dog1_float = convert_image_dtype(dog1_int, dtype=torch.float) dog2_float = convert_image_dtype(dog2_int, dtype=torch.float) batch = torch.stack([dog1_float, dog2_float]) model = fasterrcnn_resnet50_fpn(pretrained=True, progress=False) model = model.eval() outputs = model(batch) print(outputs) ##################################### # Let's plot the boxes detected by our model. We will only plot the boxes with a # score greater than a given threshold. threshold = .8 dogs_with_boxes = [
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)))
from torchvision.utils import draw_bounding_boxes drawn_boxes = draw_bounding_boxes(img, boxes, colors="red") show(drawn_boxes) ################################### # These boxes can now directly be used by detection models in torchvision. # Here is demo with a Faster R-CNN model loaded from # :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` from torchvision.models.detection import fasterrcnn_resnet50_fpn model = fasterrcnn_resnet50_fpn(pretrained=True, progress=False) print(img.size()) img = F.convert_image_dtype(img, torch.float) target = {} target["boxes"] = boxes target["labels"] = labels = torch.ones((masks.size(0), ), dtype=torch.int64) detection_outputs = model(img.unsqueeze(0), [target]) #################################### # Converting Segmentation Dataset to Detection Dataset # ---------------------------------------------------- # # With this utility it becomes very simple to convert a segmentation dataset to a detection dataset. # With this we can now use a segmentation dataset to train a detection model. # One can similarly convert panoptic dataset to detection dataset. # Here is an example where we re-purpose the dataset from the # `PenFudan Detection Tutorial <https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html>`_.
def _transform(self, input: Any, params: Dict[str, Any]) -> Any: if type(input) is features.Image: output = convert_image_dtype(input, dtype=self.dtype) return features.Image.new_like(input, output, dtype=self.dtype) else: return input