Init project

app.py

@@ -0,0 +1,15 @@
+import numpy as np
+import gradio as gr
+from model import SegmentationTool
+
+seg_tool = SegmentationTool()
+
+
+def segment(input_img):
+    mask_image, transparent_mask_image, image, items, room = seg_tool.get_mask(input_img)
+    return mask_image
+
+
+demo = gr.Interface(segment, gr.Image(), "image")
+if __name__ == "__main__":
+    demo.launch()

model.py

@@ -0,0 +1,107 @@
+import cv2
+from PIL import Image
+import numpy as np
+from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
+from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
+
+
+class SegmentationTool:
+
+    def __init__(self,
+                 segmentation_version='nvidia/segformer-b5-finetuned-ade-640-640'):
+
+        self.segmentation_version = segmentation_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)
+
+    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, 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=None):
+        if mask_items is None:
+            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]
+        mask_items = [8] # windowpane
+
+        if 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(prediction, mask_items)
+        transparent_mask_image = self._save_transparent_mask(image, prediction, mask_items)
+        return mask_image, transparent_mask_image, image, items, room

requirements.txt

@@ -0,0 +1,4 @@
+numpy
+Pillow
+opencv-python
+transformers