Dataset and Code for "Grasp-Anything: Large-scale Grasp Dataset from Foundation Models."
import argparse
import logging
import time

import numpy as np
import torch.utils.data

from hardware.device import get_device
from inference.post_process import post_process_output
from utils.data import get_dataset
from utils.dataset_processing import evaluation, grasp
from utils.visualisation.plot import save_results

logging.basicConfig(level=logging.INFO)


def parse_args():
    parser = argparse.ArgumentParser(description='Evaluate networks')

    # Network
    parser.add_argument('--network', metavar='N', type=str, nargs='+',
                        help='Path to saved networks to evaluate')
    parser.add_argument('--input-size', type=int, default=224,
                        help='Input image size for the network')

    # Dataset
    parser.add_argument('--dataset', type=str,
                        help='Dataset Name ("cornell" or "jaquard")')
    parser.add_argument('--dataset-path', type=str,
                        help='Path to dataset')
    parser.add_argument('--use-depth', type=int, default=1,
                        help='Use Depth image for evaluation (1/0)')
    parser.add_argument('--use-rgb', type=int, default=1,
                        help='Use RGB image for evaluation (1/0)')
    parser.add_argument('--augment', action='store_true',
                        help='Whether data augmentation should be applied')
    parser.add_argument('--split', type=float, default=0.01,
                        help='Fraction of data for training (remainder is validation)')
    parser.add_argument('--ds-shuffle', action='store_true', default=False,
                        help='Shuffle the dataset')
    parser.add_argument('--ds-rotate', type=float, default=0.0,
                        help='Shift the start point of the dataset to use a different test/train split')
    parser.add_argument('--num-workers', type=int, default=8,
                        help='Dataset workers')

    # Evaluation
    parser.add_argument('--n-grasps', type=int, default=1,
                        help='Number of grasps to consider per image')
    parser.add_argument('--iou-threshold', type=float, default=0.25,
                        help='Threshold for IOU matching')
    parser.add_argument('--iou-eval', action='store_true',
                        help='Compute success based on IoU metric.')
    parser.add_argument('--jacquard-output', action='store_true',
                        help='Jacquard-dataset style output')

    # Misc.
    parser.add_argument('--vis', action='store_true',
                        help='Visualise the network output')
    parser.add_argument('--cpu', dest='force_cpu', action='store_true', default=False,
                        help='Force code to run in CPU mode')
    parser.add_argument('--random-seed', type=int, default=123,
                        help='Random seed for numpy')
    parser.add_argument('--seen', type=int, default=1,
                        help='Flag for using seen classes, only work for Grasp-Anything dataset') 

    args = parser.parse_args()

    if args.jacquard_output and args.dataset != 'jacquard':
        raise ValueError('--jacquard-output can only be used with the --dataset jacquard option.')
    if args.jacquard_output and args.augment:
        raise ValueError('--jacquard-output can not be used with data augmentation.')

    return args


if __name__ == '__main__':
    args = parse_args()

    # Get the compute device
    device = get_device(args.force_cpu)

    # Load Dataset
    logging.info('Loading {} Dataset...'.format(args.dataset.title()))
    Dataset = get_dataset(args.dataset)
    test_dataset = Dataset(args.dataset_path,
                           output_size=args.input_size,
                           ds_rotate=args.ds_rotate,
                           random_rotate=args.augment,
                           random_zoom=args.augment,
                           include_depth=args.use_depth,
                           include_rgb=args.use_rgb,
                           seen=args.seen)

    indices = list(range(test_dataset.length))
    split = int(np.floor(args.split * test_dataset.length))
    if args.ds_shuffle:
        np.random.seed(args.random_seed)
        np.random.shuffle(indices)
    val_indices = indices[split:]
    val_sampler = torch.utils.data.sampler.SubsetRandomSampler(val_indices)
    logging.info('Validation size: {}'.format(len(val_indices)))

    test_data = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=1,
        num_workers=args.num_workers,
        sampler=val_sampler
    )
    logging.info('Done')

    for network in args.network:
        logging.info('\nEvaluating model {}'.format(network))

        # Load Network
        net = torch.load(network)

        results = {'correct': 0, 'failed': 0}

        if args.jacquard_output:
            jo_fn = network + '_jacquard_output.txt'
            with open(jo_fn, 'w') as f:
                pass

        start_time = time.time()

        with torch.no_grad():
            for idx, (x, y, didx, rot, zoom) in enumerate(test_data):
                xc = x.to(device)
                yc = [yi.to(device) for yi in y]
                lossd = net.compute_loss(xc, yc)

                q_img, ang_img, width_img = post_process_output(lossd['pred']['pos'], lossd['pred']['cos'],
                                                                lossd['pred']['sin'], lossd['pred']['width'])

                if args.iou_eval:
                    s = evaluation.calculate_iou_match(q_img, ang_img, test_data.dataset.get_gtbb(didx, rot, zoom),
                                                       no_grasps=args.n_grasps,
                                                       grasp_width=width_img,
                                                       threshold=args.iou_threshold
                                                       )
                    if s:
                        results['correct'] += 1
                    else:
                        results['failed'] += 1

                if args.jacquard_output:
                    grasps = grasp.detect_grasps(q_img, ang_img, width_img=width_img, no_grasps=1)
                    with open(jo_fn, 'a') as f:
                        for g in grasps:
                            f.write(test_data.dataset.get_jname(didx) + '\n')
                            f.write(g.to_jacquard(scale=1024 / 300) + '\n')

                if args.vis:
                    save_results(
                        rgb_img=test_data.dataset.get_rgb(didx, rot, zoom, normalise=False),
                        depth_img=test_data.dataset.get_depth(didx, rot, zoom),
                        grasp_q_img=q_img,
                        grasp_angle_img=ang_img,
                        no_grasps=args.n_grasps,
                        grasp_width_img=width_img
                    )

        avg_time = (time.time() - start_time) / len(test_data)
        logging.info('Average evaluation time per image: {}ms'.format(avg_time * 1000))

        if args.iou_eval:
            logging.info('IOU Results: %d/%d = %f' % (results['correct'],
                                                      results['correct'] + results['failed'],
                                                      results['correct'] / (results['correct'] + results['failed'])))

        if args.jacquard_output:
            logging.info('Jacquard output saved to {}'.format(jo_fn))

        del net
        torch.cuda.empty_cache()

evaluate.py
andvg3/Grasp-Anything
def get_network(network_name):
    network_name = network_name.lower()
    # Original GR-ConvNet
    if network_name == 'grconvnet':
        from .grconvnet import GenerativeResnet
        return GenerativeResnet
    # Configurable GR-ConvNet with multiple dropouts
    elif network_name == 'grconvnet2':
        from .grconvnet2 import GenerativeResnet
        return GenerativeResnet
    # Configurable GR-ConvNet with dropout at the end
    elif network_name == 'grconvnet3':
        from .grconvnet3 import GenerativeResnet
        return GenerativeResnet
    # Inverted GR-ConvNet
    elif network_name == 'grconvnet4':
        from .grconvnet4 import GenerativeResnet
        return GenerativeResnet
    elif network_name == 'ragt':
        from .ragt.ragt import RAGT
        return RAGT
    elif network_name == 'ggcnn':
        from .ggcnn import GGCNN
        return GGCNN
    else:
        raise NotImplementedError('Network {} is not implemented'.format(network_name))

inference/models/__init__.py
andvg3/Grasp-Anything
import torch.nn as nn
import torch.nn.functional as F

from inference.models.grasp_model import GraspModel

filter_sizes = [32, 16, 8, 8, 16, 32]
kernel_sizes = [9, 5, 3, 3, 5, 9]
strides = [3, 2, 2, 2, 2, 3]


class GGCNN(GraspModel):
    """
    GG-CNN
    Equivalient to the Keras Model used in the RSS Paper (https://arxiv.org/abs/1804.05172)
    """
    def __init__(self, input_channels=4, output_channels=1, channel_size=32, dropout=False, prob=0.0):
        super().__init__()
        self.conv1 = nn.Conv2d(input_channels, filter_sizes[0], kernel_sizes[0], stride=strides[0], padding=3)
        self.conv2 = nn.Conv2d(filter_sizes[0], filter_sizes[1], kernel_sizes[1], stride=strides[1], padding=2)
        self.conv3 = nn.Conv2d(filter_sizes[1], filter_sizes[2], kernel_sizes[2], stride=strides[2], padding=1)
        self.convt1 = nn.ConvTranspose2d(filter_sizes[2], filter_sizes[3], kernel_sizes[3], stride=strides[3], padding=1, output_padding=1)
        self.convt2 = nn.ConvTranspose2d(filter_sizes[3], filter_sizes[4], kernel_sizes[4], stride=strides[4], padding=2, output_padding=1)
        self.convt3 = nn.ConvTranspose2d(filter_sizes[4], filter_sizes[5], kernel_sizes[5], stride=strides[5], padding=5, output_padding=1)

        self.pos_output = nn.Conv2d(filter_sizes[5], 1, kernel_size=2)
        self.cos_output = nn.Conv2d(filter_sizes[5], 1, kernel_size=2)
        self.sin_output = nn.Conv2d(filter_sizes[5], 1, kernel_size=2)
        self.width_output = nn.Conv2d(filter_sizes[5], 1, kernel_size=2)

        for m in self.modules():
            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
                nn.init.xavier_uniform_(m.weight, gain=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.convt1(x))
        x = F.relu(self.convt2(x))
        x = F.relu(self.convt3(x))

        pos_output = self.pos_output(x)
        cos_output = self.cos_output(x)
        sin_output = self.sin_output(x)
        width_output = self.width_output(x)

        return pos_output, cos_output, sin_output, width_output

    def compute_loss(self, xc, yc):
        y_pos, y_cos, y_sin, y_width = yc
        pos_pred, cos_pred, sin_pred, width_pred = self(xc)

        p_loss = F.mse_loss(pos_pred, y_pos)
        cos_loss = F.mse_loss(cos_pred, y_cos)
        sin_loss = F.mse_loss(sin_pred, y_sin)
        width_loss = F.mse_loss(width_pred, y_width)

        return {
            'loss': p_loss + cos_loss + sin_loss + width_loss,
            'losses': {
                'p_loss': p_loss,
                'cos_loss': cos_loss,
                'sin_loss': sin_loss,
                'width_loss': width_loss
            },
            'pred': {
                'pos': pos_pred,
                'cos': cos_pred,
                'sin': sin_pred,
                'width': width_pred
            }
        }
inference/models/ggcnn.py
andvg3/Grasp-Anything
import numpy as np


anchor_thetas = [x*0.2094 for x in range(15)]
# anchor宽
anchor_w = 85.72
# anchor高
anchor_h = 19.15
# 每个grid cell 的anchor数
num_anchors = 3
# 输出层下采样次数
times_of_down_sampling = 5
# 输入图像尺寸
img_size = 416
# 防止角度偏移量为0
Anchor_eps = 0.000001


field_of_grid_cell = 2 ** times_of_down_sampling
num_grid_cell = int(img_size / field_of_grid_cell)
theta_margin = 180 / num_anchors


# if __name__ == '__main__':
#     print(field_of_grid_cell)
#     print(anchor_thetas)
#     print(3//0.2094)
#     a = np.arange(16).reshape((4, 4))
#     print(33.2%16.4)
#     a = np.arange(10*26*26*15*6).reshape((10, 26, 26, 15, 6))
#     b = []
#     for i in a:
#         b.append(i)
#     # b = np.array(b)
#     print(type(b[0]))
#     # print(b.shape)
#     c = np.arange(16).reshape((4, 4))
#     d = []
#     for i in c:
#         d.append(i)
#     print(d)
#     print(np.array(d))

inference/models/ragt/Anchor.py
andvg3/Grasp-Anything
"""
original code from apple:
https://github.com/apple/ml-cvnets/blob/main/cvnets/models/classification/mobilevit.py
"""

from typing import Optional, Tuple, Union, Dict
import math
import torch
import torch.nn as nn
from torch import Tensor
from torch.nn import functional as F

from inference.models.ragt.transformer import TransformerEncoder
from inference.models.ragt.model_config import get_config

from inference.models.ragt.Anchor import num_anchors


def make_divisible(
    v: Union[float, int],
    divisor: Optional[int] = 8,
    min_value: Optional[Union[float, int]] = None,
) -> Union[float, int]:
    """
    This function is taken from the original tf repo.
    It ensures that all layers have a channel number that is divisible by 8
    It can be seen here:
    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
    :param v:
    :param divisor:
    :param min_value:
    :return:
    """
    if min_value is None:
        min_value = divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_v < 0.9 * v:
        new_v += divisor
    return new_v


class ConvLayer(nn.Module):
    """
    Applies a 2D convolution over an input

    Args:
        in_channels (int): :math:`C_{in}` from an expected input of size :math:`(N, C_{in}, H_{in}, W_{in})`
        out_channels (int): :math:`C_{out}` from an expected output of size :math:`(N, C_{out}, H_{out}, W_{out})`
        kernel_size (Union[int, Tuple[int, int]]): Kernel size for convolution.
        stride (Union[int, Tuple[int, int]]): Stride for convolution. Default: 1
        groups (Optional[int]): Number of groups in convolution. Default: 1
        bias (Optional[bool]): Use bias. Default: ``False``
        use_norm (Optional[bool]): Use normalization layer after convolution. Default: ``True``
        use_act (Optional[bool]): Use activation layer after convolution (or convolution and normalization).
                                Default: ``True``

    Shape:
        - Input: :math:`(N, C_{in}, H_{in}, W_{in})`
        - Output: :math:`(N, C_{out}, H_{out}, W_{out})`

    .. note::
        For depth-wise convolution, `groups=C_{in}=C_{out}`.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[int, Tuple[int, int]],
        stride: Optional[Union[int, Tuple[int, int]]] = 1,
        groups: Optional[int] = 1,
        bias: Optional[bool] = False,
        use_norm: Optional[bool] = True,
        use_act: Optional[bool] = True,
    ) -> None:
        super().__init__()

        if isinstance(kernel_size, int):
            kernel_size = (kernel_size, kernel_size)

        if isinstance(stride, int):
            stride = (stride, stride)

        assert isinstance(kernel_size, Tuple)
        assert isinstance(stride, Tuple)

        padding = (
            int((kernel_size[0] - 1) / 2),
            int((kernel_size[1] - 1) / 2),
        )

        block = nn.Sequential()

        conv_layer = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            groups=groups,
            padding=padding,
            bias=bias
        )

        block.add_module(name="conv", module=conv_layer)

        if use_norm:
            norm_layer = nn.BatchNorm2d(num_features=out_channels, momentum=0.1)
            block.add_module(name="norm", module=norm_layer)

        if use_act:
            act_layer = nn.SiLU()
            block.add_module(name="act", module=act_layer)

        self.block = block

    def forward(self, x: Tensor) -> Tensor:
        return self.block(x)


class InvertedResidual(nn.Module):
    """
    This class implements the inverted residual block, as described in `MobileNetv2 <https://arxiv.org/abs/1801.04381>`_ paper

    Args:
        in_channels (int): :math:`C_{in}` from an expected input of size :math:`(N, C_{in}, H_{in}, W_{in})`
        out_channels (int): :math:`C_{out}` from an expected output of size :math:`(N, C_{out}, H_{out}, W_{out)`
        stride (int): Use convolutions with a stride. Default: 1
        expand_ratio (Union[int, float]): Expand the input channels by this factor in depth-wise conv
        skip_connection (Optional[bool]): Use skip-connection. Default: True

    Shape:
        - Input: :math:`(N, C_{in}, H_{in}, W_{in})`
        - Output: :math:`(N, C_{out}, H_{out}, W_{out})`

    .. note::
        If `in_channels =! out_channels` and `stride > 1`, we set `skip_connection=False`

    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        stride: int,
        expand_ratio: Union[int, float],
        skip_connection: Optional[bool] = True,
    ) -> None:
        assert stride in [1, 2]
        hidden_dim = make_divisible(int(round(in_channels * expand_ratio)), 8)

        super().__init__()

        block = nn.Sequential()
        if expand_ratio != 1:
            block.add_module(
                name="exp_1x1",
                module=ConvLayer(
                    in_channels=in_channels,
                    out_channels=hidden_dim,
                    kernel_size=1
                ),
            )

        block.add_module(
            name="conv_3x3",
            module=ConvLayer(
                in_channels=hidden_dim,
                out_channels=hidden_dim,
                stride=stride,
                kernel_size=3,
                groups=hidden_dim
            ),
        )

        block.add_module(
            name="red_1x1",
            module=ConvLayer(
                in_channels=hidden_dim,
                out_channels=out_channels,
                kernel_size=1,
                use_act=False,
                use_norm=True,
            ),
        )

        self.block = block
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.exp = expand_ratio
        self.stride = stride
        self.use_res_connect = (
            self.stride == 1 and in_channels == out_channels and skip_connection
        )

    def forward(self, x: Tensor, *args, **kwargs) -> Tensor:
        if self.use_res_connect:
            return x + self.block(x)
        else:
            return self.block(x)


class MobileViTBlock(nn.Module):
    """
    This class defines the `MobileViT block <https://arxiv.org/abs/2110.02178?context=cs.LG>`_

    Args:
        opts: command line arguments
        in_channels (int): :math:`C_{in}` from an expected input of size :math:`(N, C_{in}, H, W)`
        transformer_dim (int): Input dimension to the transformer unit
        ffn_dim (int): Dimension of the FFN block
        n_transformer_blocks (int): Number of transformer blocks. Default: 2
        head_dim (int): Head dimension in the multi-head attention. Default: 32
        attn_dropout (float): Dropout in multi-head attention. Default: 0.0
        dropout (float): Dropout rate. Default: 0.0
        ffn_dropout (float): Dropout between FFN layers in transformer. Default: 0.0
        patch_h (int): Patch height for unfolding operation. Default: 8
        patch_w (int): Patch width for unfolding operation. Default: 8
        transformer_norm_layer (Optional[str]): Normalization layer in the transformer block. Default: layer_norm
        conv_ksize (int): Kernel size to learn local representations in MobileViT block. Default: 3
        no_fusion (Optional[bool]): Do not combine the input and output feature maps. Default: False
    """

    def __init__(
        self,
        in_channels: int,
        transformer_dim: int,
        ffn_dim: int,
        n_transformer_blocks: int = 2,
        head_dim: int = 32,
        attn_dropout: float = 0.0,
        dropout: float = 0.0,
        ffn_dropout: float = 0.0,
        patch_h: int = 8,
        patch_w: int = 8,
        conv_ksize: Optional[int] = 3,
        *args,
        **kwargs
    ) -> None:
        super().__init__()

        conv_3x3_in = ConvLayer(
            in_channels=in_channels,
            out_channels=in_channels,
            kernel_size=conv_ksize,
            stride=1
        )
        conv_1x1_in = ConvLayer(
            in_channels=in_channels,
            out_channels=transformer_dim,
            kernel_size=1,
            stride=1,
            use_norm=False,
            use_act=False
        )

        conv_1x1_out = ConvLayer(
            in_channels=transformer_dim,
            out_channels=in_channels,
            kernel_size=1,
            stride=1
        )
        conv_3x3_out = ConvLayer(
            in_channels=2 * in_channels,
            out_channels=in_channels,
            kernel_size=conv_ksize,
            stride=1
        )

        self.local_rep = nn.Sequential()
        self.local_rep.add_module(name="conv_3x3", module=conv_3x3_in)
        self.local_rep.add_module(name="conv_1x1", module=conv_1x1_in)

        assert transformer_dim % head_dim == 0
        num_heads = transformer_dim // head_dim

        global_rep = [
            TransformerEncoder(
                embed_dim=transformer_dim,
                ffn_latent_dim=ffn_dim,
                num_heads=num_heads,
                attn_dropout=attn_dropout,
                dropout=dropout,
                ffn_dropout=ffn_dropout
            )
            for _ in range(n_transformer_blocks)
        ]
        global_rep.append(nn.LayerNorm(transformer_dim))
        self.global_rep = nn.Sequential(*global_rep)

        self.conv_proj = conv_1x1_out
        self.fusion = conv_3x3_out

        self.patch_h = patch_h
        self.patch_w = patch_w
        self.patch_area = self.patch_w * self.patch_h

        self.cnn_in_dim = in_channels
        self.cnn_out_dim = transformer_dim
        self.n_heads = num_heads
        self.ffn_dim = ffn_dim
        self.dropout = dropout
        self.attn_dropout = attn_dropout
        self.ffn_dropout = ffn_dropout
        self.n_blocks = n_transformer_blocks
        self.conv_ksize = conv_ksize

    def unfolding(self, x: Tensor) -> Tuple[Tensor, Dict]:
        patch_w, patch_h = self.patch_w, self.patch_h
        patch_area = patch_w * patch_h
        batch_size, in_channels, orig_h, orig_w = x.shape

        new_h = int(math.ceil(orig_h / self.patch_h) * self.patch_h)
        new_w = int(math.ceil(orig_w / self.patch_w) * self.patch_w)

        interpolate = False
        if new_w != orig_w or new_h != orig_h:
            # Note: Padding can be done, but then it needs to be handled in attention function.
            x = F.interpolate(x, size=(new_h, new_w), mode="bilinear", align_corners=False)
            interpolate = True

        # number of patches along width and height
        num_patch_w = new_w // patch_w  # n_w
        num_patch_h = new_h // patch_h  # n_h
        num_patches = num_patch_h * num_patch_w  # N

        # [B, C, H, W] -> [B * C * n_h, p_h, n_w, p_w]
        x = x.reshape(batch_size * in_channels * num_patch_h, patch_h, num_patch_w, patch_w)
        # [B * C * n_h, p_h, n_w, p_w] -> [B * C * n_h, n_w, p_h, p_w]
        x = x.transpose(1, 2)
        # [B * C * n_h, n_w, p_h, p_w] -> [B, C, N, P] where P = p_h * p_w and N = n_h * n_w
        x = x.reshape(batch_size, in_channels, num_patches, patch_area)
        # [B, C, N, P] -> [B, P, N, C]
        x = x.transpose(1, 3)
        # [B, P, N, C] -> [BP, N, C]
        x = x.reshape(batch_size * patch_area, num_patches, -1)

        info_dict = {
            "orig_size": (orig_h, orig_w),
            "batch_size": batch_size,
            "interpolate": interpolate,
            "total_patches": num_patches,
            "num_patches_w": num_patch_w,
            "num_patches_h": num_patch_h,
        }

        return x, info_dict

    def folding(self, x: Tensor, info_dict: Dict) -> Tensor:
        n_dim = x.dim()
        assert n_dim == 3, "Tensor should be of shape BPxNxC. Got: {}".format(
            x.shape
        )
        # [BP, N, C] --> [B, P, N, C]
        x = x.contiguous().view(
            info_dict["batch_size"], self.patch_area, info_dict["total_patches"], -1
        )

        batch_size, pixels, num_patches, channels = x.size()
        num_patch_h = info_dict["num_patches_h"]
        num_patch_w = info_dict["num_patches_w"]

        # [B, P, N, C] -> [B, C, N, P]
        x = x.transpose(1, 3)
        # [B, C, N, P] -> [B*C*n_h, n_w, p_h, p_w]
        x = x.reshape(batch_size * channels * num_patch_h, num_patch_w, self.patch_h, self.patch_w)
        # [B*C*n_h, n_w, p_h, p_w] -> [B*C*n_h, p_h, n_w, p_w]
        x = x.transpose(1, 2)
        # [B*C*n_h, p_h, n_w, p_w] -> [B, C, H, W]
        x = x.reshape(batch_size, channels, num_patch_h * self.patch_h, num_patch_w * self.patch_w)
        if info_dict["interpolate"]:
            x = F.interpolate(
                x,
                size=info_dict["orig_size"],
                mode="bilinear",
                align_corners=False,
            )
        return x

    def forward(self, x: Tensor) -> Tensor:
        res = x

        fm = self.local_rep(x)

        # convert feature map to patches
        patches, info_dict = self.unfolding(fm)

        # learn global representations
        for transformer_layer in self.global_rep:
            patches = transformer_layer(patches)

        # [B x Patch x Patches x C] -> [B x C x Patches x Patch]
        fm = self.folding(x=patches, info_dict=info_dict)

        fm = self.conv_proj(fm)

        fm = self.fusion(torch.cat((res, fm), dim=1))
        return fm


# 将四倍下采样特征与五倍下采样特征融合
# 将模型改为单分支输出预测
# S版本，4倍下采样输出128通道，5倍下采样输出160通道
class ResidualConvBlock(nn.Module):
    def __init__(self,
                 in_channels: int,
                 middle_channels: int,
                 ) -> None:
        super().__init__()
        self.res_conv_block = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=middle_channels,
                      kernel_size=(3, 3), stride=(1, 1), padding=1,
                      bias=False),
            nn.BatchNorm2d(num_features=middle_channels),
            nn.SiLU(),
            nn.Conv2d(in_channels=middle_channels, out_channels=in_channels,
                      kernel_size=(3, 3), stride=(1, 1), padding=1,
                      bias=False),
            nn.BatchNorm2d(num_features=in_channels),
        )
        
    def forward(self, x: Tensor) -> Tensor:
        x = F.silu(self.res_conv_block(x) + x)
        return x


class ConvSet(nn.Module):
    def __init__(self,
                 in_channels: int,
                 middle_channels: int,
                 out_channels: int,
                 ) -> None:
        super(ConvSet, self).__init__()
        self.convset = nn.Sequential()
        self.convset.add_module(name='res_conv_block1',
                                module=ResidualConvBlock(in_channels=in_channels, middle_channels=middle_channels))
        self.convset.add_module(name='res_conv_block2',
                                module=ResidualConvBlock(in_channels=in_channels, middle_channels=middle_channels))
        self.convset.add_module(name='res_conv_block3',
                                module=ResidualConvBlock(in_channels=in_channels, middle_channels=middle_channels))
        self.convset.add_module(name='res_conv_block4',
                                module=ResidualConvBlock(in_channels=in_channels, middle_channels=middle_channels))
        self.convset.add_module(name='conv_3x3',
                                module=nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
                                                 kernel_size=(3, 3), stride=(1, 1), padding=1,
                                                 bias=False))
        self.convset.add_module(name='bn', module=nn.BatchNorm2d(num_features=out_channels))
        self.convset.add_module(name='SiLU', module=nn.SiLU())

    def forward(self, x: Tensor) -> Tensor:
        x = self.convset(x)
        return x


class ConvDownSampling(nn.Module):
    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 ) -> None:
        super().__init__()
        self.convdownsampling = nn.Sequential()
        # 3x3卷积
        self.convdownsampling.add_module(name='conv1_3x3', module=nn.Conv2d(in_channels=in_channels,
                                                                            out_channels=2*in_channels,
                                                                            kernel_size=(3, 3),
                                                                            stride=(2, 2),
                                                                            padding=1,
                                                                            bias=False))
        self.convdownsampling.add_module(name='bn1', module=nn.BatchNorm2d(num_features=2*in_channels))
        self.convdownsampling.add_module(name='SiLU1', module=nn.SiLU())
        # 3x3卷积
        self.convdownsampling.add_module(name='conv2_3x3', module=nn.Conv2d(in_channels=2*in_channels,
                                                                            out_channels=out_channels,
                                                                            kernel_size=(3, 3),
                                                                            stride=(1, 1),
                                                                            padding=1,
                                                                            bias=False))
        self.convdownsampling.add_module(name='bn2', module=nn.BatchNorm2d(num_features=out_channels))
        self.convdownsampling.add_module(name='SiLU2', module=nn.SiLU())

    def forward(self, x: Tensor) -> Tensor:
        x = self.convdownsampling(x)
        return x


class PositionalEncoding(nn.Module):
    def __init__(self,
                 d_model: int,
                 max_len: int,
                 ) -> None:
        super(PositionalEncoding, self).__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        # # pe: [max_len, d_model] -> [1, max_len, d_model] -> [max_len, 1, d_model]
        # pe = pe.unsqueeze(0).transpose(0, 1)
        # pe: [max_len, d_model] -> [1, max_len, d_model]
        pe = pe.unsqueeze(0)
        # pe.requires_grad = False
        self.register_buffer('pe', pe)

    def forward(self, x: Tensor) -> Tensor:
        # # x: [max_len, N, d_model]
        # return x + self.pe[:x.size(0), :]
        # x: [N, max_len, d_model]
        # broadcast: pe: [1, max_len, d_model] -> [N, max_len, d_model]
        return x + self.pe[:, :x.size(1), :]


class GlobalTransformerEncoder(nn.Module):
    def __init__(self,
                 # in_channels: int,
                 d_model: int,
                 n_head: int,
                 ffn_dim: int,
                 encoder_layers: int,
                 ) -> None:
        super().__init__()

        # if in_channels != d_model:
        #     # 调整维度到d_model
        #     self.conv1_1x1 = nn.Conv2d(in_channels=in_channels, out_channels=d_model,
        #                             kernel_size=(1, 1), stride=(1, 1), padding=0)
        #     # 调整维度到in_channels
        #     self.conv2_1x1 = nn.Conv2d(in_channels=d_model, out_channels=in_channels,
        #                             kernel_size=(1, 1), stride=(1, 1), padding=0)
            
        # 位置编码
        self.position_encoding = PositionalEncoding(d_model=d_model, max_len=169)
        # Transformer Encoder全局表征
        self.transformer_encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=n_head,
                                                                    dim_feedforward=ffn_dim,
                                                                    batch_first=True,
                                                                    norm_first=True)
        self.transformer_encoder = nn.TransformerEncoder(self.transformer_encoder_layer,
                                                         num_layers=encoder_layers, norm=nn.LayerNorm(d_model))
        
    def sequentialize(self, x: Tensor) -> Tuple[Tensor, dict]:
        N, C, H, W = x.shape
        # [N C H W] -> [N H W C]
        x = x.permute(0, 2, 3, 1)
        # [N H W C] -> [N H*W C]
        x = x.reshape(N, H*W, C)
        shape_dict = {
            'origin_N': N,
            'origin_C': C,
            'origin_H': H,
            'origin_W': W,
        }
        return x, shape_dict

    def unsequentialize(self, x: Tensor, dim_dict: dict) -> Tensor:
        # [N H*W C] -> [N H W C]
        x = x.contiguous().view(dim_dict['origin_N'], dim_dict['origin_H'], dim_dict['origin_W'], dim_dict['origin_C'])
        # [N H W C] -> [N C H W]
        x = x.permute(0, 3, 1, 2)
        return x

    def forward(self, x: Tensor) -> Tensor:
        # x = self.conv1_1x1(x)
        x, shape_dict = self.sequentialize(x)
        x = self.position_encoding(x)
        x = self.transformer_encoder(x)
        x = self.unsequentialize(x, dim_dict=shape_dict)
        # x = self.conv2_1x1(x)
        # x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
        return x


# class TransposeConv(nn.Module):
#     def __init__(self,
#                  in_channels,
#                  out_channels
#                  ):
#         super(TransposeConv, self).__init__()
#         self.transpose_conv = nn.Sequential()
#         self.transpose_conv.add_module(name='conv_3x3', module=nn.Conv2d(in_channels=in_channels,
#                                                                          out_channels=in_channels,
#                                                                          kernel_size=(3, 3), stride=(1, 1), padding=1,
#                                                                          bias=False))
#         self.transpose_conv.add_module(name='bn1', module=nn.BatchNorm2d(in_channels))
#         self.transpose_conv.add_module(name='SiLU1', module=nn.SiLU())
#         self.transpose_conv.add_module(name='conv_1x1', module=nn.Conv2d(in_channels=in_channels,
#                                                                          out_channels=out_channels,
#                                                                          kernel_size=(1, 1), stride=(1, 1), padding=0,
#                                                                          bias=False))
#         self.transpose_conv.add_module(name='bn2', module=nn.BatchNorm2d(out_channels))
#         self.transpose_conv.add_module(name='SiLU2', module=nn.SiLU())
#         self.transpose_conv.add_module(name='transpose_conv', module=nn.ConvTranspose2d(in_channels=out_channels,
#                                                                                         out_channels=out_channels,
#                                                                                         kernel_size=(2, 2),
#                                                                                         stride=(2, 2), padding=(0, 0),
#                                                                                         output_padding=(0, 0)))
#         self.transpose_conv.add_module(name='bn3', module=nn.BatchNorm2d(out_channels))
#         self.transpose_conv.add_module(name='SiLU3', module=nn.SiLU())

