fffiloni/Video-Matting-Anything

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3 4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7import numpy as np8import torch9from torch.nn import functional as F10from torchvision.transforms.functional import resize, to_pil_image  # type: ignore11 12from copy import deepcopy13from typing import Tuple14 15 16class ResizeLongestSide:17    """18    Resizes images to the longest side 'target_length', as well as provides19    methods for resizing coordinates and boxes. Provides methods for20    transforming both numpy array and batched torch tensors.21    """22 23    def __init__(self, target_length: int) -> None:24        self.target_length = target_length25 26    def apply_image(self, image: np.ndarray) -> np.ndarray:27        """28        Expects a numpy array with shape HxWxC in uint8 format.29        """30        target_size = self.get_preprocess_shape(image.shape[0], image.shape[1], self.target_length)31        return np.array(resize(to_pil_image(image), target_size))32 33    def apply_coords(self, coords: np.ndarray, original_size: Tuple[int, ...]) -> np.ndarray:34        """35        Expects a numpy array of length 2 in the final dimension. Requires the36        original image size in (H, W) format.37        """38        old_h, old_w = original_size39        new_h, new_w = self.get_preprocess_shape(40            original_size[0], original_size[1], self.target_length41        )42        coords = deepcopy(coords).astype(float)43        coords[..., 0] = coords[..., 0] * (new_w / old_w)44        coords[..., 1] = coords[..., 1] * (new_h / old_h)45        return coords46 47    def apply_boxes(self, boxes: np.ndarray, original_size: Tuple[int, ...]) -> np.ndarray:48        """49        Expects a numpy array shape Bx4. Requires the original image size50        in (H, W) format.51        """52        boxes = self.apply_coords(boxes.reshape(-1, 2, 2), original_size)53        return boxes.reshape(-1, 4)54 55    def apply_image_torch(self, image: torch.Tensor) -> torch.Tensor:56        """57        Expects batched images with shape BxCxHxW and float format. This58        transformation may not exactly match apply_image. apply_image is59        the transformation expected by the model.60        """61        # Expects an image in BCHW format. May not exactly match apply_image.62        target_size = self.get_preprocess_shape(image.shape[2], image.shape[3], self.target_length)63        return F.interpolate(64            image, target_size, mode="bilinear", align_corners=False, antialias=True65        )66 67    def apply_coords_torch(68        self, coords: torch.Tensor, original_size: Tuple[int, ...]69    ) -> torch.Tensor:70        """71        Expects a torch tensor with length 2 in the last dimension. Requires the72        original image size in (H, W) format.73        """74        old_h, old_w = original_size75        new_h, new_w = self.get_preprocess_shape(76            original_size[0], original_size[1], self.target_length77        )78        coords = deepcopy(coords).to(torch.float)79        coords[..., 0] = coords[..., 0] * (new_w / old_w)80        coords[..., 1] = coords[..., 1] * (new_h / old_h)81        return coords82 83    def apply_boxes_torch(84        self, boxes: torch.Tensor, original_size: Tuple[int, ...]85    ) -> torch.Tensor:86        """87        Expects a torch tensor with shape Bx4. Requires the original image88        size in (H, W) format.89        """90        boxes = self.apply_coords_torch(boxes.reshape(-1, 2, 2), original_size)91        return boxes.reshape(-1, 4)92 93    @staticmethod94    def get_preprocess_shape(oldh: int, oldw: int, long_side_length: int) -> Tuple[int, int]:95        """96        Compute the output size given input size and target long side length.97        """98        scale = long_side_length * 1.0 / max(oldh, oldw)99        newh, neww = oldh * scale, oldw * scale100        neww = int(neww + 0.5)101        newh = int(newh + 0.5)102        return (newh, neww)103