AideepImage/interior-design

1import logging2from typing import List, Tuple, Dict3 4import streamlit as st5import torch6import gc7import numpy as np8from PIL import Image9 10from transformers import AutoImageProcessor, UperNetForSemanticSegmentation11 12from palette import ade_palette13 14LOGGING = logging.getLogger(__name__)15 16 17def flush():18    gc.collect()19    torch.cuda.empty_cache()20 21@st.cache_resource(max_entries=5)22def get_segmentation_pipeline() -> Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]:23    """Method to load the segmentation pipeline24    Returns:25        Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]: segmentation pipeline26    """27    image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")28    image_segmentor = UperNetForSemanticSegmentation.from_pretrained(29        "openmmlab/upernet-convnext-small")30    return image_processor, image_segmentor31 32 33@torch.inference_mode()34@torch.autocast('cuda')35def segment_image(image: Image) -> Image:36    """Method to segment image37    Args:38        image (Image): input image39    Returns:40        Image: segmented image41    """42    image_processor, image_segmentor = get_segmentation_pipeline()43    pixel_values = image_processor(image, return_tensors="pt").pixel_values44    with torch.no_grad():45        outputs = image_segmentor(pixel_values)46 47    seg = image_processor.post_process_semantic_segmentation(48        outputs, target_sizes=[image.size[::-1]])[0]49    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)50    palette = np.array(ade_palette())51    for label, color in enumerate(palette):52        color_seg[seg == label, :] = color53    color_seg = color_seg.astype(np.uint8)54    seg_image = Image.fromarray(color_seg).convert('RGB')55    return seg_image