model.py

from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib import cm

import cv2
from PIL import Image
import numpy as np

import torch
from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
from diffusers import StableDiffusionInpaintPipeline


class VirtualStagingToolV2():

    def __init__(self,
                 segmentation_version='openmmlab/upernet-convnext-tiny',
                 diffusion_version="stabilityai/stable-diffusion-2-inpainting"
                 ):

        self.segmentation_version = segmentation_version
        self.diffusion_version = diffusion_version

        if segmentation_version == "openmmlab/upernet-convnext-tiny":
            self.feature_extractor = AutoImageProcessor.from_pretrained(self.segmentation_version)
            self.segmentation_model = UperNetForSemanticSegmentation.from_pretrained(self.segmentation_version)
        elif segmentation_version == "nvidia/segformer-b5-finetuned-ade-640-640":
            self.feature_extractor = SegformerFeatureExtractor.from_pretrained(self.segmentation_version)
            self.segmentation_model = SegformerForSemanticSegmentation.from_pretrained(self.segmentation_version)

        self.diffution_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
            self.diffusion_version,
            torch_dtype=torch.float16,
        )
        self.diffution_pipeline = self.diffution_pipeline.to("cuda")

    def _predict(self, image):
        inputs = self.feature_extractor(images=image, return_tensors="pt")
        outputs = self.segmentation_model(**inputs)
        prediction = \
        self.feature_extractor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
        return prediction

    def _save_mask(self, img, prediction_array, mask_items=[]):
        mask = np.zeros_like(prediction_array, dtype=np.uint8)

        mask[np.isin(prediction_array, mask_items)] = 0
        mask[~np.isin(prediction_array, mask_items)] = 255

        buffer_size = 10

        # Dilate the binary image
        kernel = np.ones((buffer_size, buffer_size), np.uint8)
        dilated_image = cv2.dilate(mask, kernel, iterations=1)

        # Subtract the original binary image
        buffer_area = dilated_image - mask

        # Apply buffer area to the original image
        mask = cv2.bitwise_or(mask, buffer_area)

        #     # # Create a PIL Image object from the mask
        mask_image = Image.fromarray(mask, mode='L')
        # display(mask_image)

        # mask_image = mask_image.resize((512, 512))
        # mask_image.save(".tmp/mask_1.png", "PNG")
        # img = img.resize((512, 512))
        # img.save(".tmp/input_1.png", "PNG")
        return mask_image

    def _save_transparent_mask(self, img, prediction_array, mask_items=[]):
        mask = np.array(img)
        mask[~np.isin(prediction_array, mask_items), :] = 255
        mask_image = Image.fromarray(mask).convert('RGBA')

        # Set the transparency of the pixels corresponding to object 1 to 0 (fully transparent)
        mask_data = mask_image.getdata()
        mask_data = [(r, g, b, 0) if r == 255 else (r, g, b, 255) for (r, g, b, a) in mask_data]
        mask_image.putdata(mask_data)

        return mask_image

    def get_mask(self, image_path=None, image=None):
        if image_path:
            image = Image.open(image_path)
        else:
            if not image:
                raise ValueError("no image provided")

        # display(image)
        prediction = self._predict(image)

        label_ids = np.unique(prediction)

        mask_items = [0, 3, 5, 8, 14]

        if 1 in label_ids or 25 in label_ids:
            mask_items = [1, 2, 4, 25, 32]
            room = 'backyard'
        elif 73 in label_ids or 50 in label_ids or 61 in label_ids:
            mask_items = [0, 3, 5, 8, 14, 50, 61, 71, 73, 118, 124, 129
                          ]
            room = 'kitchen'
        elif 37 in label_ids or 65 in label_ids or (27 in label_ids and 47 in label_ids and 70 in label_ids):
            mask_items = [0, 3, 5, 8, 14, 27, 65]
            room = 'bathroom'
        elif 7 in label_ids:
            room = 'bedroom'
        elif 23 in label_ids or 49 in label_ids:
            mask_items = [0, 3, 5, 8, 14, 49]
            room = 'living room'
        elif 15 in label_ids and 19 in label_ids:
            room = 'dining room'
        else:
            room ='room'
        label_ids_without_mask = [i for i in label_ids if i not in mask_items]

        items = [self.segmentation_model.config.id2label[i] for i in label_ids_without_mask]

        mask_image = self._save_mask(image, prediction, mask_items)
        transparent_mask_image = self._save_transparent_mask(image, prediction, mask_items)
        return mask_image, transparent_mask_image, image, items, room

    def _edit_image(self, init_image, mask_image, prompt,  # height, width,
                    number_images=1):

        init_image = init_image.resize((512, 512)).convert("RGB")
        mask_image = mask_image.resize((512, 512)).convert("RGB")

        output_images = self.diffution_pipeline(
            prompt=prompt, image=init_image, mask_image=mask_image,
            # width=width, height=height,
            num_images_per_prompt=number_images).images
        # display(output_image)
        return output_images

    def virtual_stage(self, image_path=None, image=None, style=None,
                      color_preference=None, additional_info=None, number_images=1):
        mask_image, transparent_mask_image, init_image, items, room = self.get_mask(image_path, image)
        if not style:
            raise ValueError('style not provided.')


        if room == 'kitchen':
            items = [i for i in items if i in ['cabinet', 'shelf', 'counter', 'countertop', 'stool']]
        elif room == 'bedroom':
            items = [i for i in items if i in ['bed ', 'table', 'chest of drawers', 'desk', 'armchair', 'wardrobe']]
        elif room == 'bathroom':
            items = [i for i in items if
                     i in ['shower', 'bathtub', 'screen door', 'cabinet']]
        elif room == 'living room':
            items = [i for i in items if
                     i in ['table', 'sofa', 'chest of drawers', 'armchair', 'cabinet', 'coffee table']]
        elif room == 'dining room':
            items = [i for i in items if i in ['table', 'chair', 'cabinet']]

        items = ', '.join(items)

        if room == 'backyard':
            prompt = f'Realistic, high resolution, {room} with {style}'
        else:
            prompt = f'Realistic {items}, high resolution, in the {style} style {room}'

        if color_preference:
            prompt = f"{prompt} in {color_preference}"

        if additional_info:
            prompt = f'{prompt}. {additional_info}'
        print(prompt)

        output_images = self._edit_image(init_image, mask_image, prompt, number_images)

        final_output_images = []
        for output_image in output_images:
            output_image = output_image.resize(init_image.size)
            final_output_images.append(output_image)
        return final_output_images, transparent_mask_image