#     def forward(self, x):
#         x = self.transpose_conv(x)
#         return x


# class SpatialPyramidPoolingFast(nn.Module):
#     def __init__(
#             self,
#             in_channels
#     ):
#         super(SpatialPyramidPoolingFast, self).__init__()
#         # 降维1*1卷积
#         self.conv1_1_1 = nn.Conv2d(in_channels=in_channels, out_channels=in_channels//4,
#                                    kernel_size=(1, 1), stride=(1, 1),
#                                    bias=False)
#         self.bn1 = nn.BatchNorm2d(num_features=in_channels//4)
#         self.SiLU1 = nn.SiLU()
#         # 最大池化
#         self.MaxPool_5_5 = nn.MaxPool2d(kernel_size=(5, 5), stride=(1, 1), padding=2)
#         self.MaxPool1_3_3 = nn.MaxPool2d(kernel_size=(3, 3), stride=(1, 1), padding=1)
#         self.MaxPool2_3_3 = nn.MaxPool2d(kernel_size=(3, 3), stride=(1, 1), padding=1)
#         # 特征融合
#         self.conv2_1_1 = nn.Conv2d(in_channels=in_channels, out_channels=in_channels,
#                                    kernel_size=(1, 1), stride=(1, 1),
#                                    bias=False)
#         self.bn2 = nn.BatchNorm2d(num_features=in_channels)
#         self.SiLU2 = nn.SiLU()

#     def forward(self, x):
#         x1 = self.conv1_1_1(x)
#         x1 = self.bn1(x1)
#         x1 = self.SiLU1(x1)
#         x2 = self.MaxPool_5_5(x1)
#         x3 = self.MaxPool1_3_3(x2)
#         x4 = self.MaxPool2_3_3(x3)
#         x = torch.concat([x1, x2, x3, x4], dim=1)
#         x = self.conv2_1_1(x)
#         x = self.bn2(x)
#         x = self.SiLU2(x)
#         return x


class Detector(nn.Module):
    def __init__(self,
                 in_channels: int,
                 middle_channels: int,
                 out_channels: int,
                 ) -> None:
        super(Detector, self).__init__()
        self.detector = nn.Sequential()
        self.detector.add_module(name='res_conv_block1',
                                 module=ResidualConvBlock(in_channels=in_channels, middle_channels=middle_channels))
        self.detector.add_module(name='res_conv_block2',
                                 module=ResidualConvBlock(in_channels=in_channels, middle_channels=middle_channels))
        self.detector.add_module(name='res_conv_block3',
                                 module=ResidualConvBlock(in_channels=in_channels, middle_channels=middle_channels))
        self.detector.add_module(name='res_conv_block4',
                                 module=ResidualConvBlock(in_channels=in_channels, middle_channels=middle_channels))
        self.detector.add_module(name='res_conv_block5',
                                 module=ResidualConvBlock(in_channels=in_channels, middle_channels=middle_channels))
        # 输出
        self.detector.add_module(name='conv2_1x1', module=nn.Conv2d(in_channels=in_channels,
                                                                    out_channels=out_channels,
                                                                    kernel_size=(1, 1), stride=(1, 1), padding=0,
                                                                    bias=True))

    def forward(self, x: Tensor) -> Tensor:
        x = self.detector(x)
        return x


class MobileViT(nn.Module):
    """
    This class implements the `MobileViT architecture <https://arxiv.org/abs/2110.02178?context=cs.LG>`_
    """
    def __init__(self, model_cfg: Dict, num_classes: int = 1000):
        super().__init__()

        image_channels = 3
        out_channels = 16

        self.conv_1 = ConvLayer(
            in_channels=image_channels,
            out_channels=out_channels,
            kernel_size=3,
            stride=2
        )

        self.layer_1, out_channels = self._make_layer(input_channel=out_channels, cfg=model_cfg["layer1"])
        self.layer_2, out_channels = self._make_layer(input_channel=out_channels, cfg=model_cfg["layer2"])
        self.layer_3, out_channels = self._make_layer(input_channel=out_channels, cfg=model_cfg["layer3"])
        self.layer_4, out_channels = self._make_layer(input_channel=out_channels, cfg=model_cfg["layer4"])
        self.layer_5, out_channels = self._make_layer(input_channel=out_channels, cfg=model_cfg["layer5"])

        # exp_channels = min(model_cfg["last_layer_exp_factor"] * out_channels, 960)
        # self.conv_1x1_exp = ConvLayer(
        #     in_channels=out_channels,
        #     out_channels=exp_channels,
        #     kernel_size=1
        # )
        #
        # self.classifier = nn.Sequential()
        # self.classifier.add_module(name="global_pool", module=nn.AdaptiveAvgPool2d(1))
        # self.classifier.add_module(name="flatten", module=nn.Flatten())
        # if 0.0 < model_cfg["cls_dropout"] < 1.0:
        #     self.classifier.add_module(name="dropout", module=nn.Dropout(p=model_cfg["cls_dropout"]))
        # self.classifier.add_module(name="fc", module=nn.Linear(in_features=exp_channels, out_features=num_classes))

        # 信息融合、检测模块
        self.conv_downsampling1 = ConvDownSampling(in_channels=96, out_channels=128)
        self.convset1 = ConvSet(in_channels=256, middle_channels=128, out_channels=120)
        self.conv_downsampling2 = ConvDownSampling(in_channels=120, out_channels=160)
        self.convset2 = ConvSet(in_channels=320, middle_channels=160, out_channels=240)
        self.transformer_encoder = GlobalTransformerEncoder(d_model=240, n_head=4, ffn_dim=480, encoder_layers=4)
        self.detector = Detector(in_channels=240, middle_channels=120, out_channels=num_anchors*6)

        # weight init
        self.apply(self.init_parameters)

    def _make_layer(self, input_channel, cfg: Dict) -> Tuple[nn.Sequential, int]:
        block_type = cfg.get("block_type", "mobilevit")
        if block_type.lower() == "mobilevit":
            return self._make_mit_layer(input_channel=input_channel, cfg=cfg)
        else:
            return self._make_mobilenet_layer(input_channel=input_channel, cfg=cfg)

    @staticmethod
    def _make_mobilenet_layer(input_channel: int, cfg: Dict) -> Tuple[nn.Sequential, int]:
        output_channels = cfg.get("out_channels")
        num_blocks = cfg.get("num_blocks", 2)
        expand_ratio = cfg.get("expand_ratio", 4)
        block = []

        for i in range(num_blocks):
            stride = cfg.get("stride", 1) if i == 0 else 1

            layer = InvertedResidual(
                in_channels=input_channel,
                out_channels=output_channels,
                stride=stride,
                expand_ratio=expand_ratio
            )
            block.append(layer)
            input_channel = output_channels

        return nn.Sequential(*block), input_channel

    @staticmethod
    def _make_mit_layer(input_channel: int, cfg: Dict) -> Tuple[nn.Sequential, int]:
        stride = cfg.get("stride", 1)
        block = []

        if stride == 2:
            layer = InvertedResidual(
                in_channels=input_channel,
                out_channels=cfg.get("out_channels"),
                stride=stride,
                expand_ratio=cfg.get("mv_expand_ratio", 4)
            )

            block.append(layer)
            input_channel = cfg.get("out_channels")

        transformer_dim = cfg["transformer_channels"]
        ffn_dim = cfg.get("ffn_dim")
        num_heads = cfg.get("num_heads", 4)
        head_dim = transformer_dim // num_heads

        if transformer_dim % head_dim != 0:
            raise ValueError("Transformer input dimension should be divisible by head dimension. "
                             "Got {} and {}.".format(transformer_dim, head_dim))

        block.append(MobileViTBlock(
            in_channels=input_channel,
            transformer_dim=transformer_dim,
            ffn_dim=ffn_dim,
            n_transformer_blocks=cfg.get("transformer_blocks", 1),
            patch_h=cfg.get("patch_h", 2),
            patch_w=cfg.get("patch_w", 2),
            dropout=cfg.get("dropout", 0.1),
            ffn_dropout=cfg.get("ffn_dropout", 0.0),
            attn_dropout=cfg.get("attn_dropout", 0.1),
            head_dim=head_dim,
            conv_ksize=3
        ))

        return nn.Sequential(*block), input_channel

    @staticmethod
    def init_parameters(m):
        if isinstance(m, nn.Conv2d):
            if m.weight is not None:
                nn.init.kaiming_normal_(m.weight, mode="fan_out")
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d)):
            if m.weight is not None:
                nn.init.ones_(m.weight)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, (nn.Linear,)):
            if m.weight is not None:
                nn.init.trunc_normal_(m.weight, mean=0.0, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        else:
            pass

    def forward(self, x: Tensor) -> Tensor:
        x = self.conv_1(x)
        x = self.layer_1(x)
        x = self.layer_2(x)

        x = self.layer_3(x)  # 下采样率8，输出维度：52x52x96
        x1 = self.layer_4(x)  # 下采样率16，输出维度：26x26x128
        x2 = self.layer_5(x1)  # 下采样率32，输出维度：13x13x160

        x = self.conv_downsampling1(x)
        x = torch.concat([x, x1], dim=1)
        x = self.convset1(x)
        x = self.conv_downsampling2(x)
        x = torch.concat([x, x2], dim=1)
        x = self.convset2(x)
        x = self.transformer_encoder(x)
        x = self.detector(x)
        return x


def mobile_vit_xx_small(num_classes: int = 1000):
    # pretrain weight link
    # https://docs-assets.developer.apple.com/ml-research/models/cvnets/classification/mobilevit_xxs.pt
    config = get_config("xx_small")
    m = MobileViT(config, num_classes=num_classes)
    return m


def mobile_vit_x_small(num_classes: int = 1000):
    # pretrain weight link
    # https://docs-assets.developer.apple.com/ml-research/models/cvnets/classification/mobilevit_xs.pt
    config = get_config("x_small")
    m = MobileViT(config, num_classes=num_classes)
    return m


def mobile_vit_small(num_classes: int = 1000):
    # pretrain weight link
    # https://docs-assets.developer.apple.com/ml-research/models/cvnets/classification/mobilevit_s.pt
    config = get_config("small")
    m = MobileViT(config, num_classes=num_classes)
    return m


def get_model():
    model = mobile_vit_small()
    return model

inference/models/ragt/mobile_vit.py
andvg3/Grasp-Anything
def get_config(mode: str = "xxs") -> dict:
    if mode == "xx_small":
        mv2_exp_mult = 2
        config = {
            "layer1": {
                "out_channels": 16,
                "expand_ratio": mv2_exp_mult,
                "num_blocks": 1,
                "stride": 1,
                "block_type": "mv2",
            },
            "layer2": {
                "out_channels": 24,
                "expand_ratio": mv2_exp_mult,
                "num_blocks": 3,
                "stride": 2,
                "block_type": "mv2",
            },
            "layer3": {  # 28x28
                "out_channels": 48,
                "transformer_channels": 64,
                "ffn_dim": 128,
                "transformer_blocks": 2,
                "patch_h": 2,  # 8,
                "patch_w": 2,  # 8,
                "stride": 2,
                "mv_expand_ratio": mv2_exp_mult,
                "num_heads": 4,
                "block_type": "mobilevit",
            },
            "layer4": {  # 14x14
                "out_channels": 64,
                "transformer_channels": 80,
                "ffn_dim": 160,
                "transformer_blocks": 4,
                "patch_h": 2,  # 4,
                "patch_w": 2,  # 4,
                "stride": 2,
                "mv_expand_ratio": mv2_exp_mult,
                "num_heads": 4,
                "block_type": "mobilevit",
            },
            "layer5": {  # 7x7
                "out_channels": 80,
                "transformer_channels": 96,
                "ffn_dim": 192,
                "transformer_blocks": 3,
                "patch_h": 2,
                "patch_w": 2,
                "stride": 2,
                "mv_expand_ratio": mv2_exp_mult,
                "num_heads": 4,
                "block_type": "mobilevit",
            },
            "last_layer_exp_factor": 4,
            "cls_dropout": 0.1
        }
    elif mode == "x_small":
        mv2_exp_mult = 4
        config = {
            "layer1": {
                "out_channels": 32,
                "expand_ratio": mv2_exp_mult,
                "num_blocks": 1,
                "stride": 1,
                "block_type": "mv2",
            },
            "layer2": {
                "out_channels": 48,
                "expand_ratio": mv2_exp_mult,
                "num_blocks": 3,
                "stride": 2,
                "block_type": "mv2",
            },
            "layer3": {  # 28x28
                "out_channels": 64,
                "transformer_channels": 96,
                "ffn_dim": 192,
                "transformer_blocks": 2,
                "patch_h": 2,
                "patch_w": 2,
                "stride": 2,
                "mv_expand_ratio": mv2_exp_mult,
                "num_heads": 4,
                "block_type": "mobilevit",
            },
            "layer4": {  # 14x14
                "out_channels": 80,
                "transformer_channels": 120,
                "ffn_dim": 240,
                "transformer_blocks": 4,
                "patch_h": 2,
                "patch_w": 2,
                "stride": 2,
                "mv_expand_ratio": mv2_exp_mult,
                "num_heads": 4,
                "block_type": "mobilevit",
            },
            "layer5": {  # 7x7
                "out_channels": 96,
                "transformer_channels": 144,
                "ffn_dim": 288,
                "transformer_blocks": 3,
                "patch_h": 2,
                "patch_w": 2,
                "stride": 2,
                "mv_expand_ratio": mv2_exp_mult,
                "num_heads": 4,
                "block_type": "mobilevit",
            },
            "last_layer_exp_factor": 4,
            "cls_dropout": 0.1
        }
    elif mode == "small":
        mv2_exp_mult = 4
        config = {
            "layer1": {
                "out_channels": 32,
                "expand_ratio": mv2_exp_mult,
                "num_blocks": 1,
                "stride": 1,
                "block_type": "mv2",
            },
            "layer2": {
                "out_channels": 64,
                "expand_ratio": mv2_exp_mult,
                "num_blocks": 3,
                "stride": 2,
                "block_type": "mv2",
            },
            "layer3": {  # 28x28
                "out_channels": 96,
                "transformer_channels": 144,
                "ffn_dim": 288,
                "transformer_blocks": 2,
                "patch_h": 2,
                "patch_w": 2,
                "stride": 2,
                "mv_expand_ratio": mv2_exp_mult,
                "num_heads": 4,
                "block_type": "mobilevit",
            },
            "layer4": {  # 14x14
                "out_channels": 128,
                "transformer_channels": 192,
                "ffn_dim": 384,
                "transformer_blocks": 4,
                "patch_h": 2,
                "patch_w": 2,
                "stride": 2,
                "mv_expand_ratio": mv2_exp_mult,
                "num_heads": 4,
                "block_type": "mobilevit",
            },
            "layer5": {  # 7x7
                "out_channels": 160,
                "transformer_channels": 240,
                "ffn_dim": 480,
                "transformer_blocks": 3,
                "patch_h": 2,
                "patch_w": 2,
                "stride": 2,
                "mv_expand_ratio": mv2_exp_mult,
                "num_heads": 4,
                "block_type": "mobilevit",
            },
            "last_layer_exp_factor": 4,
            "cls_dropout": 0.1
        }
    else:
        raise NotImplementedError

    for k in ["layer1", "layer2", "layer3", "layer4", "layer5"]:
        config[k].update({"dropout": 0.1, "ffn_dropout": 0.0, "attn_dropout": 0.0})

    return config

inference/models/ragt/model_config.py
andvg3/Grasp-Anything
import torch.nn as nn
import torch.nn.functional as F

from inference.models.grasp_model import GraspModel, ResidualBlock
from .mobile_vit import get_model


class RAGT(GraspModel):

    def __init__(self, input_channels=4, output_channels=1, channel_size=18, dropout=False, prob=0.0):
        super(RAGT, self).__init__()
        self.mobile_vit = get_model()

        # Upsampling layers to increase spatial dimensions
        self.upsample_layers = nn.Sequential(
            nn.Upsample(scale_factor=33, mode='bilinear', align_corners=False),
            nn.ReLU()
        )

        self.pos_output = nn.Conv2d(in_channels=channel_size, out_channels=output_channels, kernel_size=2)
        self.cos_output = nn.Conv2d(in_channels=channel_size, out_channels=output_channels, kernel_size=2)
        self.sin_output = nn.Conv2d(in_channels=channel_size, out_channels=output_channels, kernel_size=2)
        self.width_output = nn.Conv2d(in_channels=channel_size, out_channels=output_channels, kernel_size=2)

        self.dropout = dropout
        self.dropout_pos = nn.Dropout(p=prob)
        self.dropout_cos = nn.Dropout(p=prob)
        self.dropout_sin = nn.Dropout(p=prob)
        self.dropout_wid = nn.Dropout(p=prob)

        for m in self.modules():
            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
                nn.init.xavier_uniform_(m.weight, gain=1)

    def forward(self, x_in):
        x = self.mobile_vit(x_in)
        x = self.upsample_layers(x)
        x = x[:,:,:225, :225]

        if self.dropout:
            pos_output = self.pos_output(self.dropout_pos(x))
            cos_output = self.cos_output(self.dropout_cos(x))
            sin_output = self.sin_output(self.dropout_sin(x))
            width_output = self.width_output(self.dropout_wid(x))
        else:
            pos_output = self.pos_output(x)
            cos_output = self.cos_output(x)
            sin_output = self.sin_output(x)
            width_output = self.width_output(x)

        return pos_output, cos_output, sin_output, width_output

inference/models/ragt/ragt.py
andvg3/Grasp-Anything
from typing import Optional

import torch
import torch.nn as nn
from torch import Tensor


class MultiHeadAttention(nn.Module):
    """
    This layer applies a multi-head self- or cross-attention as described in
    `Attention is all you need <https://arxiv.org/abs/1706.03762>`_ paper

    Args:
        embed_dim (int): :math:`C_{in}` from an expected input of size :math:`(N, P, C_{in})`
        num_heads (int): Number of heads in multi-head attention
        attn_dropout (float): Attention dropout. Default: 0.0
        bias (bool): Use bias or not. Default: ``True``

    Shape:
        - Input: :math:`(N, P, C_{in})` where :math:`N` is batch size, :math:`P` is number of patches,
        and :math:`C_{in}` is input embedding dim
        - Output: same shape as the input

    """

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        attn_dropout: float = 0.0,
        bias: bool = True,
        *args,
        **kwargs
    ) -> None:
        super().__init__()
        if embed_dim % num_heads != 0:
            raise ValueError(
                "Embedding dim must be divisible by number of heads in {}. Got: embed_dim={} and num_heads={}".format(
                    self.__class__.__name__, embed_dim, num_heads
                )
            )

        self.qkv_proj = nn.Linear(in_features=embed_dim, out_features=3 * embed_dim, bias=bias)

        self.attn_dropout = nn.Dropout(p=attn_dropout)
        self.out_proj = nn.Linear(in_features=embed_dim, out_features=embed_dim, bias=bias)

        self.head_dim = embed_dim // num_heads
        self.scaling = self.head_dim ** -0.5
        self.softmax = nn.Softmax(dim=-1)
        self.num_heads = num_heads
        self.embed_dim = embed_dim

    def forward(self, x_q: Tensor) -> Tensor:
        # [N, P, C]
        b_sz, n_patches, in_channels = x_q.shape

        # self-attention
        # [N, P, C] -> [N, P, 3C] -> [N, P, 3, h, c] where C = hc
        qkv = self.qkv_proj(x_q).reshape(b_sz, n_patches, 3, self.num_heads, -1)

        # [N, P, 3, h, c] -> [N, h, 3, P, C]
        qkv = qkv.transpose(1, 3).contiguous()

        # [N, h, 3, P, C] -> [N, h, P, C] x 3
        query, key, value = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]

        query = query * self.scaling

        # [N h, P, c] -> [N, h, c, P]
        key = key.transpose(-1, -2)

        # QK^T
        # [N, h, P, c] x [N, h, c, P] -> [N, h, P, P]
        attn = torch.matmul(query, key)
        attn = self.softmax(attn)
        attn = self.attn_dropout(attn)

        # weighted sum
        # [N, h, P, P] x [N, h, P, c] -> [N, h, P, c]
        out = torch.matmul(attn, value)

        # [N, h, P, c] -> [N, P, h, c] -> [N, P, C]
        out = out.transpose(1, 2).reshape(b_sz, n_patches, -1)
        out = self.out_proj(out)

        return out


class TransformerEncoder(nn.Module):
    """
    This class defines the pre-norm `Transformer encoder <https://arxiv.org/abs/1706.03762>`_
    Args:
        embed_dim (int): :math:`C_{in}` from an expected input of size :math:`(N, P, C_{in})`
        ffn_latent_dim (int): Inner dimension of the FFN
        num_heads (int) : Number of heads in multi-head attention. Default: 8
        attn_dropout (float): Dropout rate for attention in multi-head attention. Default: 0.0
        dropout (float): Dropout rate. Default: 0.0
        ffn_dropout (float): Dropout between FFN layers. Default: 0.0

    Shape:
        - Input: :math:`(N, P, C_{in})` where :math:`N` is batch size, :math:`P` is number of patches,
        and :math:`C_{in}` is input embedding dim
        - Output: same shape as the input
    """

    def __init__(
        self,
        embed_dim: int,
        ffn_latent_dim: int,
        num_heads: Optional[int] = 8,
        attn_dropout: Optional[float] = 0.0,
        dropout: Optional[float] = 0.0,
        ffn_dropout: Optional[float] = 0.0,
        *args,
        **kwargs
    ) -> None:

        super().__init__()

        attn_unit = MultiHeadAttention(
            embed_dim,
            num_heads,
            attn_dropout=attn_dropout,
            bias=True
        )

        self.pre_norm_mha = nn.Sequential(
            nn.LayerNorm(embed_dim),
            attn_unit,
            nn.Dropout(p=dropout)
        )

        self.pre_norm_ffn = nn.Sequential(
            nn.LayerNorm(embed_dim),
            nn.Linear(in_features=embed_dim, out_features=ffn_latent_dim, bias=True),
            nn.SiLU(),
            nn.Dropout(p=ffn_dropout),
            nn.Linear(in_features=ffn_latent_dim, out_features=embed_dim, bias=True),
            nn.Dropout(p=dropout)
        )
        self.embed_dim = embed_dim
        self.ffn_dim = ffn_latent_dim
        self.ffn_dropout = ffn_dropout
        self.std_dropout = dropout

    def forward(self, x: Tensor) -> Tensor:
        # multi-head attention
        res = x
        x = self.pre_norm_mha(x)
        x = x + res

        # feed forward network
        x = x + self.pre_norm_ffn(x)
        return x

inference/models/ragt/transformer.py
andvg3/Grasp-Anything
from ._version import version as __version__
from models.grasp_det_seg.config import load_config
import torch
import models.grasp_det_seg.models as models
from models.grasp_det_seg.algos.detection import PredictionGenerator, ProposalMatcher, DetectionLoss
from models.grasp_det_seg.algos.fpn import RPNAlgoFPN, DetectionAlgoFPN
from models.grasp_det_seg.algos.rpn import AnchorMatcher, ProposalGenerator, RPNLoss
from models.grasp_det_seg.algos.semantic_seg import SemanticSegAlgo, SemanticSegLoss
from models.grasp_det_seg.config import load_config
from models.grasp_det_seg.models.det_seg import DetSegNet, NETWORK_INPUTS
from models.grasp_det_seg.modules.fpn import FPN, FPNBody
from models.grasp_det_seg.modules.heads import RPNHead, FPNSemanticHeadDeeplab, FPNROIHead
from models.grasp_det_seg.utils.misc import norm_act_from_config, freeze_params, NORM_LAYERS, OTHER_LAYERS


def make_config(args):
    args.config = 'models/grasp_det_seg/config/defaults/det_seg_OCID.ini'
    print("Loading configuration from %s", args.config)

    conf = load_config(args.config, args.config)

    return conf

def make_model(config):
    body_config = config["body"]
    fpn_config = config["fpn"]
    rpn_config = config["rpn"]
    roi_config = config["roi"]
    sem_config = config["sem"]
    general_config = config["general"]
    classes = {"total": int(general_config["num_things"]) + int(general_config["num_stuff"]), "stuff":
        int(general_config["num_stuff"]), "thing": int(general_config["num_things"]),
               "semantic": int(general_config["num_semantic"])}
    # BN + activation
    norm_act_static, norm_act_dynamic = norm_act_from_config(body_config)

    # Create backbone
    print("Creating backbone model %s", body_config["body"])
    body_fn = models.__dict__["net_" + body_config["body"]]
    body_params = body_config.getstruct("body_params") if body_config.get("body_params") else {}
    body = body_fn(norm_act=norm_act_static, **body_params)
    if body_config.get("weights"):
        body.load_state_dict(torch.load(body_config["weights"], map_location="cpu"))

    # Freeze parameters
    for n, m in body.named_modules():
        for mod_id in range(1, body_config.getint("num_frozen") + 1):
            if ("mod%d" % mod_id) in n:
                freeze_params(m)

    body_channels = body_config.getstruct("out_channels")

    # Create FPN
    fpn_inputs = fpn_config.getstruct("inputs")
    fpn = FPN([body_channels[inp] for inp in fpn_inputs],
              fpn_config.getint("out_channels"),
              fpn_config.getint("extra_scales"),
              norm_act_static,
              fpn_config["interpolation"])
    body = FPNBody(body, fpn, fpn_inputs)

    # Create RPN
    proposal_generator = ProposalGenerator(rpn_config.getfloat("nms_threshold"),
                                           rpn_config.getint("num_pre_nms_train"),
                                           rpn_config.getint("num_post_nms_train"),
                                           rpn_config.getint("num_pre_nms_val"),
                                           rpn_config.getint("num_post_nms_val"),
                                           rpn_config.getint("min_size"))
    anchor_matcher = AnchorMatcher(rpn_config.getint("num_samples"),
                                   rpn_config.getfloat("pos_ratio"),
                                   rpn_config.getfloat("pos_threshold"),
                                   rpn_config.getfloat("neg_threshold"),
                                   rpn_config.getfloat("void_threshold"))
    rpn_loss = RPNLoss(rpn_config.getfloat("sigma"))
    rpn_algo = RPNAlgoFPN(
        proposal_generator, anchor_matcher, rpn_loss,
        rpn_config.getint("anchor_scale"), rpn_config.getstruct("anchor_ratios"),
        fpn_config.getstruct("out_strides"), rpn_config.getint("fpn_min_level"), rpn_config.getint("fpn_levels"))
    rpn_head = RPNHead(
        fpn_config.getint("out_channels"), len(rpn_config.getstruct("anchor_ratios")), 1,
        rpn_config.getint("hidden_channels"), norm_act_dynamic)

    # Create detection network
    prediction_generator = PredictionGenerator(roi_config.getfloat("nms_threshold"),
                                               roi_config.getfloat("score_threshold"),
                                               roi_config.getint("max_predictions"))
    proposal_matcher = ProposalMatcher(classes,
                                       roi_config.getint("num_samples"),
                                       roi_config.getfloat("pos_ratio"),
                                       roi_config.getfloat("pos_threshold"),
                                       roi_config.getfloat("neg_threshold_hi"),
                                       roi_config.getfloat("neg_threshold_lo"),
                                       roi_config.getfloat("void_threshold"))
    roi_loss = DetectionLoss(roi_config.getfloat("sigma"))
    roi_size = roi_config.getstruct("roi_size")
    roi_algo = DetectionAlgoFPN(
        prediction_generator, proposal_matcher, roi_loss, classes, roi_config.getstruct("bbx_reg_weights"),
        roi_config.getint("fpn_canonical_scale"), roi_config.getint("fpn_canonical_level"), roi_size,
        roi_config.getint("fpn_min_level"), roi_config.getint("fpn_levels"))
    roi_head = FPNROIHead(fpn_config.getint("out_channels"), classes, roi_size, norm_act=norm_act_dynamic)

    # Create semantic segmentation network
    sem_loss = SemanticSegLoss(ohem=sem_config.getfloat("ohem"))
    sem_algo = SemanticSegAlgo(sem_loss, classes["semantic"])
    sem_head = FPNSemanticHeadDeeplab(fpn_config.getint("out_channels"),
                                      sem_config.getint("fpn_min_level"),
                                      sem_config.getint("fpn_levels"),
                                      classes["semantic"],
                                      pooling_size=sem_config.getstruct("pooling_size"),
                                      norm_act=norm_act_static)

    return DetSegNet(body, rpn_head, roi_head, sem_head, rpn_algo, roi_algo, sem_algo, classes)
models/grasp_det_seg/__init__.py
andvg3/Grasp-Anything
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as functional

from models.grasp_det_seg.modules.losses import smooth_l1
from models.grasp_det_seg.utils.bbx import ious, calculate_shift, bbx_overlap, mask_overlap
from models.grasp_det_seg.utils.misc import Empty
from models.grasp_det_seg.utils.nms import nms
from models.grasp_det_seg.utils.parallel import PackedSequence


