Add mmflow to mmdeploy

.pre-commit-config.yaml

@@ -5,7 +5,7 @@ repos:
       - id: flake8
         args: ["--exclude=*/client/inference_pb2.py,*/client/inference_pb2_grpc.py"]
   - repo: https://github.com/PyCQA/isort
-    rev: 5.10.1
+    rev: 5.11.5
     hooks:
       - id: isort
   - repo: https://github.com/pre-commit/mirrors-yapf

configs/mmflow/raft/opticalflow_onnxruntime_static.py

@@ -1,4 +1,5 @@
-_base_ = ['../_base_/opticalflow_static.py','../../_base_/backends/onnxruntime.py']
+_base_ = [
+    '../_base_/opticalflow_static.py', '../../_base_/backends/onnxruntime.py'
+]
 
-onnx_config = dict(
-    verbose=True)
+onnx_config = dict(verbose=True)

configs/mmflow/raft/opticalflow_tensorrt_static-440x1024.py

@@ -1,4 +1,6 @@
-_base_ = ['../_base_/opticalflow_static.py', '../../_base_/backends/tensorrt.py']
+_base_ = [
+    '../_base_/opticalflow_static.py', '../../_base_/backends/tensorrt.py'
+]
 
 onnx_config = dict(input_shape=[440, 1024])
 backend_config = dict(

mmdeploy/codebase/mmflow/__init__.py

@@ -1,4 +1,2 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .deploy import *  # noqa: F401,F403
-from .models import *
-from .ops import *

mmdeploy/codebase/mmflow/deploy/__init__.py

@@ -2,4 +2,4 @@
 from .flow import Flow
 from .mmflow import MMFlow  # noqa: F401,F403
 
-__all__ = ['MMFlow', 'Flow']
+__all__ = ['MMFlow', 'Flow']

mmdeploy/codebase/mmflow/deploy/flow.py

@@ -1,16 +1,17 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import copy
 import warnings
 from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
-import copy
+
 import mmcv
 import numpy as np
 import torch
 from mmcv.parallel import collate, scatter
 from torch.utils.data import Dataset
-from mmflow.core.evaluation.metrics import eval_metrics
+
 from mmdeploy.codebase.base import BaseTask
+from mmdeploy.utils import Task, get_input_shape
 from .mmflow import MMFLOW_TASK
-from mmdeploy.utils import Task, get_input_shape, load_config
 
 
 def process_model_config(model_cfg: mmcv.Config,
@@ -46,7 +47,6 @@ def process_model_config(model_cfg: mmcv.Config,
     if 'flow_bw_gt' in cfg.pipeline[-1]['meta_keys']:
         cfg.pipeline[-1]['meta_keys'].remove('flow_bw_gt')
 
-
     return cfg
 
 
@@ -132,7 +132,7 @@ def create_input(self,
 
         test_pipeline = Compose(cfg.pipeline)
         valid_masks = []
-        datas = []
+        data_list = []
         if isinstance(imgs[0], np.ndarray) and isinstance(imgs[1], np.ndarray):
             # directly add img and valid mask
             data = dict(img1=imgs[0], img2=imgs[1])
@@ -144,10 +144,10 @@ def create_input(self,
                 img2_prefix=None,
                 valid=valid)
         data['img_fields'] = ['img1', 'img2']
-        data = test_pipeline(data)
-        datas.append(data)
+        data = test_pipeline(data_list)
+        data_list.append(data)
         valid_masks.append(valid)
-        data = collate(datas, samples_per_gpu=len(datas))
+        data = collate(data_list, samples_per_gpu=len(data_list))
 
         # just get the actual data from DataContainer
         data['img_metas'] = data['img_metas'].data[0]
@@ -271,46 +271,6 @@ def evaluate_outputs(model_cfg,
         """
         from mmcv.utils import get_logger
         logger = get_logger('test', log_file=log_file)
-        # if metrics:
-        #     for i in range(len(outputs)):
-        #         result = outputs[i]
-        #         for k in result.keys():
-        #             if img_meta.get(k + '_gt', None) is None:
-        #                 # img_meta does not have flow_bw_gt, so just check
-        #                 # the forward predicted.
-        #                 if k == 'flow_bw':
-        #                     continue
-        #                 elif k == 'flow_fw':
-        #                     batch_flow_gt.append(img_meta['flow_gt'])
-        #             else:
-        #                 batch_flow_gt.append(img_meta[k + '_gt'])
-        # 
-        #             batch_flow.append(result[k])
-        #             batch_valid.append(
-        #                 img_meta.get('valid', np.ones_like(result[k][..., 0])))
-        #     results = dataset.evaluate(outputs, metrics, metric_options)
-        #     if json_file is not None:
-        #         mmcv.dump(results, json_file, indent=4)
-        # 
-        #     for k, v in results.items():
-        #         logger.info(f'{k} : {v:.2f}')
-        # else:
-        #     logger.warning('Evaluation metrics are not specified.')
-        #     scores = np.vstack(outputs)
-        #     pred_score = np.max(scores, axis=1)
-        #     pred_label = np.argmax(scores, axis=1)
-        #     pred_class = [dataset.CLASSES[lb] for lb in pred_label]
-        #     results = {
-        #         'pred_score': pred_score,
-        #         'pred_label': pred_label,
-        #         'pred_class': pred_class
-        #     }
-        #     if not out:
-        #         logger.info('the predicted result for the first element is '
-        #                     f'pred_score = {pred_score[0]:.2f}, '
-        #                     f'pred_label = {pred_label[0]} '
-        #                     f'and pred_class = {pred_class[0]}. '
-        #                     'Specify --out to save all results to files.')
         if out:
             logger.debug(f'writing results to {out}')
             mmcv.dump(outputs, out)

mmdeploy/codebase/mmflow/deploy/mmflow.py

@@ -1,10 +1,11 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from typing import Optional,Sequence, Union
+from typing import Optional, Sequence, Union
 
 import mmcv
 import torch
 from mmcv.utils import Registry
 from torch.utils.data import DataLoader, Dataset
+
 from mmdeploy.codebase.base import CODEBASE, BaseTask, MMCodebase
 from mmdeploy.utils import Codebase, get_task_type, load_config
 
@@ -107,9 +108,9 @@ def build_dataloader(dataset: Dataset,
         """
         from mmflow.datasets import build_dataloader as build_dataloader_mmflow
         return build_dataloader_mmflow(dataset, samples_per_gpu,
-                                       workers_per_gpu,sample_ratio, num_gpus, dist,
-                                       shuffle, seed,
-                                       persistent_workers, **kwargs)
+                                       workers_per_gpu, sample_ratio, num_gpus,
+                                       dist, shuffle, seed, persistent_workers,
+                                       **kwargs)
 
     @staticmethod
     def single_gpu_test(model: torch.nn.Module,
@@ -130,6 +131,6 @@ def single_gpu_test(model: torch.nn.Module,
             list: The prediction results.
         """
         from mmflow.apis import single_gpu_test
-        show_dir  = out_dir
-        outputs = single_gpu_test(model, data_loader,out_dir,show_dir)
+        show_dir = out_dir
+        outputs = single_gpu_test(model, data_loader, out_dir, show_dir)
         return outputs

mmdeploy/codebase/mmflow/deploy/mmflow_model.py

@@ -1,12 +1,11 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-import os.path as osp
-from typing import Dict, List, Optional, Sequence, Union
-from numpy import ndarray
+from typing import List, Optional, Sequence, Union
+
 import mmcv
 import numpy as np
 import torch
 from mmcv.utils import Registry
-# from mmflow.core import psnr, ssim, tensor2img
+from numpy import ndarray
 
 from mmdeploy.codebase.base import BaseBackendModel
 from mmdeploy.utils import (Backend, get_backend, get_codebase_config,
@@ -23,7 +22,7 @@ def __build_backend_model(cls_name: str, registry: Registry, *args, **kwargs):
 
 @__BACKEND_MODEL.register_module('end2end')
 class End2EndModel(BaseBackendModel):
-    """End to end model for inference of optical flow
+    """End to end model for inference of optical flow.
 
     Args:
         backend (Backend): The backend enum, specifying backend type.
@@ -44,7 +43,6 @@ def __init__(self,
                  **kwargs):
         super().__init__(deploy_cfg=deploy_cfg)
         self.deploy_cfg = deploy_cfg
-        # self.test_cfg = model_cfg.test_cfg
         self._init_wrapper(
             backend=backend,
             backend_files=backend_files,
@@ -63,17 +61,16 @@ def _init_wrapper(self, backend: Backend, backend_files: Sequence[str],
             deploy_cfg=self.deploy_cfg,
             **kwargs)
 
-    def forward(self,
-                imgs: torch.Tensor,
-                *args, **kwargs) -> Sequence[ndarray]:
+    def forward(self, imgs: torch.Tensor, *args,
+                **kwargs) -> Sequence[ndarray]:
         """Run test inference for restorer.
 
         We want forward() to output an image or a evaluation result.
         When test_mode is set, the output is evaluation result. Otherwise
         it is an image.
 
         Args:
-            lq (torch.Tensor): The input low-quality image of the model.
+            imgs (torch.Tensor): The input low-quality image of the model.
             test_mode (bool): When test_mode is set, the output is evaluation
                 result. Otherwise it is an image. Default to `False`.
             *args: Other arguments.
@@ -94,10 +91,7 @@ def forward_test(self, imgs: torch.Tensor, *args, **kwargs):
         """Run inference for restorer to generate evaluation result.
 
         Args:
-            lq (torch.Tensor): The input low-quality image of the model.
-            gt (torch.Tensor): The ground truth of input image. Defaults to
-                `None`.
-            meta (List[Dict]): The meta infomations of MMEditing.
+            imgs (torch.Tensor): The input low-quality image of the model.
             save_path (str): Path to save image. Default: None.
             *args: Other arguments.
             **kwargs: Other key-pair arguments.
@@ -126,8 +120,6 @@ def evaluate(self, output: Union[torch.Tensor, np.ndarray],
 
         if isinstance(output, np.ndarray):
             output = torch.from_numpy(output)
-        # output = tensor2img(output)
-        # gt = tensor2img(gt)
 
         eval_result = dict()
         for metric in self.test_cfg.metrics:
@@ -191,22 +183,13 @@ def forward(self,
         Returns:
             list | dict: High resolution image or a evaluation results.
         """
-        img = tensor2img(lq)
-        output = self.wrapper.invoke(img)
-        if test_mode:
-            output = torch.from_numpy(output)
-            output = output.permute(2, 0, 1)
-            output = output / 255.
-            results = self.test_post_process([output], lq, gt)
-            return results
-        else:
-            return [output]
+        pass
 
 
 def build_mmflow_model(model_files: Sequence[str],
-                                 model_cfg: Union[str, mmcv.Config],
-                                 deploy_cfg: Union[str, mmcv.Config],
-                                 device: str, **kwargs):
+                       model_cfg: Union[str, mmcv.Config],
+                       deploy_cfg: Union[str,
+                                         mmcv.Config], device: str, **kwargs):
     model_cfg = load_config(model_cfg)[0]
     deploy_cfg = load_config(deploy_cfg)[0]
 

mmdeploy/codebase/mmflow/models/__init__.py

@@ -1,5 +1 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from . import decoders
-
-
-

mmdeploy/codebase/mmflow/models/decoders/__init__.py

@@ -1,3 +1 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from . import raft_decoder
-

mmdeploy/codebase/mmflow/models/decoders/raft_decoder.py

@@ -1,47 +1,23 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-import math
-from typing import Dict, Optional, Sequence, Union
+from typing import Dict, Sequence
 
 import numpy as np
 import torch
 
 from mmdeploy.core import FUNCTION_REWRITER
 
-@FUNCTION_REWRITER.register_rewriter(
-    func_name='mmflow.models.decoders.raft_decoder.CorrelationPyramid.forward',backend='default')
-def raft_decoders_forward(ctx,self, feat1: torch.Tensor,
-            feat2: torch.Tensor) -> Sequence[torch.Tensor]:
-    """Forward function for Correlation pyramid.
-
-    Args:
-        feat1 (Tensor): The feature from first input image.
-        feat2 (Tensor): The feature from second input image.
-
-    Returns:
-        Sequence[Tensor]: The list of correlation which is pooled using
-            average pooling with kernel sizes {1, 2, 4, 8}.
-    """
-    N, C, H, W = feat1.shape
-    corr = torch.matmul(
-        feat1.view(N, C, -1).permute(0, 2, 1),
-        feat2.view(N, C, -1)).view(N, H, W, H, W)
-    corr = corr.reshape(N * H * W, 1, H, W) / torch.sqrt(C.clone().detach().float())
-    corr_pyramid = [corr]
-    for _ in range(self.num_levels - 1):
-        _corr = self.pool(corr_pyramid[-1])
-        corr_pyramid.append(_corr)
-
-    return corr_pyramid
 
 @FUNCTION_REWRITER.register_rewriter(
     func_name='mmflow.models.decoders.raft_decoder.RAFTDecoder.forward_test')
-def raft_decoder_forward_test(ctx,self,
-                 feat1: torch.Tensor,
-                 feat2: torch.Tensor,
-                 flow: torch.Tensor,
-                 h_feat: torch.Tensor,
-                 cxt_feat: torch.Tensor,
-                 img_metas=None) -> Sequence[Dict[str, np.ndarray]]:
+def raft_decoder_forward_test(
+        ctx,
+        self,
+        feat1: torch.Tensor,
+        feat2: torch.Tensor,
+        flow: torch.Tensor,
+        h_feat: torch.Tensor,
+        cxt_feat: torch.Tensor,
+        img_metas=None) -> Sequence[Dict[str, np.ndarray]]:
     """Forward function when model training.
 
     Args:
@@ -61,7 +37,7 @@ def raft_decoder_forward_test(ctx,self,
 
     flow_result = flow_pred[-1]
     # flow maps with the shape [H, W, 2]
-    flow_result = flow_result.permute(0, 2, 3, 1).squeeze(dim=0).clone().detach()
+    flow_result = flow_result.permute(0, 2, 3, 1)
     # unravel batch dim
     flow_result = [dict(flow=flow_result)]
-    return self.get_flow(flow_result, img_metas=img_metas)
+    return self.get_flow(flow_result, img_metas=img_metas)