mirror of
https://github.com/modelscope/modelscope.git
synced 2025-12-16 16:27:45 +01:00
373 lines
16 KiB
Python
373 lines
16 KiB
Python
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||
import os
|
||
import unittest
|
||
from os import path as osp
|
||
|
||
import cv2
|
||
from PIL import Image
|
||
|
||
from modelscope.models import Model
|
||
from modelscope.outputs import OutputKeys
|
||
from modelscope.pipelines import pipeline
|
||
from modelscope.utils.constant import Tasks
|
||
from modelscope.utils.cv.image_utils import created_boxed_image
|
||
from modelscope.utils.test_utils import test_level
|
||
|
||
|
||
class OfaTasksTest(unittest.TestCase):
|
||
|
||
def setUp(self) -> None:
|
||
self.output_dir = 'unittest_output'
|
||
os.makedirs(self.output_dir, exist_ok=True)
|
||
|
||
def save_img(self, image_in, box, image_out):
|
||
cv2.imwrite(
|
||
osp.join(self.output_dir, image_out),
|
||
created_boxed_image(image_in, box))
|
||
|
||
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
||
def test_run_with_image_captioning_with_model(self):
|
||
model = Model.from_pretrained('damo/ofa_image-caption_coco_large_en')
|
||
img_captioning = pipeline(
|
||
task=Tasks.image_captioning,
|
||
model=model,
|
||
)
|
||
image = 'data/test/images/image_captioning.png'
|
||
result = img_captioning(image)
|
||
print(result[OutputKeys.CAPTION])
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_image_captioning_with_name(self):
|
||
img_captioning = pipeline(
|
||
Tasks.image_captioning,
|
||
model='damo/ofa_image-caption_coco_large_en')
|
||
result = img_captioning('data/test/images/image_captioning.png')
|
||
print(result[OutputKeys.CAPTION])
|
||
|
||
img_captioning.model.num_return_sequences = 2
|
||
result = img_captioning('data/test/images/image_captioning.png')
|
||
print(result[OutputKeys.CAPTION])
|
||
|
||
# test batch infer
|
||
img_captioning.model.num_return_sequences = 1
|
||
results = img_captioning(
|
||
[{
|
||
'image': 'data/test/images/image_captioning.png'
|
||
} for _ in range(3)],
|
||
batch_size=2)
|
||
for r in results:
|
||
print(r[OutputKeys.CAPTION])
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_ocr_recognize_with_name(self):
|
||
ocr_recognize = pipeline(
|
||
Tasks.ocr_recognition,
|
||
model='damo/ofa_ocr-recognition_scene_base_zh')
|
||
result = ocr_recognize('data/test/images/image_ocr_recognition.jpg')
|
||
print(result[OutputKeys.TEXT])
|
||
# test batch infer
|
||
results = ocr_recognize(
|
||
['data/test/images/image_ocr_recognition.jpg' for _ in range(3)],
|
||
batch_size=2)
|
||
for r in results:
|
||
print(r[OutputKeys.TEXT])
|
||
|
||
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
||
def test_run_with_image_classification_with_model(self):
|
||
model = Model.from_pretrained(
|
||
'damo/ofa_image-classification_imagenet_large_en')
|
||
ofa_pipe = pipeline(Tasks.image_classification, model=model)
|
||
image = 'data/test/images/image_classification.png'
|
||
result = ofa_pipe(image)
|
||
print(result)
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_image_classification_with_name(self):
|
||
ofa_pipe = pipeline(
|
||
Tasks.image_classification,
|
||
model='damo/ofa_image-classification_imagenet_large_en')
|
||
image = 'data/test/images/image_classification.png'
|
||
result = ofa_pipe(image)
|
||
print(result)
|
||
|
||
# test batch infer
|
||
image = ['data/test/images/image_classification.png' for _ in range(3)]
|
||
results = ofa_pipe(image, batch_size=2)
|
||
for r in results:
|
||
print(r[OutputKeys.LABELS], r[OutputKeys.SCORES])
|
||
|
||
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
||
def test_run_with_summarization_with_model(self):
|
||
model = Model.from_pretrained(
|
||
'damo/ofa_summarization_gigaword_large_en')
|
||
ofa_pipe = pipeline(Tasks.text_summarization, model=model)
|
||
text = 'five-time world champion michelle kwan withdrew' + \
|
||
'from the #### us figure skating championships on wednesday ,' + \
|
||
' but will petition us skating officials for the chance to ' + \
|
||
'compete at the #### turin olympics .'
|
||
input = {'text': text}
|
||
result = ofa_pipe(input)
|
||
print(result)
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_summarization_with_name(self):
|
||
ofa_pipe = pipeline(
|
||
Tasks.text_summarization,
|
||
model='damo/ofa_summarization_gigaword_large_en')
|
||
text = 'five-time world champion michelle kwan withdrew' + \
|
||
'from the #### us figure skating championships on wednesday ,' + \
|
||
' but will petition us skating officials for the chance to ' + \
|
||
'compete at the #### turin olympics .'
|
||
input = {'text': text}
|
||
result = ofa_pipe(input)
|
||
print(result)
|
||
|
||
# test for return multiple sequences
|
||
ofa_pipe.model.num_return_sequences = 2
|
||
result = ofa_pipe(input)
|
||
print(result)
|
||
# test batch infer
|
||
ofa_pipe.model.num_return_sequences = 1
|
||
input = [{'text': text} for _ in range(3)]
|
||
results = ofa_pipe(input, batch_size=2)
|
||
for r in results:
|
||
print(r[OutputKeys.TEXT])
|
||
|
||
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
||
def test_run_with_text_classification_with_model(self):
|
||
model = Model.from_pretrained(
|
||
'damo/ofa_text-classification_mnli_large_en')
|
||
ofa_pipe = pipeline(Tasks.text_classification, model=model)
|
||
text = 'One of our number will carry out your instructions minutely.'
|
||
text2 = 'A member of my team will execute your orders with immense precision.'
|
||
result = ofa_pipe((text, text2))
|
||
result = ofa_pipe({'text': text, 'text2': text2})
|
||
print(result)
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_text_classification_with_name(self):
|
||
ofa_pipe = pipeline(
|
||
Tasks.text_classification,
|
||
model='damo/ofa_text-classification_mnli_large_en')
|
||
text = 'One of our number will carry out your instructions minutely.'
|
||
text2 = 'A member of my team will execute your orders with immense precision.'
|
||
result = ofa_pipe((text, text2))
|
||
print(result)
|
||
# test batch infer
|
||
inputs = [(text, text2) for _ in range(3)]
|
||
results = ofa_pipe(inputs, batch_size=2)
|
||
for r in results:
|
||
print(r[OutputKeys.LABELS], r[OutputKeys.SCORES])
|
||
|
||
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
||
def test_run_with_visual_entailment_with_model(self):
|
||
model = Model.from_pretrained(
|
||
'damo/ofa_visual-entailment_snli-ve_large_en')
|
||
ofa_pipe = pipeline(Tasks.visual_entailment, model=model)
|
||
image = 'data/test/images/dogs.jpg'
|
||
text = 'there are two birds.'
|
||
input = {'image': image, 'text': text}
|
||
result = ofa_pipe(input)
|
||
print(result)
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_visual_entailment_with_name(self):
|
||
ofa_pipe = pipeline(
|
||
Tasks.visual_entailment,
|
||
model='damo/ofa_visual-entailment_snli-ve_large_en')
|
||
image = 'data/test/images/dogs.jpg'
|
||
text = 'there are two birds.'
|
||
input = {'image': image, 'text': text}
|
||
result = ofa_pipe(input)
|
||
print(result)
|
||
# test batch infer
|
||
input = [{'image': image, 'text': text} for _ in range(3)]
|
||
results = ofa_pipe(input, batch_size=2)
|
||
for r in results:
|
||
print(r[OutputKeys.LABELS], r[OutputKeys.SCORES])
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_visual_grounding_with_model(self):
|
||
model = Model.from_pretrained(
|
||
'damo/ofa_visual-grounding_refcoco_large_en')
|
||
ofa_pipe = pipeline(Tasks.visual_grounding, model=model)
|
||
image = 'data/test/images/visual_grounding.png'
|
||
text = 'a blue turtle-like pokemon with round head'
|
||
input = {'image': image, 'text': text}
|
||
result = ofa_pipe(input)
|
||
print(result)
|
||
image_name = image.split('/')[-2]
|
||
self.save_img(
|
||
image,
|
||
result[OutputKeys.BOXES][0], # just one box
|
||
osp.join('large_en_model_' + image_name + '.png'))
|
||
|
||
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
||
def test_run_with_visual_grounding_with_name(self):
|
||
ofa_pipe = pipeline(
|
||
Tasks.visual_grounding,
|
||
model='damo/ofa_visual-grounding_refcoco_large_en')
|
||
image = 'data/test/images/visual_grounding.png'
|
||
text = 'a blue turtle-like pokemon with round head'
|
||
input = {'image': image, 'text': text}
|
||
result = ofa_pipe(input)
|
||
print(result)
|
||
image_name = image.split('/')[-2]
|
||
self.save_img(image, result[OutputKeys.BOXES][0],
|
||
osp.join('large_en_name_' + image_name + '.png'))
|
||
# test batch infer
|
||
result = ofa_pipe([input for _ in range(3)], batch_size=2)
|
||
print(result)
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_visual_grounding_zh_with_name(self):
|
||
model = 'damo/ofa_visual-grounding_refcoco_large_zh'
|
||
ofa_pipe = pipeline(Tasks.visual_grounding, model=model)
|
||
image = 'data/test/images/visual_grounding.png'
|
||
text = '一个圆头的蓝色宝可梦'
|
||
input = {'image': image, 'text': text}
|
||
result = ofa_pipe(input)
|
||
print(result)
|
||
image_name = image.split('/')[-1]
|
||
self.save_img(image, result[OutputKeys.BOXES][0],
|
||
osp.join('large_zh_name_' + image_name))
|
||
|
||
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
||
def test_run_with_visual_question_answering_with_model(self):
|
||
model = Model.from_pretrained(
|
||
'damo/ofa_visual-question-answering_pretrain_large_en')
|
||
ofa_pipe = pipeline(Tasks.visual_question_answering, model=model)
|
||
image = 'data/test/images/visual_question_answering.png'
|
||
text = 'what is grown on the plant?'
|
||
input = {'image': image, 'text': text}
|
||
result = ofa_pipe(input)
|
||
print(result)
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_visual_question_answering_with_name(self):
|
||
model = 'damo/ofa_visual-question-answering_pretrain_large_en'
|
||
ofa_pipe = pipeline(Tasks.visual_question_answering, model=model)
|
||
image = 'data/test/images/visual_question_answering.png'
|
||
text = 'what is grown on the plant?'
|
||
input = {'image': image, 'text': text}
|
||
result = ofa_pipe(input)
|
||
print(result)
|
||
|
||
# test batch infer
|
||
result = ofa_pipe([input for _ in range(3)], batch_size=2)
|
||
print(result)
|
||
|
||
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
||
def test_run_with_image_captioning_distilled_with_model(self):
|
||
model = Model.from_pretrained(
|
||
'damo/ofa_image-caption_coco_distilled_en')
|
||
img_captioning = pipeline(
|
||
task=Tasks.image_captioning,
|
||
model=model,
|
||
)
|
||
image_path = 'data/test/images/image_captioning.png'
|
||
image = Image.open(image_path)
|
||
result = img_captioning(image)
|
||
print(result[OutputKeys.CAPTION])
|
||
|
||
# test batch infer
|
||
print(img_captioning([image for _ in range(3)], batch_size=2))
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_visual_entailment_distilled_model_with_name(self):
|
||
ofa_pipe = pipeline(
|
||
Tasks.visual_entailment,
|
||
model='damo/ofa_visual-entailment_snli-ve_distilled_v2_en')
|
||
image = 'data/test/images/dogs.jpg'
|
||
text = 'there are two birds.'
|
||
input = {'image': image, 'text': text}
|
||
result = ofa_pipe(input)
|
||
print(result)
|
||
|
||
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
||
def test_run_with_visual_grounding_distilled_model_with_model(self):
|
||
model = Model.from_pretrained(
|
||
'damo/ofa_visual-grounding_refcoco_distilled_en')
|
||
ofa_pipe = pipeline(Tasks.visual_grounding, model=model)
|
||
image = 'data/test/images/visual_grounding.png'
|
||
text = 'a blue turtle-like pokemon with round head'
|
||
input = {'image': image, 'text': text}
|
||
result = ofa_pipe(input)
|
||
print(result)
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_text_to_image_synthesis_with_name(self):
|
||
model = 'damo/ofa_text-to-image-synthesis_coco_large_en'
|
||
ofa_pipe = pipeline(Tasks.text_to_image_synthesis, model=model)
|
||
ofa_pipe.model.generator.beam_size = 2
|
||
example = {'text': 'a bear in the water.'}
|
||
result = ofa_pipe(example)
|
||
result[OutputKeys.OUTPUT_IMGS][0].save('result.png')
|
||
print(f'Output written to {osp.abspath("result.png")}')
|
||
|
||
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
||
def test_run_with_text_to_image_synthesis_with_model(self):
|
||
model = Model.from_pretrained(
|
||
'damo/ofa_text-to-image-synthesis_coco_large_en')
|
||
ofa_pipe = pipeline(Tasks.text_to_image_synthesis, model=model)
|
||
ofa_pipe.model.generator.beam_size = 2
|
||
example = {'text': 'a bear in the water.'}
|
||
result = ofa_pipe(example)
|
||
result[OutputKeys.OUTPUT_IMG].save('result.png')
|
||
print(f'Output written to {osp.abspath("result.png")}')
|
||
|
||
@unittest.skipUnless(
|
||
test_level() >= 1,
|
||
'skip test in current test level, model has no text2phone_dict.txt')
|
||
def test_run_with_asr_with_name(self):
|
||
model = 'damo/ofa_mmspeech_pretrain_base_zh'
|
||
ofa_pipe = pipeline(Tasks.auto_speech_recognition, model=model)
|
||
example = {'wav': 'data/test/audios/asr_example_ofa.wav'}
|
||
result = ofa_pipe(example)
|
||
print(result[OutputKeys.TEXT])
|
||
# test batch infer
|
||
result = ofa_pipe([example for _ in range(3)], batch_size=2)
|
||
for r in result:
|
||
print(r[OutputKeys.TEXT])
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_sudoku_with_name(self):
|
||
model = 'damo/ofa_sudoku_kaggle_large'
|
||
ofa_pipe = pipeline(Tasks.sudoku, model=model)
|
||
# the valid num is 1-9，and use 0 represents the empty block
|
||
# the separator of column is ` : `, and the separator of row is ` | `
|
||
example = '5 : 3 : 0 : 0 : 7 : 0 : 0 : 0 : 0 | \
|
||
6 : 0 : 0 : 1 : 9 : 5 : 0 : 0 : 0 | \
|
||
0 : 9 : 8 : 0 : 0 : 0 : 0 : 6 : 0 | \
|
||
8 : 0 : 0 : 0 : 6 : 0 : 0 : 0 : 3 | \
|
||
4 : 0 : 0 : 8 : 0 : 3 : 0 : 0 : 1 | \
|
||
7 : 0 : 0 : 0 : 2 : 0 : 0 : 0 : 6 | \
|
||
0 : 6 : 0 : 0 : 0 : 0 : 2 : 8 : 0 | \
|
||
0 : 0 : 0 : 4 : 1 : 9 : 0 : 0 : 5 | \
|
||
0 : 0 : 0 : 0 : 8 : 0 : 0 : 7 : 9'
|
||
|
||
result = ofa_pipe(example)
|
||
print(result[OutputKeys.TEXT])
|
||
# test batch infer
|
||
result = ofa_pipe([example for _ in range(3)], batch_size=2)
|
||
for r in result:
|
||
print(r[OutputKeys.TEXT])
|
||
|
||
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
|
||
def test_run_with_text2sql_with_name(self):
|
||
model = 'damo/ofa_text2sql_spider_large_en'
|
||
ofa_pipe = pipeline(Tasks.text2sql, model=model)
|
||
text = 'Show all book categories and the number of books in each category.'
|
||
database = 'culture_company' # optional, default `culture_company`
|
||
example = {'text': text, 'database': database}
|
||
result = ofa_pipe(example)
|
||
print(result[OutputKeys.TEXT])
|
||
# test batch infer
|
||
result = ofa_pipe([example for _ in range(3)], batch_size=2)
|
||
for r in result:
|
||
print(r[OutputKeys.TEXT])
|
||
|
||
|
||
if __name__ == '__main__':
|
||
unittest.main()
|