class PredictionGenerator(nn.Module):
    """Perform NMS-based selection of detections

    Parameters
    ----------
    nms_threshold : float
        IoU threshold for the class-specific NMS
    score_threshold : float
        Minimum class probability for a detection to be kept
    max_predictions : int
        Maximum number of detections to keep for each image
    """

    def __init__(self,
                 nms_threshold,
                 score_threshold,
                 max_predictions):
        super(PredictionGenerator, self).__init__()
        self.nms_threshold = nms_threshold
        self.score_threshold = score_threshold
        self.max_predictions = max_predictions

    @staticmethod
    def _proposals_for_img(proposals, proposals_idx, roi_cls_logits, roi_bbx_logits, img_it):
        relevant = proposals_idx == img_it
        if relevant.any():
            return proposals[relevant], roi_cls_logits[relevant], roi_bbx_logits[relevant]
        else:
            return None, None, None

    def forward(self, boxes, scores):
        """Perform NMS-based selection of detections

        Parameters
        ----------
        boxes : sequence of torch.Tensor
            Sequence of N tensors of class-specific bounding boxes with shapes M_i x C x 4, entries can be None
        scores : sequence of torch.Tensor
            Sequence of N tensors of class probabilities with shapes M_i x (C + 1), entries can be None

        Returns
        -------
        bbx_pred : PackedSequence
            A sequence of N tensors of bounding boxes with shapes S_i x 4, entries are None for images in which no
            detection can be kept according to the selection parameters
        cls_pred : PackedSequence
            A sequence of N tensors of thing class predictions with shapes S_i, entries are None for images in which no
            detection can be kept according to the selection parameters
        obj_pred : PackedSequence
            A sequence of N tensors of detection confidences with shapes S_i, entries are None for images in which no
            detection can be kept according to the selection parameters
        """

        bbx_pred, cls_pred, obj_pred = [], [], []
        for bbx_i, obj_i in zip(boxes, scores):
            try:
                if bbx_i is None or obj_i is None:
                    raise Empty

                # Do NMS separately for each class
                bbx_pred_i, cls_pred_i, obj_pred_i = [], [], []
                for cls_id, (bbx_cls_i, obj_cls_i) in enumerate(zip(torch.unbind(bbx_i, dim=1),
                                                                    torch.unbind(obj_i, dim=1)[1:])):
                    # Filter out low-scoring predictions
                    idx = obj_cls_i > self.score_threshold
                    if not idx.any().item():
                        continue
                    bbx_cls_i = bbx_cls_i[idx]
                    obj_cls_i = obj_cls_i[idx]

                    # Filter out empty predictions
                    idx = (bbx_cls_i[:, 2] > bbx_cls_i[:, 0]) & (bbx_cls_i[:, 3] > bbx_cls_i[:, 1])
                    if not idx.any().item():
                        continue
                    bbx_cls_i = bbx_cls_i[idx]
                    obj_cls_i = obj_cls_i[idx]

                    # Do NMS
                    idx = nms(bbx_cls_i.contiguous(), obj_cls_i.contiguous(), threshold=self.nms_threshold, n_max=-1)
                    if idx.numel() == 0:
                        continue
                    bbx_cls_i = bbx_cls_i[idx]
                    obj_cls_i = obj_cls_i[idx]

                    # Save remaining outputs
                    bbx_pred_i.append(bbx_cls_i)
                    cls_pred_i.append(bbx_cls_i.new_full((bbx_cls_i.size(0),), cls_id, dtype=torch.long))
                    obj_pred_i.append(obj_cls_i)

                # Compact predictions from the classes
                if len(bbx_pred_i) == 0:
                    raise Empty
                bbx_pred_i = torch.cat(bbx_pred_i, dim=0)
                cls_pred_i = torch.cat(cls_pred_i, dim=0)
                obj_pred_i = torch.cat(obj_pred_i, dim=0)

                # Do post-NMS selection (if needed)
                if bbx_pred_i.size(0) > self.max_predictions:
                    _, idx = obj_pred_i.topk(self.max_predictions)
                    bbx_pred_i = bbx_pred_i[idx]
                    cls_pred_i = cls_pred_i[idx]
                    obj_pred_i = obj_pred_i[idx]

                # Save results
                bbx_pred.append(bbx_pred_i)
                cls_pred.append(cls_pred_i)
                obj_pred.append(obj_pred_i)
            except Empty:
                bbx_pred.append(None)
                cls_pred.append(None)
                obj_pred.append(None)

        return PackedSequence(bbx_pred), PackedSequence(cls_pred), PackedSequence(obj_pred)


class ProposalMatcher:
    """Match proposals to ground truth boxes

    Parameters
    ----------
    classes : dict
        Dictionary with the number of classes in the dataset -- expected keys: "total", "stuff", "thing"
    num_samples : int
        Maximum number of non-void proposals to keep for each image
    pos_ratio : float
        Fraction of `num_samples` reserved for positive proposals
    pos_threshold : float
        Minimum IoU threshold to mark a proposal as positive
    neg_threshold_hi : float
        Maximum IoU threshold to mark a proposal as negative / background
    neg_threshold_lo : float
        Minimum IoU threshold to mark a proposal as negative / background
    void_threshold : float
        If not zero, minimum overlap threshold with void regions to mark a proposal as void
    """

    def __init__(self,
                 classes,
                 num_samples=128,
                 pos_ratio=0.25,
                 pos_threshold=0.5,
                 neg_threshold_hi=0.5,
                 neg_threshold_lo=0.0,
                 void_threshold=0.):
        self.num_stuff = classes["stuff"]
        self.num_samples = num_samples
        self.pos_ratio = pos_ratio
        self.pos_threshold = pos_threshold
        self.neg_threshold_hi = neg_threshold_hi
        self.neg_threshold_lo = neg_threshold_lo
        self.void_threshold = void_threshold

    def _subsample(self, pos_idx, neg_idx):
        num_pos = int(self.num_samples * self.pos_ratio)
        pos_idx = torch.nonzero(pos_idx).view(-1)
        if pos_idx.numel() > 0:
            rand_selection = np.random.permutation(pos_idx.numel()).astype(np.int64)
            rand_selection = torch.from_numpy(rand_selection).to(pos_idx.device)

            num_pos = min(num_pos, pos_idx.numel())
            pos_idx = pos_idx[rand_selection[:num_pos]]
        else:
            num_pos = 0
            pos_idx = torch.tensor((), dtype=torch.long, device=pos_idx.device)

        num_neg = self.num_samples - num_pos
        neg_idx = torch.nonzero(neg_idx).view(-1)
        if neg_idx.numel() > 0:
            rand_selection = np.random.permutation(neg_idx.numel()).astype(np.int64)
            rand_selection = torch.from_numpy(rand_selection).to(neg_idx.device)

            num_neg = min(num_neg, neg_idx.numel())
            neg_idx = neg_idx[rand_selection[:num_neg]]
        else:
            neg_idx = torch.tensor((), dtype=torch.long, device=neg_idx.device)

        return pos_idx, neg_idx

    def __call__(self,
                 proposals,
                 bbx,
                 cat,
                 iscrowd):
        """Match proposals to ground truth boxes

        Parameters
        ----------
        proposals : PackedSequence
            A sequence of N tensors with shapes P_i x 4 containing bounding box proposals, entries can be None
        bbx : sequence of torch.Tensor
            A sequence of N tensors with shapes K_i x 4 containing ground truth bounding boxes, entries can be None
        cat : sequence of torch.Tensor
            A sequence of N tensors with shapes K_i containing ground truth instance -> category mappings, entries can
            be None
        iscrowd : sequence of torch.Tensor
            Sequence of N tensors of ground truth crowd regions (shapes H_i x W_i), or ground truth crowd bounding boxes
            (shapes K_i x 4), entries can be None

        Returns
        -------
        out_proposals : PackedSequence
            A sequence of N tensors with shapes S_i x 4 containing the non-void bounding box proposals, entries are None
            for images that do not contain any non-void proposal
        match : PackedSequence
            A sequence of matching results with shape S_i, with the following semantic:
              - match[i, j] == -1: the j-th anchor in image i is negative
              - match[i, j] == k, k >= 0: the j-th anchor in image i is matched to bbx[i][k] (bbox = gt)
        """
        out_proposals = []
        match = []

        for proposals_i, bbx_i_ in zip(proposals, bbx):
            bbx_i = bbx_i_[:, [0, 1, 3, 4]]

            try:
                # Append proposals to ground truth bounding boxes before proceeding
                if bbx_i is not None and proposals_i is not None:
                    proposals_i = torch.cat([bbx_i, proposals_i], dim=0)

                elif bbx_i is not None:
                    proposals_i = bbx_i
                else:
                    raise Empty

                if proposals_i.size(0) == 0:
                    raise Empty

                # Find positives and negatives based on IoU
                if bbx_i is not None:
                    iou = ious(proposals_i, bbx_i)
                    best_iou, best_gt = iou.max(dim=1)

                    pos_idx = best_iou >= self.pos_threshold
                    neg_idx = (best_iou >= self.neg_threshold_lo) & (best_iou < self.neg_threshold_hi)
                else:
                    # No ground truth boxes: all proposals that are non-void are negative
                    pos_idx = proposals_i.new_zeros(proposals_i.size(0), dtype=torch.uint8)
                    neg_idx = proposals_i.new_ones(proposals_i.size(0), dtype=torch.uint8)

                # Check that there are still some non-voids and do sub-sampling
                if not pos_idx.any().item() and not neg_idx.any().item():
                    raise Empty
                pos_idx, neg_idx = self._subsample(pos_idx, neg_idx)

                # Gather selected proposals
                out_proposals_i = proposals_i[torch.cat([pos_idx, neg_idx])]

                # Save matching
                match_i = out_proposals_i.new_full((out_proposals_i.size(0),), -1, dtype=torch.long)
                match_i[:pos_idx.numel()] = best_gt[pos_idx]

                # Save to output
                out_proposals.append(out_proposals_i)
                match.append(match_i)
            except Empty:
                out_proposals.append(None)
                match.append(None)

        return PackedSequence(out_proposals), PackedSequence(match)


class DetectionLoss:
    """Detection loss"""

    def __init__(self, sigma):
        self.sigma = sigma

    def bbx_loss(self, bbx_logits, bbx_lbl, num_non_void):
        bbx_logits = bbx_logits.view(-1, 4)
        bbx_lbl = bbx_lbl.view(-1, 4)

        bbx_loss = smooth_l1(bbx_logits, bbx_lbl, self.sigma).sum(dim=-1).sum()
        bbx_loss /= num_non_void
        return bbx_loss

    def __call__(self, cls_logits, bbx_logits, cls_lbl, bbx_lbl):
        """Detection loss
        """
        # Get contiguous view of the labels
        cls_lbl, _ = cls_lbl.contiguous
        bbx_lbl, _ = bbx_lbl.contiguous

        # Classification loss
        cls_loss = functional.cross_entropy(cls_logits, cls_lbl)

        # Regression loss
        positives = cls_lbl > 0
        num_non_void = cls_lbl.numel()
        if positives.any().item():
            cls_lbl = cls_lbl[positives]
            bbx_logits = bbx_logits[positives]

            idx = torch.arange(0, bbx_logits.size(0), dtype=torch.long, device=bbx_logits.device)
            bbx_loss = self.bbx_loss(bbx_logits[idx, cls_lbl - 1], bbx_lbl, num_non_void)
        else:
            bbx_loss = bbx_logits.sum() * 0

        return cls_loss, bbx_loss


class DetectionAlgo:
    """Base class for detection algorithms
    """

    def __init__(self, classes, bbx_reg_weights):
        self.num_stuff = classes["stuff"]
        self.bbx_reg_weights = bbx_reg_weights

    @staticmethod
    def _split_and_clip(boxes, scores, index, valid_size):
        boxes_out, scores_out = [], []
        for img_id, valid_size_i in enumerate(valid_size):
            idx = index == img_id
            if idx.any().item():
                boxes_i = boxes[idx]
                boxes_i[:, :, [0, 2]] = torch.clamp(boxes_i[:, :, [0, 2]], min=0, max=valid_size_i[0])
                boxes_i[:, :, [1, 3]] = torch.clamp(boxes_i[:, :, [1, 3]], min=0, max=valid_size_i[1])

                boxes_out.append(boxes_i)
                scores_out.append(scores[idx])
            else:
                boxes_out.append(None)
                scores_out.append(None)

        return boxes_out, scores_out

    def _match_to_lbl(self, proposals, bbx, cat, match):
        cls_lbl = []
        bbx_lbl = []

        for i, (proposals_i, bbx_i_, match_i) in enumerate(zip(proposals, bbx, match)):
            bbx_i = bbx_i_[:, [0, 1, 3, 4]]
            cat_i = bbx_i_[:, 5].long()
            if match_i is not None:
                pos = match_i >= 0

                # Objectness labels
                cls_lbl_i = proposals_i.new_zeros(proposals_i.size(0), dtype=torch.long)
                cls_lbl_i[pos] = cat_i[match_i[pos]] + 1 - self.num_stuff

                # Bounding box regression labels
                if pos.any().item():
                    bbx_lbl_i = calculate_shift(proposals_i[pos], bbx_i[match_i[pos]])
                    bbx_lbl_i *= bbx_lbl_i.new(self.bbx_reg_weights)
                else:
                    bbx_lbl_i = None

                cls_lbl.append(cls_lbl_i)
                bbx_lbl.append(bbx_lbl_i)
            else:
                cls_lbl.append(None)
                bbx_lbl.append(None)

        return PackedSequence(cls_lbl), PackedSequence(bbx_lbl)

    def training(self, head, x, proposals, bbx, cat, iscrowd, img_size):
        """Given input features, proposals and ground truth compute detection losses
        """
        raise NotImplementedError()

    def inference(self, head, x, proposals, valid_size, img_size):
        """Given input features compute detection predictions
        """
        raise NotImplementedError()

models/grasp_det_seg/algos/detection.py
andvg3/Grasp-Anything
import torch
# from inplace_abn import active_group, set_active_group

from models.grasp_det_seg.utils.bbx import shift_boxes
from models.grasp_det_seg.utils.misc import Empty
from models.grasp_det_seg.utils.parallel import PackedSequence
from models.grasp_det_seg.utils.roi_sampling import roi_sampling
from .detection import DetectionAlgo
from .rpn import RPNAlgo


class RPNAlgoFPN(RPNAlgo):
    """RPN algorithm for FPN-based region proposal networks

    Parameters
    ----------
    proposal_generator : RPNProposalGenerator
    anchor_matcher : RPNAnchorMatcher
    loss : RPNLoss
    anchor_scale : float
        Anchor scale factor, this is multiplied by the RPN stride at each level to determine the actual anchor sizes
    anchor_ratios : sequence of float
        Anchor aspect ratios
    anchor_strides: sequence of int
        Effective strides of the RPN outputs at each FPN level
    min_level : int
        First FPN level to work on
    levels : int
        Number of FPN levels to work on
    """

    def __init__(self,
                 proposal_generator,
                 anchor_matcher,
                 loss,
                 anchor_scale,
                 anchor_ratios,
                 anchor_strides,
                 min_level,
                 levels):
        super(RPNAlgoFPN, self).__init__((anchor_scale,), anchor_ratios)
        self.proposal_generator = proposal_generator
        self.anchor_matcher = anchor_matcher
        self.loss = loss
        self.min_level = min_level
        self.levels = levels

        # Cache per-cell anchors
        self.anchor_strides = anchor_strides[min_level:min_level + levels]
        self.anchors = [self._base_anchors(stride) for stride in self.anchor_strides]

    @staticmethod
    def _get_logits(head, x):
        obj_logits, bbx_logits, h, w = [], [], [], []
        for x_i in x:
            obj_logits_i, bbx_logits_i = head(x_i)
            h_i, w_i = (int(s) for s in obj_logits_i.shape[-2:])

            obj_logits_i = obj_logits_i.permute(0, 2, 3, 1).contiguous().view(obj_logits_i.size(0), -1)
            bbx_logits_i = bbx_logits_i.permute(0, 2, 3, 1).contiguous().view(bbx_logits_i.size(0), -1, 4)

            obj_logits.append(obj_logits_i)
            bbx_logits.append(bbx_logits_i)
            h.append(h_i)
            w.append(w_i)

        return torch.cat(obj_logits, dim=1), torch.cat(bbx_logits, dim=1), h, w

    def _inference(self, obj_logits, bbx_logits, anchors, valid_size, training):
        # Compute shifted boxes
        boxes = shift_boxes(anchors, bbx_logits)

        # Clip boxes to their image sizes
        for i, (height, width) in enumerate(valid_size):
            boxes[i, :, [0, 2]] = boxes[i, :, [0, 2]].clamp(min=0, max=height)
            boxes[i, :, [1, 3]] = boxes[i, :, [1, 3]].clamp(min=0, max=width)

        return self.proposal_generator(boxes, obj_logits, training)

    def training(self, head, x, bbx, iscrowd, valid_size, training=True, do_inference=False):
        # Calculate logits for the levels that we need
        x = x[self.min_level:self.min_level + self.levels]
        obj_logits, bbx_logits, h, w = self._get_logits(head, x)

        with torch.no_grad():
            # Compute anchors for each scale and merge them
            anchors = []
            for h_i, w_i, stride_i, anchors_i in zip(h, w, self.anchor_strides, self.anchors):
                anchors.append(self._shifted_anchors(
                    anchors_i, stride_i, h_i, w_i, bbx_logits.dtype, bbx_logits.device))
            anchors = torch.cat(anchors, dim=0)
            # obj_lbl: binary class label for each anchor (being an object or not)
            # bbx_lbl: coordinates for each bbx with pos object_lbl
            match = self.anchor_matcher(anchors, bbx, iscrowd, valid_size)
            obj_lbl, bbx_lbl = self._match_to_lbl(anchors, bbx, match)

        # Compute losses
        obj_loss, bbx_loss = self.loss(obj_logits, bbx_logits, obj_lbl, bbx_lbl)

        # Optionally, also run inference
        if do_inference:
            with torch.no_grad():
                proposals = self._inference(obj_logits, bbx_logits, anchors, valid_size, training)
        else:
            proposals = None

        return obj_loss, bbx_loss, proposals

    def inference(self, head, x, valid_size, training):
        # Calculate logits for the levels that we need
        x = x[self.min_level:self.min_level + self.levels]
        obj_logits, bbx_logits, h, w = self._get_logits(head, x)

        # Compute anchors for each scale and merge them
        anchors = []
        for h_i, w_i, stride_i, anchors_i in zip(h, w, self.anchor_strides, self.anchors):
            anchors.append(self._shifted_anchors(
                anchors_i, stride_i, h_i, w_i, bbx_logits.dtype, bbx_logits.device))
        anchors = torch.cat(anchors, dim=0)

        return self._inference(obj_logits, bbx_logits, anchors, valid_size, training)


class DetectionAlgoFPN(DetectionAlgo):
    """Detection algorithm for FPN networks
    """

    def __init__(self,
                 prediction_generator,
                 proposal_matcher,
                 loss,
                 classes,
                 bbx_reg_weights,
                 canonical_scale,
                 canonical_level,
                 roi_size,
                 min_level,
                 levels):
        super(DetectionAlgoFPN, self).__init__(classes, bbx_reg_weights)
        self.prediction_generator = prediction_generator
        self.proposal_matcher = proposal_matcher
        self.loss = loss
        self.canonical_scale = canonical_scale
        self.canonical_level = canonical_level
        self.roi_size = roi_size
        self.min_level = min_level
        self.levels = levels

    def _target_level(self, boxes):
        scales = (boxes[:, 2:] - boxes[:, :2]).prod(dim=-1).sqrt()
        target_level = torch.floor(self.canonical_level + torch.log2(scales / self.canonical_scale + 1e-6))
        return target_level.clamp(min=self.min_level, max=self.min_level + self.levels - 1)

    def _rois(self, x, proposals, proposals_idx, img_size):
        stride = proposals.new([fs / os for fs, os in zip(x.shape[-2:], img_size)])
        proposals = (proposals - 0.5) * stride.repeat(2) + 0.5
        return roi_sampling(x, proposals, proposals_idx, self.roi_size)

    def _head(self, head, x, proposals, proposals_idx, img_size):
        # Find target levels
        target_level = self._target_level(proposals)

        # Sample rois
        rois = x[0].new_zeros(proposals.size(0), x[0].size(1), self.roi_size[0], self.roi_size[1])
        for level_i, x_i in enumerate(x):
            idx = target_level == (level_i + self.min_level)
            if idx.any().item():
                rois[idx] = self._rois(x_i, proposals[idx], proposals_idx[idx], img_size)

        # Run head
        return head(rois)

    def training(self, head, x, proposals, bbx, cat, iscrowd, img_size):
        x = x[self.min_level:self.min_level + self.levels]

        try:
            if proposals.all_none:
                raise Empty

            with torch.no_grad():
                # Match proposals to ground truth
                proposals, match = self.proposal_matcher(proposals, bbx, cat, iscrowd)
                cls_lbl, bbx_lbl = self._match_to_lbl(proposals, bbx, cat, match)

            if proposals.all_none:
                raise Empty

            # Run head
            # set_active_group(head, active_group(True))
            proposals, proposals_idx = proposals.contiguous
            cls_logits, bbx_logits = self._head(head, x, proposals, proposals_idx, img_size)

            # Calculate loss
            cls_loss, bbx_loss = self.loss(cls_logits, bbx_logits, cls_lbl, bbx_lbl)
        except Empty:
            # active_group(False)
            cls_loss = bbx_loss = sum(x_i.sum() for x_i in x) * 0

        return cls_loss, bbx_loss

    def inference(self, head, x, proposals, valid_size, img_size):
        x = x[self.min_level:self.min_level + self.levels]

        if not proposals.all_none:
            # Run head on the given proposals
            # set_active_group(head, active_group(True))
            proposals, proposals_idx = proposals.contiguous
            cls_logits, bbx_logits = self._head(head, x, proposals, proposals_idx, img_size)
            
            # Shift the proposals according to the logits
            
            bbx_reg_weights = x[0].new(self.bbx_reg_weights)
            boxes = shift_boxes(proposals.unsqueeze(1), bbx_logits / bbx_reg_weights)
            scores = torch.softmax(cls_logits, dim=1)
            

            # Split boxes and scores by image, clip to valid size
            
            boxes, scores = self._split_and_clip(boxes, scores, proposals_idx, valid_size)
            bbx_pred, cls_pred, obj_pred = self.prediction_generator(boxes, scores)
        else:
            # active_group(False)
            bbx_pred = PackedSequence([None for _ in range(x[0].size(0))])
            cls_pred = PackedSequence([None for _ in range(x[0].size(0))])
            obj_pred = PackedSequence([None for _ in range(x[0].size(0))])

        return bbx_pred, cls_pred, obj_pred
models/grasp_det_seg/algos/fpn.py
andvg3/Grasp-Anything
from math import sqrt

import torch
import torch.nn as nn
import torch.nn.functional as functional

from models.grasp_det_seg.modules.losses import smooth_l1
from models.grasp_det_seg.utils.bbx import ious, calculate_shift
from models.grasp_det_seg.utils.misc import Empty
from models.grasp_det_seg.utils.nms import nms
from models.grasp_det_seg.utils.parallel import PackedSequence

CHUNK_SIZE = 16


class ProposalGenerator:
    """Perform NMS-based selection of proposals

    Parameters
    ----------
    nms_threshold : float
        Intersection over union threshold for the NMS
    num_pre_nms_train : int
        Number of top-scoring proposals to feed to NMS, training mode
    num_post_nms_train : int
        Number of top-scoring proposal to keep after NMS, training mode
    num_pre_nms_val : int
        Number of top-scoring proposals to feed to NMS, validation mode
    num_post_nms_val : int
        Number of top-scoring proposal to keep after NMS, validation mode
    min_size : int
        Minimum size for proposals, discard anything with a side smaller than this
    """

    def __init__(self,
                 nms_threshold=0.7,
                 num_pre_nms_train=12000,
                 num_post_nms_train=2000,
                 num_pre_nms_val=6000,
                 num_post_nms_val=300,
                 min_size=0):
        super(ProposalGenerator, self).__init__()
        self.nms_threshold = nms_threshold
        self.num_pre_nms_train = num_pre_nms_train
        self.num_post_nms_train = num_post_nms_train
        self.num_pre_nms_val = num_pre_nms_val
        self.num_post_nms_val = num_post_nms_val
        self.min_size = min_size

    def __call__(self, boxes, scores, training):
        """Perform NMS-based selection of proposals
        """
        if training:
            num_pre_nms = self.num_pre_nms_train
            num_post_nms = self.num_post_nms_train
        else:
            num_pre_nms = self.num_pre_nms_val
            num_post_nms = self.num_post_nms_val

        proposals = []
        for bbx_i, obj_i in zip(boxes, scores):
            try:
                # Optional size pre-selection
                if self.min_size > 0:
                    bbx_size = bbx_i[:, 2:] - bbx_i[:, :2]
                    valid = (bbx_size[:, 0] >= self.min_size) & (bbx_size[:, 1] >= self.min_size)

                    if valid.any().item():
                        bbx_i, obj_i = bbx_i[valid], obj_i[valid]
                    else:
                        raise Empty

                # Score pre-selection
                obj_i, idx = obj_i.topk(min(obj_i.size(0), num_pre_nms))
                bbx_i = bbx_i[idx]

                # NMS
                idx = nms(bbx_i, obj_i, self.nms_threshold, num_post_nms)
                if idx.numel() == 0:
                    raise Empty
                bbx_i = bbx_i[idx]

                proposals.append(bbx_i)
            except Empty:
                proposals.append(None)

        return PackedSequence(proposals)


class AnchorMatcher:
    """Match anchors to ground truth boxes
    """

    def __init__(self,
                 num_samples=256,
                 pos_ratio=.5,
                 pos_threshold=.7,
                 neg_threshold=.3,
                 void_threshold=0.):
        self.num_samples = num_samples
        self.pos_ratio = pos_ratio
        self.pos_threshold = pos_threshold
        self.neg_threshold = neg_threshold
        self.void_threshold = void_threshold

    def _subsample(self, match):
        num_pos = int(self.num_samples * self.pos_ratio)
        pos_idx = torch.nonzero(match >= 0).view(-1)
        if pos_idx.numel() > num_pos:
            rand_selection = torch.randperm(pos_idx.numel(), dtype=torch.long, device=match.device)[num_pos:]
            match[pos_idx[rand_selection]] = -2
        else:
            num_pos = pos_idx.numel()

        num_neg = self.num_samples - num_pos
        neg_idx = torch.nonzero(match == -1).view(-1)
        if neg_idx.numel() > num_neg:
            rand_selection = torch.randperm(neg_idx.numel(), dtype=torch.long, device=match.device)[num_neg:]
            match[neg_idx[rand_selection]] = -2

    @staticmethod
    def _is_inside(bbx, valid_size):
        p0y, p0x, p1y, p1x = bbx[:, 0], bbx[:, 1], bbx[:, 2], bbx[:, 3]
        return (p0y >= 0) & (p0x >= 0) & (p1y <= valid_size[0]) & (p1x <= valid_size[1])

    def __call__(self, anchors, bbx, iscrowd, valid_size):
        """Match anchors to ground truth boxes
        """
        match = []
        for bbx_i_, valid_size_i in zip(bbx, valid_size):
            bbx_i = bbx_i_[:,[0,1,3,4]]

            # Default labels: everything is void
            match_i = anchors.new_full((anchors.size(0),), -2, dtype=torch.long)

            try:
                # Find anchors that are entirely within the original image area
                valid = self._is_inside(anchors, valid_size_i)

                if not valid.any().item():
                    raise Empty

                valid_anchors = anchors[valid]

                if bbx_i is not None:
                    max_a2g_iou = bbx_i.new_zeros(valid_anchors.size(0))
                    max_a2g_idx = bbx_i.new_full((valid_anchors.size(0),), -1, dtype=torch.long)
                    max_g2a_iou = []
                    max_g2a_idx = []

                    # Calculate assignments iteratively to save memory
                    for j, bbx_i_j in enumerate(torch.split(bbx_i, CHUNK_SIZE, dim=0)):
                        iou = ious(valid_anchors, bbx_i_j)

                        # Anchor -> GT
                        iou_max, iou_idx = iou.max(dim=1)
                        replace_idx = iou_max > max_a2g_iou

                        max_a2g_idx[replace_idx] = iou_idx[replace_idx] + j * CHUNK_SIZE
                        max_a2g_iou[replace_idx] = iou_max[replace_idx]

                        # GT -> Anchor
                        max_g2a_iou_j, max_g2a_idx_j = iou.transpose(0, 1).max(dim=1)
                        max_g2a_iou.append(max_g2a_iou_j)
                        max_g2a_idx.append(max_g2a_idx_j)

                        del iou

                    max_g2a_iou = torch.cat(max_g2a_iou, dim=0)
                    max_g2a_idx = torch.cat(max_g2a_idx, dim=0)

                    a2g_pos = max_a2g_iou >= self.pos_threshold
                    a2g_neg = max_a2g_iou < self.neg_threshold
                    g2a_pos = max_g2a_iou > 0

                    valid_match = valid_anchors.new_full((valid_anchors.size(0),), -2, dtype=torch.long)
                    valid_match[a2g_pos] = max_a2g_idx[a2g_pos]
                    valid_match[a2g_neg] = -1
                    valid_match[max_g2a_idx[g2a_pos]] = g2a_pos.nonzero().squeeze()
                else:
                    # No ground truth boxes for this image: everything that is not void is negative
                    valid_match = valid_anchors.new_full((valid_anchors.size(0),), -1, dtype=torch.long)

                # Subsample positives and negatives
                self._subsample(valid_match)

                match_i[valid] = valid_match
            except Empty:
                pass

            match.append(match_i)

        return torch.stack(match, dim=0)


class RPNLoss:
    """RPN loss function

    Parameters
    ----------
    sigma : float
        "bandwidth" parameter of the smooth-L1 loss used for bounding box regression
    """

    def __init__(self, sigma):
        self.sigma = sigma

    def bbx_loss(self, bbx_logits, bbx_lbl, num_non_void):
        bbx_logits = bbx_logits.view(-1, 4)
        bbx_lbl = bbx_lbl.view(-1, 4)

        bbx_loss = smooth_l1(bbx_logits, bbx_lbl, self.sigma).sum(dim=-1).sum()
        bbx_loss *= torch.clamp(1 / num_non_void, max=1.)
        return bbx_loss

    def __call__(self, obj_logits, bbx_logits, obj_lbl, bbx_lbl):
        """RPN loss function
        """
        # Get contiguous view of the labels
        positives = obj_lbl == 1
        non_void = obj_lbl != -1
        num_non_void = non_void.float().sum()

        # Objectness loss
        obj_loss = functional.binary_cross_entropy_with_logits(
            obj_logits, positives.float(), non_void.float(), reduction="sum")
        obj_loss *= torch.clamp(1. / num_non_void, max=1.)

        # Bounding box regression loss
        if positives.any().item():
            bbx_logits = bbx_logits[positives.unsqueeze(-1).expand_as(bbx_logits)]
            bbx_lbl = bbx_lbl[positives.unsqueeze(-1).expand_as(bbx_lbl)]
            bbx_loss = self.bbx_loss(bbx_logits, bbx_lbl, num_non_void)
        else:
            bbx_loss = bbx_logits.sum() * 0

        return obj_loss.mean(), bbx_loss.mean()


class RPNAlgo:
    """Base class for RPN algorithms

    Parameters
    ----------
    anchor_scales : sequence of float
        Anchor scale factors, these will be multiplied by the RPN stride to determine the actual anchor sizes
    anchor_ratios : sequence of float
        Anchor aspect ratios
    """

    def __init__(self, anchor_scales, anchor_ratios):
        self.anchor_scales = anchor_scales
        self.anchor_ratios = anchor_ratios

    def _base_anchors(self, stride):
        # Pre-generate per-cell anchors
        anchors = []
        center = stride / 2.
        for scale in self.anchor_scales:
            for ratio in self.anchor_ratios:
                h = stride * scale * sqrt(ratio)
                w = stride * scale * sqrt(1. / ratio)

                anchor = (
                    center - h / 2.,
                    center - w / 2.,
                    center + h / 2.,
                    center + w / 2.
                )
                anchors.append(anchor)

        return anchors

    @staticmethod
    def _shifted_anchors(anchors, stride, height, width, dtype=torch.float32, device="cpu"):
        grid_y = torch.arange(0, stride * height, stride, dtype=dtype, device=device)
        grid_x = torch.arange(0, stride * width, stride, dtype=dtype, device=device)
        grid = torch.stack([grid_y.view(-1, 1).repeat(1, width), grid_x.view(1, -1).repeat(height, 1)], dim=-1)

        anchors = torch.tensor(anchors, dtype=dtype, device=device)
        shifted_anchors = anchors.view(1, 1, -1, 4) + grid.repeat(1, 1, 2).unsqueeze(2)
        return shifted_anchors.view(-1, 4)

    @staticmethod
    def _match_to_lbl(anchors, bbx, match):
        pos, neg = match >= 0, match == -1

        # Objectness labels from matching tensor
        obj_lbl = torch.full_like(match, -1)
        obj_lbl[neg] = 0
        obj_lbl[pos] = 1

        # Bounding box regression labels from matching tensor
        bbx_lbl = anchors.new_zeros(len(bbx), anchors.size(0), anchors.size(1))
        for i, (pos_i, bbx_i_, match_i) in enumerate(zip(pos, bbx, match)):
            bbx_i = bbx_i_[:,[0,1,3,4]]
            if pos_i.any():
                bbx_lbl[i, pos_i] = calculate_shift(anchors[pos_i], bbx_i[match_i[pos_i]])

        return obj_lbl, bbx_lbl

    def training(self, head, x, bbx, iscrowd, valid_size, training=True, do_inference=False):
        """Given input features and ground truth compute losses and, optionally, predictions
        """
        raise NotImplementedError()

    def inference(self, head, x, valid_size, training):
        """Given input features compute object proposals
        """
        raise NotImplementedError()

