SamsungLabs / ritm_interactive_segmentation

Reviving Iterative Training with Mask Guidance for Interactive Segmentation
MIT License
632 stars 125 forks source link

Is your models are deployable in onnx format? #31

Open sbmalik opened 2 years ago

sbmalik commented 2 years ago

I wanted to run these models on some other platforms and needed tha onnx version. Kindly guide me if they are available?

FrancescoSaverioZuppichini commented 2 years ago

I've tried to export it but there is a mishape error caused by this line https://github.com/saic-vul/ritm_interactive_segmentation/blob/708182de6e3a73da4034579935aef4c19ceae80a/isegm/model/ops.py#L77

Ehteshamciitwah commented 2 years ago

I've tried to export it but there is a mishape error caused by this line

https://github.com/saic-vul/ritm_interactive_segmentation/blob/708182de6e3a73da4034579935aef4c19ceae80a/isegm/model/ops.py#L77

I had the same issue. but after removing that coords[invalid_points, :, :, :] = 1e6 . i successfully converted and load the onnx model Thank you

FrancescoSaverioZuppichini commented 2 years ago

but I need you need to set the invalid points

Ehteshamciitwah commented 2 years ago

Hello, i just want to use inference in onnx. So i do some changing in ops.py file

FrancescoSaverioZuppichini commented 2 years ago

Not sure I understood, I advice you to take your time to write a better-explained message.

Without the line you have removed, the output shouldn't be the same.

bfan1256 commented 1 year ago

Is there any updates on this? @Ehteshamciitwah were you able to get a functional ONNX model?

Ehteshamciitwah commented 1 year ago

yes. I changed it successfully to onnx format. During conversion, first, I use else condition making self.cpu_mode=False. https://github.com/saic-vul/ritm_interactive_segmentation/blob/708182de6e3a73da4034579935aef4c19ceae80a/isegm/model/ops.py#L56 second, i removed the line https://github.com/saic-vul/ritm_interactive_segmentation/blob/708182de6e3a73da4034579935aef4c19ceae80a/isegm/model/ops.py#L77

After changing, get_coord_features code is: ''' def get_coord_features(self, points, batchsize, rows, cols):

    num_points=torch.div(points.shape[1] , 2, rounding_mode='floor') #int 1
    points = points.view(-1, points.size(2))
    points, points_order = torch.split(points, [2, 1], dim=1)

    invalid_points = torch.max(points, dim=1, keepdim=False)[0] < 0      #tensor([False,  True, False,  True], device='cuda:0')
    row_array = torch.arange(start=0, end=rows, step=1, dtype=torch.float32, device=points.device)
    col_array = torch.arange(start=0, end=cols, step=1, dtype=torch.float32, device=points.device)

    coord_rows, coord_cols = torch.meshgrid(row_array, col_array) #400x400 each
    coords = torch.stack((coord_rows, coord_cols), dim=0).unsqueeze(0).repeat(points.size(0), 1, 1, 1)

    add_xy = (points * self.spatial_scale).view(points.size(0), points.size(1), 1, 1)   #4x2x1x1
    coords.add_(-add_xy)
    if not self.use_disks:
        coords.div_(self.norm_radius * self.spatial_scale)
    coords.mul_(coords)

    coords[:, 0] += coords[:, 1]
    coords = coords[:, :1]

    # coords[invalid_points, :, :, :] = 1e6

    coords = coords.view(-1, num_points, 1, rows, cols)
    coords = coords.min(dim=1)[0]  # -> (bs * num_masks * 2) x 1 x h x w
    coords = coords.view(-1, 2, rows, cols)  #2x2x400x400

    if self.use_disks:
        coords = (coords <= (self.norm_radius * self.spatial_scale) ** 2).float()
    else:
        coords.sqrt_().mul_(2).tanh_()

    return coords    #2x2x400x400

'''

Caooc commented 1 year ago

是的。我成功地将其更改为onnx格式。在转换过程中,首先,我使用 else 条件制作 self.cpu_mode=False。https://github.com/saic-vul/ritm_interactive_segmentation/blob/708182de6e3a73da4034579935aef4c19ceae80a/isegm/model/ops.py#L56 秒,我删除了该行 https://github.com/saic-vul/ritm_interactive_segmentation/blob/708182de6e3a73da4034579935aef4c19ceae80a/isegm/model/ops.py#L77

