modelscope/tests/pipelines/test_crowd_counting.py

# Copyright (c) Alibaba, Inc. and its affiliates.
import unittest

import cv2
import numpy as np
from PIL import Image

from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.utils.cv.heatmap import numpy_to_cv2img
from modelscope.utils.logger import get_logger
from modelscope.utils.test_utils import test_level

logger = get_logger()


class CrowdCountingTest(unittest.TestCase):

    def setUp(self) -> None:
        self.input_location = 'data/test/images/crowd_counting.jpg'
        self.model_id = 'damo/cv_hrnet_crowd-counting_dcanet'

    def save_result(self, result):
        print('scores:', result[OutputKeys.SCORES])
        vis_img = result[OutputKeys.OUTPUT_IMG]
        vis_img = numpy_to_cv2img(vis_img)
        cv2.imwrite('result.jpg', vis_img)

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_crowd_counting(self):
        crowd_counting = pipeline(Tasks.crowd_counting, model=self.model_id)
        result = crowd_counting(self.input_location)
        if result:
            self.save_result(result)
        else:
            raise ValueError('process error')

    @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
    def test_crowd_counting_with_image(self):
        crowd_counting = pipeline(Tasks.crowd_counting, model=self.model_id)
        img = Image.open(self.input_location)
        result = crowd_counting(img)
        if result:
            self.save_result(result)
        else:
            raise ValueError('process error')

    @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
    def test_crowd_counting_with_default_task(self):
        crowd_counting = pipeline(Tasks.crowd_counting)
        result = crowd_counting(self.input_location)
        if result:
            self.save_result(result)
        else:
            raise ValueError('process error')


if __name__ == '__main__':
    unittest.main()