models/grasp_det_seg/algos/rpn.py
andvg3/Grasp-Anything
from math import ceil

import torch
import torch.nn.functional as functional

from models.grasp_det_seg.utils.parallel import PackedSequence
from models.grasp_det_seg.utils.sequence import pack_padded_images


class SemanticSegLoss:
    """Semantic segmentation loss

    Parameters
    ----------
    ohem : float or None
        Online hard example mining fraction, or `None` to disable OHEM
    ignore_index : int
        Index of the void class
    """

    def __init__(self, ohem=None, ignore_index=255):
        if ohem is not None and (ohem <= 0 or ohem > 1):
            raise ValueError("ohem should be in (0, 1]")
        self.ohem = ohem
        self.ignore_index = ignore_index

    def __call__(self, sem_logits, sem):
        """Compute the semantic segmentation loss
        """
        sem_loss = []
        for sem_logits_i, sem_i in zip(sem_logits, sem):
            sem_loss_i = functional.cross_entropy(
                sem_logits_i.unsqueeze(0), sem_i.unsqueeze(0), ignore_index=self.ignore_index, reduction="none")
            sem_loss_i = sem_loss_i.view(-1)

            if self.ohem is not None and self.ohem != 1:
                top_k = int(ceil(sem_loss_i.numel() * self.ohem))
                if top_k != sem_loss_i.numel():
                    sem_loss_i, _ = sem_loss_i.topk(top_k)

            sem_loss.append(sem_loss_i.mean())

        return sum(sem_loss) / len(sem_logits)


class SemanticSegAlgo:
    """Semantic segmentation algorithm
    """

    def __init__(self, loss, num_classes, ignore_index=255):
        self.loss = loss
        self.num_classes = num_classes
        self.ignore_index = ignore_index

    @staticmethod
    def _pack_logits(sem_logits, valid_size, img_size):
        sem_logits = functional.interpolate(sem_logits, size=img_size, mode="bilinear", align_corners=False)
        return pack_padded_images(sem_logits, valid_size)

    def _confusion_matrix(self, sem_pred, sem):
        confmat = sem[0].new_zeros(self.num_classes * self.num_classes, dtype=torch.float)

        for sem_pred_i, sem_i in zip(sem_pred, sem):
            valid = sem_i != self.ignore_index
            if valid.any():
                sem_pred_i = sem_pred_i[valid]
                sem_i = sem_i[valid]

                confmat.index_add_(
                    0, sem_i.view(-1) * self.num_classes + sem_pred_i.view(-1), confmat.new_ones(sem_i.numel()))

        return confmat.view(self.num_classes, self.num_classes)

    @staticmethod
    def _logits(head, x, valid_size, img_size):
        sem_logits, sem_feats = head(x)
        return sem_logits,SemanticSegAlgo._pack_logits(sem_logits, valid_size, img_size), sem_feats

    def training(self, head, x, sem, valid_size, img_size):
        """Given input features and ground truth compute semantic segmentation loss, confusion matrix and prediction
        """
        # Compute logits and prediction
        sem_logits_low_res, sem_logits, sem_feats = self._logits(head, x, valid_size, img_size)
        sem_pred = PackedSequence([sem_logits_i.max(dim=0)[1] for sem_logits_i in sem_logits])
        sem_pred_low_res = PackedSequence([sem_logits_low_res_i.max(dim=0)[1].float() for sem_logits_low_res_i in sem_logits_low_res])

        # Compute loss and confusion matrix
        sem_loss = self.loss(sem_logits, sem)
        conf_mat = self._confusion_matrix(sem_pred, sem)

        return sem_loss, conf_mat, sem_pred,sem_logits,sem_logits_low_res,sem_pred_low_res,sem_feats

    def inference(self, head, x, valid_size, img_size):
        """Given input features compute semantic segmentation prediction
        """
        sem_logits_low_res, sem_logits, sem_feats = self._logits(head, x, valid_size, img_size)
        sem_pred = PackedSequence([sem_logits_i.max(dim=0)[1] for sem_logits_i in sem_logits])
        sem_pred_low_res = PackedSequence([sem_logits_low_res_i.max(dim=0)[1].float() for sem_logits_low_res_i in sem_logits_low_res])

        return sem_pred, sem_feats, sem_pred_low_res


def confusion_matrix(sem_pred, sem, num_classes, ignore_index=255):
    confmat = sem_pred.new_zeros(num_classes * num_classes, dtype=torch.float)

    valid = sem != ignore_index
    if valid.any():
        sem_pred = sem_pred[valid]
        sem = sem[valid]

        confmat.index_add_(0, sem.view(-1) * num_classes + sem_pred.view(-1), confmat.new_ones(sem.numel()))

    return confmat.view(num_classes, num_classes)

models/grasp_det_seg/algos/semantic_seg.py
andvg3/Grasp-Anything
import ast
import configparser
from os import path, listdir

_CONVERTERS = {
    "struct": ast.literal_eval
}

def load_config(config_file, defaults_file):
    parser = configparser.ConfigParser(allow_no_value=True, converters=_CONVERTERS)
    parser.read([defaults_file, config_file])
    return parser

models/grasp_det_seg/config/config.py
andvg3/Grasp-Anything
import numpy as np

cls_names = {
    'background'      : '0',
    'apple'           : '1',
    'ball'            : '2',
    'banana'          : '3',
    'bell_pepper'     : '4',
    'binder'          : '5',
    'bowl'            : '6',
    'cereal_box'      : '7',
    'coffee_mug'      : '8',
    'flashlight'      : '9',
    'food_bag'        : '10',
    'food_box'        : '11',
    'food_can'        : '12',
    'glue_stick'      : '13',
    'hand_towel'      : '14',
    'instant_noodles' : '15',
    'keyboard'        : '16',
    'kleenex'         : '17',
    'lemon'           : '18',
    'lime'            : '19',
    'marker'          : '20',
    'orange'          : '21',
    'peach'           : '22',
    'pear'            : '23',
    'potato'          : '24',
    'shampoo'         : '25',
    'soda_can'        : '26',
    'sponge'          : '27',
    'stapler'         : '28',
    'tomato'          : '29',
    'toothpaste'      : '30',
    'unknown'         : '31'
    }

colors = {
    '0': np.array([0, 0, 0]),
    '1': np.array([ 211, 47, 47 ]),
    '2': np.array([  0, 255, 0]),
    '3': np.array([123, 31, 162]),
    '4': np.array([ 81, 45, 168 ]),
    '5': np.array([ 48, 63, 159 ]),
    '6': np.array([25, 118, 210]),
    '7': np.array([ 2, 136, 209 ]),
    '8': np.array([  153, 51, 102 ]),
    '9': np.array([ 0, 121, 107 ]),
    '10': np.array([ 56, 142, 60 ]),
    '11': np.array([  104, 159, 56  ]),
    '12': np.array([ 175, 180, 43 ]),
    '13': np.array([  251, 192, 45  ]),
    '14': np.array([  255, 160, 0 ]),
    '15': np.array([ 245, 124, 0 ]),
    '16': np.array([ 230, 74, 25 ]),
    '17': np.array([  93, 64, 55 ]),
    '18': np.array([ 97, 97, 97 ]),
    '19': np.array([  84, 110, 122  ]),
    '20': np.array([ 255, 255, 102]),
    '21': np.array([ 0, 151, 167 ]),
    '22': np.array([  153, 255, 102  ]),
    '23': np.array([  51, 255, 102 ]),
    '24': np.array([  0, 255, 255  ]),
    '25': np.array([  255, 255, 255  ]),
    '26': np.array([  255, 204, 204  ]),
    '27': np.array([  153, 102, 0 ]),
    '28': np.array([  204, 255, 204  ]),
    '29': np.array([ 204, 255, 0  ]),
    '30': np.array([  255, 0, 255  ]),
    '31': np.array([ 194, 24, 91 ]),
}

colors_list = list(colors.values())
cls_list = list(cls_names.keys())
models/grasp_det_seg/data_OCID/OCID_class_dict.py
andvg3/Grasp-Anything
from .dataset import OCIDDataset, OCIDTestDataset
from .misc import iss_collate_fn, read_boxes_from_file, prepare_frcnn_format
from .transform import OCIDTransform, OCIDTestTransform
models/grasp_det_seg/data_OCID/__init__.py
andvg3/Grasp-Anything
from os import path
import cv2
import numpy as np
import torch.utils.data as data
import os
from PIL import Image


class OCIDDataset(data.Dataset):
    """OCID_grasp dataset for grasp detection and semantic segmentation
    """

    def __init__(self, data_path, root_dir, split_name, transform):
        super(OCIDDataset, self).__init__()
        self.data_path = data_path
        self.root_dir = root_dir
        self.split_name = split_name
        self.transform = transform

        self._images = self._load_split()

    def _load_split(self):
        with open(path.join(self.data_path, self.split_name + ".txt"), "r") as fid:
            images = [x.strip() for x in fid.readlines()]

        return images

    def _load_item(self, item):
        seq_path, im_name = item.split(',')
        sample_path = os.path.join(self.root_dir, seq_path)
        img_path = os.path.join(sample_path, 'rgb', im_name)
        mask_path = os.path.join(sample_path, 'seg_mask_labeled_combi', im_name)
        anno_path = os.path.join(sample_path, 'Annotations', im_name[:-4] + '.txt')
        img_bgr = cv2.imread(img_path)
        img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

        with open(anno_path, "r") as f:
            points_list = []
            boxes_list = []
            for count, line in enumerate(f):
                line = line.rstrip()
                [x, y] = line.split(' ')

                x = float(x)
                y = float(y)

                pt = (x, y)
                points_list.append(pt)

                if len(points_list) == 4:
                    boxes_list.append(points_list)
                    points_list = []

        msk = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
        box_arry = np.asarray(boxes_list)
        return img, msk, box_arry

    @property
    def categories(self):
        """Category names"""
        return self._meta["categories"]

    @property
    def num_categories(self):
        """Number of categories"""
        return len(self.categories)

    @property
    def num_stuff(self):
        """Number of "stuff" categories"""
        return self._meta["num_stuff"]

    @property
    def num_thing(self):
        """Number of "thing" categories"""
        return self.num_categories - self.num_stuff

    @property
    def original_ids(self):
        """Original class id of each category"""
        return self._meta["original_ids"]

    @property
    def palette(self):
        """Default palette to be used when color-coding semantic labels"""
        return np.array(self._meta["palette"], dtype=np.uint8)

    @property
    def img_sizes(self):
        """Size of each image of the dataset"""
        return [img_desc["size"] for img_desc in self._images]

    @property
    def img_categories(self):
        """Categories present in each image of the dataset"""
        return [img_desc["cat"] for img_desc in self._images]

    @property
    def get_images(self):
        """Categories present in each image of the dataset"""
        return self._images

    def __len__(self):
        return len(self._images)

    def __getitem__(self, item):
        im_rgb, msk, bbox_infos = self._load_item(item)

        rec, im_size = self.transform(im_rgb, msk, bbox_infos)

        rec["abs_path"] = item
        rec["root_path"] = self.root_dir
        rec["im_size"] = im_size
        return rec

    def get_raw_image(self, idx):
        """Load a single, unmodified image with given id from the dataset"""
        img_file = path.join(self._img_dir, idx)
        if path.exists(img_file + ".png"):
            img_file = img_file + ".png"
        elif path.exists(img_file + ".jpg"):
            img_file = img_file + ".jpg"
        else:
            raise IOError("Cannot find any image for id {} in {}".format(idx, self._img_dir))

        return Image.open(img_file)

    def get_image_desc(self, idx):
        """Look up an image descriptor given the id"""
        matching = [img_desc for img_desc in self._images if img_desc["id"] == idx]
        if len(matching) == 1:
            return matching[0]
        else:
            raise ValueError("No image found with id %s" % idx)


class OCIDTestDataset(data.Dataset):

    def __init__(self, data_path, root_dir, split_name, transform):
        super(OCIDTestDataset, self).__init__()
        self.data_path = data_path
        self.root_dir = root_dir
        self.split_name = split_name
        self.transform = transform

        self._images = self._load_split()

    def _load_split(self):
        with open(path.join(self.data_path, self.split_name + ".txt"), "r") as fid:
            images = [x.strip() for x in fid.readlines()]
        return images

    @property
    def img_sizes(self):
        """Size of each image of the dataset"""
        return [img_desc["size"] for img_desc in self._images]

    @property
    def get_images(self):
        """Categories present in each image of the dataset"""
        return self._images

    def __len__(self):
        return len(self._images)

    def __getitem__(self, item):
        seq_path, im_name = item.split(',')
        sample_path = os.path.join(self.root_dir, seq_path)
        img_path = os.path.join(sample_path, 'rgb', im_name)
        img_bgr = cv2.imread(img_path)
        im_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

        img_, im_size = self.transform(im_rgb)

        return {"img": img_,
                "root_path": self.root_dir,
                "abs_path": item,
                "im_size": im_size
                }

models/grasp_det_seg/data_OCID/dataset.py
andvg3/Grasp-Anything
import torch
import numpy as np

from grasp_det_seg.utils.parallel import PackedSequence


def iss_collate_fn(items):
    """Collate function for ISS batches"""
    out = {}
    if len(items) > 0:
        for key in items[0]:
            out[key] = [item[key] for item in items]
            if isinstance(items[0][key], torch.Tensor):
                out[key] = PackedSequence(out[key])
    return out

def prepare_frcnn_format(boxes,im_size):
    boxes_ary = np.asarray(boxes)

    boxes_ary = np.swapaxes(boxes_ary, 1, 2)
    xy_ctr = np.sum(boxes_ary, axis=2) / 4
    x_ctr = xy_ctr[:, 0]
    y_ctr = xy_ctr[:, 1]
    width = np.sqrt(np.sum((boxes_ary[:, :, 0] - boxes_ary[:, :, 1]) ** 2, axis=1))
    height = np.sqrt(np.sum((boxes_ary[:, :, 1] - boxes_ary[:, :, 2]) ** 2, axis=1))

    theta = np.zeros((boxes_ary.shape[0]), dtype=np.int)
    theta = np.arctan((boxes_ary[:, 1, 1] - boxes_ary[:, 1, 0]) / (boxes_ary[:, 0, 0] - boxes_ary[:, 0, 1]))
    b = np.arctan((boxes_ary[:, 1, 0] - boxes_ary[:, 1, 1]) / (boxes_ary[:, 0, 1] - boxes_ary[:, 0, 0]))
    theta[np.where(boxes_ary[:, 0, 0] <= boxes_ary[:, 0, 1])] = b[np.where(boxes_ary[:, 0, 0] <= boxes_ary[:, 0, 1])]

    # used for fasterrcnn loss
    x_min = x_ctr - width / 2
    x_max = x_ctr + width / 2
    y_min = y_ctr - height / 2
    y_max = y_ctr + height / 2

    x_coords = np.vstack((x_min, x_max))
    y_coords = np.vstack((y_min, y_max))

    mat = np.asarray((np.all(x_coords > im_size[1], axis=0), np.all(x_coords < 0, axis=0),
                      np.all(y_coords > im_size[0], axis=0), np.all(y_coords < 0, axis=0)))

    fail = np.any(mat, axis=0)
    correct_idx = np.where(fail == False)
    theta_deg = np.rad2deg(theta) + 90
    cls = (np.round((theta_deg) / (180 / 18))).astype(int)
    cls[np.where(cls == 18)] = 0

    ret_value = (boxes_ary[correct_idx], theta_deg[correct_idx],cls[correct_idx])
    return ret_value

def read_boxes_from_file(gt_path,delta_xy):
    with open(gt_path)as f:
        points_list = []
        box_list = []
        for count, line in enumerate(f):
            line = line.rstrip()
            [x, y] = line.split(' ')
            x = float(x) - int(delta_xy[0])
            y = float(y) - int(delta_xy[1])

            pt = (x, y)
            points_list.append(pt)

            if len(points_list) == 4:
                box_list.append(points_list)
                points_list = []
    return box_list

models/grasp_det_seg/data_OCID/misc.py
andvg3/Grasp-Anything
import math

import torch
from torch import distributed
from torch.utils.data.sampler import Sampler


class ARBatchSampler(Sampler):
    def __init__(self, data_source, batch_size, drop_last=False, epoch=0):
        super(ARBatchSampler, self).__init__(data_source)
        self.data_source = data_source
        self.batch_size = batch_size
        self.drop_last = drop_last
        self._epoch = epoch

        # Split images by orientation
        self.img_sets = self.data_source.get_images

    def _split_images(self, indices):
        # returns lists of [im_id, aspect_ratio]

        img_sizes = self.data_source.img_sizes
        img_sets = [[], []]
        for img_id in indices:
            aspect_ratio = img_sizes[img_id][0] / img_sizes[img_id][1]
            if aspect_ratio < 1:
                img_sets[0].append({"id": img_id, "ar": aspect_ratio})
            else:
                img_sets[1].append({"id": img_id, "ar": aspect_ratio})

        return img_sets

    def _generate_batches(self):
        g = torch.Generator()
        g.manual_seed(self._epoch)

        self.img_sets = [self.img_sets[i] for i in torch.randperm(len(self.img_sets), generator=g)]

        batches = []
        leftover = []
        batch = []
        for img in self.img_sets:
            batch.append(img)
            if len(batch) == self.batch_size:
                batches.append(batch)
                batch = []
        leftover += batch

        if not self.drop_last:
            batch = []
            for img in leftover:
                batch.append(img)
                if len(batch) == self.batch_size:
                    batches.append(batch)
                    batch = []

            if len(batch) != 0:
                batches.append(batch)

        return batches

    def set_epoch(self, epoch):
        self._epoch = epoch

    def __len__(self):
        if self.drop_last:
            return len(self.img_sets) // self.batch_size
        else:
            return (len(self.img_sets) + self.batch_size - 1) // self.batch_size


    def __iter__(self):
        batches = self._generate_batches()
        for batch in batches:
            batch = sorted(batch, key=lambda i: i["ar"])
            batch = [i["id"] for i in batch]
            yield batch


class DistributedARBatchSampler(ARBatchSampler):
    def __init__(self, data_source, batch_size, num_replicas=None, rank=None, drop_last=False, epoch=0):
        super(DistributedARBatchSampler, self).__init__(data_source, batch_size, drop_last, epoch)

        # Automatically get world size and rank if not provided
        if num_replicas is None:
            num_replicas = distributed.get_world_size()
        if rank is None:
            rank = distributed.get_rank()

        self.num_replicas = num_replicas
        self.rank = rank

        tot_batches = super(DistributedARBatchSampler, self).__len__()
        self.num_batches = int(math.ceil(tot_batches / self.num_replicas))

    def __len__(self):
        return self.num_batches

    def __iter__(self):
        batches = self._generate_batches()

        g = torch.Generator()
        g.manual_seed(self._epoch)
        indices = list(torch.randperm(len(batches), generator=g))

        # add extra samples to make it evenly divisible
        indices += indices[:(self.num_batches * self.num_replicas - len(indices))]
        assert len(indices) == self.num_batches * self.num_replicas

        # subsample
        offset = self.num_batches * self.rank
        indices = indices[offset:offset + self.num_batches]
        assert len(indices) == self.num_batches

        for idx in indices:
            yield batches[idx]

models/grasp_det_seg/data_OCID/sampler.py
andvg3/Grasp-Anything
import random
import scipy
import numpy as np
import torch
from PIL import Image
import cv2
from torchvision.transforms import functional as tfn


class OCIDTransform:
    """Transformer function for OCID_grasp dataset
    """

    def __init__(self,
                 shortest_size,
                 longest_max_size,
                 rgb_mean=None,
                 rgb_std=None,
                 random_flip=False,
                 random_scale=None,
                 rotate_and_scale=False):
        self.shortest_size = shortest_size
        self.longest_max_size = longest_max_size
        self.rgb_mean = rgb_mean
        self.rgb_std = rgb_std
        self.random_flip = random_flip
        self.random_scale = random_scale
        self.rotate_and_scale = rotate_and_scale

    def _adjusted_scale(self, in_width, in_height, target_size):
        min_size = min(in_width, in_height)
        max_size = max(in_width, in_height)
        scale = target_size / min_size

        if int(max_size * scale) > self.longest_max_size:
            scale = self.longest_max_size / max_size

        return scale

    @staticmethod
    def _random_flip(img, msk):
        if random.random() < 0.5:
            img = img.transpose(Image.FLIP_LEFT_RIGHT)
            msk = [m.transpose(Image.FLIP_LEFT_RIGHT) for m in msk]
            return img, msk
        else:
            return img, msk

    def _random_target_size(self):
        if len(self.random_scale) == 2:
            target_size = random.uniform(self.shortest_size * self.random_scale[0],
                                         self.shortest_size * self.random_scale[1])
        else:
            target_sizes = [self.shortest_size * scale for scale in self.random_scale]
            target_size = random.choice(target_sizes)
        return int(target_size)

    def _normalize_image(self, img):
        if self.rgb_mean is not None:
            img.sub_(img.new(self.rgb_mean).view(-1, 1, 1))
        if self.rgb_std is not None:
            img.div_(img.new(self.rgb_std).view(-1, 1, 1))
        return img

    @staticmethod
    def _Rotate2D(pts, cnt, ang):
        ang = np.deg2rad(ang)
        return scipy.dot(pts - cnt,
                         scipy.array([[scipy.cos(ang), scipy.sin(ang)], [-scipy.sin(ang), scipy.cos(ang)]])) + cnt

    @staticmethod
    def _prepare_frcnn_format(boxes, im_size):
        A = boxes
        xy_ctr = np.sum(A, axis=2) / 4
        x_ctr = xy_ctr[:, 0]
        y_ctr = xy_ctr[:, 1]
        width = np.sqrt(np.sum((A[:, :, 0] - A[:, :, 1]) ** 2, axis=1))
        height = np.sqrt(np.sum((A[:, :, 1] - A[:, :, 2]) ** 2, axis=1))

        theta = np.zeros((A.shape[0]), dtype=np.int)

        theta = np.arctan((A[:, 1, 1] - A[:, 1, 0]) / (A[:, 0, 0] - A[:, 0, 1]))
        b = np.arctan((A[:, 1, 0] - A[:, 1, 1]) / (A[:, 0, 1] - A[:, 0, 0]))
        theta[np.where(A[:, 0, 0] <= A[:, 0, 1])] = b[np.where(A[:, 0, 0] <= A[:, 0, 1])]

        # used for fasterrcnn loss
        x_min = x_ctr - width / 2
        x_max = x_ctr + width / 2
        y_min = y_ctr - height / 2
        y_max = y_ctr + height / 2

        x_coords = np.vstack((x_min, x_max))
        y_coords = np.vstack((y_min, y_max))

        mat = np.asarray((np.all(x_coords > im_size[1], axis=0), np.all(x_coords < 0, axis=0),
                          np.all(y_coords > im_size[0], axis=0), np.all(y_coords < 0, axis=0)))

        fail = np.any(mat, axis=0)
        correct_idx = np.where(fail == False)
        theta_deg = np.rad2deg(theta) + 90
        cls = (np.round((theta_deg) / (180 / 18))).astype(int)
        cls[np.where(cls == 18)] = 0

        if np.any(cls) > 17:
            assert False

        ret_value = (
        x_min[correct_idx], y_min[correct_idx], theta_deg[correct_idx], x_max[correct_idx], y_max[correct_idx],
        cls[correct_idx])
        return ret_value

    def _rotateAndScale(self, img, msk, all_boxes):
        im_size = [self.shortest_size, self.longest_max_size]

        img_pad = cv2.copyMakeBorder(img, 200, 200, 200, 200, borderType=cv2.BORDER_REPLICATE)
        msk_pad = cv2.copyMakeBorder(msk, 200, 200, 200, 200, borderType=cv2.BORDER_CONSTANT)

        (oldY, oldX, chan) = img_pad.shape  # note: numpy uses (y,x) convention but most OpenCV functions use (x,y)

        theta = float(np.random.randint(360) - 1)
        dx = np.random.randint(101) - 51
        dy = np.random.randint(101) - 51

        M = cv2.getRotationMatrix2D(center=(oldX / 2, oldY / 2), angle=theta,
                                    scale=1.0)  # rotate about center of image.

        # choose a new image size.
        newX, newY = oldX, oldY
        # include this if you want to prevent corners being cut off
        r = np.deg2rad(theta)
        newX, newY = (abs(np.sin(r) * newY) + abs(np.cos(r) * newX), abs(np.sin(r) * newX) + abs(np.cos(r) * newY))

        # Find the translation that moves the result to the center of that region.
        (tx, ty) = ((newX - oldX) / 2, (newY - oldY) / 2)
        M[0, 2] += tx
        M[1, 2] += ty

        imgRotate = cv2.warpAffine(img_pad, M, dsize=(int(newX), int(newY)))
        mskRotate = cv2.warpAffine(msk_pad, M, dsize=(int(newX), int(newY)))

        imgRotateCrop = imgRotate[
                        int(imgRotate.shape[0] / 2 - (im_size[0] / 2)) - dx:int(
                            imgRotate.shape[0] / 2 + (im_size[0] / 2)) - dx,
                        int(imgRotate.shape[1] / 2 - (im_size[1] / 2)) - dy:int(
                            imgRotate.shape[1] / 2 + (im_size[1] / 2)) - dy, :]
        mskRotateCrop = mskRotate[
                        int(mskRotate.shape[0] / 2 - (im_size[0] / 2)) - dx:int(
                            mskRotate.shape[0] / 2 + (im_size[0] / 2)) - dx,
                        int(mskRotate.shape[1] / 2 - (im_size[1] / 2)) - dy:int(
                            mskRotate.shape[1] / 2 + (im_size[1] / 2)) - dy]

        bbsInShift = np.zeros_like(all_boxes)
        bbsInShift[:, 0, :] = all_boxes[:, 0, :] - (im_size[1] / 2)
        bbsInShift[:, 1, :] = all_boxes[:, 1, :] - (im_size[0] / 2)
        R = np.array([[np.cos(theta / 180 * np.pi), -np.sin(theta / 180 * np.pi)],
                      [np.sin(theta / 180 * np.pi), np.cos(theta / 180 * np.pi)]])
        R_all = np.expand_dims(R, axis=0)  #
        R_all = np.repeat(R_all, all_boxes.shape[0], axis=0)
        bbsInShift = np.swapaxes(bbsInShift, 1, 2)

        bbsRotated = np.dot(bbsInShift, R_all.T)
        bbsRotated = bbsRotated[:, :, :, 0]
        bbsRotated = np.swapaxes(bbsRotated, 1, 2)
        bbsInShiftBack = np.asarray(bbsRotated)
        bbsInShiftBack[:, 0, :] = (bbsRotated[:, 0, :] + (im_size[1] / 2) + dy)
        bbsInShiftBack[:, 1, :] = (bbsRotated[:, 1, :] + (im_size[0] / 2) + dx)

        return imgRotateCrop, mskRotateCrop, bbsInShiftBack

    def __call__(self, img_, msk_, bbox_infos_):
        im_size = [self.shortest_size, self.longest_max_size]
        bbox_infos_ = np.swapaxes(bbox_infos_, 1, 2)

        x_min = int(img_.shape[0] / 2 - int(im_size[0] / 2))
        x_max = int(img_.shape[0] / 2 + int(im_size[0] / 2))
        y_min = int(img_.shape[1] / 2 - int(im_size[1] / 2))
        y_max = int(img_.shape[1] / 2 + int(im_size[1] / 2))

        new_origin = np.array([[y_min], [x_min]])

        img = img_[x_min:x_max, y_min:y_max, :]

        msk = msk_[x_min:x_max, y_min:y_max]

        bbox_infos_ = bbox_infos_ - new_origin
        bbox_infos = np.copy(bbox_infos_)

        if self.rotate_and_scale:
            img, msk, bbox_transformed = self._rotateAndScale(img, msk, bbox_infos_)
            bbox_infos = bbox_transformed
        # Random flip
        if self.random_flip:
            img, msk = self._random_flip(img, msk)

        # Adjust scale, possibly at random
        if self.random_scale is not None:
            target_size = self._random_target_size()
        else:
            target_size = self.shortest_size

        ret = self._prepare_frcnn_format(bbox_infos, im_size)
        (x1, y1, theta, x2, y2, cls) = ret
        if len(cls) == 0:
            print('NO valid boxes after augmentation, switch to gt values')
            ret = self._prepare_frcnn_format(bbox_infos_, im_size)
            img = img_[x_min:x_max, y_min:y_max, :]

            msk = msk_[x_min:x_max, y_min:y_max]

        bbox_infos = np.asarray(ret).T
        bbox_infos = bbox_infos.astype(np.float32)

        # Image transformations
        img = tfn.to_tensor(img)
        img = self._normalize_image(img)

        # Label transformations
        msk = np.stack([np.array(m, dtype=np.int32, copy=False) for m in msk], axis=0)

        # Convert labels to torch and extract bounding boxes
        msk = torch.from_numpy(msk.astype(np.long))

        bbx = torch.from_numpy(np.asarray(bbox_infos)).contiguous()
        if bbox_infos.shape[1] != 6:
            assert False

        return dict(img=img, msk=msk, bbx=bbx), im_size


class OCIDTestTransform:
    """Transformer function for OCID_grasp dataset, used at test time
    """

    def __init__(self,
                 shortest_size,
                 longest_max_size,
                 rgb_mean=None,
                 rgb_std=None):
        self.longest_max_size = longest_max_size
        self.shortest_size = shortest_size
        self.rgb_mean = rgb_mean
        self.rgb_std = rgb_std

    def _adjusted_scale(self, in_width, in_height):
        min_size = min(in_width, in_height)
        scale = self.shortest_size / min_size
        return scale

    def _normalize_image(self, img):
        if self.rgb_mean is not None:
            img.sub_(img.new(self.rgb_mean).view(-1, 1, 1))
        if self.rgb_std is not None:
            img.div_(img.new(self.rgb_std).view(-1, 1, 1))
        return img

    def __call__(self, img):
        im_size = [self.shortest_size, self.longest_max_size]

        x_min = int(img.shape[0] / 2 - int(im_size[0] / 2))
        x_max = int(img.shape[0] / 2 + int(im_size[0] / 2))
        y_min = int(img.shape[1] / 2 - int(im_size[1] / 2))
        y_max = int(img.shape[1] / 2 + int(im_size[1] / 2))

        img = img[x_min:x_max, y_min:y_max, :]

        # Image transformations
        img = tfn.to_tensor(img)
        img = self._normalize_image(img)

        return img, im_size

models/grasp_det_seg/data_OCID/transform.py
andvg3/Grasp-Anything
from collections import OrderedDict

import torch
import torch.nn as nn
import numpy as np

from models.grasp_det_seg.utils.sequence import pad_packed_images
from inference.models.grasp_model import GraspModel
from utils.dataset_processing.grasp import Grasp
from utils.dataset_processing.grasp import GraspRectangles

def _gr_text_to_no(l, offset=(0, 0)):
    """
    Transform a single point from a Cornell file line to a pair of ints.
    :param l: Line from Cornell grasp file (str)
    :param offset: Offset to apply to point positions
    :return: Point [y, x]
    """
    x, y = l.split()
    return [int(round(float(y))) - offset[0], int(round(float(x))) - offset[1]]

NETWORK_INPUTS = ["img", "msk", "bbx"]

