EQUES/suspicious-behavior-detector

1from typing import Dict, List, Optional, Tuple, Union2 3import torch4import torchvision5from torch import nn, Tensor6from torchvision import ops7from torchvision.transforms import functional as F, InterpolationMode, transforms as T8 9 10def _flip_coco_person_keypoints(kps, width):11    flip_inds = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]12    flipped_data = kps[:, flip_inds]13    flipped_data[..., 0] = width - flipped_data[..., 0]14    # Maintain COCO convention that if visibility == 0, then x, y = 015    inds = flipped_data[..., 2] == 016    flipped_data[inds] = 017    return flipped_data18 19 20class Compose:21    def __init__(self, transforms):22        self.transforms = transforms23 24    def __call__(self, image, target):25        for t in self.transforms:26            image, target = t(image, target)27        return image, target28 29class ToTensor(object):30    def __call__(self, image, target):31        image = F.to_tensor(image)32        return image, target33 34class RandomHorizontalFlip(T.RandomHorizontalFlip):35    def forward(36        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None37    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:38        if torch.rand(1) < self.p:39            image = F.hflip(image)40            if target is not None:41                _, _, width = F.get_dimensions(image)42                target["boxes"][:, [0, 2]] = width - target["boxes"][:, [2, 0]]43                if "masks" in target:44                    target["masks"] = target["masks"].flip(-1)45                if "keypoints" in target:46                    keypoints = target["keypoints"]47                    keypoints = _flip_coco_person_keypoints(keypoints, width)48                    target["keypoints"] = keypoints49        return image, target50 51 52class PILToTensor(nn.Module):53    def forward(54        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None55    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:56        image = F.pil_to_tensor(image)57        return image, target58 59 60class ConvertImageDtype(nn.Module):61    def __init__(self, dtype: torch.dtype) -> None:62        super().__init__()63        self.dtype = dtype64 65    def forward(66        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None67    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:68        image = F.convert_image_dtype(image, self.dtype)69        return image, target70 71 72class RandomIoUCrop(nn.Module):73    def __init__(74        self,75        min_scale: float = 0.3,76        max_scale: float = 1.0,77        min_aspect_ratio: float = 0.5,78        max_aspect_ratio: float = 2.0,79        sampler_options: Optional[List[float]] = None,80        trials: int = 40,81    ):82        super().__init__()83        # Configuration similar to https://github.com/weiliu89/caffe/blob/ssd/examples/ssd/ssd_coco.py#L89-L17484        self.min_scale = min_scale85        self.max_scale = max_scale86        self.min_aspect_ratio = min_aspect_ratio87        self.max_aspect_ratio = max_aspect_ratio88        if sampler_options is None:89            sampler_options = [0.0, 0.1, 0.3, 0.5, 0.7, 0.9, 1.0]90        self.options = sampler_options91        self.trials = trials92 93    def forward(94        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None95    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:96        if target is None:97            raise ValueError("The targets can't be None for this transform.")98 99        if isinstance(image, torch.Tensor):100            if image.ndimension() not in {2, 3}:101                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")102            elif image.ndimension() == 2:103                image = image.unsqueeze(0)104 105        _, orig_h, orig_w = F.get_dimensions(image)106 107        while True:108            # sample an option109            idx = int(torch.randint(low=0, high=len(self.options), size=(1,)))110            min_jaccard_overlap = self.options[idx]111            if min_jaccard_overlap >= 1.0:  # a value larger than 1 encodes the leave as-is option112                return image, target113 114            for _ in range(self.trials):115                # check the aspect ratio limitations116                r = self.min_scale + (self.max_scale - self.min_scale) * torch.rand(2)117                new_w = int(orig_w * r[0])118                new_h = int(orig_h * r[1])119                aspect_ratio = new_w / new_h120                if not (self.min_aspect_ratio <= aspect_ratio <= self.max_aspect_ratio):121                    continue122 123                # check for 0 area crops124                r = torch.rand(2)125                left = int((orig_w - new_w) * r[0])126                top = int((orig_h - new_h) * r[1])127                right = left + new_w128                bottom = top + new_h129                if left == right or top == bottom:130                    continue131 132                # check for any valid boxes with centers within the crop area133                cx = 0.5 * (target["boxes"][:, 0] + target["boxes"][:, 2])134                cy = 0.5 * (target["boxes"][:, 1] + target["boxes"][:, 3])135                is_within_crop_area = (left < cx) & (cx < right) & (top < cy) & (cy < bottom)136                if not is_within_crop_area.any():137                    continue138 139                # check at least 1 box with jaccard limitations140                boxes = target["boxes"][is_within_crop_area]141                ious = torchvision.ops.boxes.box_iou(142                    boxes, torch.tensor([[left, top, right, bottom]], dtype=boxes.dtype, device=boxes.device)143                )144                if ious.max() < min_jaccard_overlap:145                    continue146 147                # keep only valid boxes and perform cropping148                target["boxes"] = boxes149                target["labels"] = target["labels"][is_within_crop_area]150                target["boxes"][:, 0::2] -= left151                target["boxes"][:, 1::2] -= top152                target["boxes"][:, 0::2].clamp_(min=0, max=new_w)153                target["boxes"][:, 1::2].clamp_(min=0, max=new_h)154                image = F.crop(image, top, left, new_h, new_w)155 156                return image, target157 158 159class RandomZoomOut(nn.Module):160    def __init__(161        self, fill: Optional[List[float]] = None, side_range: Tuple[float, float] = (1.0, 4.0), p: float = 0.5162    ):163        super().__init__()164        if fill is None:165            fill = [0.0, 0.0, 0.0]166        self.fill = fill167        self.side_range = side_range168        if side_range[0] < 1.0 or side_range[0] > side_range[1]:169            raise ValueError(f"Invalid canvas side range provided {side_range}.")170        self.p = p171 172    @torch.jit.unused173    def _get_fill_value(self, is_pil):174        # type: (bool) -> int175        # We fake the type to make it work on JIT176        return tuple(int(x) for x in self.fill) if is_pil else 0177 178    def forward(179        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None180    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:181        if isinstance(image, torch.Tensor):182            if image.ndimension() not in {2, 3}:183                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")184            elif image.ndimension() == 2:185                image = image.unsqueeze(0)186 187        if torch.rand(1) >= self.p:188            return image, target189 190        _, orig_h, orig_w = F.get_dimensions(image)191 192        r = self.side_range[0] + torch.rand(1) * (self.side_range[1] - self.side_range[0])193        canvas_width = int(orig_w * r)194        canvas_height = int(orig_h * r)195 196        r = torch.rand(2)197        left = int((canvas_width - orig_w) * r[0])198        top = int((canvas_height - orig_h) * r[1])199        right = canvas_width - (left + orig_w)200        bottom = canvas_height - (top + orig_h)201 202        if torch.jit.is_scripting():203            fill = 0204        else:205            fill = self._get_fill_value(F._is_pil_image(image))206 207        image = F.pad(image, [left, top, right, bottom], fill=fill)208        if isinstance(image, torch.Tensor):209            # PyTorch's pad supports only integers on fill. So we need to overwrite the colour210            v = torch.tensor(self.fill, device=image.device, dtype=image.dtype).view(-1, 1, 1)211            image[..., :top, :] = image[..., :, :left] = image[..., (top + orig_h) :, :] = image[212                ..., :, (left + orig_w) :213            ] = v214 215        if target is not None:216            target["boxes"][:, 0::2] += left217            target["boxes"][:, 1::2] += top218 219        return image, target220 221 222class RandomPhotometricDistort(nn.Module):223    def __init__(224        self,225        contrast: Tuple[float, float] = (0.5, 1.5),226        saturation: Tuple[float, float] = (0.5, 1.5),227        hue: Tuple[float, float] = (-0.05, 0.05),228        brightness: Tuple[float, float] = (0.875, 1.125),229        p: float = 0.5,230    ):231        super().__init__()232        self._brightness = T.ColorJitter(brightness=brightness)233        self._contrast = T.ColorJitter(contrast=contrast)234        self._hue = T.ColorJitter(hue=hue)235        self._saturation = T.ColorJitter(saturation=saturation)236        self.p = p237 238    def forward(239        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None240    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:241        if isinstance(image, torch.Tensor):242            if image.ndimension() not in {2, 3}:243                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")244            elif image.ndimension() == 2:245                image = image.unsqueeze(0)246 247        r = torch.rand(7)248 249        if r[0] < self.p:250            image = self._brightness(image)251 252        contrast_before = r[1] < 0.5253        if contrast_before:254            if r[2] < self.p:255                image = self._contrast(image)256 257        if r[3] < self.p:258            image = self._saturation(image)259 260        if r[4] < self.p:261            image = self._hue(image)262 263        if not contrast_before:264            if r[5] < self.p:265                image = self._contrast(image)266 267        if r[6] < self.p:268            channels, _, _ = F.get_dimensions(image)269            permutation = torch.randperm(channels)270 271            is_pil = F._is_pil_image(image)272            if is_pil:273                image = F.pil_to_tensor(image)274                image = F.convert_image_dtype(image)275            image = image[..., permutation, :, :]276            if is_pil:277                image = F.to_pil_image(image)278 279        return image, target280 281 282class ScaleJitter(nn.Module):283    """Randomly resizes the image and its bounding boxes  within the specified scale range.284    The class implements the Scale Jitter augmentation as described in the paper285    `"Simple Copy-Paste is a Strong Data Augmentation Method for Instance Segmentation" <https://arxiv.org/abs/2012.07177>`_.286 287    Args:288        target_size (tuple of ints): The target size for the transform provided in (height, weight) format.289        scale_range (tuple of ints): scaling factor interval, e.g (a, b), then scale is randomly sampled from the290            range a <= scale <= b.291        interpolation (InterpolationMode): Desired interpolation enum defined by292            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``.293    """294 295    def __init__(296        self,297        target_size: Tuple[int, int],298        scale_range: Tuple[float, float] = (0.1, 2.0),299        interpolation: InterpolationMode = InterpolationMode.BILINEAR,300    ):301        super().__init__()302        self.target_size = target_size303        self.scale_range = scale_range304        self.interpolation = interpolation305 306    def forward(307        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None308    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:309        if isinstance(image, torch.Tensor):310            if image.ndimension() not in {2, 3}:311                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")312            elif image.ndimension() == 2:313                image = image.unsqueeze(0)314 315        _, orig_height, orig_width = F.get_dimensions(image)316 317        scale = self.scale_range[0] + torch.rand(1) * (self.scale_range[1] - self.scale_range[0])318        r = min(self.target_size[1] / orig_height, self.target_size[0] / orig_width) * scale319        new_width = int(orig_width * r)320        new_height = int(orig_height * r)321 322        image = F.resize(image, [new_height, new_width], interpolation=self.interpolation)323 324        if target is not None:325            target["boxes"][:, 0::2] *= new_width / orig_width326            target["boxes"][:, 1::2] *= new_height / orig_height327            if "masks" in target:328                target["masks"] = F.resize(329                    target["masks"], [new_height, new_width], interpolation=InterpolationMode.NEAREST330                )331 332        return image, target333 334 335class FixedSizeCrop(nn.Module):336    def __init__(self, size, fill=0, padding_mode="constant"):337        super().__init__()338        size = tuple(T._setup_size(size, error_msg="Please provide only two dimensions (h, w) for size."))339        self.crop_height = size[0]340        self.crop_width = size[1]341        self.fill = fill  # TODO: Fill is currently respected only on PIL. Apply tensor patch.342        self.padding_mode = padding_mode343 344    def _pad(self, img, target, padding):345        # Taken from the functional_tensor.py pad346        if isinstance(padding, int):347            pad_left = pad_right = pad_top = pad_bottom = padding348        elif len(padding) == 1:349            pad_left = pad_right = pad_top = pad_bottom = padding[0]350        elif len(padding) == 2:351            pad_left = pad_right = padding[0]352            pad_top = pad_bottom = padding[1]353        else:354            pad_left = padding[0]355            pad_top = padding[1]356            pad_right = padding[2]357            pad_bottom = padding[3]358 359        padding = [pad_left, pad_top, pad_right, pad_bottom]360        img = F.pad(img, padding, self.fill, self.padding_mode)361        if target is not None:362            target["boxes"][:, 0::2] += pad_left363            target["boxes"][:, 1::2] += pad_top364            if "masks" in target:365                target["masks"] = F.pad(target["masks"], padding, 0, "constant")366 367        return img, target368 369    def _crop(self, img, target, top, left, height, width):370        img = F.crop(img, top, left, height, width)371        if target is not None:372            boxes = target["boxes"]373            boxes[:, 0::2] -= left374            boxes[:, 1::2] -= top375            boxes[:, 0::2].clamp_(min=0, max=width)376            boxes[:, 1::2].clamp_(min=0, max=height)377 378            is_valid = (boxes[:, 0] < boxes[:, 2]) & (boxes[:, 1] < boxes[:, 3])379 380            target["boxes"] = boxes[is_valid]381            target["labels"] = target["labels"][is_valid]382            if "masks" in target:383                target["masks"] = F.crop(target["masks"][is_valid], top, left, height, width)384 385        return img, target386 387    def forward(self, img, target=None):388        _, height, width = F.get_dimensions(img)389        new_height = min(height, self.crop_height)390        new_width = min(width, self.crop_width)391 392        if new_height != height or new_width != width:393            offset_height = max(height - self.crop_height, 0)394            offset_width = max(width - self.crop_width, 0)395 396            r = torch.rand(1)397            top = int(offset_height * r)398            left = int(offset_width * r)399 400            img, target = self._crop(img, target, top, left, new_height, new_width)401 402        pad_bottom = max(self.crop_height - new_height, 0)403        pad_right = max(self.crop_width - new_width, 0)404        if pad_bottom != 0 or pad_right != 0:405            img, target = self._pad(img, target, [0, 0, pad_right, pad_bottom])406 407        return img, target408 409 410class RandomShortestSize(nn.Module):411    def __init__(412        self,413        min_size: Union[List[int], Tuple[int], int],414        max_size: int,415        interpolation: InterpolationMode = InterpolationMode.BILINEAR,416    ):417        super().__init__()418        self.min_size = [min_size] if isinstance(min_size, int) else list(min_size)419        self.max_size = max_size420        self.interpolation = interpolation421 422    def forward(423        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None424    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:425        _, orig_height, orig_width = F.get_dimensions(image)426 427        min_size = self.min_size[torch.randint(len(self.min_size), (1,)).item()]428        r = min(min_size / min(orig_height, orig_width), self.max_size / max(orig_height, orig_width))429 430        new_width = int(orig_width * r)431        new_height = int(orig_height * r)432 433        image = F.resize(image, [new_height, new_width], interpolation=self.interpolation)434 435        if target is not None:436            target["boxes"][:, 0::2] *= new_width / orig_width437            target["boxes"][:, 1::2] *= new_height / orig_height438            if "masks" in target:439                target["masks"] = F.resize(440                    target["masks"], [new_height, new_width], interpolation=InterpolationMode.NEAREST441                )442 443        return image, target444 445 446def _copy_paste(447    image: torch.Tensor,448    target: Dict[str, Tensor],449    paste_image: torch.Tensor,450    paste_target: Dict[str, Tensor],451    blending: bool = True,452    resize_interpolation: F.InterpolationMode = F.InterpolationMode.BILINEAR,453) -> Tuple[torch.Tensor, Dict[str, Tensor]]:454 455    # Random paste targets selection:456    num_masks = len(paste_target["masks"])457 458    if num_masks < 1:459        # Such degerante case with num_masks=0 can happen with LSJ460        # Let's just return (image, target)461        return image, target462 463    # We have to please torch script by explicitly specifying dtype as torch.long464    random_selection = torch.randint(0, num_masks, (num_masks,), device=paste_image.device)465    random_selection = torch.unique(random_selection).to(torch.long)466 467    paste_masks = paste_target["masks"][random_selection]468    paste_boxes = paste_target["boxes"][random_selection]469    paste_labels = paste_target["labels"][random_selection]470 471    masks = target["masks"]472 473    # We resize source and paste data if they have different sizes474    # This is something we introduced here as originally the algorithm works475    # on equal-sized data (for example, coming from LSJ data augmentations)476    size1 = image.shape[-2:]477    size2 = paste_image.shape[-2:]478    if size1 != size2:479        paste_image = F.resize(paste_image, size1, interpolation=resize_interpolation)480        paste_masks = F.resize(paste_masks, size1, interpolation=F.InterpolationMode.NEAREST)481        # resize bboxes:482        ratios = torch.tensor((size1[1] / size2[1], size1[0] / size2[0]), device=paste_boxes.device)483        paste_boxes = paste_boxes.view(-1, 2, 2).mul(ratios).view(paste_boxes.shape)484 485    paste_alpha_mask = paste_masks.sum(dim=0) > 0486 487    if blending:488        paste_alpha_mask = F.gaussian_blur(489            paste_alpha_mask.unsqueeze(0),490            kernel_size=(5, 5),491            sigma=[492                2.0,493            ],494        )495 496    # Copy-paste images:497    image = (image * (~paste_alpha_mask)) + (paste_image * paste_alpha_mask)498 499    # Copy-paste masks:500    masks = masks * (~paste_alpha_mask)501    non_all_zero_masks = masks.sum((-1, -2)) > 0502    masks = masks[non_all_zero_masks]503 504    # Do a shallow copy of the target dict505    out_target = {k: v for k, v in target.items()}506 507    out_target["masks"] = torch.cat([masks, paste_masks])508 509    # Copy-paste boxes and labels510    boxes = ops.masks_to_boxes(masks)511    out_target["boxes"] = torch.cat([boxes, paste_boxes])512 513    labels = target["labels"][non_all_zero_masks]514    out_target["labels"] = torch.cat([labels, paste_labels])515 516    # Update additional optional keys: area and iscrowd if exist517    if "area" in target:518        out_target["area"] = out_target["masks"].sum((-1, -2)).to(torch.float32)519 520    if "iscrowd" in target and "iscrowd" in paste_target:521        # target['iscrowd'] size can be differ from mask size (non_all_zero_masks)522        # For example, if previous transforms geometrically modifies masks/boxes/labels but523        # does not update "iscrowd"524        if len(target["iscrowd"]) == len(non_all_zero_masks):525            iscrowd = target["iscrowd"][non_all_zero_masks]526            paste_iscrowd = paste_target["iscrowd"][random_selection]527            out_target["iscrowd"] = torch.cat([iscrowd, paste_iscrowd])528 529    # Check for degenerated boxes and remove them530    boxes = out_target["boxes"]531    degenerate_boxes = boxes[:, 2:] <= boxes[:, :2]532    if degenerate_boxes.any():533        valid_targets = ~degenerate_boxes.any(dim=1)534 535        out_target["boxes"] = boxes[valid_targets]536        out_target["masks"] = out_target["masks"][valid_targets]537        out_target["labels"] = out_target["labels"][valid_targets]538 539        if "area" in out_target:540            out_target["area"] = out_target["area"][valid_targets]541        if "iscrowd" in out_target and len(out_target["iscrowd"]) == len(valid_targets):542            out_target["iscrowd"] = out_target["iscrowd"][valid_targets]543 544    return image, out_target545 546 547class SimpleCopyPaste(torch.nn.Module):548    def __init__(self, blending=True, resize_interpolation=F.InterpolationMode.BILINEAR):549        super().__init__()550        self.resize_interpolation = resize_interpolation551        self.blending = blending552 553    def forward(554        self, images: List[torch.Tensor], targets: List[Dict[str, Tensor]]555    ) -> Tuple[List[torch.Tensor], List[Dict[str, Tensor]]]:556        torch._assert(557            isinstance(images, (list, tuple)) and all([isinstance(v, torch.Tensor) for v in images]),558            "images should be a list of tensors",559        )560        torch._assert(561            isinstance(targets, (list, tuple)) and len(images) == len(targets),562            "targets should be a list of the same size as images",563        )564        for target in targets:565            # Can not check for instance type dict with inside torch.jit.script566            # torch._assert(isinstance(target, dict), "targets item should be a dict")567            for k in ["masks", "boxes", "labels"]:568                torch._assert(k in target, f"Key {k} should be present in targets")569                torch._assert(isinstance(target[k], torch.Tensor), f"Value for the key {k} should be a tensor")570 571        # images = [t1, t2, ..., tN]572        # Let's define paste_images as shifted list of input images573        # paste_images = [t2, t3, ..., tN, t1]574        # FYI: in TF they mix data on the dataset level575        images_rolled = images[-1:] + images[:-1]576        targets_rolled = targets[-1:] + targets[:-1]577 578        output_images: List[torch.Tensor] = []579        output_targets: List[Dict[str, Tensor]] = []580 581        for image, target, paste_image, paste_target in zip(images, targets, images_rolled, targets_rolled):582            output_image, output_data = _copy_paste(583                image,584                target,585                paste_image,586                paste_target,587                blending=self.blending,588                resize_interpolation=self.resize_interpolation,589            )590            output_images.append(output_image)591            output_targets.append(output_data)592 593        return output_images, output_targets594 595    def __repr__(self) -> str:596        s = f"{self.__class__.__name__}(blending={self.blending}, resize_interpolation={self.resize_interpolation})"597        return s