| 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.no_grad()
|
| 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) |