class DetSegNet(GraspModel):
    def __init__(self,
                 body,
                 rpn_head,
                 roi_head,
                 sem_head,
                 rpn_algo,
                 detection_algo,
                 semantic_seg_algo,
                 classes):
        super(DetSegNet, self).__init__()
        self.num_stuff = classes["stuff"]

        # Modules
        self.body = body
        self.rpn_head = rpn_head
        self.roi_head = roi_head
        self.sem_head = sem_head

        # Algorithms
        self.rpn_algo = rpn_algo
        self.detection_algo = detection_algo
        self.semantic_seg_algo = semantic_seg_algo

    def _prepare_inputs(self, msk, cat, iscrowd, bbx):
        cat_out, iscrowd_out, bbx_out, ids_out, sem_out = [], [], [], [], []
        for msk_i, cat_i, iscrowd_i, bbx_i in zip(msk, cat, iscrowd, bbx):
            msk_i = msk_i.squeeze(0)
            thing = (cat_i >= self.num_stuff) & (cat_i != 255)
            valid = thing & ~iscrowd_i

            if valid.any().item():
                cat_out.append(cat_i[valid])
                bbx_out.append(bbx_i[valid])
                ids_out.append(torch.nonzero(valid))
            else:
                cat_out.append(None)
                bbx_out.append(None)
                ids_out.append(None)

            if iscrowd_i.any().item():
                iscrowd_i = iscrowd_i & thing
                iscrowd_out.append(iscrowd_i[msk_i])
            else:
                iscrowd_out.append(None)

            sem_out.append(cat_i[msk_i])

        return cat_out, iscrowd_out, bbx_out, ids_out, sem_out

    def _convert_to_bb(self, bbx_pred, angle_pred):
        grs = []
        # bbxs = torch.cat(tuple(bbx_pred), dim=0)
        # angle_preds = torch.cat(tuple(angle_pred), dim=0)
        bbxs = bbx_pred[0]
        angle_preds = angle_pred[0]
        for i, bbx in enumerate(bbxs):
            x, y, w, h = bbx.tolist()
            theta = angle_preds[i].item()
            grs.append(Grasp(np.array([y, x]), -theta / 180.0 * np.pi, w, h).as_gr)
        grs = GraspRectangles(grs)
        # grs = grs.scale(scale)
        return grs
    
    def _numpy_to_torch(self, s):
        if len(s.shape) == 2:
            return torch.from_numpy(np.expand_dims(s, 0).astype(np.float32))
        else:
            return torch.from_numpy(s.astype(np.float32))


    def forward(self, img, msk=None, cat=None, iscrowd=None, bbx=None, do_loss=False, do_prediction=True):
        # Pad the input images
        output_size = img.shape[-1]
        device = img.device
        img.requires_grad_(True)
        img, valid_size = pad_packed_images(img)
        img_size = img.shape[-2:]

        # Convert ground truth to the internal format
        if do_loss:
            sem, _ = pad_packed_images(msk)
            msk, _ = pad_packed_images(msk)

        # Run network body
        x = self.body(img)

        # RPN part
        if do_loss:
            obj_loss, bbx_loss, proposals = self.rpn_algo.training(
                self.rpn_head, x, bbx, iscrowd, valid_size, training=self.training, do_inference=True)
        elif do_prediction:
            proposals = self.rpn_algo.inference(self.rpn_head, x, valid_size, self.training)
            obj_loss, bbx_loss = None, None
        else:
            obj_loss, bbx_loss, proposals = None, None, None

        # ROI part
        if do_loss:
            roi_cls_loss, roi_bbx_loss = self.detection_algo.training(
                self.roi_head, x, proposals, bbx, cat, iscrowd, img_size)
        else:
            roi_cls_loss, roi_bbx_loss = None, None
        if do_prediction:
            bbx_pred, cls_pred, obj_pred = self.detection_algo.inference(
                self.roi_head, x, proposals, valid_size, img_size)
        else:
            bbx_pred, cls_pred, obj_pred = None, None, None

        # Segmentation part
        # if do_loss:
        #     sem_loss, conf_mat, sem_pred,sem_logits,sem_logits_low_res, sem_pred_low_res, sem_feats  =\
        #         self.semantic_seg_algo.training(self.sem_head, x, sem, valid_size, img_size)
        # elif do_prediction:
        #     sem_pred,sem_feats,_ = self.semantic_seg_algo.inference(self.sem_head, x, valid_size, img_size)
        #     sem_loss, conf_mat = None, None
        # else:
        #     sem_loss, conf_mat, sem_pred, sem_feats = None, None, None, None

        grs = self._convert_to_bb(bbx_pred, angle_pred=obj_pred)
        pos_img, ang_img, width_img = grs.draw((output_size, output_size))
        pos = self._numpy_to_torch(pos_img).to(device)
        cos = self._numpy_to_torch(np.cos(2 * ang_img)).to(device)
        sin = self._numpy_to_torch(np.sin(2 * ang_img)).to(device)
        width = self._numpy_to_torch(width_img).to(device)

        pos.requires_grad_(True)
        cos.requires_grad_(True)
        sin.requires_grad_(True)
        width.requires_grad_(True)
        return pos, cos, sin, width
        # Prepare outputs
        # loss = OrderedDict([
        #     ("obj_loss", obj_loss),
        #     ("bbx_loss", bbx_loss),
        #     ("roi_cls_loss", roi_cls_loss),
        #     ("roi_bbx_loss", roi_bbx_loss),
        #     ("sem_loss", sem_loss)
        # ])
        # pred = OrderedDict([
        #     ("bbx_pred", bbx_pred),
        #     ("cls_pred", cls_pred),
        #     ("obj_pred", obj_pred),
        #     ("sem_pred", sem_pred)
        # ])
        # conf = OrderedDict([
        #     ("sem_conf", conf_mat)
        # ])
        # return loss, pred, conf

models/grasp_det_seg/models/det_seg.py
andvg3/Grasp-Anything
import sys
from collections import OrderedDict
from functools import partial

import torch.nn as nn
# from inplace_abn import ABN

from models.grasp_det_seg.modules.misc import GlobalAvgPool2d
from models.grasp_det_seg.modules.residual import ResidualBlock
from models.grasp_det_seg.utils.misc import try_index


class ResNet(nn.Module):
    """Standard residual network

    Parameters
    ----------
    structure : list of int
        Number of residual blocks in each of the four modules of the network
    bottleneck : bool
        If `True` use "bottleneck" residual blocks with 3 convolutions, otherwise use standard blocks
    norm_act : callable or list of callable
        Function to create normalization / activation Module. If a list is passed it should have four elements, one for
        each module of the network
    classes : int
        If not `0` also include global average pooling and a fully-connected layer with `classes` outputs at the end
        of the network
    dilation : int or list of int
        List of dilation factors for the four modules of the network, or `1` to ignore dilation
    dropout : list of float or None
        If present, specifies the amount of dropout to apply in the blocks of each of the four modules of the network
    caffe_mode : bool
        If `True`, use bias in the first convolution for compatibility with the Caffe pretrained models
    """

    def __init__(self,
                 structure,
                 bottleneck,
                 norm_act=nn.BatchNorm2d,
                 classes=0,
                 dilation=1,
                 dropout=None,
                 caffe_mode=False):
        super(ResNet, self).__init__()
        self.structure = structure
        self.bottleneck = bottleneck
        self.dilation = dilation
        self.dropout = dropout
        self.caffe_mode = caffe_mode

        if len(structure) != 4:
            raise ValueError("Expected a structure with four values")
        if dilation != 1 and len(dilation) != 4:
            raise ValueError("If dilation is not 1 it must contain four values")

        # Initial layers
        layers = [
            ("conv1", nn.Conv2d(3, 64, 7, stride=2, padding=3, bias=caffe_mode)),
            ("bn1", nn.BatchNorm2d(64))
        ]
        if try_index(dilation, 0) == 1:
            layers.append(("pool1", nn.MaxPool2d(3, stride=2, padding=1)))
        self.mod1 = nn.Sequential(OrderedDict(layers))

        # Groups of residual blocks
        in_channels = 64
        if self.bottleneck:
            channels = (64, 64, 256)
        else:
            channels = (64, 64)
        for mod_id, num in enumerate(structure):
            mod_dropout = None
            if self.dropout is not None:
                if self.dropout[mod_id] is not None:
                    mod_dropout = partial(nn.Dropout, p=self.dropout[mod_id])

            # Create blocks for module
            blocks = []
            for block_id in range(num):
                stride, dil = self._stride_dilation(dilation, mod_id, block_id)
                blocks.append((
                    "block%d" % (block_id + 1),
                    ResidualBlock(in_channels, channels, norm_act=norm_act,
                                  stride=stride, dilation=dil, dropout=mod_dropout)
                ))

                # Update channels and p_keep
                in_channels = channels[-1]

            # Create module
            self.add_module("mod%d" % (mod_id + 2), nn.Sequential(OrderedDict(blocks)))

            # Double the number of channels for the next module
            channels = [c * 2 for c in channels]

        # Pooling and predictor
        if classes != 0:
            self.classifier = nn.Sequential(OrderedDict([
                ("avg_pool", GlobalAvgPool2d()),
                ("fc", nn.Linear(in_channels, classes))
            ]))

    @staticmethod
    def _stride_dilation(dilation, mod_id, block_id):
        d = try_index(dilation, mod_id)
        s = 2 if d == 1 and block_id == 0 and mod_id > 0 else 1
        return s, d

    def forward(self, x):
        outs = OrderedDict()

        outs["mod1"] = self.mod1(x)
        outs["mod2"] = self.mod2(outs["mod1"])
        outs["mod3"] = self.mod3(outs["mod2"])
        outs["mod4"] = self.mod4(outs["mod3"])
        outs["mod5"] = self.mod5(outs["mod4"])

        if hasattr(self, "classifier"):
            outs["classifier"] = self.classifier(outs["mod5"])

        return outs


_NETS = {
    "18": {"structure": [2, 2, 2, 2], "bottleneck": False},
    "34": {"structure": [3, 4, 6, 3], "bottleneck": False},
    "50": {"structure": [3, 4, 6, 3], "bottleneck": True},
    "101": {"structure": [3, 4, 23, 3], "bottleneck": True},
    "152": {"structure": [3, 8, 36, 3], "bottleneck": True},
}

__all__ = []
for name, params in _NETS.items():
    net_name = "net_resnet" + name
    setattr(sys.modules[__name__], net_name, partial(ResNet, **params))
    __all__.append(net_name)

models/grasp_det_seg/models/resnet.py
andvg3/Grasp-Anything
from collections import OrderedDict

import torch.nn as nn
import torch.nn.functional as functional
# from inplace_abn import ABN


class FPN(nn.Module):
    """Feature Pyramid Network module

    Parameters
    ----------
    in_channels : sequence of int
        Number of feature channels in each of the input feature levels
    out_channels : int
        Number of output feature channels (same for each level)
    extra_scales : int
        Number of extra low-resolution scales
    norm_act : callable
        Function to create normalization + activation modules
    interpolation : str
        Interpolation mode to use when up-sampling, see `torch.nn.functional.interpolate`
    """

    def __init__(self, in_channels, out_channels=256, extra_scales=0, norm_act=nn.BatchNorm2d, interpolation="nearest"):
        super(FPN, self).__init__()
        self.interpolation = interpolation

        # Lateral connections and output convolutions
        self.lateral = nn.ModuleList([
            self._make_lateral(channels, out_channels, norm_act) for channels in in_channels
        ])
        self.output = nn.ModuleList([
            self._make_output(out_channels, norm_act) for _ in in_channels
        ])

        if extra_scales > 0:
            self.extra = nn.ModuleList([
                self._make_extra(in_channels[-1] if i == 0 else out_channels, out_channels, norm_act)
                for i in range(extra_scales)
            ])

        self.reset_parameters()

    def reset_parameters(self):
        gain = nn.init.calculate_gain("relu", self.lateral[0].bn.weight)
        for mod in self.modules():
            if isinstance(mod, nn.Conv2d):
                nn.init.xavier_normal_(mod.weight, gain)
            elif isinstance(mod, nn.BatchNorm2d):
                nn.init.constant_(mod.weight, 1.)
            if hasattr(mod, "bias") and mod.bias is not None:
                nn.init.constant_(mod.bias, 0.)

    @staticmethod
    def _make_lateral(input_channels, hidden_channels, norm_act):
        return nn.Sequential(OrderedDict([
            ("conv", nn.Conv2d(input_channels, hidden_channels, 1, bias=False)),
            ("bn", nn.BatchNorm2d(hidden_channels))
        ]))

    @staticmethod
    def _make_output(channels, norm_act):
        return nn.Sequential(OrderedDict([
            ("conv", nn.Conv2d(channels, channels, 3, padding=1, bias=False)),
            ("bn", nn.BatchNorm2d(channels))
        ]))

    @staticmethod
    def _make_extra(input_channels, out_channels, norm_act):
        return nn.Sequential(OrderedDict([
            ("conv", nn.Conv2d(input_channels, out_channels, 3, stride=2, padding=1, bias=False)),
            ("bn", norm_act(out_channels))
        ]))

    def forward(self, xs):
        """Feature Pyramid Network module

        Parameters
        ----------
        xs : sequence of torch.Tensor
            The input feature maps, tensors with shapes N x C_i x H_i x W_i

        Returns
        -------
        ys : sequence of torch.Tensor
            The output feature maps, tensors with shapes N x K x H_i x W_i
        """
        ys = []
        interp_params = {"mode": self.interpolation}
        if self.interpolation == "bilinear":
            interp_params["align_corners"] = False

        # Build pyramid
        for x_i, lateral_i in zip(xs[::-1], self.lateral[::-1]):
            x_i = lateral_i(x_i)
            if len(ys) > 0:
                x_i = x_i + functional.interpolate(ys[0], size=x_i.shape[-2:], **interp_params)
            ys.insert(0, x_i)

        # Compute outputs
        ys = [output_i(y_i) for y_i, output_i in zip(ys, self.output)]

        # Compute extra outputs if necessary
        if hasattr(self, "extra"):
            y = xs[-1]
            for extra_i in self.extra:
                y = extra_i(y)
                ys.append(y)

        return ys


class FPNBody(nn.Module):
    """Wrapper for a backbone network and an FPN module

    Parameters
    ----------
    backbone : torch.nn.Module
        Backbone network, which takes a batch of images and produces a dictionary of intermediate features
    fpn : torch.nn.Module
        FPN module, which takes a list of intermediate features and produces a list of outputs
    fpn_inputs : iterable
        An iterable producing the names of the intermediate features to take from the backbone's output and pass
        to the FPN
    """

    def __init__(self, backbone, fpn, fpn_inputs=()):
        super(FPNBody, self).__init__()
        self.fpn_inputs = fpn_inputs

        self.backbone = backbone
        self.fpn = fpn

    def forward(self, x):
        x = self.backbone(x)
        xs = [x[fpn_input] for fpn_input in self.fpn_inputs]
        return self.fpn(xs)

models/grasp_det_seg/modules/fpn.py
andvg3/Grasp-Anything
from .fpn import FPNROIHead, FPNSemanticHeadDeeplab
from .rpn import RPNHead

models/grasp_det_seg/modules/heads/__init__.py
andvg3/Grasp-Anything
from collections import OrderedDict

import torch
import torch.nn as nn
import torch.nn.functional as functional
# from inplace_abn import ABN
from models.grasp_det_seg.utils.misc import try_index

class FPNROIHead(nn.Module):
    """ROI head module for FPN
    """

    def __init__(self, in_channels, classes, roi_size, hidden_channels=1024, norm_act=nn.BatchNorm2d):
        super(FPNROIHead, self).__init__()

        self.fc = nn.Sequential(OrderedDict([
            ("fc1", nn.Linear(int(roi_size[0] * roi_size[1] * in_channels / 4), hidden_channels, bias=False)),
            ("bn1",nn.BatchNorm1d(hidden_channels)),
            ("fc2", nn.Linear(hidden_channels, hidden_channels, bias=False)),
            ("bn2", nn.BatchNorm1d(hidden_channels))
        ]))
        self.roi_cls = nn.Linear(hidden_channels, classes["thing"] + 1)
        self.roi_bbx = nn.Linear(hidden_channels, classes["thing"] * 4)

        self.reset_parameters()

    def reset_parameters(self):
        gain = nn.init.calculate_gain("relu", self.fc.bn1.weight)

        for name, mod in self.named_modules():
            if isinstance(mod, nn.Linear):
                if "roi_cls" in name:
                    nn.init.xavier_normal_(mod.weight, .01)
                elif "roi_bbx" in name:
                    nn.init.xavier_normal_(mod.weight, .001)
                else:
                    nn.init.xavier_normal_(mod.weight, gain)
            elif isinstance(mod, nn.BatchNorm2d):
                nn.init.constant_(mod.weight, 1.)

            if hasattr(mod, "bias") and mod.bias is not None:
                nn.init.constant_(mod.bias, 0.)

    def forward(self, x):
        """ROI head module for FPN
        """
        x = functional.avg_pool2d(x, 2)

        # Run head
        x = self.fc(x.view(x.size(0), -1))
        return self.roi_cls(x), self.roi_bbx(x).view(x.size(0), -1, 4)

class FPNSemanticHeadDeeplab(nn.Module):
    """Semantic segmentation head for FPN-style networks, extending Deeplab v3 for FPN bodies"""

    class _MiniDL(nn.Module):
        def __init__(self, in_channels, out_channels, dilation, pooling_size, norm_act):
            super(FPNSemanticHeadDeeplab._MiniDL, self).__init__()
            self.pooling_size = pooling_size

            self.conv1_3x3 = nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False)
            self.conv1_dil = nn.Conv2d(in_channels, out_channels, 3, dilation=dilation, padding=dilation, bias=False)
            self.conv1_glb = nn.Conv2d(in_channels, out_channels, 1, bias=False)
            self.bn1 = nn.BatchNorm2d(out_channels * 3)

            self.conv2 = nn.Conv2d(out_channels * 3, out_channels, 1, bias=False)
            self.bn2 = nn.BatchNorm2d(out_channels)

        def _global_pooling(self, x):
            pooling_size = (min(try_index(self.pooling_size, 0), x.shape[2]),
                            min(try_index(self.pooling_size, 1), x.shape[3]))
            padding = (
                (pooling_size[1] - 1) // 2,
                (pooling_size[1] - 1) // 2 if pooling_size[1] % 2 == 1 else (pooling_size[1] - 1) // 2 + 1,
                (pooling_size[0] - 1) // 2,
                (pooling_size[0] - 1) // 2 if pooling_size[0] % 2 == 1 else (pooling_size[0] - 1) // 2 + 1
            )

            pool = functional.avg_pool2d(x, pooling_size, stride=1)
            pool = functional.pad(pool, pad=padding, mode="replicate")
            return pool

        def forward(self, x):
            x = torch.cat([
                self.conv1_3x3(x),
                self.conv1_dil(x),
                self.conv1_glb(self._global_pooling(x)),
            ], dim=1)
            x = self.bn1(x)
            x = self.conv2(x)
            x = self.bn2(x)
            return x

    def __init__(self,
                 in_channels,
                 min_level,
                 levels,
                 num_classes,
                 hidden_channels=128,
                 dilation=6,
                 pooling_size=(64, 64),
                 norm_act=nn.BatchNorm2d,
                 interpolation="bilinear"):
        super(FPNSemanticHeadDeeplab, self).__init__()
        self.min_level = min_level
        self.levels = levels
        self.interpolation = interpolation

        self.output = nn.ModuleList([
            self._MiniDL(in_channels, hidden_channels, dilation, pooling_size, norm_act) for _ in range(levels)
        ])
        self.conv_sem = nn.Conv2d(hidden_channels * levels, num_classes, 1)

        self.reset_parameters()

    def reset_parameters(self):
        gain = nn.init.calculate_gain("relu", self.output[0].bn1.weight)
        for name, mod in self.named_modules():
            if isinstance(mod, nn.Conv2d):
                if "conv_sem" not in name:
                    nn.init.xavier_normal_(mod.weight, gain)
                else:
                    nn.init.xavier_normal_(mod.weight, .1)
            elif isinstance(mod, nn.BatchNorm2d):
                nn.init.constant_(mod.weight, 1.)
            if hasattr(mod, "bias") and mod.bias is not None:
                nn.init.constant_(mod.bias, 0.)

    def forward(self, xs):
        xs = xs[self.min_level:self.min_level + self.levels]

        ref_size = xs[0].shape[-2:]
        interp_params = {"mode": self.interpolation}
        if self.interpolation == "bilinear":
            interp_params["align_corners"] = False

        for i, output in enumerate(self.output):
            xs[i] = output(xs[i])
            if i > 0:
                xs[i] = functional.interpolate(xs[i], size=ref_size, **interp_params)

        xs_feats = torch.cat(xs, dim=1)
        xs = self.conv_sem(xs_feats)

        return xs,xs_feats

models/grasp_det_seg/modules/heads/fpn.py
andvg3/Grasp-Anything
import torch.nn as nn

# from inplace_abn import ABN

class RPNHead(nn.Module):
    """RPN head module

    Parameters
    ----------
    in_channels : int
        Number of channels in the input feature map
    num_anchors : int
        Number of anchors predicted at each spatial location
    stride : int
        Stride of the internal convolutions
    hidden_channels : int
        Number of channels in the internal intermediate feature map
    norm_act : callable
        Function to create normalization + activation modules
    """

    def __init__(self, in_channels, num_anchors, stride=1, hidden_channels=255, norm_act=nn.BatchNorm2d):
        super(RPNHead, self).__init__()

        self.conv1 = nn.Conv2d(in_channels, hidden_channels, 3, padding=1, stride=stride, bias=False)
        self.bn1 = nn.BatchNorm2d(hidden_channels)
        self.conv_obj = nn.Conv2d(hidden_channels, num_anchors, 1)
        self.conv_bbx = nn.Conv2d(hidden_channels, num_anchors * 4, 1)

        self.reset_parameters()

    def reset_parameters(self):
        activation = "relu"
        activation_param = self.bn1.weight

        # Hidden convolution
        gain = nn.init.calculate_gain(activation, activation_param)
        nn.init.xavier_normal_(self.conv1.weight, gain)
        self.bn1.reset_parameters()

        # Classifiers
        for m in [self.conv_obj, self.conv_bbx]:
            nn.init.xavier_normal_(m.weight, .01)
            nn.init.constant_(m.bias, 0)

    def forward(self, x):
        """RPN head module
        """
        x = self.conv1(x)
        x = self.bn1(x)
        return self.conv_obj(x), self.conv_bbx(x)

models/grasp_det_seg/modules/heads/rpn.py
andvg3/Grasp-Anything
import torch

from models.grasp_det_seg.utils.parallel import PackedSequence


def smooth_l1(x1, x2, sigma):
    """Smooth L1 loss"""
    sigma2 = sigma ** 2

    diff = x1 - x2
    abs_diff = diff.abs()

    mask = (abs_diff.detach() < (1. / sigma2)).float()
    return mask * (sigma2 / 2.) * diff ** 2 + (1 - mask) * (abs_diff - 0.5 / sigma2)


def ohem_loss(loss, ohem=None):
    if isinstance(loss, torch.Tensor):
        loss = loss.view(loss.size(0), -1)
        if ohem is None:
            return loss.mean()

        top_k = min(max(int(ohem * loss.size(1)), 1), loss.size(1))
        if top_k != loss.size(1):
            loss, _ = loss.topk(top_k, dim=1)

        return loss.mean()
    elif isinstance(loss, PackedSequence):
        if ohem is None:
            return sum(loss_i.mean() for loss_i in loss) / len(loss)

        loss_out = loss.data.new_zeros(())
        for loss_i in loss:
            loss_i = loss_i.view(-1)

            top_k = min(max(int(ohem * loss_i.numel()), 1), loss_i.numel())
            if top_k != loss_i.numel():
                loss_i, _ = loss_i.topk(top_k, dim=0)

            loss_out += loss_i.mean()

        return loss_out / len(loss)

models/grasp_det_seg/modules/losses.py
andvg3/Grasp-Anything
import torch
import torch.nn as nn
import torch.nn.functional as functional

# from inplace_abn import ABN


class GlobalAvgPool2d(nn.Module):
    """Global average pooling over the input's spatial dimensions"""

    def __init__(self):
        super(GlobalAvgPool2d, self).__init__()

    def forward(self, inputs):
        in_size = inputs.size()
        return inputs.view((in_size[0], in_size[1], -1)).mean(dim=2)


class Interpolate(nn.Module):
    """nn.Module wrapper to nn.functional.interpolate"""

    def __init__(self, size=None, scale_factor=None, mode="nearest", align_corners=None):
        super(Interpolate, self).__init__()
        self.size = size
        self.scale_factor = scale_factor
        self.mode = mode
        self.align_corners = align_corners

    def forward(self, x):
        return functional.interpolate(x, self.size, self.scale_factor, self.mode, self.align_corners)


class ActivatedAffine(nn.BatchNorm2d):
    """Drop-in replacement for ABN which performs inference-mode BN + activation"""

    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, activation="leaky_relu",
                 activation_param=0.01):
        super(ActivatedAffine, self).__init__(num_features, eps, momentum, affine, activation, activation_param)

    @staticmethod
    def _broadcast_shape(x):
        out_size = []
        for i, s in enumerate(x.size()):
            if i != 1:
                out_size.append(1)
            else:
                out_size.append(s)
        return out_size

    def forward(self, x):
        inv_var = torch.rsqrt(self.running_var + self.eps)
        if self.affine:
            alpha = self.weight * inv_var
            beta = self.bias - self.running_mean * alpha
        else:
            alpha = inv_var
            beta = - self.running_mean * alpha

        x.mul_(alpha.view(self._broadcast_shape(x)))
        x.add_(beta.view(self._broadcast_shape(x)))

        if self.activation == "relu":
            return functional.relu(x, inplace=True)
        elif self.activation == "leaky_relu":
            return functional.leaky_relu(x, negative_slope=self.activation_param, inplace=True)
        elif self.activation == "elu":
            return functional.elu(x, alpha=self.activation_param, inplace=True)
        elif self.activation == "identity":
            return x
        else:
            raise RuntimeError("Unknown activation function {}".format(self.activation))


class ActivatedGroupNorm(nn.BatchNorm2d):
    """GroupNorm + activation function compatible with the ABN interface"""

    def __init__(self, num_channels, num_groups, eps=1e-5, affine=True, activation="leaky_relu", activation_param=0.01):
        super(ActivatedGroupNorm, self).__init__(num_channels, eps, affine=affine, activation=activation,
                                                 activation_param=activation_param)
        self.num_groups = num_groups

        # Delete running mean and var since they are not used here
        delattr(self, "running_mean")
        delattr(self, "running_var")

    def reset_parameters(self):
        if self.affine:
            nn.init.constant_(self.weight, 1)
            nn.init.constant_(self.bias, 0)

    def forward(self, x):
        x = functional.group_norm(x, self.num_groups, self.weight, self.bias, self.eps)

        if self.activation == "relu":
            return functional.relu(x, inplace=True)
        elif self.activation == "leaky_relu":
            return functional.leaky_relu(x, negative_slope=self.activation_param, inplace=True)
        elif self.activation == "elu":
            return functional.elu(x, alpha=self.activation_param, inplace=True)
        elif self.activation == "identity":
            return x
        else:
            raise RuntimeError("Unknown activation function {}".format(self.activation))

models/grasp_det_seg/modules/misc.py
andvg3/Grasp-Anything
from collections import OrderedDict

import torch.nn as nn
import torch.nn.functional as functional
# from inplace_abn import ABN


class ResidualBlock(nn.Module):
    """Configurable residual block

    Parameters
    ----------
    in_channels : int
        Number of input channels.
    channels : list of int
        Number of channels in the internal feature maps. Can either have two or three elements: if three construct
        a residual block with two `3 x 3` convolutions, otherwise construct a bottleneck block with `1 x 1`, then
        `3 x 3` then `1 x 1` convolutions.
    stride : int
        Stride of the first `3 x 3` convolution
    dilation : int
        Dilation to apply to the `3 x 3` convolutions.
    groups : int
        Number of convolution groups. This is used to create ResNeXt-style blocks and is only compatible with
        bottleneck blocks.
    norm_act : callable
        Function to create normalization / activation Module.
    dropout: callable
        Function to create Dropout Module.
    """

    def __init__(self,
                 in_channels,
                 channels,
                 stride=1,
                 dilation=1,
                 groups=1,
                 norm_act=nn.BatchNorm2d,
                 dropout=None):
        super(ResidualBlock, self).__init__()

        # Check parameters for inconsistencies
        if len(channels) != 2 and len(channels) != 3:
            raise ValueError("channels must contain either two or three values")
        if len(channels) == 2 and groups != 1:
            raise ValueError("groups > 1 are only valid if len(channels) == 3")

        is_bottleneck = len(channels) == 3
        need_proj_conv = stride != 1 or in_channels != channels[-1]

        if not is_bottleneck:
            bn2 = norm_act(channels[1])
            bn2.activation = "identity"
            layers = [
                ("conv1", nn.Conv2d(in_channels, channels[0], 3, stride=stride, padding=dilation, bias=False,
                                    dilation=dilation)),
                ("bn1", norm_act(channels[0])),
                ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=1, padding=dilation, bias=False,
                                    dilation=dilation)),
                ("bn2", bn2)
            ]
            if dropout is not None:
                layers = layers[0:2] + [("dropout", dropout())] + layers[2:]
        else:
            bn3 = nn.BatchNorm2d(channels[2])
            # bn3.activation = "identity"
            layers = [
                ("conv1", nn.Conv2d(in_channels, channels[0], 1, stride=1, padding=0, bias=False)),
                ("bn1", nn.BatchNorm2d(channels[0])),
                ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=stride, padding=dilation, bias=False,
                                    groups=groups, dilation=dilation)),
                ("bn2", nn.BatchNorm2d(channels[1])),
                ("conv3", nn.Conv2d(channels[1], channels[2], 1, stride=1, padding=0, bias=False)),
                ("bn3", bn3)
            ]
            if dropout is not None:
                layers = layers[0:4] + [("dropout", dropout())] + layers[4:]
        self.convs = nn.Sequential(OrderedDict(layers))

        if need_proj_conv:
            self.proj_conv = nn.Conv2d(in_channels, channels[-1], 1, stride=stride, padding=0, bias=False)
            self.proj_bn = nn.BatchNorm2d(channels[-1])
            self.proj_bn.activation = "identity"

    def forward(self, x):
        if hasattr(self, "proj_conv"):
            residual = self.proj_conv(x)
            residual = self.proj_bn(residual)
        else:
            residual = x

        x = self.convs(x) + residual

        return x

        # if self.convs.bn1.activation == "relu":
        #     return functional.relu(x, inplace=True)
        # elif self.convs.bn1.activation == "leaky_relu":
        #     return functional.leaky_relu(x, negative_slope=self.convs.bn1.activation_param, inplace=True)
        # elif self.convs.bn1.activation == "elu":
        #     return functional.elu(x, alpha=self.convs.bn1.activation_param, inplace=True)
        # elif self.convs.bn1.activation == "identity":
        #     return x
        # else:
        #     raise RuntimeError("Unknown activation function {}".format(self.activation))

models/grasp_det_seg/modules/residual.py
andvg3/Grasp-Anything
from math import log
import math
import numpy as np
import torch

# from . import _backend

__all__ = [
    "extract_boxes",
    "shift_boxes",
    "shift_boxes_rotation",
    "calculate_shift",
    "calculate_shift_rotation",
    "corners_to_center_scale",
    "center_scale_to_corners",
    "invert_roi_bbx",
    "ious",
    "mask_overlap",
    "bbx_overlap"
]


