| | import cv2 |
| | import numpy as np |
| | import PIL.Image |
| | import torch |
| | from controlnet_aux.util import HWC3, ade_palette |
| | from transformers import AutoImageProcessor, UperNetForSemanticSegmentation |
| |
|
| | from cv_utils import resize_image |
| |
|
| |
|
| | class ImageSegmentor: |
| | def __init__(self): |
| | self.image_processor = AutoImageProcessor.from_pretrained( |
| | 'openmmlab/upernet-convnext-small') |
| | self.image_segmentor = UperNetForSemanticSegmentation.from_pretrained( |
| | 'openmmlab/upernet-convnext-small') |
| |
|
| | @torch.inference_mode() |
| | def __call__(self, image: np.ndarray, **kwargs) -> PIL.Image.Image: |
| | detect_resolution = kwargs.pop('detect_resolution', 512) |
| | image_resolution = kwargs.pop('image_resolution', 512) |
| | image = HWC3(image) |
| | image = resize_image(image, resolution=detect_resolution) |
| | image = PIL.Image.fromarray(image) |
| |
|
| | pixel_values = self.image_processor(image, |
| | return_tensors='pt').pixel_values |
| | outputs = self.image_segmentor(pixel_values) |
| | seg = self.image_processor.post_process_semantic_segmentation( |
| | outputs, target_sizes=[image.size[::-1]])[0] |
| | color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) |
| | for label, color in enumerate(ade_palette()): |
| | color_seg[seg == label, :] = color |
| | color_seg = color_seg.astype(np.uint8) |
| |
|
| | color_seg = resize_image(color_seg, |
| | resolution=image_resolution, |
| | interpolation=cv2.INTER_NEAREST) |
| | return PIL.Image.fromarray(color_seg) |
| |
|
