Files
modelscope/tests/pipelines/test_dense_optical_flow_estimation.py
pipikk 802f90d2e1 feat: add raft model (#702)
* first version

* fix minor bugs

* add test images

* TorchModel & docstr

* modify pipeline implementation

* minor

* add datasets

* add extractor

* add update

* add augmentor

* add flow_viz

* add frame_utils

* add utils

* add raft_model

* add dense_optical_flow_estimation_pipeline

* add image_utils

* add test_dense_optical_flow_estimation

* test

* update cv/__init__

* [3]update cv/__init__

* update submodule data/test

* correct yapf

* fix bugs

* move test data

* update submodule

* minor

* update submodule

---------

Co-authored-by: kejie <kejie.qkj@alibaba-inc.com>
2024-04-08 15:27:20 +08:00

40 lines
1.3 KiB
Python

# Copyright (c) Alibaba, Inc. and its affiliates.
import unittest
import cv2
import numpy as np
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.utils.test_utils import test_level
class DenseOpticalFlowEstimationTest(unittest.TestCase):
def setUp(self) -> None:
self.task = 'dense-optical-flow-estimation'
self.model_id = 'Damo_XR_Lab/cv_raft_dense-optical-flow_things'
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
def test_dense_optical_flow_estimation(self):
input_location = [[
'data/test/images/dense_flow1.png',
'data/test/images/dense_flow2.png',
# 'modelscope/models/cv/dense_optical_flow_estimation/data/test/images/dense_flow1.png',
# 'modelscope/models/cv/dense_optical_flow_estimation/data/test/images/dense_flow2.png'
]]
estimator = pipeline(
Tasks.dense_optical_flow_estimation, model=self.model_id)
result = estimator(input_location)
# flow = result[0][OutputKeys.FLOWS]
flow_vis = result[0][OutputKeys.FLOWS_COLOR]
cv2.imwrite('result.jpg', flow_vis)
print('test_dense_optical_flow_estimation DONE')
if __name__ == '__main__':
unittest.main()