更改后,get_coord_features代码为:“”定义get_coord_features(self,points,batchsize,rows,cols):

    num_points=torch.div(points.shape[1] , 2, rounding_mode='floor') #int 1
    points = points.view(-1, points.size(2))
    points, points_order = torch.split(points, [2, 1], dim=1)

    invalid_points = torch.max(points, dim=1, keepdim=False)[0] < 0      #tensor([False,  True, False,  True], device='cuda:0')
    row_array = torch.arange(start=0, end=rows, step=1, dtype=torch.float32, device=points.device)
    col_array = torch.arange(start=0, end=cols, step=1, dtype=torch.float32, device=points.device)

    coord_rows, coord_cols = torch.meshgrid(row_array, col_array) #400x400 each
    coords = torch.stack((coord_rows, coord_cols), dim=0).unsqueeze(0).repeat(points.size(0), 1, 1, 1)

    add_xy = (points * self.spatial_scale).view(points.size(0), points.size(1), 1, 1)   #4x2x1x1
    coords.add_(-add_xy)
    if not self.use_disks:
        coords.div_(self.norm_radius * self.spatial_scale)
    coords.mul_(coords)

    coords[:, 0] += coords[:, 1]
    coords = coords[:, :1]

    # coords[invalid_points, :, :, :] = 1e6

    coords = coords.view(-1, num_points, 1, rows, cols)
    coords = coords.min(dim=1)[0]  # -> (bs * num_masks * 2) x 1 x h x w
    coords = coords.view(-1, 2, rows, cols)  #2x2x400x400

    if self.use_disks:
        coords = (coords <= (self.norm_radius * self.spatial_scale) ** 2).float()
    else:
        coords.sqrt_().mul_(2).tanh_()

    return coords    #2x2x400x400

'''

Hi, can you provide the code to export onnx?

Ehteshamciitwah commented 1 year ago

Kindly update the code according to your requirements:

import torch from isegm.utils import exp from isegm.inference import utils import argparse def parse_args(): parser = argparse.ArgumentParser()

parser.add_argument('--checkpoint',default='coco_lvis_h18_itermask',type=str, required=False,
                    help='The path to the checkpoint. '
                         'This can be a relative path (relative to cfg.INTERACTIVE_MODELS_PATH) '
                         'or an absolute path. The file extension can be omitted.')
parser.add_argument('--gpu', type=int, default=0,
                    help='Id of GPU to use.')
parser.add_argument('--image', type=str, default='./3096.jpg',
                    help='path to image')
parser.add_argument('--cpu', action='store_true', default=False,
                    help='Use only CPU for inference.')

parser.add_argument('--limit-longest-size', type=int, default=800,
                    help='If the largest side of an image exceeds this value, '
                         'it is resized so that its largest side is equal to this value.')

parser.add_argument('--cfg', type=str, default="config.yml",
                    help='The path to the config file.')

args = parser.parse_args()
# parser.parse_args()
if args.cpu:
    args.device =torch.device('cpu')
else:
    args.device = torch.device(f'cuda:{args.gpu}')
cfg = exp.load_config_file(args.cfg, return_edict=True)

return args, cfg

def main(): args, cfg = parse_args() torch.backends.cudnn.deterministic = True checkpoint_path = utils.find_checkpoint(cfg.INTERACTIVE_MODELS_PATH, args.checkpoint) model = utils.load_is_model(checkpoint_path, args.device, cpu_dist_maps=True) point_nd=torch.tensor([[[239.0, 75.0, 0.0], [ -1, -1, -1]]],device='cuda')#tensor(2,2,3)

image=torch.randn(1,4,400,400,device='cuda')
dummy_input=(image,point_nd)
input_names = ["image","points"]
output_names = ["output"]
torch.onnx.export(
  model,
  dummy_input,
  "ritm.onnx",
  verbose=True,
  opset_version=13,
  input_names=input_names,
  output_names=output_names,
  dynamic_axes={
      'image':{2:'height',3:'width'},
      'points':{1:"custom"},
          }
        )

if name == 'main': main()

duxuan11 commented 1 year ago

是的。我成功地将其更改为onnx格式。在转换过程中,首先,我使用 else 条件制作 self.cpu_mode=False。

May I ask why the image has 4 channels? Can you send me a copy of the Onnx inference code?

duxuan11 commented 1 year ago

Hello, i just want to use inference in onnx. So i do some changing in ops.py file

  • use else condition instead of if in ops.py file
  • remove coords[invalid_points, :, :, :] = 1e6
  • run demo.py file to test and it works well without any problem
  • Then i change it into onnx and test onnx with their Ui by replacing https://github.com/saic-vul/ritm_interactive_segmentation/blob/708182de6e3a73da4034579935aef4c19ceae80a/isegm/inference/predictors/base.py#L80 with onnx run. So i think invalid points may have not much use. for onnx conversion you need two dummy inputs as tuples. you need to use dynamic axes during export.
    def _get_prediction(self, image_nd, clicks_lists, is_image_changed):
    points_nd = self.get_points_nd(clicks_lists)
    return self.net(image_nd, points_nd)['instances']

    I modified this code as follows

    def _get_prediction(self, image_nd, clicks_lists, is_image_changed):
    points_nd = self.get_points_nd(clicks_lists)
    #return self.net(image_nd, points_nd)['instances']
    #print(image_nd.shape)
    image_nd = np.asarray(image_nd).astype(np.float32)
    point_nd=np.array([[[239.0, 75.0, 0.0],[ -1, -1, -1]]]).astype(np.float32)
    sess = rt.InferenceSession('ritm.onnx')
    pred_onnx = sess.run([output_name], { "image": image,"points": point_nd})
    return pred_onnx

    I found the image_nd is (2,4,256,256) size, but the onnx model input of image is size of (1, 4,height, width). So how to modify this code ?