run_react.py

import cv2
from PIL import Image
import numpy as np
import os
import sys

from animeinsseg import AnimeInsSeg, AnimeInstances
from animeinsseg.anime_instances import get_color


ckpt = r'models/AnimeInstanceSegmentation/rtmdetl_e60.ckpt'

mask_thres = 0.3
instance_thres = 0.3
refine_kwargs = {'refine_method': 'refinenet_isnet'} # set to None if not using refinenet
# refine_kwargs = None

# imgp = 'examples/mafuyu.png'
imgp = sys.argv[1]
img = cv2.imread(imgp)
img_name = os.path.basename(imgp).split('.')[0]


net = AnimeInsSeg(ckpt, mask_thr=mask_thres, refine_kwargs=refine_kwargs)
instances: AnimeInstances = net.infer(
    img,
    output_type='numpy',
    pred_score_thr=instance_thres
)

x = 0.5 ## 取りたい画像の大きさによって変える
drawed = img.copy()
im_h, im_w = img.shape[:2]
img_with_alpha = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)

# instances.bboxes, instances.masks will be None, None if no obj is detected

nb_object = len(list(zip(instances.bboxes, instances.masks)))

for ii, (xywh, mask) in enumerate(zip(instances.bboxes, instances.masks)):
    color = get_color(ii)

    mask_alpha = 0.5
    linewidth = max(round(sum(img.shape) / 2 * 0.003), 2)

    # draw bbox
    p1, p2 = (int(xywh[0]), int(xywh[1])), (int(xywh[2] + xywh[0]), int(xywh[3] + xywh[1]))
    # cv2.rectangle(drawed, p1, p2, color, thickness=linewidth, lineType=cv2.LINE_AA)

    threshold_w = x * (int(im_w)/nb_object)
    threshold_h = x * (int(im_h)/nb_object)


    if int(xywh[2]) < threshold_w  and  int(xywh[3]) < threshold_h:
        print("too small")
        continue

    ## draw mask and cutout
    mask_uint8 = mask.astype(np.uint8) * 255
    # masked_img = cv2.bitwise_and(img, img, mask = mask)##アルファチャンネル使わないときはこれ
    img_with_alpha[:, :, 3] = mask_uint8
    cv2.imwrite(f'output_images/masked_seg_{ii}_{img_name}.png', img_with_alpha)
    
    # draw mask
    p = mask.astype(np.float32)
    blend_mask = np.full((im_h, im_w, 3), color, dtype=np.float32)
    alpha_msk = (mask_alpha * p)[..., None]
    alpha_ori = 1 - alpha_msk
    drawed = drawed * alpha_ori + alpha_msk * blend_mask



###元実装
    drawed = drawed.astype(np.uint8)
    # Image.fromarray(drawed[..., ::-1]).show()

output_image_path = f'output_images/seg_{img_name}.png'
print(output_image_path)

# # save
Image.fromarray(drawed[..., ::-1]).save(output_image_path)