def extract_boxes(mask, num_instances):
    """Calculate bounding boxes from instance segmentation mask

    Parameters
    ----------
    mask : torch.Tensor
        A tensor with shape H x W containing an instance segmentation mask
    num_instances : int
        The number of instances to look for

    Returns
    -------
    bbx : torch.Tensor
        A tensor with shape `num_instances` x 4 containing the coordinates of the bounding boxes in "corners" form

    """
    if mask.ndimension() == 2:
        mask = mask.unsqueeze(0)
    return _backend.extract_boxes(mask, num_instances)


def shift_boxes(bbx, shift, dim=-1, scale_clip=log(1000. / 16.)):
    """Shift bounding boxes using the faster r-CNN formulas

    Each 4-vector of `bbx` and `shift` contain, respectively, bounding box coordiantes in "corners" form and shifts
    in the form `(dy, dx, dh, dw)`. The output is calculated according to the Faster r-CNN formulas:

        y_out = y_in + h_in * dy
        x_out = x_in + w_in * dx
        h_out = h_in * exp(dh)
        w_out = w_in * exp(dw)

    Parameters
    ----------
    bbx : torch.Tensor
        A tensor of bounding boxes with shape N_0 x ... x N_i = 4 x ... x N_n
    shift : torch.Tensor
        A tensor of shifts with shape N_0 x ... x N_i = 4 x ... x N_n
    dim : int
        The dimension i of the input tensors which contains the bounding box coordinates and the shifts
    scale_clip : float
        Maximum scale shift value to avoid exp overflow

    Returns
    -------
    bbx_out : torch.Tensor
        A tensor of shifted bounding boxes with shape N_0 x ... x N_i = 4 x ... x N_n

    """
    yx_in, hw_in = corners_to_center_scale(*bbx.split(2, dim=dim))
    dyx, dhw = shift.split(2, dim=dim)

    yx_out = yx_in + hw_in * dyx
    hw_out = hw_in * dhw.clamp(max=scale_clip).exp()

    return torch.cat(center_scale_to_corners(yx_out, hw_out), dim=dim)

def shift_boxes_rotation(bbx,theta, shift, dim=-1, scale_clip=log(1000. / 16.)):
    """Shift bounding boxes using the faster r-CNN formulas

    Each 4-vector of `bbx` and `shift` contain, respectively, bounding box coordiantes in "corners" form and shifts
    in the form `(dy, dx, dh, dw)`. The output is calculated according to the Faster r-CNN formulas:

        y_out = y_in + h_in * dy
        x_out = x_in + w_in * dx
        h_out = h_in * exp(dh)
        w_out = w_in * exp(dw)

    Parameters
    ----------
    bbx : torch.Tensor
        A tensor of bounding boxes with shape N_0 x ... x N_i = 4 x ... x N_n
    shift : torch.Tensor
        A tensor of shifts with shape N_0 x ... x N_i = 4 x ... x N_n
    dim : int
        The dimension i of the input tensors which contains the bounding box coordinates and the shifts
    scale_clip : float
        Maximum scale shift value to avoid exp overflow

    Returns
    -------
    bbx_out : torch.Tensor
        A tensor of shifted bounding boxes with shape N_0 x ... x N_i = 4 x ... x N_n

    """
    # convert degree to rad
    theta_ = (theta * torch.Tensor([math.pi]).float().to('cuda:0')) / 180.


    yx_in, hw_in = corners_to_center_scale(*bbx.split(2, dim=dim))
    y_in,x_in = yx_in.split(1,dim=dim)
    h_in,w_in = hw_in.split(1,dim=dim)
    dyx, dhw,_ = shift.split((2,2,1), dim=dim)

    dy, dx, dh,dw, dtheta = shift.split((1,1,1,1,1), dim=dim)

    pred_ctr_x = dx * w_in * torch.cos(theta_.unsqueeze(1)) - dy * h_in * torch.sin(theta_.unsqueeze(1)) + x_in
    pred_ctr_y = dx * w_in * torch.sin(theta_.unsqueeze(1)) + dy * h_in * torch.cos(theta_.unsqueeze(1)) + y_in
    pred_w = torch.exp(dw.clamp(max=scale_clip)) * w_in
    pred_h = torch.exp(dh.clamp(max=scale_clip)) * h_in

    pred_angle = (torch.Tensor([math.pi]).float().to('cuda:0')) * dtheta + theta_.unsqueeze(1)#[:, np.newaxis]
    #pred_angle = pred_angle % (torch.Tensor([math.pi]).float().to('cuda:0'))
    pred_angle = torch.fmod(pred_angle,torch.Tensor([math.pi]).float().to('cuda:0')) * (180./torch.Tensor([math.pi]).float().to('cuda:0'))
    #torch.fmod(theta_gt - cls_pred_i, torch.Tensor([math.pi]).float().to('cuda:0'))
    yx_out_ = yx_in + hw_in * dyx
    hw_out_ = hw_in * dhw.clamp(max=scale_clip).exp()
    yx_out = torch.cat((pred_ctr_y,pred_ctr_x),dim=dim)
    hw_out = torch.cat((pred_h,pred_w),dim=dim)

    return torch.cat(center_scale_to_corners(yx_out, hw_out), dim=dim),pred_angle


def calculate_shift(bbx0, bbx1, dim=-1, eps=1e-5):
    """Calculate shift parameters between bounding boxes using the faster r-CNN formulas

    Each 4-vector of `bbx0` and `bbx1` contains bounding box coordiantes in "corners" form. The output is calculated
    according to the Faster r-CNN formulas:

        dy = (y1 - y0) / h0
        dx = (x1 - x0) / w0
        dh = log(h1 / h0)
        dw = log(w1 / w0)

    Parameters
    ----------
    bbx0 : torch.Tensor
        A tensor of source bounding boxes with shape N_0 x ... x N_i = 4 x ... x N_n
    bbx1 : torch.Tensor
        A tensor of target bounding boxes with shape N_0 x ... x N_i = 4 x ... x N_n
    dim : int
        The dimension `i` of the input tensors which contains the bounding box coordinates
    eps : float
        Small number used to avoid overflow

    Returns
    -------
    shift : torch.Tensor
        A tensor of calculated shifts from `bbx0` to `bbx1` with shape N_0 x ... x N_i = 4 x ... x N_n

    """
    # 0 -> anchor ; 1 -> gt
    yx0, hw0 = corners_to_center_scale(*bbx0.split(2, dim=dim))
    yx1, hw1 = corners_to_center_scale(*bbx1.split(2, dim=dim))

    hw0 = hw0.clamp(min=eps)

    dyx = (yx1 - yx0) / hw0
    dhw = (hw1 / hw0).log()

    return torch.cat([dyx, dhw], dim=dim)

def calculate_shift_rotation(bbx0, bbx1,cls_pred_i,theta_gt, dim=-1, eps=1e-5):
    """Calculate shift parameters between bounding boxes using the faster r-CNN formulas

    Each 4-vector of `bbx0` and `bbx1` contains bounding box coordiantes in "corners" form. The output is calculated
    according to the Faster r-CNN formulas:

        dy = (y1 - y0) / h0
        dx = (x1 - x0) / w0
        dh = log(h1 / h0)
        dw = log(w1 / w0)

    Parameters
    ----------
    bbx0 : torch.Tensor
        A tensor of source bounding boxes with shape N_0 x ... x N_i = 4 x ... x N_n
    bbx1 : torch.Tensor
        A tensor of target bounding boxes with shape N_0 x ... x N_i = 4 x ... x N_n
    dim : int
        The dimension `i` of the input tensors which contains the bounding box coordinates
    eps : float
        Small number used to avoid overflow

    Returns
    -------
    shift : torch.Tensor
        A tensor of calculated shifts from `bbx0` to `bbx1` with shape N_0 x ... x N_i = 4 x ... x N_n

    """
    # 0 -> anchor ; 1 -> gt
    yx0, hw0 = corners_to_center_scale(*bbx0.split(2, dim=dim))
    yx1, hw1 = corners_to_center_scale(*bbx1.split(2, dim=dim))

    hw0 = hw0.clamp(min=eps)

    # convert degree to rad
    cls_pred_i_ = (cls_pred_i * torch.Tensor([math.pi]).float().to('cuda:0')) / 180.
    theta_gt_ = (theta_gt * torch.Tensor([math.pi]).float().to('cuda:0')) / 180.
    #cls_pred_i_ = cls_pred_i
    #theta_gt_ = theta_gt


    #dyx = (yx1 - yx0) / hw0
    #dyx = (yx1 - yx0)
    #tx_mat = [torch.cos(cls_pred_i), torch.sin(cls_pred_i)]
    #ty_mat = [torch.cos(cls_pred_i), -torch.sin(cls_pred_i)]
    dx = (1/hw0[:,1]) * ((yx1[:,1] - yx0[:,1]) * torch.cos(cls_pred_i_) + (yx1[:,0] - yx0[:,0]) * torch.sin(cls_pred_i_))
    dy = (1/hw0[:,0]) * ((yx1[:,0] - yx0[:,0]) * torch.cos(cls_pred_i_) - (yx1[:,1] - yx0[:,1]) * torch.sin(cls_pred_i_))
    #t_theta = torch.Tensor([1/2*math.pi]).float().to('cuda:0') * \
    #          torch.fmod(theta_gt-cls_pred_i,torch.Tensor([2*math.pi]).float().to('cuda:0'))
    t_theta = torch.Tensor([1/math.pi]).float().to('cuda:0') * \
              torch.fmod(theta_gt_-cls_pred_i_,torch.Tensor([math.pi]).float().to('cuda:0'))
    dhw = (hw1 / hw0).log()
    dyx = torch.cat([dy.unsqueeze(1),dx.unsqueeze(1)],dim =dim)

    return torch.cat([dyx, dhw,t_theta.unsqueeze(1)], dim=dim)

def corners_to_center_scale(p0, p1):
    """Convert bounding boxes from "corners" form to "center+scale" form"""
    yx = 0.5 * (p0 + p1)
    hw = p1 - p0
    return yx, hw


def center_scale_to_corners(yx, hw):
    """Convert bounding boxes from "center+scale" form to "corners" form"""
    hw_half = 0.5 * hw
    p0 = yx - hw_half
    p1 = yx + hw_half
    return p0, p1


def invert_roi_bbx(bbx, roi_size, img_size):
    """Compute bbx coordinates to perform inverse roi sampling"""
    bbx_size = bbx[:, 2:] - bbx[:, :2]
    return torch.cat([
        -bbx.new(roi_size) * bbx[:, :2] / bbx_size,
        bbx.new(roi_size) * (bbx.new(img_size) - bbx[:, :2]) / bbx_size
    ], dim=1)


def ious(bbx0, bbx1):
    """Calculate intersection over union between sets of bounding boxes

    Parameters
    ----------
    bbx0 : torch.Tensor
        A tensor of bounding boxes in "corners" form with shape N x 4
    bbx1 : torch.Tensor
        A tensor of bounding boxes in "corners" form with shape M x 4

    Returns
    -------
    iou : torch.Tensor
        A tensor with shape N x M containing the IoUs between all pairs of bounding boxes in bbx0 and bbx1
    """
    bbx0_tl, bbx0_br = bbx0.unsqueeze(dim=1).split(2, -1)
    bbx1_tl, bbx1_br = bbx1.unsqueeze(dim=0).split(2, -1)

    # Intersection coordinates
    int_tl = torch.max(bbx0_tl, bbx1_tl)
    int_br = torch.min(bbx0_br, bbx1_br)

    intersection = (int_br - int_tl).clamp(min=0).prod(dim=-1)
    bbx0_area = (bbx0_br - bbx0_tl).prod(dim=-1)
    bbx1_area = (bbx1_br - bbx1_tl).prod(dim=-1)
    return intersection / (bbx0_area + bbx1_area - intersection)


def mask_overlap(bbx, mask):
    """Calculate overlap between a set of bounding boxes and a mask

    Parameters
    ----------
    bbx : torch.Tensor
        A tensor of bounding boxes in "corners" form with shape N x 4
    mask : torch.Tensor
        A binary tensor with shape H x W

    Returns
    -------
    overlap : torch.Tensor
        A tensor with shape N containing the proportion of non-zero pixels in each box
    """
    # Compute integral image of the mask
    int_mask = bbx.new_zeros((mask.size(0) + 1, mask.size(1) + 1))
    int_mask[1:, 1:] = mask > 0
    int_mask = int_mask.cumsum(0).cumsum(1)

    count = _backend.mask_count(bbx, int_mask)
    area = (bbx[:, 2:] - bbx[:, :2]).prod(dim=1)

    return count / area


def bbx_overlap(bbx0, bbx1):
    """Calculate intersection over area between two sets of bounding boxes

    Intersection over area is defined as:
        area(inter(bbx0, bbx1)) / area(bbx0)

    Parameters
    ----------
    bbx0 : torch.Tensor
        A tensor of bounding boxes in "corners" form with shape N x 4
    bbx1 : torch.Tensor
        A tensor of bounding boxes in "corners" form with shape M x 4

    Returns
    -------
    ratios : torch.Tensor
        A tensor with shape N x M containing the intersection over areas between all pairs of bounding boxes
    """
    bbx0_tl, bbx0_br = bbx0.unsqueeze(dim=1).split(2, -1)
    bbx1_tl, bbx1_br = bbx1.unsqueeze(dim=0).split(2, -1)

    # Intersection coordinates
    int_tl = torch.max(bbx0_tl, bbx1_tl)
    int_br = torch.min(bbx0_br, bbx1_br)

    intersection = (int_br - int_tl).clamp(min=0).prod(dim=-1)
    bbx0_area = (bbx0_br - bbx0_tl).prod(dim=-1)

    return intersection / bbx0_area

models/grasp_det_seg/utils/bbx/bbx.py
andvg3/Grasp-Anything
import logging
from math import log10
from os import path

from .meters import AverageMeter

_NAME = "GraspDetSeg_CNN"


def _current_total_formatter(current, total):
    width = int(log10(total)) + 1
    return ("[{:" + str(width) + "}/{:" + str(width) + "}]").format(current, total)


def init(log_dir, name):
    logger = logging.getLogger(_NAME)
    logger.setLevel(logging.DEBUG)

    # Set console logging
    console_handler = logging.StreamHandler()
    console_formatter = logging.Formatter(fmt="%(asctime)s %(message)s", datefmt="%H:%M:%S")
    console_handler.setFormatter(console_formatter)
    console_handler.setLevel(logging.DEBUG)
    logger.addHandler(console_handler)

    # Setup file logging
    file_handler = logging.FileHandler(path.join(log_dir, name + ".log"), mode="w")
    file_formatter = logging.Formatter(fmt="%(levelname).1s %(asctime)s %(message)s", datefmt="%y-%m-%d %H:%M:%S")
    file_handler.setFormatter(file_formatter)
    file_handler.setLevel(logging.INFO)
    logger.addHandler(file_handler)


def get_logger():
    return logging.getLogger(_NAME)


def iteration(summary, phase, global_step, epoch, num_epochs, step, num_steps, values, multiple_lines=False):
    logger = get_logger()

    # Build message and write summary
    msg = _current_total_formatter(epoch, num_epochs) + " " + _current_total_formatter(step, num_steps)
    for k, v in values.items():
        if isinstance(v, AverageMeter):
            msg += "\n" if multiple_lines else "" + "\t{}={:.3f} ({:.3f})".format(k, v.value.item(), v.mean.item())
            if summary is not None:
                summary.add_scalar("{}/{}".format(phase, k), v.value.item(), global_step)
        else:
            msg += "\n" if multiple_lines else "" + "\t{}={:.3f}".format(k, v)
            if summary is not None:
                summary.add_scalar("{}/{}".format(phase, k), v, global_step)

    # Write log
    logger.info(msg)

models/grasp_det_seg/utils/logging.py
andvg3/Grasp-Anything
import io
from collections import OrderedDict
from functools import partial

import torch
import torch.distributed as dist
import torch.nn as nn
# from inplace_abn import InPlaceABN, InPlaceABNSync, ABN

from models.grasp_det_seg.modules.misc import ActivatedAffine, ActivatedGroupNorm
from . import scheduler as lr_scheduler

NORM_LAYERS = [nn.BatchNorm2d, nn.BatchNorm2d, nn.BatchNorm2d, nn.BatchNorm2d, nn.GroupNorm]
OTHER_LAYERS = [nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.ConvTranspose1d, nn.ConvTranspose2d, nn.ConvTranspose3d]


class Empty(Exception):
    """Exception to facilitate handling of empty predictions, annotations etc."""
    pass


def try_index(scalar_or_list, i):
    try:
        return scalar_or_list[i]
    except TypeError:
        return scalar_or_list


def config_to_string(config):
    with io.StringIO() as sio:
        config.write(sio)
        config_str = sio.getvalue()
    return config_str


def scheduler_from_config(scheduler_config, optimizer, epoch_length):
    assert scheduler_config["type"] in ("linear", "step", "poly", "multistep")

    params = scheduler_config.getstruct("params")

    if scheduler_config["type"] == "linear":
        if scheduler_config["update_mode"] == "batch":
            count = epoch_length * scheduler_config.getint("epochs")
        else:
            count = scheduler_config.getint("epochs")

        beta = float(params["from"])
        alpha = float(params["to"] - beta) / count

        scheduler = lr_scheduler.LambdaLR(optimizer, lambda it: it * alpha + beta)
    elif scheduler_config["type"] == "step":
        scheduler = lr_scheduler.StepLR(optimizer, params["step_size"], params["gamma"])
    elif scheduler_config["type"] == "poly":
        if scheduler_config["update_mode"] == "batch":
            count = epoch_length * scheduler_config.getint("epochs")
        else:
            count = scheduler_config.getint("epochs")
        scheduler = lr_scheduler.LambdaLR(optimizer, lambda it: (1 - float(it) / count) ** params["gamma"])
    elif scheduler_config["type"] == "multistep":
        scheduler = lr_scheduler.MultiStepLR(optimizer, params["milestones"], params["gamma"])
    else:
        raise ValueError("Unrecognized scheduler type {}, valid options: 'linear', 'step', 'poly', 'multistep'"
                         .format(scheduler_config["type"]))

    if scheduler_config.getint("burn_in_steps") != 0:
        scheduler = lr_scheduler.BurnInLR(scheduler,
                                          scheduler_config.getint("burn_in_steps"),
                                          scheduler_config.getfloat("burn_in_start"))

    return scheduler


def norm_act_from_config(body_config):
    """Make normalization + activation function from configuration

    Available normalization modes are:
      - `bn`: Standard In-Place Batch Normalization
      - `syncbn`: Synchronized In-Place Batch Normalization
      - `syncbn+bn`: Synchronized In-Place Batch Normalization in the "static" part of the network, Standard In-Place
        Batch Normalization in the "dynamic" parts
      - `gn`: Group Normalization
      - `syncbn+gn`: Synchronized In-Place Batch Normalization in the "static" part of the network, Group Normalization
        in the "dynamic" parts
      - `off`: No normalization (preserve scale and bias parameters)

    The "static" part of the network includes the backbone, FPN and semantic segmentation components, while the
    "dynamic" part of the network includes the RPN, detection and instance segmentation components. Note that this
    distinction is due to historical reasons and for back-compatibility with the CVPR2019 pre-trained models.

    Parameters
    ----------
    body_config
        Configuration object containing the following fields: `normalization_mode`, `activation`, `activation_slope`
        and `gn_groups`

    Returns
    -------
    norm_act_static : callable
        Function that returns norm_act modules for the static parts of the network
    norm_act_dynamic : callable
        Function that returns norm_act modules for the dynamic parts of the network
    """
    mode = body_config["normalization_mode"]
    activation = body_config["activation"]
    slope = body_config.getfloat("activation_slope")
    groups = body_config.getint("gn_groups")

    if mode == "bn":
        norm_act_static = norm_act_dynamic = partial(nn.BatchNorm2d, activation=activation, activation_param=slope)
    elif mode == "syncbn":
        norm_act_static = norm_act_dynamic = partial(nn.BatchNorm2d, activation=activation, activation_param=slope)
    elif mode == "syncbn+bn":
        norm_act_static = partial(nn.BatchNorm2d, activation=activation, activation_param=slope)
        norm_act_dynamic = partial(nn.BatchNorm2d, activation=activation, activation_param=slope)
    elif mode == "gn":
        norm_act_static = norm_act_dynamic = partial(
            ActivatedGroupNorm, num_groups=groups, activation=activation, activation_param=slope)
    elif mode == "syncbn+gn":
        norm_act_static = partial(nn.BatchNorm2d, activation=activation, activation_param=slope)
        norm_act_dynamic = partial(ActivatedGroupNorm, num_groups=groups, activation=activation, activation_param=slope)
    elif mode == "off":
        norm_act_static = norm_act_dynamic = partial(ActivatedAffine, activation=activation, activation_param=slope)
    else:
        raise ValueError("Unrecognized normalization_mode {}, valid options: 'bn', 'syncbn', 'syncbn+bn', 'gn', "
                         "'syncbn+gn', 'off'".format(mode))

    return norm_act_static, norm_act_dynamic


def freeze_params(module):
    """Freeze all parameters of the given module"""
    for p in module.parameters():
        p.requires_grad_(False)


def all_reduce_losses(losses):
    """Coalesced mean all reduce over a dictionary of 0-dimensional tensors"""
    names, values = [], []
    for k, v in losses.items():
        names.append(k)
        values.append(v)

    # Peform the actual coalesced all_reduce
    values = torch.cat([v.view(1) for v in values], dim=0)
    dist.all_reduce(values, dist.ReduceOp.SUM)
    values.div_(dist.get_world_size())
    values = torch.chunk(values, values.size(0), dim=0)

    # Reconstruct the dictionary
    return OrderedDict((k, v.view(())) for k, v in zip(names, values))

models/grasp_det_seg/utils/misc.py
andvg3/Grasp-Anything
from typing import Tuple

import torch
import torchvision
from torch import Tensor
from torchvision.extension import _assert_has_ops

# from ..utils import _log_api_usage_once
# from ._box_convert import _box_cxcywh_to_xyxy, _box_xywh_to_xyxy, _box_xyxy_to_cxcywh, _box_xyxy_to_xywh
# from ._utils import _upcast

from torchvision.ops.boxes import box_iou

def nms(bboxes: torch.Tensor, scores: torch.Tensor, threshold: float=0.5, n_max=-1) -> torch.Tensor:
    order = torch.argsort(-scores)
    indices = torch.arange(bboxes.shape[0], device=bboxes.device)
    keep = torch.ones_like(indices, dtype=torch.bool, device=bboxes.device)
    for i in indices:
        if keep[i]:
            bbox = bboxes[order[i]]
            iou = box_iou(bbox[None,...],(bboxes[order[i + 1:]]) * keep[i + 1:][...,None])
            overlapped = torch.nonzero(iou > threshold)
            keep[overlapped + i + 1] = 0
    return order[keep]

# def nms(bbx, scores, threshold=0.5, n_max=-1):
#     """Perform non-maxima suppression

#     Select up to n_max bounding boxes from bbx, giving priorities to bounding boxes with greater scores. Each selected
#     bounding box suppresses all other not yet selected boxes that intersect it by more than the given threshold.

#     Parameters
#     ----------
#     bbx : torch.Tensor
#         A tensor of bounding boxes with shape N x 4
#     scores : torch.Tensor
#         A tensor of bounding box scores with shape N
#     threshold : float
#         The minimum iou value for a pair of bounding boxes to be considered a match
#     n_max : int
#         Maximum number of bounding boxes to select. If n_max <= 0, keep all surviving boxes

#     Returns
#     -------
#     selection : torch.Tensor
#         A tensor with the indices of the selected boxes

#     """
#     selection = _backend.nms(bbx, scores, threshold, n_max)
#     return selection.to(device=bbx.device)

models/grasp_det_seg/utils/nms/nms.py
andvg3/Grasp-Anything
import torch
import torch.nn.functional as F
import torch.autograd as autograd
from torch.autograd.function import once_differentiable

# from . import _backend

# _INTERPOLATION = {"bilinear": _backend.Interpolation.Bilinear, "nearest": _backend.Interpolation.Nearest}
# _PADDING = {"zero": _backend.PaddingMode.Zero, "border": _backend.PaddingMode.Border}


class ROISampling(autograd.Function):
    @staticmethod
    def forward(ctx, x, bbx, idx, roi_size, interpolation, padding, valid_mask):
        ctx.save_for_backward(bbx, idx)
        ctx.input_shape = (x.size(0), x.size(2), x.size(3))
        ctx.valid_mask = valid_mask

        try:
            ctx.interpolation = _INTERPOLATION[interpolation]
        except KeyError:
            raise ValueError("Unknown interpolation {}".format(interpolation))
        try:
            ctx.padding = _PADDING[padding]
        except KeyError:
            raise ValueError("Unknown padding {}".format(padding))

        y, mask = _backend.roi_sampling_forward(x, bbx, idx, roi_size, ctx.interpolation, ctx.padding, valid_mask)

        if not torch.is_floating_point(x):
            ctx.mark_non_differentiable(y)
        if valid_mask:
            ctx.mark_non_differentiable(mask)
            return y, mask
        else:
            return y

    @staticmethod
    @once_differentiable
    def backward(ctx, *args):
        if ctx.valid_mask:
            dy, _ = args
        else:
            dy = args[0]

        assert torch.is_floating_point(dy), "ROISampling.backward is only defined for floating point types"
        bbx, idx = ctx.saved_tensors

        dx = _backend.roi_sampling_backward(dy, bbx, idx, ctx.input_shape, ctx.interpolation, ctx.padding)
        return dx, None, None, None, None, None, None


def roi_sampling(x, bbx, idx, roi_size, interpolation='bilinear', padding_mode='zeros'):
    """
    Performs RoI sampling on the input feature map `x` based on the bounding boxes `bbx`.

    Arguments:
        x (torch.Tensor): Input feature map with shape (N, C, H, W).
        bbx (torch.Tensor): RoI bounding boxes with shape (M, 4), where M is the number of RoIs,
                            and each RoI is represented by (xmin, ymin, xmax, ymax).
        idx (torch.Tensor): The corresponding indices of the RoIs in the input feature map `x`.
                            Should be a 1D tensor of integers with shape (M,).
        roi_size (int): The output size of the RoI.
        interpolation (str): The interpolation method to use for resizing. Default is 'bilinear'.
        padding_mode (str): The padding mode to use for resizing. Default is 'zeros'.

    Returns:
        torch.Tensor: Output RoI feature map with shape (M, C, roi_size, roi_size).
    """
    num_rois = bbx.size(0)
    output = []

    for i in range(num_rois):
        roi_idx = int(idx[i])
        roi = bbx[i]
        xmin, ymin, xmax, ymax = roi

        # Calculate RoI coordinates in integer format
        roi_xmin = int(torch.round(xmin).item())
        roi_ymin = int(torch.round(ymin).item())
        roi_xmax = int(torch.round(xmax).item())
        roi_ymax = int(torch.round(ymax).item())

        # Crop the RoI region from the input feature map
        roi_feature = x[roi_idx:roi_idx + 1, :, roi_ymin:roi_ymax, roi_xmin:roi_xmax]

        # Resize the RoI to the desired output size using RoIAlign
        roi_feature = F.adaptive_avg_pool2d(roi_feature, roi_size)

        output.append(roi_feature)

    return torch.cat(output, dim=0)

# def roi_sampling(x, bbx, idx, roi_size, interpolation="bilinear", padding="border", valid_mask=False):
#     """Sample ROIs from a batch of images using bi-linear interpolation

#     ROIs are sampled from the input by bi-linear interpolation, using the following equations to transform from
#     ROI coordinates to image coordinates:

#         y_img = y0 + y_roi / h_roi * (y1 - y0),     for y_roi in range(0, h_roi)
#         x_img = x0 + x_roi / w_roi * (x1 - x0),     for x_roi in range(0, w_roi)

#     where `(h_roi, w_roi)` is the shape of the ROI and `(y0, x0, y1, x1)` are its bounding box coordinates on the image

#     Parameters
#     ----------
#     x : torch.Tensor
#         A tensor with shape N x C x H x W containing a batch of images to sample from
#     bbx : torch.Tensor
#         A tensor with shape K x 4 containing the bounding box coordinates of the ROIs in "corners" format
#     idx : torch.Tensor
#         A tensor with shape K containing the batch indices of the image each ROI should be sampled from
#     roi_size : tuple of int
#         The size `(h_roi, w_roi)` of the output ROIs
#     interpolation : str
#         Sampling mode, one of "bilinear" or "nearest"
#     padding : str
#         Padding mode, one of "border" or "zero"
#     valid_mask : bool
#         If `True` also return a mask tensor that indicates which points of the outputs where sampled from within the
#         valid region of the input

#     Returns
#     -------
#     y : torch.Tensor
#         A tensor with shape K x C x h_roi x w_roi containing the sampled ROIs
#     mask : torch.Tensor
#         Optional output returned only when valid_mask is `True`: a mask tensor with shape K x h_roi x w_roi, whose
#         entries are `!= 0` where the corresponding location in `y` was sampled from within the limits of the input image
#     """
#     return ROISampling.apply(x, bbx, idx, roi_size, interpolation, padding, valid_mask)


__all__ = ["roi_sampling"]

models/grasp_det_seg/utils/roi_sampling/functions.py
andvg3/Grasp-Anything
import os
import glob
import shutil
import pickle

def copy_files_with_structure(source_dir, destination_dir, filter_func=None):
    # Create the destination directory if it doesn't exist
    if not os.path.exists(destination_dir):
        os.makedirs(destination_dir)

    for root, _, files in os.walk(source_dir):
        # Get the relative path of the current directory within the source directory
        relative_path = os.path.relpath(root, source_dir)
        # Create the corresponding directory structure in the destination directory
        dest_dir = os.path.join(destination_dir, relative_path)

        # Create the directory in the destination directory
        if not os.path.exists(dest_dir):
            os.makedirs(dest_dir)

        # Copy each file in the current directory to the corresponding directory in the destination directory
        for file in files:
            source_file_path = os.path.join(root, file)
            dest_file_path = os.path.join(dest_dir, file)

            # Use the filter function if provided to check if the file should be copied
            if filter_func is None or filter_func(source_file_path):
                shutil.copy2(source_file_path, dest_file_path)  # Use shutil.copy2 to preserve metadata

# Example usage:
def seen_filter_cornell(file_path):
    idx = file_path.split('/')[-1].split('.')[0]
    idx = idx[:7]
    return idx in seen_filters

def unseen_filter_cornell(file_path):
    idx = file_path.split('/')[-1].split('.')[0]
    idx = idx[:7]
    return idx in unseen_filters

cornell_path = 'data/cornell'
cornell_filter_path = 'split/cornell'
source_directory = "data/cornell"
seen_destination_directory = "data/cornell_seen"
unseen_destination_directory = "data/cornell_unseen"

with open(os.path.join(cornell_filter_path, 'seen.obj'), 'rb') as f:
    seen_filters = pickle.load(f)

with open(os.path.join(cornell_filter_path, 'unseen.obj'), 'rb') as f:
    unseen_filters = pickle.load(f)

copy_files_with_structure(source_directory, seen_destination_directory, filter_func=seen_filter_cornell)
copy_files_with_structure(source_directory, unseen_destination_directory, filter_func=unseen_filter_cornell)

split/control_manage_cornell.py
andvg3/Grasp-Anything
import os
import glob
import shutil
import pickle

