490 lines
21 KiB
Python
490 lines
21 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import os
|
|
import unittest
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
from scepter.modules.annotator.registry import ANNOTATORS
|
|
from scepter.modules.utils.config import Config
|
|
from scepter.modules.utils.distribute import we
|
|
from scepter.modules.utils.file_system import FS
|
|
|
|
|
|
class AnnotatorTest(unittest.TestCase):
|
|
def setUp(self):
|
|
print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
|
|
_ = FS.init_fs_client(Config(cfg_dict={
|
|
'NAME': 'ModelscopeFs',
|
|
'TEMP_DIR': './cache/data'
|
|
},
|
|
load=False),
|
|
overwrite=False)
|
|
image_path = 'asset/images/sunflower.jpeg'
|
|
image = Image.open(image_path)
|
|
if image.mode != 'RGB':
|
|
image = image.convert('RGB')
|
|
self.image = np.array(image)
|
|
|
|
self.save_dir = './cache/save_data/images'
|
|
if not os.path.exists(self.save_dir):
|
|
os.makedirs(self.save_dir)
|
|
|
|
def tearDown(self):
|
|
super().tearDown()
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_canny(self):
|
|
# canny
|
|
canny_dict = {
|
|
'NAME': 'CannyAnnotator',
|
|
'LOW_THRESHOLD': 100,
|
|
'HIGH_THRESHOLD': 200
|
|
}
|
|
canny_anno = Config(cfg_dict=canny_dict, load=False)
|
|
canny_ins = ANNOTATORS.build(canny_anno).to(we.device_id)
|
|
canny_image = canny_ins(self.image)
|
|
print("canny's shape:", canny_image.shape)
|
|
Image.fromarray(canny_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_canny.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_canny_random(self):
|
|
# canny
|
|
canny_dict = {
|
|
'NAME': 'CannyAnnotator',
|
|
'LOW_THRESHOLD': 100,
|
|
'HIGH_THRESHOLD': 200,
|
|
'RANDOM_CFG': {
|
|
'PROBA': 1.0,
|
|
'MIN_LOW_THRESHOLD': 50,
|
|
'MAX_LOW_THRESHOLD': 100,
|
|
'MIN_HIGH_THRESHOLD': 200,
|
|
'MAX_HIGH_THRESHOLD': 350
|
|
}
|
|
}
|
|
canny_anno = Config(cfg_dict=canny_dict, load=False)
|
|
canny_ins = ANNOTATORS.build(canny_anno).to(we.device_id)
|
|
canny_image = canny_ins(self.image)
|
|
print("canny's shape:", canny_image.shape)
|
|
Image.fromarray(canny_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_canny_random.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_hed(self):
|
|
# hed
|
|
hed_dict = {
|
|
'NAME':
|
|
'HedAnnotator',
|
|
'PRETRAINED_MODEL':
|
|
'ms://iic/scepter_scedit@annotator/ckpts/ControlNetHED.pth'
|
|
}
|
|
hed_anno = Config(cfg_dict=hed_dict, load=False)
|
|
hed_ins = ANNOTATORS.build(hed_anno).to(we.device_id)
|
|
hed_image = hed_ins(self.image)
|
|
print("hed's shape:", hed_image.shape)
|
|
Image.fromarray(hed_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_hed.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_openpose(self):
|
|
# openpose
|
|
openpose_dict = {
|
|
'NAME':
|
|
'OpenposeAnnotator',
|
|
'BODY_MODEL_PATH':
|
|
'ms://iic/scepter_scedit@annotator/ckpts/body_pose_model.pth',
|
|
'HAND_MODEL_PATH':
|
|
'ms://iic/scepter_scedit@annotator/ckpts/hand_pose_model.pth'
|
|
}
|
|
openpose_anno = Config(cfg_dict=openpose_dict, load=False)
|
|
openpose_ins = ANNOTATORS.build(openpose_anno).to(we.device_id)
|
|
openpose_image = openpose_ins(self.image)
|
|
print("openpose's shape:", openpose_image.shape)
|
|
Image.fromarray(openpose_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_openpose.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_midas(self):
|
|
# midas
|
|
midas_dict = {
|
|
'NAME': 'MidasDetector',
|
|
'PRETRAINED_MODEL':
|
|
'ms://iic/scepter_scedit@annotator/ckpts/dpt_hybrid-midas-501f0c75.pt',
|
|
'A': 6.2,
|
|
'BG_TH': 0.1
|
|
}
|
|
midas_anno = Config(cfg_dict=midas_dict, load=False)
|
|
midas_ins = ANNOTATORS.build(midas_anno).to(we.device_id)
|
|
midas_image = midas_ins(self.image)
|
|
print("midas's shape:", midas_image.shape)
|
|
Image.fromarray(midas_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_midas.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_mlsd(self):
|
|
# mlsd
|
|
mlsd_dict = {
|
|
'NAME': 'MLSDdetector',
|
|
'PRETRAINED_MODEL':
|
|
'ms://iic/scepter_scedit@annotator/ckpts/mlsd_large_512_fp32.pth',
|
|
'THR_V': 0.1,
|
|
'THR_D': 0.1
|
|
}
|
|
mlsd_anno = Config(cfg_dict=mlsd_dict, load=False)
|
|
mlsd_ins = ANNOTATORS.build(mlsd_anno).to(we.device_id)
|
|
mlsd_image = mlsd_ins(self.image)
|
|
print("mlsd's shape:", mlsd_image.shape)
|
|
Image.fromarray(mlsd_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_mlsd.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_color(self):
|
|
# color
|
|
color_dict = {'NAME': 'ColorAnnotator', 'RATIO': 64}
|
|
color_anno = Config(cfg_dict=color_dict, load=False)
|
|
color_ins = ANNOTATORS.build(color_anno).to(we.device_id)
|
|
color_image = color_ins(self.image)
|
|
print("color's shape:", color_image.shape)
|
|
Image.fromarray(color_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_color.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_color_random(self):
|
|
# color
|
|
color_dict = {
|
|
'NAME': 'ColorAnnotator',
|
|
'RATIO': 64,
|
|
'RANDOM_CFG': {
|
|
'PROBA': 1.0,
|
|
# 'MIN_RATIO': 64,
|
|
# 'MAX_RATIO': 128
|
|
'CHOICE_RATIO': [32, 64, 128]
|
|
}
|
|
}
|
|
color_anno = Config(cfg_dict=color_dict, load=False)
|
|
color_ins = ANNOTATORS.build(color_anno).to(we.device_id)
|
|
color_image = color_ins(self.image)
|
|
print("color's shape:", color_image.shape)
|
|
Image.fromarray(color_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_color_random.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_multi(self):
|
|
# multi annotators
|
|
canny_dict = {
|
|
'NAME': 'CannyAnnotator',
|
|
'LOW_THRESHOLD': 100,
|
|
'HIGH_THRESHOLD': 200,
|
|
'INPUT_KEYS': ['img'],
|
|
'OUTPUT_KEYS': ['canny_img']
|
|
}
|
|
hed_dict = {
|
|
'NAME': 'HedAnnotator',
|
|
'PRETRAINED_MODEL':
|
|
'ms://iic/scepter_scedit@annotator/ckpts/ControlNetHED.pth',
|
|
'INPUT_KEYS': ['img'],
|
|
'OUTPUT_KEYS': ['hed_img']
|
|
}
|
|
openpose_dict = {
|
|
'NAME': 'OpenposeAnnotator',
|
|
'BODY_MODEL_PATH':
|
|
'ms://iic/scepter_scedit@annotator/ckpts/body_pose_model.pth',
|
|
'HAND_MODEL_PATH':
|
|
'ms://iic/scepter_scedit@annotator/ckpts/hand_pose_model.pth',
|
|
'INPUT_KEYS': ['img'],
|
|
'OUTPUT_KEYS': ['openpose_img']
|
|
}
|
|
midas_dict = {
|
|
'NAME': 'MidasDetector',
|
|
'PRETRAINED_MODEL':
|
|
'ms://iic/scepter_scedit@annotator/ckpts/dpt_hybrid-midas-501f0c75.pt',
|
|
'INPUT_KEYS': ['img'],
|
|
'OUTPUT_KEYS': ['midas_img']
|
|
}
|
|
mlsd_dict = {
|
|
'NAME': 'MLSDdetector',
|
|
'PRETRAINED_MODEL':
|
|
'ms://iic/scepter_scedit@annotator/ckpts/mlsd_large_512_fp32.pth',
|
|
'INPUT_KEYS': ['img'],
|
|
'OUTPUT_KEYS': ['mlsd_img']
|
|
}
|
|
color_dict = {'NAME': 'ColorAnnotator', 'RATIO': 64}
|
|
general_dict = {
|
|
'NAME':
|
|
'GeneralAnnotator',
|
|
'ANNOTATORS': [
|
|
canny_dict, hed_dict, openpose_dict, midas_dict, mlsd_dict,
|
|
color_dict
|
|
]
|
|
}
|
|
general_anno = Config(cfg_dict=general_dict, load=False)
|
|
general_ins = ANNOTATORS.build(general_anno).to(we.device_id)
|
|
output_image = general_ins({'img': self.image})
|
|
for key, save_image in output_image.items():
|
|
Image.fromarray(save_image).save(
|
|
os.path.join(self.save_dir, f'sunflower_multi_{key}.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_processor(self):
|
|
from scepter.modules.annotator.utils import AnnotatorProcessor
|
|
anno_processor = AnnotatorProcessor(anno_type='hed')
|
|
output_image = anno_processor.run(self.image, 'hed')
|
|
Image.fromarray(output_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_processor_hed.png'))
|
|
|
|
anno_processor = AnnotatorProcessor(
|
|
anno_type=['canny', 'color', 'depth'])
|
|
output_image = anno_processor.run(self.image, 'color')
|
|
Image.fromarray(output_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_processor_color.png'))
|
|
|
|
output_image = anno_processor.run(self.image, ['canny', 'depth'])
|
|
for key, save_image in output_image.items():
|
|
Image.fromarray(save_image).save(
|
|
os.path.join(self.save_dir, f'sunflower_processor_{key}.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_doodle(self):
|
|
doodle_dict = {
|
|
'NAME': 'DoodleAnnotator', 'PROCESSOR_TYPE': 'pidinet_sketch',
|
|
'PROCESSOR_CFG': [
|
|
{'NAME': 'PiDiAnnotator',
|
|
'PRETRAINED_MODEL': 'ms://iic/scepter_annotator@annotator/ckpts/table5_pidinet.pth'},
|
|
{'NAME': 'SketchAnnotator',
|
|
'PRETRAINED_MODEL': 'ms://iic/scepter_annotator@annotator/ckpts/sketch_simplification_gan.pth'}
|
|
]
|
|
}
|
|
doodle_anno = Config(cfg_dict=doodle_dict, load=False)
|
|
doodle_ins = ANNOTATORS.build(doodle_anno).to(we.device_id)
|
|
doodle_image = doodle_ins(self.image)
|
|
print("doodle's shape:", doodle_image.shape)
|
|
Image.fromarray(doodle_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_doodle.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_gray(self):
|
|
gray_dict = {'NAME': 'GrayAnnotator'}
|
|
gray_anno = Config(cfg_dict=gray_dict, load=False)
|
|
gray_ins = ANNOTATORS.build(gray_anno).to(we.device_id)
|
|
gray_image = gray_ins(self.image)
|
|
print("gray's shape:", gray_image.shape)
|
|
Image.fromarray(gray_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_gray.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_drawing(self):
|
|
cont_dict = {'NAME': 'InfoDrawContourAnnotator', 'INPUT_NC': 3, 'OUTPUT_NC': 1, 'N_RESIDUAL_BLOCKS': 3,
|
|
'SIGMOID': True,
|
|
'PRETRAINED_MODEL': 'ms://iic/scepter_annotator@annotator/ckpts/informative_drawing_contour_style.pth'}
|
|
cont_anno = Config(cfg_dict=cont_dict, load=False)
|
|
cont_ins = ANNOTATORS.build(cont_anno).to(we.device_id)
|
|
cont_image = cont_ins(self.image)
|
|
print("cont's shape:", cont_image.shape)
|
|
Image.fromarray(cont_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_drawing_contour_style.png'))
|
|
|
|
cont_dict = {'NAME': 'InfoDrawAnimeAnnotator', 'INPUT_NC': 3, 'OUTPUT_NC': 1, 'N_RESIDUAL_BLOCKS': 3,
|
|
'SIGMOID': True,
|
|
'PRETRAINED_MODEL': 'ms://iic/scepter_annotator@annotator/ckpts/informative_drawing_anime_style.pth'}
|
|
cont_anno = Config(cfg_dict=cont_dict, load=False)
|
|
cont_ins = ANNOTATORS.build(cont_anno).to(we.device_id)
|
|
cont_image = cont_ins(self.image)
|
|
print("cont's shape:", cont_image.shape)
|
|
Image.fromarray(cont_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_drawing_anime_style.png'))
|
|
|
|
cont_dict = {'NAME': 'InfoDrawOpenSketchAnnotator', 'INPUT_NC': 3, 'OUTPUT_NC': 1, 'N_RESIDUAL_BLOCKS': 3,
|
|
'SIGMOID': True,
|
|
'PRETRAINED_MODEL': 'ms://iic/scepter_annotator@annotator/ckpts/informative_drawing_opensketch_style.pth'}
|
|
cont_anno = Config(cfg_dict=cont_dict, load=False)
|
|
cont_ins = ANNOTATORS.build(cont_anno).to(we.device_id)
|
|
cont_image = cont_ins(self.image)
|
|
print("cont's shape:", cont_image.shape)
|
|
Image.fromarray(cont_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_drawing_opensketch_style.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_outpainting(self):
|
|
outpaint_dict = {'NAME': 'OutpaintingAnnotator', 'RETURN_SOURCE': False}
|
|
outpaint_anno = Config(cfg_dict=outpaint_dict, load=False)
|
|
outpaint_ins = ANNOTATORS.build(outpaint_anno).to(we.device_id)
|
|
outpaint_image = outpaint_ins(self.image)
|
|
print("outpaint's shape:", outpaint_image.shape)
|
|
Image.fromarray(outpaint_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_outpaint.png'))
|
|
|
|
outpaint_dict = {'NAME': 'OutpaintingAnnotator',
|
|
'RANDOM_CFG': {'DIRECTION_RANGE': ['left', 'right', 'up'], 'RATIO_RANGE': [0.2, 0.8]}}
|
|
outpaint_anno = Config(cfg_dict=outpaint_dict, load=False)
|
|
outpaint_ins = ANNOTATORS.build(outpaint_anno).to(we.device_id)
|
|
outpaint_image = outpaint_ins(self.image, return_mask=True)
|
|
print("outpaint image's shape:", outpaint_image['image'].shape)
|
|
print("outpaint mask's shape:", outpaint_image['mask'].shape)
|
|
Image.fromarray(outpaint_image['image']).save(
|
|
os.path.join(self.save_dir, 'sunflower_outpaint_rand_image.png'))
|
|
Image.fromarray(outpaint_image['mask']).save(
|
|
os.path.join(self.save_dir, 'sunflower_outpaint_rand_mask.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_inpainting(self):
|
|
inpaint_dict = {'NAME': 'InpaintingAnnotator'}
|
|
inpaint_anno = Config(cfg_dict=inpaint_dict, load=False)
|
|
inpaint_ins = ANNOTATORS.build(inpaint_anno).to(we.device_id)
|
|
mask = np.zeros_like(self.image)
|
|
mask = cv2.rectangle(mask, (0, 0), (150, 150), (255, 255, 255), -1)
|
|
mask = mask[:, :, 0] # one channel format
|
|
inpaint_image = inpaint_ins(self.image, mask=mask, return_mask=True)
|
|
print("inpaint image's shape:", inpaint_image['image'].shape)
|
|
print("inpaint mask's shape:", inpaint_image['mask'].shape)
|
|
Image.fromarray(inpaint_image['image']).save(
|
|
os.path.join(self.save_dir, 'sunflower_inpaint_image.png'))
|
|
Image.fromarray(inpaint_image['mask']).save(
|
|
os.path.join(self.save_dir, 'sunflower_inpaint_mask.png'))
|
|
|
|
inpaint_dict = {'NAME': 'InpaintingAnnotator'}
|
|
inpaint_anno = Config(cfg_dict=inpaint_dict, load=False)
|
|
inpaint_ins = ANNOTATORS.build(inpaint_anno).to(we.device_id)
|
|
inpaint_image = inpaint_ins(self.image, return_mask=True)
|
|
print("inpaint image's shape:", inpaint_image['image'].shape)
|
|
print("inpaint mask's shape:", inpaint_image['mask'].shape)
|
|
Image.fromarray(inpaint_image['image']).save(
|
|
os.path.join(self.save_dir, 'sunflower_inpaint_image_2.png'))
|
|
Image.fromarray(inpaint_image['mask']).save(
|
|
os.path.join(self.save_dir, 'sunflower_inpaint_mask_2.png'))
|
|
|
|
inpaint_dict = {'NAME': 'InpaintingAnnotator',
|
|
'MASK_CFG': {"irregular_proba": 0.5,
|
|
"irregular_kwargs": {"min_times": 4,
|
|
"max_times": 10,
|
|
"max_width": 150,
|
|
"max_angle": 4,
|
|
"max_len": 200},
|
|
"box_proba": 0.5,
|
|
"box_kwargs": {"margin": 0,
|
|
"bbox_min_size": 50,
|
|
"bbox_max_size": 150,
|
|
"max_times": 5,
|
|
"min_times": 1}
|
|
}
|
|
}
|
|
inpaint_anno = Config(cfg_dict=inpaint_dict, load=False)
|
|
inpaint_ins = ANNOTATORS.build(inpaint_anno).to(we.device_id)
|
|
inpaint_image = inpaint_ins(self.image, return_mask=True)
|
|
print("inpaint image's shape:", inpaint_image['image'].shape)
|
|
print("inpaint mask's shape:", inpaint_image['mask'].shape)
|
|
Image.fromarray(inpaint_image['image']).save(
|
|
os.path.join(self.save_dir, 'sunflower_inpaint_image_3.png'))
|
|
Image.fromarray(inpaint_image['mask']).save(
|
|
os.path.join(self.save_dir, 'sunflower_inpaint_mask_3.png'))
|
|
|
|
inpaint_image = inpaint_ins(self.image, return_mask=True, mask_color=255)
|
|
print("inpaint image's shape:", inpaint_image['image'].shape)
|
|
print("inpaint mask's shape:", inpaint_image['mask'].shape)
|
|
Image.fromarray(inpaint_image['image']).save(
|
|
os.path.join(self.save_dir, 'sunflower_inpaint_image_4.png'))
|
|
Image.fromarray(inpaint_image['mask']).save(
|
|
os.path.join(self.save_dir, 'sunflower_inpaint_mask_4.png'))
|
|
|
|
inpaint_image = inpaint_ins(self.image, return_mask=True, mask_color=255, return_invert=False)
|
|
print("inpaint image's shape:", inpaint_image['image'].shape)
|
|
print("inpaint mask's shape:", inpaint_image['mask'].shape)
|
|
Image.fromarray(inpaint_image['image']).save(
|
|
os.path.join(self.save_dir, 'sunflower_inpaint_image_5.png'))
|
|
Image.fromarray(inpaint_image['mask']).save(
|
|
os.path.join(self.save_dir, 'sunflower_inpaint_mask_5.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_deg(self):
|
|
deg_dict = {'NAME': 'DegradationAnnotator'}
|
|
deg_anno = Config(cfg_dict=deg_dict, load=False)
|
|
deg_ins = ANNOTATORS.build(deg_anno).to(we.device_id)
|
|
deg_image = deg_ins(self.image)
|
|
print("deg's shape:", deg_image.shape)
|
|
Image.fromarray(deg_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_deg.png'))
|
|
|
|
deg_dict = {
|
|
'NAME': 'DegradationAnnotator',
|
|
'RANDOM_DEGRADATION': True,
|
|
'PARAMS': {
|
|
'gaussian_noise': {},
|
|
'resize': {'scale': [0.4, 0.8]},
|
|
'jpeg': {'jpeg_level': [25, 75]},
|
|
'gaussian_blur': {'kernel_size': [7, 9, 11, 13, 15], 'sigma': [0.9, 1.8]}
|
|
}
|
|
}
|
|
deg_anno = Config(cfg_dict=deg_dict, load=False)
|
|
deg_ins = ANNOTATORS.build(deg_anno).to(we.device_id)
|
|
deg_image = deg_ins(self.image)
|
|
print("deg's shape:", deg_image.shape)
|
|
Image.fromarray(deg_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_deg_2.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_seg(self):
|
|
seg_dict = {
|
|
'NAME': 'ESAMAnnotator',
|
|
'PRETRAINED_MODEL': 'ms://iic/scepter_annotator@annotator/ckpts/efficient_sam_vits.pt',
|
|
'SAVE_MODE': 'P',
|
|
'GRID_SIZE': 32,
|
|
}
|
|
seg_anno = Config(cfg_dict=seg_dict, load=False)
|
|
seg_ins = ANNOTATORS.build(seg_anno).to(we.device_id)
|
|
seg_image = seg_ins(self.image)
|
|
print("seg's shape:", seg_image.shape)
|
|
Image.fromarray(seg_image).save(
|
|
os.path.join(self.save_dir, 'sunflower_esam_seg.png'))
|
|
|
|
seg_dict = {
|
|
'NAME': 'ESAMAnnotator',
|
|
'PRETRAINED_MODEL': 'ms://iic/scepter_annotator@annotator/ckpts/efficient_sam_vits.pt',
|
|
'SAVE_MODE': 'P',
|
|
'GRID_SIZE': 32,
|
|
'USE_DOMINANT_COLOR': True,
|
|
'RETURN_MASK': True
|
|
}
|
|
seg_anno = Config(cfg_dict=seg_dict, load=False)
|
|
seg_ins = ANNOTATORS.build(seg_anno).to(we.device_id)
|
|
seg_image = seg_ins(self.image)
|
|
print("seg image's shape:", seg_image['image'].shape)
|
|
Image.fromarray(seg_image['image']).save(
|
|
os.path.join(self.save_dir, 'sunflower_esam_seg_dominant_image.png'))
|
|
print("seg mask's shape:", seg_image['mask'].shape)
|
|
Image.fromarray(seg_image['mask']).save(
|
|
os.path.join(self.save_dir, 'sunflower_esam_seg_dominant_mask.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_samdraw(self):
|
|
sam_dict = {
|
|
'NAME': 'SAMAnnotatorDraw',
|
|
'TASK_TYPE': 'input_box',
|
|
'SAM_MODEL': 'vit_b',
|
|
'PRETRAINED_MODEL': 'ms://iic/scepter_annotator@annotator/ckpts/sam_vit_b_01ec64.pth'
|
|
}
|
|
sam_anno = Config(cfg_dict=sam_dict, load=False)
|
|
sam_ins = ANNOTATORS.build(sam_anno).to(we.device_id)
|
|
sam_res = sam_ins(self.image, input_box=[0, 0, 200, 200], task_type='input_box', multimask_output=False)
|
|
Image.fromarray(sam_res['mask']).save(os.path.join(self.save_dir, f'sunflower_sam_mask.png'))
|
|
|
|
@unittest.skip('')
|
|
def test_annotator_lama(self):
|
|
lama_dict = {
|
|
'NAME': 'LamaAnnotator',
|
|
'PRETRAINED_MODEL': 'ms://iic/cv_fft_inpainting_lama/'
|
|
}
|
|
lama_anno = Config(cfg_dict=lama_dict, load=False)
|
|
lama_ins = ANNOTATORS.build(lama_anno).to(we.device_id)
|
|
mask = np.zeros_like(self.image)
|
|
mask = cv2.rectangle(mask, (0, 0), (150, 150), (255, 255, 255), -1)
|
|
mask = mask[:, :, 0]
|
|
lama_res = lama_ins(self.image, mask)
|
|
print("lama's shape:", lama_res.shape)
|
|
Image.fromarray(lama_res).save(os.path.join(self.save_dir, f'sunflower_lama_mask2.png'))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|