def copy_files_with_structure(source_dir, destination_dir, filter_func=None):
    # Create the destination directory if it doesn't exist
    if not os.path.exists(destination_dir):
        os.makedirs(destination_dir)

    for root, _, files in os.walk(source_dir):
        # Get the relative path of the current directory within the source directory
        relative_path = os.path.relpath(root, source_dir)
        # Create the corresponding directory structure in the destination directory
        dest_dir = os.path.join(destination_dir, relative_path)

        # Create the directory in the destination directory
        if not os.path.exists(dest_dir):
            os.makedirs(dest_dir)

        # Copy each file in the current directory to the corresponding directory in the destination directory
        for file in files:
            source_file_path = os.path.join(root, file)
            dest_file_path = os.path.join(dest_dir, file)

            # Use the filter function if provided to check if the file should be copied
            if filter_func is None or filter_func(source_file_path):
                shutil.copy2(source_file_path, dest_file_path)  # Use shutil.copy2 to preserve metadata

# Example usage:
def seen_filter_jacquard(file_path):
    idx = file_path.split('/')[-1].split('.')[0]
    idx = "_".join(idx.split('_')[:2])
    return idx in seen_filters

def unseen_filter_jacquard(file_path):
    idx = file_path.split('/')[-1].split('.')[0]
    idx = "_".join(idx.split('_')[:2])
    return idx in unseen_filters

jacquard_path = 'data/jacquard'
jacquard_filter_path = 'split/jacquard'
source_directory = "data/jacquard"
seen_destination_directory = "data/jacquard_seen"
unseen_destination_directory = "data/jacquard_unseen"

with open(os.path.join(jacquard_filter_path, 'seen.obj'), 'rb') as f:
    seen_filters = pickle.load(f)

with open(os.path.join(jacquard_filter_path, 'unseen.obj'), 'rb') as f:
    unseen_filters = pickle.load(f)

copy_files_with_structure(source_directory, seen_destination_directory, filter_func=seen_filter_jacquard)
copy_files_with_structure(source_directory, unseen_destination_directory, filter_func=unseen_filter_jacquard)

split/control_manage_jacquard.py
andvg3/Grasp-Anything
import argparse
import datetime
import json
import logging
import os
import sys

import cv2
import numpy as np
import tensorboardX
import torch
import torch.optim as optim
import torch.utils.data
from torchsummary import summary

from hardware.device import get_device
from inference.models import get_network
from inference.post_process import post_process_output
from utils.data import get_dataset
from utils.dataset_processing import evaluation
from utils.visualisation.gridshow import gridshow


def parse_args():
    parser = argparse.ArgumentParser(description='Train network')

    # Network
    parser.add_argument('--network', type=str, default='grconvnet3',
                        help='Network name in inference/models')
    parser.add_argument('--input-size', type=int, default=224,
                        help='Input image size for the network')
    parser.add_argument('--use-depth', type=int, default=1,
                        help='Use Depth image for training (1/0)')
    parser.add_argument('--use-rgb', type=int, default=1,
                        help='Use RGB image for training (1/0)')
    parser.add_argument('--use-dropout', type=int, default=1,
                        help='Use dropout for training (1/0)')
    parser.add_argument('--dropout-prob', type=float, default=0.1,
                        help='Dropout prob for training (0-1)')
    parser.add_argument('--channel-size', type=int, default=32,
                        help='Internal channel size for the network')
    parser.add_argument('--iou-threshold', type=float, default=0.25,
                        help='Threshold for IOU matching')

    # Datasets
    parser.add_argument('--dataset', type=str,
                        help='Dataset Name ("cornell" or "jaquard")')
    parser.add_argument('--dataset-path', type=str,
                        help='Path to dataset')
    parser.add_argument('--split', type=float, default=0.9,
                        help='Fraction of data for training (remainder is validation)')
    parser.add_argument('--ds-shuffle', action='store_true', default=False,
                        help='Shuffle the dataset')
    parser.add_argument('--ds-rotate', type=float, default=0.0,
                        help='Shift the start point of the dataset to use a different test/train split')
    parser.add_argument('--num-workers', type=int, default=8,
                        help='Dataset workers')

    # Training
    parser.add_argument('--batch-size', type=int, default=8,
                        help='Batch size')
    parser.add_argument('--epochs', type=int, default=50,
                        help='Training epochs')
    parser.add_argument('--batches-per-epoch', type=int, default=1000,
                        help='Batches per Epoch')
    parser.add_argument('--optim', type=str, default='adam',
                        help='Optmizer for the training. (adam or SGD)')

    # Logging etc.
    parser.add_argument('--description', type=str, default='',
                        help='Training description')
    parser.add_argument('--logdir', type=str, default='logs/',
                        help='Log directory')
    parser.add_argument('--vis', action='store_true',
                        help='Visualise the training process')
    parser.add_argument('--cpu', dest='force_cpu', action='store_true', default=False,
                        help='Force code to run in CPU mode')
    parser.add_argument('--random-seed', type=int, default=123,
                        help='Random seed for numpy')
    parser.add_argument('--seen', type=int, default=1,
                        help='Flag for using seen classes, only work for Grasp-Anything dataset') 

    args = parser.parse_args()
    return args


def validate(net, device, val_data, iou_threshold):
    """
    Run validation.
    :param net: Network
    :param device: Torch device
    :param val_data: Validation Dataset
    :param iou_threshold: IoU threshold
    :return: Successes, Failures and Losses
    """
    net.eval()

    results = {
        'correct': 0,
        'failed': 0,
        'loss': 0,
        'losses': {

        }
    }

    ld = len(val_data)

    with torch.no_grad():
        for x, y, didx, rot, zoom_factor in val_data:
            xc = x.to(device)
            yc = [yy.to(device) for yy in y]
            lossd = net.compute_loss(xc, yc)

            loss = lossd['loss']

            results['loss'] += loss.item() / ld
            for ln, l in lossd['losses'].items():
                if ln not in results['losses']:
                    results['losses'][ln] = 0
                results['losses'][ln] += l.item() / ld

            q_out, ang_out, w_out = post_process_output(lossd['pred']['pos'], lossd['pred']['cos'],
                                                        lossd['pred']['sin'], lossd['pred']['width'])

            s = evaluation.calculate_iou_match(q_out,
                                               ang_out,
                                               val_data.dataset.get_gtbb(didx, rot, zoom_factor),
                                               no_grasps=1,
                                               grasp_width=w_out,
                                               threshold=iou_threshold
                                               )

            if s:
                results['correct'] += 1
            else:
                results['failed'] += 1

    return results


def train(epoch, net, device, train_data, optimizer, batches_per_epoch, vis=False):
    """
    Run one training epoch
    :param epoch: Current epoch
    :param net: Network
    :param device: Torch device
    :param train_data: Training Dataset
    :param optimizer: Optimizer
    :param batches_per_epoch:  Data batches to train on
    :param vis:  Visualise training progress
    :return:  Average Losses for Epoch
    """
    results = {
        'loss': 0,
        'losses': {
        }
    }

    net.train()

    batch_idx = 0
    # Use batches per epoch to make training on different sized datasets (cornell/jacquard) more equivalent.
    while batch_idx <= batches_per_epoch:
        for x, y, _, _, _ in train_data:
            batch_idx += 1
            if batch_idx >= batches_per_epoch:
                break

            xc = x.to(device)
            yc = [yy.to(device) for yy in y]
            lossd = net.compute_loss(xc, yc)

            loss = lossd['loss']

            if batch_idx % 100 == 0:
                logging.info('Epoch: {}, Batch: {}, Loss: {:0.4f}'.format(epoch, batch_idx, loss.item()))

            results['loss'] += loss.item()
            for ln, l in lossd['losses'].items():
                if ln not in results['losses']:
                    results['losses'][ln] = 0
                results['losses'][ln] += l.item()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # Display the images
            if vis:
                imgs = []
                n_img = min(4, x.shape[0])
                for idx in range(n_img):
                    imgs.extend([x[idx,].numpy().squeeze()] + [yi[idx,].numpy().squeeze() for yi in y] + [
                        x[idx,].numpy().squeeze()] + [pc[idx,].detach().cpu().numpy().squeeze() for pc in
                                                      lossd['pred'].values()])
                gridshow('Display', imgs,
                         [(xc.min().item(), xc.max().item()), (0.0, 1.0), (0.0, 1.0), (-1.0, 1.0),
                          (0.0, 1.0)] * 2 * n_img,
                         [cv2.COLORMAP_BONE] * 10 * n_img, 10)
                cv2.waitKey(2)

    results['loss'] /= batch_idx
    for l in results['losses']:
        results['losses'][l] /= batch_idx

    return results


def run():
    args = parse_args()

    # Set-up output directories
    dt = datetime.datetime.now().strftime('%y%m%d_%H%M')
    net_desc = '{}_{}'.format(dt, '_'.join(args.description.split()))

    save_folder = os.path.join(args.logdir, net_desc)
    if not os.path.exists(save_folder):
        os.makedirs(save_folder)
    tb = tensorboardX.SummaryWriter(save_folder)

    # Save commandline args
    if args is not None:
        params_path = os.path.join(save_folder, 'commandline_args.json')
        with open(params_path, 'w') as f:
            json.dump(vars(args), f)

    # Initialize logging
    logging.root.handlers = []
    logging.basicConfig(
        level=logging.INFO,
        filename="{0}/{1}.log".format(save_folder, 'log'),
        format='[%(asctime)s] {%(pathname)s:%(lineno)d} %(levelname)s - %(message)s',
        datefmt='%H:%M:%S'
    )
    # set up logging to console
    console = logging.StreamHandler()
    console.setLevel(logging.DEBUG)
    # set a format which is simpler for console use
    formatter = logging.Formatter('%(name)-12s: %(levelname)-8s %(message)s')
    console.setFormatter(formatter)
    # add the handler to the root logger
    logging.getLogger('').addHandler(console)

    # Get the compute device
    device = get_device(args.force_cpu)

    # Load Dataset
    logging.info('Loading {} Dataset...'.format(args.dataset.title()))
    Dataset = get_dataset(args.dataset)
    dataset = Dataset(args.dataset_path,
                      output_size=args.input_size,
                      ds_rotate=args.ds_rotate,
                      random_rotate=True,
                      random_zoom=True,
                      include_depth=args.use_depth,
                      include_rgb=args.use_rgb,
                      seen=args.seen)
    logging.info('Dataset size is {}'.format(dataset.length))

    # Creating data indices for training and validation splits
    indices = list(range(dataset.length))
    split = int(np.floor(args.split * dataset.length))
    if args.ds_shuffle:
        np.random.seed(args.random_seed)
        np.random.shuffle(indices)
    train_indices, val_indices = indices[:split], indices[split:]
    logging.info('Training size: {}'.format(len(train_indices)))
    logging.info('Validation size: {}'.format(len(val_indices)))

    # Creating data samplers and loaders
    train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_indices)
    val_sampler = torch.utils.data.sampler.SubsetRandomSampler(val_indices)

    train_data = torch.utils.data.DataLoader(
        dataset,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        sampler=train_sampler
    )
    val_data = torch.utils.data.DataLoader(
        dataset,
        batch_size=1,
        num_workers=args.num_workers,
        sampler=val_sampler
    )
    logging.info('Done')

    # Load the network
    logging.info('Loading Network...')
    input_channels = 1 * args.use_depth + 3 * args.use_rgb
    network = get_network(args.network)
    net = network(
        input_channels=input_channels,
        dropout=args.use_dropout,
        prob=args.dropout_prob,
        channel_size=args.channel_size
    )

    net = net.to(device)
    logging.info('Done')

    if args.optim.lower() == 'adam':
        optimizer = optim.Adam(net.parameters())
    elif args.optim.lower() == 'sgd':
        optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)
    else:
        raise NotImplementedError('Optimizer {} is not implemented'.format(args.optim))

    # Print model architecture.
    # summary(net, (input_channels, args.input_size, args.input_size))
    # f = open(os.path.join(save_folder, 'arch.txt'), 'w')
    # sys.stdout = f
    # summary(net, (input_channels, args.input_size, args.input_size))
    # sys.stdout = sys.__stdout__
    # f.close()

    best_iou = 0.0
    for epoch in range(args.epochs):
        logging.info('Beginning Epoch {:02d}'.format(epoch))
        train_results = train(epoch, net, device, train_data, optimizer, args.batches_per_epoch, vis=args.vis)

        # Log training losses to tensorboard
        tb.add_scalar('loss/train_loss', train_results['loss'], epoch)
        for n, l in train_results['losses'].items():
            tb.add_scalar('train_loss/' + n, l, epoch)

        # Run Validation
        logging.info('Validating...')
        test_results = validate(net, device, val_data, args.iou_threshold)
        logging.info('%d/%d = %f' % (test_results['correct'], test_results['correct'] + test_results['failed'],
                                     test_results['correct'] / (test_results['correct'] + test_results['failed'])))

        # Log validation results to tensorbaord
        tb.add_scalar('loss/IOU', test_results['correct'] / (test_results['correct'] + test_results['failed']), epoch)
        tb.add_scalar('loss/val_loss', test_results['loss'], epoch)
        for n, l in test_results['losses'].items():
            tb.add_scalar('val_loss/' + n, l, epoch)

        # Save best performing network
        iou = test_results['correct'] / (test_results['correct'] + test_results['failed'])
        if iou > best_iou or epoch == 0 or (epoch % 10) == 0:
            torch.save(net, os.path.join(save_folder, 'epoch_%02d_iou_%0.2f' % (epoch, iou)))
            best_iou = iou


if __name__ == '__main__':
    run()

train_network.py
andvg3/Grasp-Anything
import argparse
import datetime
import json
import logging
import os
import sys

import cv2
import numpy as np
import tensorboardX
import torch
import torch.optim as optim
import torch.distributed as dist
import torch.utils.data
from torchsummary import summary

from hardware.device import get_device
from inference.models import get_network
from inference.post_process import post_process_output
from utils.data import get_dataset
from utils.dataset_processing import evaluation
from utils.visualisation.gridshow import gridshow

from models.grasp_det_seg import make_config, make_model


def parse_args():
    parser = argparse.ArgumentParser(description='Train network')

    # Network
    parser.add_argument('--network', type=str, default='grconvnet3',
                        help='Network name in inference/models')
    parser.add_argument('--input-size', type=int, default=224,
                        help='Input image size for the network')
    parser.add_argument('--use-depth', type=int, default=1,
                        help='Use Depth image for training (1/0)')
    parser.add_argument('--use-rgb', type=int, default=1,
                        help='Use RGB image for training (1/0)')
    parser.add_argument('--use-dropout', type=int, default=1,
                        help='Use dropout for training (1/0)')
    parser.add_argument('--dropout-prob', type=float, default=0.1,
                        help='Dropout prob for training (0-1)')
    parser.add_argument('--channel-size', type=int, default=32,
                        help='Internal channel size for the network')
    parser.add_argument('--iou-threshold', type=float, default=0.25,
                        help='Threshold for IOU matching')

    # Datasets
    parser.add_argument('--dataset', type=str,
                        help='Dataset Name ("cornell" or "jaquard")')
    parser.add_argument('--dataset-path', type=str,
                        help='Path to dataset')
    parser.add_argument('--split', type=float, default=0.9,
                        help='Fraction of data for training (remainder is validation)')
    parser.add_argument('--ds-shuffle', action='store_true', default=False,
                        help='Shuffle the dataset')
    parser.add_argument('--ds-rotate', type=float, default=0.0,
                        help='Shift the start point of the dataset to use a different test/train split')
    parser.add_argument('--num-workers', type=int, default=8,
                        help='Dataset workers')

    # Training
    parser.add_argument('--batch-size', type=int, default=8,
                        help='Batch size')
    parser.add_argument('--epochs', type=int, default=50,
                        help='Training epochs')
    parser.add_argument('--batches-per-epoch', type=int, default=1000,
                        help='Batches per Epoch')
    parser.add_argument('--optim', type=str, default='adam',
                        help='Optmizer for the training. (adam or SGD)')

    # Logging etc.
    parser.add_argument('--description', type=str, default='',
                        help='Training description')
    parser.add_argument('--logdir', type=str, default='logs/',
                        help='Log directory')
    parser.add_argument('--vis', action='store_true',
                        help='Visualise the training process')
    parser.add_argument('--cpu', dest='force_cpu', action='store_true', default=False,
                        help='Force code to run in CPU mode')
    parser.add_argument('--random-seed', type=int, default=123,
                        help='Random seed for numpy')
    parser.add_argument('--seen', type=int, default=1,
                        help='Flag for using seen classes, only work for Grasp-Anything dataset') 

    args = parser.parse_args()
    return args


def validate(net, device, val_data, iou_threshold):
    """
    Run validation.
    :param net: Network
    :param device: Torch device
    :param val_data: Validation Dataset
    :param iou_threshold: IoU threshold
    :return: Successes, Failures and Losses
    """
    net.eval()

    results = {
        'correct': 0,
        'failed': 0,
        'loss': 0,
        'losses': {

        }
    }

    ld = len(val_data)

    with torch.no_grad():
        for x, y, didx, rot, zoom_factor in val_data:
            xc = x.to(device)
            yc = [yy.to(device) for yy in y]
            lossd = net.compute_loss(xc, yc)

            loss = lossd['loss']

            results['loss'] += loss.item() / ld
            for ln, l in lossd['losses'].items():
                if ln not in results['losses']:
                    results['losses'][ln] = 0
                results['losses'][ln] += l.item() / ld

            q_out, ang_out, w_out = post_process_output(lossd['pred']['pos'], lossd['pred']['cos'],
                                                        lossd['pred']['sin'], lossd['pred']['width'])

            s = evaluation.calculate_iou_match(q_out,
                                               ang_out,
                                               val_data.dataset.get_gtbb(didx, rot, zoom_factor),
                                               no_grasps=1,
                                               grasp_width=w_out,
                                               threshold=iou_threshold
                                               )

            if s:
                results['correct'] += 1
            else:
                results['failed'] += 1

    return results


def train(epoch, net, device, train_data, optimizer, batches_per_epoch, vis=False):
    """
    Run one training epoch
    :param epoch: Current epoch
    :param net: Network
    :param device: Torch device
    :param train_data: Training Dataset
    :param optimizer: Optimizer
    :param batches_per_epoch:  Data batches to train on
    :param vis:  Visualise training progress
    :return:  Average Losses for Epoch
    """
    results = {
        'loss': 0,
        'losses': {
        }
    }

    net.train()

    batch_idx = 0
    # Use batches per epoch to make training on different sized datasets (cornell/jacquard) more equivalent.
    while batch_idx <= batches_per_epoch:
        for x, y, _, _, _ in train_data:
            batch_idx += 1
            if batch_idx >= batches_per_epoch:
                break

            xc = x.to(device)
            yc = [yy.to(device) for yy in y]
            lossd = net.compute_loss(xc, yc)
            loss = lossd['loss']

            if batch_idx % 100 == 0:
                logging.info('Epoch: {}, Batch: {}, Loss: {:0.4f}'.format(epoch, batch_idx, loss.item()))

            results['loss'] += loss.item()
            for ln, l in lossd['losses'].items():
                if ln not in results['losses']:
                    results['losses'][ln] = 0
                results['losses'][ln] += l.item()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # Display the images
            if vis:
                imgs = []
                n_img = min(4, x.shape[0])
                for idx in range(n_img):
                    imgs.extend([x[idx,].numpy().squeeze()] + [yi[idx,].numpy().squeeze() for yi in y] + [
                        x[idx,].numpy().squeeze()] + [pc[idx,].detach().cpu().numpy().squeeze() for pc in
                                                      lossd['pred'].values()])
                gridshow('Display', imgs,
                         [(xc.min().item(), xc.max().item()), (0.0, 1.0), (0.0, 1.0), (-1.0, 1.0),
                          (0.0, 1.0)] * 2 * n_img,
                         [cv2.COLORMAP_BONE] * 10 * n_img, 10)
                cv2.waitKey(2)

    results['loss'] /= batch_idx
    for l in results['losses']:
        results['losses'][l] /= batch_idx

    return results


def run():
    args = parse_args()
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("gloo", rank=0, world_size=1)

    # Set-up output directories
    dt = datetime.datetime.now().strftime('%y%m%d_%H%M')
    net_desc = '{}_{}'.format(dt, '_'.join(args.description.split()))

    save_folder = os.path.join(args.logdir, net_desc)
    if not os.path.exists(save_folder):
        os.makedirs(save_folder)
    tb = tensorboardX.SummaryWriter(save_folder)

    # Save commandline args
    if args is not None:
        params_path = os.path.join(save_folder, 'commandline_args.json')
        with open(params_path, 'w') as f:
            json.dump(vars(args), f)

    # Initialize logging
    logging.root.handlers = []
    logging.basicConfig(
        level=logging.INFO,
        filename="{0}/{1}.log".format(save_folder, 'log'),
        format='[%(asctime)s] {%(pathname)s:%(lineno)d} %(levelname)s - %(message)s',
        datefmt='%H:%M:%S'
    )
    # set up logging to console
    console = logging.StreamHandler()
    console.setLevel(logging.DEBUG)
    # set a format which is simpler for console use
    formatter = logging.Formatter('%(name)-12s: %(levelname)-8s %(message)s')
    console.setFormatter(formatter)
    # add the handler to the root logger
    logging.getLogger('').addHandler(console)

    # Get the compute device
    device = get_device(args.force_cpu)

    # Load Dataset
    logging.info('Loading {} Dataset...'.format(args.dataset.title()))
    Dataset = get_dataset(args.dataset)
    dataset = Dataset(args.dataset_path,
                      output_size=args.input_size,
                      ds_rotate=args.ds_rotate,
                      random_rotate=True,
                      random_zoom=True,
                      include_depth=args.use_depth,
                      include_rgb=args.use_rgb,
                      seen=args.seen)
    logging.info('Dataset size is {}'.format(dataset.length))

    # Creating data indices for training and validation splits
    indices = list(range(dataset.length))
    split = int(np.floor(args.split * dataset.length))
    if args.ds_shuffle:
        np.random.seed(args.random_seed)
        np.random.shuffle(indices)
    train_indices, val_indices = indices[:split], indices[split:]
    logging.info('Training size: {}'.format(len(train_indices)))
    logging.info('Validation size: {}'.format(len(val_indices)))

    # Creating data samplers and loaders
    train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_indices)
    val_sampler = torch.utils.data.sampler.SubsetRandomSampler(val_indices)

    train_data = torch.utils.data.DataLoader(
        dataset,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        sampler=train_sampler
    )
    val_data = torch.utils.data.DataLoader(
        dataset,
        batch_size=1,
        num_workers=args.num_workers,
        sampler=val_sampler
    )
    logging.info('Done')

    # Debugging
    logging.info('Debugging...')
    config = make_config(args)
    net = make_model(config)
    net.to(device)

    # Load the network
    logging.info('Loading Network...')
    input_channels = 1 * args.use_depth + 3 * args.use_rgb
    network = get_network(args.network)
    # net = network(
    #     input_channels=input_channels,
    #     dropout=args.use_dropout,
    #     prob=args.dropout_prob,
    #     channel_size=args.channel_size
    # )

    net = net.to(device)
    logging.info('Done')

    if args.optim.lower() == 'adam':
        optimizer = optim.Adam(net.parameters())
    elif args.optim.lower() == 'sgd':
        optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)
    else:
        raise NotImplementedError('Optimizer {} is not implemented'.format(args.optim))

    # Print model architecture.
    # summary(net, (input_channels, args.input_size, args.input_size))
    # f = open(os.path.join(save_folder, 'arch.txt'), 'w')
    # sys.stdout = f
    # summary(net, (input_channels, args.input_size, args.input_size))
    # sys.stdout = sys.__stdout__
    # f.close()

    best_iou = 0.0
    for epoch in range(args.epochs):
        logging.info('Beginning Epoch {:02d}'.format(epoch))
        train_results = train(epoch, net, device, train_data, optimizer, args.batches_per_epoch, vis=args.vis)

        # Log training losses to tensorboard
        tb.add_scalar('loss/train_loss', train_results['loss'], epoch)
        for n, l in train_results['losses'].items():
            tb.add_scalar('train_loss/' + n, l, epoch)

        # Run Validation
        logging.info('Validating...')
        test_results = validate(net, device, val_data, args.iou_threshold)
        logging.info('%d/%d = %f' % (test_results['correct'], test_results['correct'] + test_results['failed'],
                                     test_results['correct'] / (test_results['correct'] + test_results['failed'])))

        # Log validation results to tensorbaord
        tb.add_scalar('loss/IOU', test_results['correct'] / (test_results['correct'] + test_results['failed']), epoch)
        tb.add_scalar('loss/val_loss', test_results['loss'], epoch)
        for n, l in test_results['losses'].items():
            tb.add_scalar('val_loss/' + n, l, epoch)

        # Save best performing network
        iou = test_results['correct'] / (test_results['correct'] + test_results['failed'])
        if iou > best_iou or epoch == 0 or (epoch % 10) == 0:
            torch.save(net, os.path.join(save_folder, 'epoch_%02d_iou_%0.2f' % (epoch, iou)))
            best_iou = iou


if __name__ == '__main__':
    run()

train_network_grasp_det_seg.py
andvg3/Grasp-Anything
def get_dataset(dataset_name):
    if dataset_name == 'cornell':
        from .cornell_data import CornellDataset
        return CornellDataset
    elif dataset_name == 'jacquard':
        from .jacquard_data import JacquardDataset
        return JacquardDataset
    elif dataset_name == 'grasp-anything':
        from .grasp_anything_data import GraspAnythingDataset
        return GraspAnythingDataset
    elif dataset_name == 'vmrd':
        from .vmrd_data import VMRDDataset
        return VMRDDataset
    elif dataset_name == 'ocid':
        from .ocid_grasp_data import OCIDGraspDataset
        return OCIDGraspDataset
    else:
        raise NotImplementedError('Dataset Type {} is Not implemented'.format(dataset_name))

utils/data/__init__.py
andvg3/Grasp-Anything
import glob
import os
import re

import pickle
import torch

from utils.dataset_processing import grasp, image, mask
from .grasp_data import GraspDatasetBase


class GraspAnythingDataset(GraspDatasetBase):
    """
    Dataset wrapper for the Grasp-Anything dataset.
    """

    def __init__(self, file_path, ds_rotate=0, **kwargs):
        """
        :param file_path: Grasp-Anything Dataset directory.
        :param ds_rotate: If splitting the dataset, rotate the list of items by this fraction first
        :param kwargs: kwargs for GraspDatasetBase
        """
        super(GraspAnythingDataset, self).__init__(**kwargs)

        self.grasp_files = glob.glob(os.path.join(file_path, 'positive_grasp', '*.pt'))
        self.prompt_files = glob.glob(os.path.join(file_path, 'prompt', '*.pkl'))
        self.rgb_files = glob.glob(os.path.join(file_path, 'image', '*.jpg'))
        self.mask_files = glob.glob(os.path.join(file_path, 'mask', '*.npy'))

        if kwargs["seen"]:
            with open(os.path.join('split/grasp-anything/seen.obj'), 'rb') as f:
                idxs = pickle.load(f)

            self.grasp_files = list(filter(lambda x: x.split('/')[-1].split('.')[0] in idxs, self.grasp_files))
        else:
            with open(os.path.join('split/grasp-anything/unseen.obj'), 'rb') as f:
                idxs = pickle.load(f)

            self.grasp_files = list(filter(lambda x: x.split('/')[-1].split('.')[0] in idxs, self.grasp_files))

        self.grasp_files.sort()
        self.prompt_files.sort()
        self.rgb_files.sort()
        self.mask_files.sort()

        self.length = len(self.grasp_files)

        if self.length == 0:
            raise FileNotFoundError('No dataset files found. Check path: {}'.format(file_path))

        if ds_rotate:
            self.grasp_files = self.grasp_files[int(self.length * ds_rotate):] + self.grasp_files[
                                                                                 :int(self.length * ds_rotate)]
            

    def _get_crop_attrs(self, idx):
        gtbbs = grasp.GraspRectangles.load_from_grasp_anything_file(self.grasp_files[idx])
        center = gtbbs.center
        left = max(0, min(center[1] - self.output_size // 2, 416 - self.output_size))
        top = max(0, min(center[0] - self.output_size // 2, 416 - self.output_size))
        return center, left, top

    def get_gtbb(self, idx, rot=0, zoom=1.0):       
        # Jacquard try
        gtbbs = grasp.GraspRectangles.load_from_grasp_anything_file(self.grasp_files[idx], scale=self.output_size / 416.0)

        c = self.output_size // 2
        gtbbs.rotate(rot, (c, c))
        gtbbs.zoom(zoom, (c, c))

        # Cornell try
        # gtbbs = grasp.GraspRectangles.load_from_grasp_anything_file(self.grasp_files[idx])
        # center, left, top = self._get_crop_attrs(idx)
        # gtbbs.rotate(rot, center)
        # gtbbs.offset((-top, -left))
        # gtbbs.zoom(zoom, (self.output_size // 2, self.output_size // 2))
        return gtbbs

    def get_depth(self, idx, rot=0, zoom=1.0):
        depth_img = image.DepthImage.from_tiff(self.depth_files[idx])
        center, left, top = self._get_crop_attrs(idx)
        depth_img.rotate(rot, center)
        depth_img.crop((top, left), (min(480, top + self.output_size), min(640, left + self.output_size)))
        depth_img.normalise()
        depth_img.zoom(zoom)
        depth_img.resize((self.output_size, self.output_size))
        return depth_img.img

    def get_rgb(self, idx, rot=0, zoom=1.0, normalise=True):
        mask_file = self.grasp_files[idx].replace("positive_grasp", "mask").replace(".pt", ".npy")
        mask_img = mask.Mask.from_file(mask_file)
        rgb_file = re.sub(r"_\d{1}\.pt", ".jpg", self.grasp_files[idx])
        rgb_file = rgb_file.replace("positive_grasp", "image")
        rgb_img = image.Image.from_file(rgb_file)
        rgb_img = image.Image.mask_out_image(rgb_img, mask_img)

        # Jacquard try
        rgb_img.rotate(rot)
        rgb_img.zoom(zoom)
        rgb_img.resize((self.output_size, self.output_size))
        if normalise:
            rgb_img.normalise()
            rgb_img.img = rgb_img.img.transpose((2, 0, 1))
        return rgb_img.img

        # Cornell try
        # center, left, top = self._get_crop_attrs(idx)
        # rgb_img.rotate(rot, center)
        # rgb_img.crop((top, left), (min(480, top + self.output_size), min(640, left + self.output_size)))
        # rgb_img.zoom(zoom)
        # rgb_img.resize((self.output_size, self.output_size))
        # if normalise:
        #     rgb_img.normalise()
        #     rgb_img.img = rgb_img.img.transpose((2, 0, 1))
        # return rgb_img.img

utils/data/grasp_anything_data.py
andvg3/Grasp-Anything
import random

import numpy as np
import torch
import torch.utils.data


class GraspDatasetBase(torch.utils.data.Dataset):
    """
    An abstract dataset for training networks in a common format.
    """

    def __init__(self, output_size=224, include_depth=True, include_rgb=False, random_rotate=False,
                 random_zoom=False, input_only=False, seen=True):
        """
        :param output_size: Image output size in pixels (square)
        :param include_depth: Whether depth image is included
        :param include_rgb: Whether RGB image is included
        :param random_rotate: Whether random rotations are applied
        :param random_zoom: Whether random zooms are applied
        :param input_only: Whether to return only the network input (no labels)
        """
        self.output_size = output_size
        self.random_rotate = random_rotate
        self.random_zoom = random_zoom
        self.input_only = input_only
        self.include_depth = include_depth
        self.include_rgb = include_rgb

        self.grasp_files = []

        if include_depth is False and include_rgb is False:
            raise ValueError('At least one of Depth or RGB must be specified.')

    @staticmethod
    def numpy_to_torch(s):
        if len(s.shape) == 2:
            return torch.from_numpy(np.expand_dims(s, 0).astype(np.float32))
        else:
            return torch.from_numpy(s.astype(np.float32))

    def get_gtbb(self, idx, rot=0, zoom=1.0):
        raise NotImplementedError()

    def get_depth(self, idx, rot=0, zoom=1.0):
        raise NotImplementedError()

    def get_rgb(self, idx, rot=0, zoom=1.0):
        raise NotImplementedError()

    def __getitem__(self, idx):
        if self.random_rotate:
            rotations = [0, np.pi / 2, 2 * np.pi / 2, 3 * np.pi / 2]
            rot = random.choice(rotations)
        else:
            rot = 0.0

        if self.random_zoom:
            zoom_factor = np.random.uniform(0.5, 1.0)
        else:
            zoom_factor = 1.0

        # Load the depth image
        if self.include_depth:
            depth_img = self.get_depth(idx, rot, zoom_factor)

        # Load the RGB image
        if self.include_rgb:
            rgb_img = self.get_rgb(idx, rot, zoom_factor)

        # Load the grasps
        bbs = self.get_gtbb(idx, rot, zoom_factor)

        pos_img, ang_img, width_img = bbs.draw((self.output_size, self.output_size))
        width_img = np.clip(width_img, 0.0, self.output_size / 2) / (self.output_size / 2)

        if self.include_depth and self.include_rgb:
            x = self.numpy_to_torch(
                np.concatenate(
                    (np.expand_dims(depth_img, 0),
                     rgb_img),
                    0
                )
            )
        elif self.include_depth:
            x = self.numpy_to_torch(depth_img)
        elif self.include_rgb:
            x = self.numpy_to_torch(rgb_img)

        pos = self.numpy_to_torch(pos_img)
        cos = self.numpy_to_torch(np.cos(2 * ang_img))
        sin = self.numpy_to_torch(np.sin(2 * ang_img))
        width = self.numpy_to_torch(width_img)

        return x, (pos, cos, sin, width), idx, rot, zoom_factor

    def __len__(self):
        return len(self.grasp_files)

utils/data/grasp_data.py
andvg3/Grasp-Anything
import glob
import os

from utils.dataset_processing import grasp, image
from .grasp_data import GraspDatasetBase


class OCIDGraspDataset(GraspDatasetBase):
    """
    Dataset wrapper for the Cornell dataset.
    """

    def __init__(self, file_path, ds_rotate=0, **kwargs):
        """
        :param file_path: Cornell Dataset directory.
        :param ds_rotate: If splitting the dataset, rotate the list of items by this fraction first
        :param kwargs: kwargs for GraspDatasetBase
        """
        super(OCIDGraspDataset, self).__init__(**kwargs)

        self.grasp_files = glob.glob(os.path.join(file_path, '*', '*', '*', '*', '*', 'Annotations', '*.txt'))
        self.grasp_files.sort()
        self.length = len(self.grasp_files)

        if self.length == 0:
            raise FileNotFoundError('No dataset files found. Check path: {}'.format(file_path))

        if ds_rotate:
            self.grasp_files = self.grasp_files[int(self.length * ds_rotate):] + self.grasp_files[
                                                                                 :int(self.length * ds_rotate)]

    def _get_crop_attrs(self, idx):
        gtbbs = grasp.GraspRectangles.load_from_ocid_grasp_file(self.grasp_files[idx])
        center = gtbbs.center
        left = max(0, min(center[1] - self.output_size // 2, 640 - self.output_size))
        top = max(0, min(center[0] - self.output_size // 2, 480 - self.output_size))
        return center, left, top

    def get_gtbb(self, idx, rot=0, zoom=1.0):
        gtbbs = grasp.GraspRectangles.load_from_cornell_file(self.grasp_files[idx])
        center, left, top = self._get_crop_attrs(idx)
        gtbbs.rotate(rot, center)
        gtbbs.offset((-top, -left))
        gtbbs.zoom(zoom, (self.output_size // 2, self.output_size // 2))
        return gtbbs

    def get_depth(self, idx, rot=0, zoom=1.0):
        depth_img = image.DepthImage.from_tiff(self.depth_files[idx])
        center, left, top = self._get_crop_attrs(idx)
        depth_img.rotate(rot, center)
        depth_img.crop((top, left), (min(480, top + self.output_size), min(640, left + self.output_size)))
        depth_img.normalise()
        depth_img.zoom(zoom)
        depth_img.resize((self.output_size, self.output_size))
        return depth_img.img

    def get_rgb(self, idx, rot=0, zoom=1.0, normalise=True):
        rgb_file = self.grasp_files[idx].replace("Annotations", "rgb").replace("txt", "png")
        rgb_img = image.Image.from_file(rgb_file)
        center, left, top = self._get_crop_attrs(idx)
        rgb_img.rotate(rot, center)
        rgb_img.crop((top, left), (min(480, top + self.output_size), min(640, left + self.output_size)))
        rgb_img.zoom(zoom)
        rgb_img.resize((self.output_size, self.output_size))
        if normalise:
            rgb_img.normalise()
            rgb_img.img = rgb_img.img.transpose((2, 0, 1))
        return rgb_img.img

utils/data/ocid_grasp_data.py
andvg3/Grasp-Anything
import glob
import os
import re

import pickle
import torch

from utils.dataset_processing import grasp, image, mask
from .grasp_data import GraspDatasetBase


class VMRDDataset(GraspDatasetBase):
    """
    Dataset wrapper for the Grasp-Anything dataset.
    """

    def __init__(self, file_path, ds_rotate=0, **kwargs):
        """
        :param file_path: Grasp-Anything Dataset directory.
        :param ds_rotate: If splitting the dataset, rotate the list of items by this fraction first
        :param kwargs: kwargs for GraspDatasetBase
        """
        super(VMRDDataset, self).__init__(**kwargs)

        self.grasp_files = glob.glob(os.path.join(file_path, 'Grasps', '*.txt'))
        self.rgb_files = glob.glob(os.path.join(file_path, 'JPEGImages', '*.jpg'))
        self.length = len(self.grasp_files)

        if self.length == 0:
            raise FileNotFoundError('No dataset files found. Check path: {}'.format(file_path))

        if ds_rotate:
            self.grasp_files = self.grasp_files[int(self.length * ds_rotate):] + self.grasp_files[
                                                                                 :int(self.length * ds_rotate)]
            

    def _get_crop_attrs(self, idx):
        gtbbs = grasp.GraspRectangles.load_from_vmrd_file(self.grasp_files[idx])
        center = gtbbs.center
        left = max(0, min(center[1] - self.output_size // 2, 1008 - self.output_size))
        top = max(0, min(center[0] - self.output_size // 2, 756 - self.output_size))
        return center, left, top

    def get_gtbb(self, idx, rot=0, zoom=1.0):       
        # Jacquard try
        gtbbs = grasp.GraspRectangles.load_from_vmrd_file(self.grasp_files[idx])
        center, left, top = self._get_crop_attrs(idx)
        gtbbs.rotate(rot, center)
        gtbbs.offset((-top, -left))
        gtbbs.zoom(zoom, (self.output_size // 2, self.output_size // 2))
        return gtbbs

    def get_depth(self, idx, rot=0, zoom=1.0):
        depth_img = image.DepthImage.from_tiff(self.depth_files[idx])
        center, left, top = self._get_crop_attrs(idx)
        depth_img.rotate(rot, center)
        depth_img.crop((top, left), (min(480, top + self.output_size), min(640, left + self.output_size)))
        depth_img.normalise()
        depth_img.zoom(zoom)
        depth_img.resize((self.output_size, self.output_size))
        return depth_img.img

    def get_rgb(self, idx, rot=0, zoom=1.0, normalise=True):
        rgb_file = self.grasp_files[idx].replace("Grasps", "JPEGImages").replace("txt", "jpg")
        rgb_img = image.Image.from_file(rgb_file)

        # Cornell try
        center, left, top = self._get_crop_attrs(idx)
        rgb_img.rotate(rot, center)
        rgb_img.crop((top, left), (min(756, top + self.output_size), min(1008, left + self.output_size)))
        rgb_img.zoom(zoom)
        rgb_img.resize((self.output_size, self.output_size))
        if normalise:
            rgb_img.normalise()
            rgb_img.img = rgb_img.img.transpose((2, 0, 1))

        return rgb_img.img

utils/data/vmrd_data.py
andvg3/Grasp-Anything
import torch
import matplotlib.pyplot as plt
import numpy as np
from skimage.draw import polygon
from skimage.feature import peak_local_max


def _gr_text_to_no(l, offset=(0, 0)):
    """
    Transform a single point from a Cornell file line to a pair of ints.
    :param l: Line from Cornell grasp file (str)
    :param offset: Offset to apply to point positions
    :return: Point [y, x]
    """
    x, y = l.split()
    return [int(round(float(y))) - offset[0], int(round(float(x))) - offset[1]]


def _grasp_anything_format(grasp: list):
    _, x, y, w, h, theta = grasp
    # index based on row, column (y,x), and the Grasp-Anything dataset's angles are flipped around an axis.
    return Grasp(np.array([y, x]), -theta / 180.0 * np.pi, w, h).as_gr


class GraspRectangles:
    """
    Convenience class for loading and operating on sets of Grasp Rectangles.
    """

    def __init__(self, grs=None):
        if grs:
            self.grs = grs
        else:
            self.grs = []

    def __getitem__(self, item):
        return self.grs[item]

    def __iter__(self):
        return self.grs.__iter__()

    def __getattr__(self, attr):
        """
        Test if GraspRectangle has the desired attr as a function and call it.
        """
        # Fuck yeah python.
        if hasattr(GraspRectangle, attr) and callable(getattr(GraspRectangle, attr)):
            return lambda *args, **kwargs: list(map(lambda gr: getattr(gr, attr)(*args, **kwargs), self.grs))
        else:
            raise AttributeError("Couldn't find function %s in BoundingBoxes or BoundingBox" % attr)

    @classmethod
    def load_from_array(cls, arr):
        """
        Load grasp rectangles from numpy array.
        :param arr: Nx4x2 array, where each 4x2 array is the 4 corner pixels of a grasp rectangle.
        :return: GraspRectangles()
        """
        grs = []
        for i in range(arr.shape[0]):
            grp = arr[i, :, :].squeeze()
            if grp.max() == 0:
                break
            else:
                grs.append(GraspRectangle(grp))
        return cls(grs)

    @classmethod
    def load_from_cornell_file(cls, fname):
        """
        Load grasp rectangles from a Cornell dataset grasp file.
        :param fname: Path to text file.
        :return: GraspRectangles()
        """
        grs = []
        with open(fname) as f:
            while True:
                # Load 4 lines at a time, corners of bounding box.
                p0 = f.readline()
                if not p0:
                    break  # EOF
                p1, p2, p3 = f.readline(), f.readline(), f.readline()
                try:
                    gr = np.array([
                        _gr_text_to_no(p0),
                        _gr_text_to_no(p1),
                        _gr_text_to_no(p2),
                        _gr_text_to_no(p3)
                    ])

                    grs.append(GraspRectangle(gr))

                except ValueError:
                    # Some files contain weird values.
                    continue
        
        return cls(grs)

    @classmethod
    def load_from_ocid_grasp_file(cls, fname):
        """
        Load grasp rectangles from a Cornell dataset grasp file.
        :param fname: Path to text file.
        :return: GraspRectangles()
        """
        grs = []
        with open(fname) as f:
            while True:
                # Load 4 lines at a time, corners of bounding box.
                p0 = f.readline()
                if not p0:
                    break  # EOF
                p1, p2, p3 = f.readline(), f.readline(), f.readline()
                try:
                    gr = np.array([
                        _gr_text_to_no(p0),
                        _gr_text_to_no(p1),
                        _gr_text_to_no(p2),
                        _gr_text_to_no(p3)
                    ])

                    grs.append(GraspRectangle(gr))

                except ValueError:
                    # Some files contain weird values.
                    continue
        
        return cls(grs)

    @classmethod
    def load_from_vmrd_file(cls, fname):
        """
        Load grasp rectangles from a VMRD dataset grasp file.
        :param fname: Path to text file.
        :return: GraspRectangles()
        """
        grs = []
        with open(fname) as f:
            grasp_lines = f.readlines()
            for grasp_line in grasp_lines:
                # x1, y1, x2, y2, x3, y3, x4, y4 = list(map(lambda x: int(round(float(x))), grasp_line.split(' ')[:8]))
                y1, x1, y2, x2, y3, x3, y4, x4 = list(map(lambda x: int(round(float(x))), grasp_line.split(' ')[:8]))
                try:
                    gr = np.array([
                        [x1, y1],
                        [x2, y2],
                        [x3, y3],
                        [x4, y4],
                    ])

                    grs.append(GraspRectangle(gr))

                except ValueError:
                    # Some files contain weird values.
                    continue
        
        return cls(grs)

    @classmethod
    def load_from_jacquard_file(cls, fname, scale=1.0):
        """
        Load grasp rectangles from a Jacquard dataset file.
        :param fname: Path to file.
        :param scale: Scale to apply (e.g. if resizing images)
        :return: GraspRectangles()
        """
        grs = []
        with open(fname) as f:
            for l in f:
                x, y, theta, w, h = [float(v) for v in l[:-1].split(';')]
                # index based on row, column (y,x), and the Jacquard dataset's angles are flipped around an axis.
                grs.append(Grasp(np.array([y, x]), -theta / 180.0 * np.pi, w, h).as_gr)
        grs = cls(grs)
        grs.scale(scale)
        return grs

    @classmethod
    def load_from_grasp_anything_file(cls, fname, scale=1.0):
        """
        Load grasp rectangles from a Grasp-Anything dataset grasp file.
        :param fname: Path to text file.
        :return: GraspRectangles()
        """
        grs = None
        with open(fname, 'rb') as f:
            pos_grasps = torch.load(f)
        # add_fn = fname.replace("positive_grasp", "negative_grasp")
        # with open(fname, 'rb') as f:
        #     neg_grasps = torch.load(f)
        # grasps = torch.cat((pos_grasps, neg_grasps), dim=0).tolist()
        grasps = pos_grasps.tolist()
        grs = list(map(lambda x: _grasp_anything_format(x), grasps))

        grs = cls(grs)
        grs.scale(scale)
        return grs
    
    def append(self, gr):
        """
        Add a grasp rectangle to this GraspRectangles object
        :param gr: GraspRectangle
        """
        self.grs.append(gr)

    def copy(self):
        """
        :return: A deep copy of this object and all of its GraspRectangles.
        """
        new_grs = GraspRectangles()
        for gr in self.grs:
            new_grs.append(gr.copy())
        return new_grs

    def show(self, ax=None, shape=None):
        """
        Draw all GraspRectangles on a matplotlib plot.
        :param ax: (optional) existing axis
        :param shape: (optional) Plot shape if no existing axis
        """
        if ax is None:
            f = plt.figure()
            ax = f.add_subplot(1, 1, 1)
            ax.imshow(np.zeros(shape))
            ax.axis([0, shape[1], shape[0], 0])
            self.plot(ax)
            plt.show()
        else:
            self.plot(ax)

    def draw(self, shape, position=True, angle=True, width=True):
        """
        Plot all GraspRectangles as solid rectangles in a numpy array, e.g. as network training data.
        :param shape: output shape
        :param position: If True, Q output will be produced
        :param angle: If True, Angle output will be produced
        :param width: If True, Width output will be produced
        :return: Q, Angle, Width outputs (or None)
        """
        if position:
            pos_out = np.zeros(shape)
        else:
            pos_out = None
        if angle:
            ang_out = np.zeros(shape)
        else:
            ang_out = None
        if width:
            width_out = np.zeros(shape)
        else:
            width_out = None

        for gr in self.grs:
            rr, cc = gr.compact_polygon_coords(shape)
            if position:
                pos_out[rr, cc] = 1.0
            if angle:
                ang_out[rr, cc] = gr.angle
            if width:
                width_out[rr, cc] = gr.length

        return pos_out, ang_out, width_out

    def to_array(self, pad_to=0):
        """
        Convert all GraspRectangles to a single array.
        :param pad_to: Length to 0-pad the array along the first dimension
        :return: Nx4x2 numpy array
        """
        a = np.stack([gr.points for gr in self.grs])
        if pad_to:
            if pad_to > len(self.grs):
                a = np.concatenate((a, np.zeros((pad_to - len(self.grs), 4, 2))))
        return a.astype(np.int)

    @property
    def center(self):
        """
        Compute mean center of all GraspRectangles
        :return: float, mean centre of all GraspRectangles
        """
        points = [gr.points for gr in self.grs]
        return np.mean(np.vstack(points), axis=0).astype(np.int)


class GraspRectangle:
    """
    Representation of a grasp in the common "Grasp Rectangle" format.
    """

    def __init__(self, points):
        self.points = points

    def __str__(self):
        return str(self.points)

    @property
    def angle(self):
        """
        :return: Angle of the grasp to the horizontal.
        """
        dx = self.points[1, 1] - self.points[0, 1]
        dy = self.points[1, 0] - self.points[0, 0]
        return (np.arctan2(-dy, dx) + np.pi / 2) % np.pi - np.pi / 2

    @property
    def as_grasp(self):
        """
        :return: GraspRectangle converted to a Grasp
        """
        return Grasp(self.center, self.angle, self.length, self.width)

    @property
    def center(self):
        """
        :return: Rectangle center point
        """
        return self.points.mean(axis=0).astype(np.int)

    @property
    def length(self):
        """
        :return: Rectangle length (i.e. along the axis of the grasp)
        """
        dx = self.points[1, 1] - self.points[0, 1]
        dy = self.points[1, 0] - self.points[0, 0]
        return np.sqrt(dx ** 2 + dy ** 2)

    @property
    def width(self):
        """
        :return: Rectangle width (i.e. perpendicular to the axis of the grasp)
        """
        dy = self.points[2, 1] - self.points[1, 1]
        dx = self.points[2, 0] - self.points[1, 0]
        return np.sqrt(dx ** 2 + dy ** 2)

    def polygon_coords(self, shape=None):
        """
        :param shape: Output Shape
        :return: Indices of pixels within the grasp rectangle polygon.
        """
        return polygon(self.points[:, 0], self.points[:, 1], shape)

    def compact_polygon_coords(self, shape=None):
        """
        :param shape: Output shape
        :return: Indices of pixels within the centre thrid of the grasp rectangle.
        """
        return Grasp(self.center, self.angle, self.length / 3, self.width).as_gr.polygon_coords(shape)

    def iou(self, gr, angle_threshold=np.pi / 6):
        """
        Compute IoU with another grasping rectangle
        :param gr: GraspingRectangle to compare
        :param angle_threshold: Maximum angle difference between GraspRectangles
        :return: IoU between Grasp Rectangles
        """
        if abs((self.angle - gr.angle + np.pi / 2) % np.pi - np.pi / 2) > angle_threshold:
            return 0

        rr1, cc1 = self.polygon_coords()
        rr2, cc2 = polygon(gr.points[:, 0], gr.points[:, 1])

        try:
            r_max = max(rr1.max(), rr2.max()) + 1
            c_max = max(cc1.max(), cc2.max()) + 1
        except:
            return 0

        canvas = np.zeros((r_max, c_max))
        canvas[rr1, cc1] += 1
        canvas[rr2, cc2] += 1
        union = np.sum(canvas > 0)
        if union == 0:
            return 0
        intersection = np.sum(canvas == 2)
        return intersection / union

    def copy(self):
        """
        :return: Copy of self.
        """
        return GraspRectangle(self.points.copy())

    def offset(self, offset):
        """
        Offset grasp rectangle
        :param offset: array [y, x] distance to offset
        """
        self.points += np.array(offset).reshape((1, 2))

    def rotate(self, angle, center):
        """
        Rotate grasp rectangle
        :param angle: Angle to rotate (in radians)
        :param center: Point to rotate around (e.g. image center)
        """
        R = np.array(
            [
                [np.cos(-angle), np.sin(-angle)],
                [-1 * np.sin(-angle), np.cos(-angle)],
            ]
        )
        c = np.array(center).reshape((1, 2))
        self.points = ((np.dot(R, (self.points - c).T)).T + c).astype(np.int)

    def scale(self, factor):
        """
        :param factor: Scale grasp rectangle by factor
        """
        if factor == 1.0:
            return
        self.points *= factor

    def plot(self, ax, color=None):
        """
        Plot grasping rectangle.
        :param ax: Existing matplotlib axis
        :param color: matplotlib color code (optional)
        """
        points = np.vstack((self.points, self.points[0]))
        ax.plot(points[:, 1], points[:, 0], color=color)

    def zoom(self, factor, center):
        """
        Zoom grasp rectangle by given factor.
        :param factor: Zoom factor
        :param center: Zoom zenter (focus point, e.g. image center)
        """
        T = np.array(
            [
                [1 / factor, 0],
                [0, 1 / factor]
            ]
        )
        c = np.array(center).reshape((1, 2))
        self.points = ((np.dot(T, (self.points - c).T)).T + c).astype(np.int)


class Grasp:
    """
    A Grasp represented by a center pixel, rotation angle and gripper width (length)
    """

    def __init__(self, center, angle, length=60, width=30):
        self.center = center
        self.angle = angle  # Positive angle means rotate anti-clockwise from horizontal.
        self.length = length
        self.width = width

    @property
    def as_gr(self):
        """
        Convert to GraspRectangle
        :return: GraspRectangle representation of grasp.
        """
        xo = np.cos(self.angle)
        yo = np.sin(self.angle)

        y1 = self.center[0] + self.length / 2 * yo
        x1 = self.center[1] - self.length / 2 * xo
        y2 = self.center[0] - self.length / 2 * yo
        x2 = self.center[1] + self.length / 2 * xo

        return GraspRectangle(np.array(
            [
                [y1 - self.width / 2 * xo, x1 - self.width / 2 * yo],
                [y2 - self.width / 2 * xo, x2 - self.width / 2 * yo],
                [y2 + self.width / 2 * xo, x2 + self.width / 2 * yo],
                [y1 + self.width / 2 * xo, x1 + self.width / 2 * yo],
            ]
        ).astype(np.float))

    def max_iou(self, grs):
        """
        Return maximum IoU between self and a list of GraspRectangles
        :param grs: List of GraspRectangles
        :return: Maximum IoU with any of the GraspRectangles
        """
        self_gr = self.as_gr
        max_iou = 0
        for gr in grs:
            iou = self_gr.iou(gr)
            max_iou = max(max_iou, iou)
        return max_iou

    def plot(self, ax, color=None):
        """
        Plot Grasp
        :param ax: Existing matplotlib axis
        :param color: (optional) color
        """
        self.as_gr.plot(ax, color)

    def to_jacquard(self, scale=1):
        """
        Output grasp in "Jacquard Dataset Format" (https://jacquard.liris.cnrs.fr/database.php)
        :param scale: (optional) scale to apply to grasp
        :return: string in Jacquard format
        """
        # Output in jacquard format.
        return '%0.2f;%0.2f;%0.2f;%0.2f;%0.2f' % (
            self.center[1] * scale, self.center[0] * scale, -1 * self.angle * 180 / np.pi, self.length * scale,
            self.width * scale)


def detect_grasps(q_img, ang_img, width_img=None, no_grasps=1):
    """
    Detect grasps in a network output.
    :param q_img: Q image network output
    :param ang_img: Angle image network output
    :param width_img: (optional) Width image network output
    :param no_grasps: Max number of grasps to return
    :return: list of Grasps
    """
    local_max = peak_local_max(q_img, min_distance=20, threshold_abs=0.2, num_peaks=no_grasps)

    grasps = []
    for grasp_point_array in local_max:
        grasp_point = tuple(grasp_point_array)

        grasp_angle = ang_img[grasp_point]

        g = Grasp(grasp_point, grasp_angle)
        if width_img is not None:
            g.length = width_img[grasp_point]
            g.width = g.length / 2

        grasps.append(g)

    return grasps

utils/dataset_processing/grasp.py
andvg3/Grasp-Anything
import warnings

import cv2
import matplotlib.pyplot as plt
import numpy as np
from imageio import imread
import imageio
from skimage.transform import rotate, resize

warnings.filterwarnings("ignore", category=UserWarning)


class Image:
    """
    Wrapper around an image with some convenient functions.
    """

    def __init__(self, img):
        self.img = img

    def __getattr__(self, attr):
        # Pass along any other methods to the underlying ndarray
        return getattr(self.img, attr)

    @classmethod
    def from_file(cls, fname):
        return cls(imread(fname))

    def copy(self):
        """
        :return: Copy of self.
        """
        return self.__class__(self.img.copy())
    
    @classmethod
    def mask_out_image(cls, image, mask):
        # Apply the mask to the image
        masked_image = np.array(image)
        masked_image[:, :, 0] = masked_image[:, :, 0] * mask + 255 * (1-mask)
        masked_image[:, :, 1] = masked_image[:, :, 1] * mask + 255 * (1-mask)
        masked_image[:, :, 2] = masked_image[:, :, 2] * mask + 255 * (1-mask)

        return cls(imageio.core.util.Array(masked_image))

    def crop(self, top_left, bottom_right, resize=None):
        """
        Crop the image to a bounding box given by top left and bottom right pixels.
        :param top_left: tuple, top left pixel.
        :param bottom_right: tuple, bottom right pixel
        :param resize: If specified, resize the cropped image to this size
        """
        self.img = self.img[top_left[0]:bottom_right[0], top_left[1]:bottom_right[1]]
        if resize is not None:
            self.resize(resize)

    def cropped(self, *args, **kwargs):
        """
        :return: Cropped copy of the image.
        """
        i = self.copy()
        i.crop(*args, **kwargs)
        return i

    def normalise(self):
        """
        Normalise the image by converting to float [0,1] and zero-centering
        """
        self.img = self.img.astype(np.float32) / 255.0
        self.img -= self.img.mean()

    def resize(self, shape):
        """
        Resize image to shape.
        :param shape: New shape.
        """
        if self.img.shape == shape:
            return
        self.img = resize(self.img, shape, preserve_range=True).astype(self.img.dtype)

    def resized(self, *args, **kwargs):
        """
        :return: Resized copy of the image.
        """
        i = self.copy()
        i.resize(*args, **kwargs)
        return i

    def rotate(self, angle, center=None):
        """
        Rotate the image.
        :param angle: Angle (in radians) to rotate by.
        :param center: Center pixel to rotate if specified, otherwise image center is used.
        """
        if center is not None:
            center = (center[1], center[0])
        self.img = rotate(self.img, angle / np.pi * 180, center=center, mode='symmetric', preserve_range=True).astype(
            self.img.dtype)

    def rotated(self, *args, **kwargs):
        """
        :return: Rotated copy of image.
        """
        i = self.copy()
        i.rotate(*args, **kwargs)
        return i

    def show(self, ax=None, **kwargs):
        """
        Plot the image
        :param ax: Existing matplotlib axis (optional)
        :param kwargs: kwargs to imshow
        """
        if ax:
            ax.imshow(self.img, **kwargs)
        else:
            plt.imshow(self.img, **kwargs)
            plt.show()

    def zoom(self, factor):
        """
        "Zoom" the image by cropping and resizing.
        :param factor: Factor to zoom by. e.g. 0.5 will keep the center 50% of the image.
        """
        sr = int(self.img.shape[0] * (1 - factor)) // 2
        sc = int(self.img.shape[1] * (1 - factor)) // 2
        orig_shape = self.img.shape
        self.img = self.img[sr:self.img.shape[0] - sr, sc: self.img.shape[1] - sc].copy()
        self.img = resize(self.img, orig_shape, mode='symmetric', preserve_range=True).astype(self.img.dtype)

    def zoomed(self, *args, **kwargs):
        """
        :return: Zoomed copy of the image.
        """
        i = self.copy()
        i.zoom(*args, **kwargs)
        return i


class DepthImage(Image):
    def __init__(self, img):
        super().__init__(img)

    @classmethod
    def from_pcd(cls, pcd_filename, shape, default_filler=0, index=None):
        """
            Create a depth image from an unstructured PCD file.
            If index isn't specified, use euclidean distance, otherwise choose x/y/z=0/1/2
        """
        img = np.zeros(shape)
        if default_filler != 0:
            img += default_filler

        with open(pcd_filename) as f:
            for l in f.readlines():
                ls = l.split()

                if len(ls) != 5:
                    # Not a point line in the file.
                    continue
                try:
                    # Not a number, carry on.
                    float(ls[0])
                except ValueError:
                    continue

                i = int(ls[4])
                r = i // shape[1]
                c = i % shape[1]

                if index is None:
                    x = float(ls[0])
                    y = float(ls[1])
                    z = float(ls[2])

                    img[r, c] = np.sqrt(x ** 2 + y ** 2 + z ** 2)

                else:
                    img[r, c] = float(ls[index])

        return cls(img / 1000.0)

    @classmethod
    def from_tiff(cls, fname):
        return cls(imread(fname))

    def inpaint(self, missing_value=0):
        """
        Inpaint missing values in depth image.
        :param missing_value: Value to fill in teh depth image.
        """
        # cv2 inpainting doesn't handle the border properly
        # https://stackoverflow.com/questions/25974033/inpainting-depth-map-still-a-black-image-border
        self.img = cv2.copyMakeBorder(self.img, 1, 1, 1, 1, cv2.BORDER_DEFAULT)
        mask = (self.img == missing_value).astype(np.uint8)

        # Scale to keep as float, but has to be in bounds -1:1 to keep opencv happy.
        scale = np.abs(self.img).max()
        self.img = self.img.astype(np.float32) / scale  # Has to be float32, 64 not supported.
        self.img = cv2.inpaint(self.img, mask, 1, cv2.INPAINT_NS)

        # Back to original size and value range.
        self.img = self.img[1:-1, 1:-1]
        self.img = self.img * scale

    def gradients(self):
        """
        Compute gradients of the depth image using Sobel filtesr.
        :return: Gradients in X direction, Gradients in Y diretion, Magnitude of XY gradients.
        """
        grad_x = cv2.Sobel(self.img, cv2.CV_64F, 1, 0, borderType=cv2.BORDER_DEFAULT)
        grad_y = cv2.Sobel(self.img, cv2.CV_64F, 0, 1, borderType=cv2.BORDER_DEFAULT)
        grad = np.sqrt(grad_x ** 2 + grad_y ** 2)

        return DepthImage(grad_x), DepthImage(grad_y), DepthImage(grad)

    def normalise(self):
        """
        Normalise by subtracting the mean and clippint [-1, 1]
        """
        self.img = np.clip((self.img - self.img.mean()), -1, 1)


class WidthImage(Image):
    """
    A width image is one that describes the desired gripper width at each pixel.
    """

    def zoom(self, factor):
        """
        "Zoom" the image by cropping and resizing.  Also scales the width accordingly.
        :param factor: Factor to zoom by. e.g. 0.5 will keep the center 50% of the image.
        """
        super().zoom(factor)
        self.img = self.img / factor

    def normalise(self):
        """
        Normalise by mapping [0, 150] -> [0, 1]
        """
        self.img = np.clip(self.img, 0, 150.0) / 150.0

utils/dataset_processing/image.py
andvg3/Grasp-Anything
import warnings

import numpy as np

warnings.filterwarnings("ignore", category=UserWarning)

class Mask:
    """
    Wrapper around an mask with some convenient functions.
    """

    def __init__(self, mask):
        self.mask = mask

    def __getattr__(self, attr):
        # Pass along any other methods to the underlying ndarray
        return getattr(self.mask, attr)
    
    @classmethod
    def from_file(cls, fname):
        return np.load(fname, mmap_mode='r')
utils/dataset_processing/mask.py
andvg3/Grasp-Anything