How to get segmentation pth

[DukeofCambridge]

Hi, thanks for the great work.
I wonder how you get the segmentation pth file with ODISE. Seems it will only give a visualization image by default.
I can get similar pth file with Mask2Former as it provides an index image with each pixel labeled by the class index.

Is there an official guidance? Thanks.

[Zhu-Liyuan]

Hi, thank you for your interest in our work.
I didn't include ODISE because it requires a very old version of diffuser and stable diffusion, which are not compatible with ReStyle3D repo. I suggest you do the following:

• Install ODISE in a separate conda/pyvenv environment.
• Use the following scipt to get the segmentation pth file for running ReStyle3D.

odise_utils.py:

import requests
import itertools
from contextlib import ExitStack
import torch
from mask2former.data.datasets.register_ade20k_panoptic import ADE20K_150_CATEGORIES
from PIL import Image
import numpy as np

from detectron2.config import instantiate
from detectron2.data import MetadataCatalog
from detectron2.data import detection_utils as utils
from detectron2.data import transforms as T
from detectron2.config import LazyCall as L
from detectron2.data.datasets.builtin_meta import COCO_CATEGORIES
from detectron2.evaluation import inference_context
from detectron2.utils.env import seed_all_rng
from detectron2.utils.logger import setup_logger
from detectron2.utils.visualizer import ColorMode, Visualizer, random_color

from odise import model_zoo
from odise.checkpoint import ODISECheckpointer
from odise.config import instantiate_odise
from odise.data import get_openseg_labels
from odise.modeling.wrapper import OpenPanopticInference

setup_logger()
logger = setup_logger(name="odise")

COCO_THING_CLASSES = [
    label
    for idx, label in enumerate(get_openseg_labels("coco_panoptic", True))
    if COCO_CATEGORIES[idx]["isthing"] == 1
]
COCO_THING_COLORS = [c["color"] for c in COCO_CATEGORIES if c["isthing"] == 1]
COCO_STUFF_CLASSES = [
    label
    for idx, label in enumerate(get_openseg_labels("coco_panoptic", True))
    if COCO_CATEGORIES[idx]["isthing"] == 0
]
COCO_STUFF_COLORS = [c["color"] for c in COCO_CATEGORIES if c["isthing"] == 0]

ADE_THING_CLASSES = [
    label
    for idx, label in enumerate(get_openseg_labels("ade20k_150", True))
    if ADE20K_150_CATEGORIES[idx]["isthing"] == 1
]
ADE_THING_COLORS = [c["color"] for c in ADE20K_150_CATEGORIES if c["isthing"] == 1]
ADE_STUFF_CLASSES = [
    label
    for idx, label in enumerate(get_openseg_labels("ade20k_150", True))
    if ADE20K_150_CATEGORIES[idx]["isthing"] == 0
]
ADE_STUFF_COLORS = [c["color"] for c in ADE20K_150_CATEGORIES if c["isthing"] == 0]

LVIS_CLASSES = get_openseg_labels("lvis_1203", True)
# use beautiful coco colors
LVIS_COLORS = list(
    itertools.islice(itertools.cycle([c["color"] for c in COCO_CATEGORIES]), len(LVIS_CLASSES))
)


class VisualizationDemo(object):
    def __init__(self, model, metadata, aug, instance_mode=ColorMode.IMAGE):
        """
        Args:
            model (nn.Module):
            metadata (MetadataCatalog): image metadata.
            instance_mode (ColorMode):
            parallel (bool): whether to run the model in different processes from visualization.
                Useful since the visualization logic can be slow.
        """
        self.model = model
        self.metadata = metadata
        self.aug = aug
        self.cpu_device = torch.device("cpu")
        self.instance_mode = instance_mode

    def predict(self, original_image):
        """
        Args:
            original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).

        Returns:
            predictions (dict):
                the output of the model for one image only.
                See :doc:`/tutorials/models` for details about the format.
        """
        height, width = original_image.shape[:2]
        aug_input = T.AugInput(original_image, sem_seg=None)
        self.aug(aug_input)
        image = aug_input.image
        image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))

        inputs = {"image": image, "height": height, "width": width}
        predictions = self.model([inputs])[0]
        return predictions

    def run_on_image(self, image):
        """
        Args:
            image (np.ndarray): an image of shape (H, W, C) (in BGR order).
                This is the format used by OpenCV.
        Returns:
            predictions (dict): the output of the model.
            vis_output (VisImage): the visualized image output.
        """
        vis_output, pred_labels = None, []
        predictions = self.predict(image)
        visualizer = Visualizer(image, self.metadata, instance_mode=self.instance_mode)
        if "panoptic_seg" in predictions:
            panoptic_seg, segments_info = predictions["panoptic_seg"]
            vis_output = visualizer.draw_panoptic_seg(
                panoptic_seg.to(self.cpu_device), segments_info
            )
        else:
            if "sem_seg" in predictions:
                vis_output = visualizer.draw_sem_seg(
                    predictions["sem_seg"].argmax(dim=0).to(self.cpu_device)
                )
            if "instances" in predictions:
                instances = predictions["instances"].to(self.cpu_device)
                vis_output = visualizer.draw_instance_predictions(predictions=instances)

        return predictions, vis_output

def add_labels_to_segments_info(predictions, metadata):
    if "panoptic_seg" in predictions:
        panoptic_seg, segments_info = predictions["panoptic_seg"]
        for segment in segments_info:
            label_id = segment["category_id"]
            isthing = segment["isthing"]
            if isthing:
                label = metadata.thing_classes[label_id]
            else:
                label = metadata.stuff_classes[label_id]
            segment["label"] = label
        return panoptic_seg, segments_info
    return None, []

def build_demo_classes_and_metadata(vocab, label_list):
    extra_classes = []

    if vocab:
        for words in vocab.split(";"):
            extra_classes.append([word.strip() for word in words.split(",")])
    extra_colors = [random_color(rgb=True, maximum=1) for _ in range(len(extra_classes))]

    demo_thing_classes = extra_classes
    demo_stuff_classes = []
    demo_thing_colors = extra_colors
    demo_stuff_colors = []

    if "COCO" in label_list:
        demo_thing_classes += COCO_THING_CLASSES
        demo_stuff_classes += COCO_STUFF_CLASSES
        demo_thing_colors += COCO_THING_COLORS
        demo_stuff_colors += COCO_STUFF_COLORS
    if "ADE" in label_list:
        demo_thing_classes += ADE_THING_CLASSES
        demo_stuff_classes += ADE_STUFF_CLASSES
        demo_thing_colors += ADE_THING_COLORS
        demo_stuff_colors += ADE_STUFF_COLORS
    if "LVIS" in label_list:
        demo_thing_classes += LVIS_CLASSES
        demo_thing_colors += LVIS_COLORS

    MetadataCatalog.pop("odise_demo_metadata", None)
    demo_metadata = MetadataCatalog.get("odise_demo_metadata")
    demo_metadata.thing_classes = [c[0] for c in demo_thing_classes]
    demo_metadata.stuff_classes = [
        *demo_metadata.thing_classes,
        *[c[0] for c in demo_stuff_classes],
    ]
    demo_metadata.thing_colors = demo_thing_colors
    demo_metadata.stuff_colors = demo_thing_colors + demo_stuff_colors
    demo_metadata.stuff_dataset_id_to_contiguous_id = {
        idx: idx for idx in range(len(demo_metadata.stuff_classes))
    }
    demo_metadata.thing_dataset_id_to_contiguous_id = {
        idx: idx for idx in range(len(demo_metadata.thing_classes))
    }

    demo_classes = demo_thing_classes + demo_stuff_classes

    return demo_classes, demo_metadata


def load_model(config_path="Panoptic/odise_label_coco_50e.py", seed=42):
    cfg = model_zoo.get_config(config_path, trained=True)

    cfg.model.overlap_threshold = 0
    seed_all_rng(seed)

    
    cfg.dataloader.test.mapper.augmentations=[
            L(T.ResizeShortestEdge)(short_edge_length=960, sample_style="choice", max_size=2560),
        ]
    dataset_cfg = cfg.dataloader.test
    
    aug = instantiate(dataset_cfg.mapper).augmentations

    model = instantiate_odise(cfg.model)
    model.to(cfg.train.device)
    ODISECheckpointer(model).load(cfg.train.init_checkpoint)

    return model, aug

def inference(image, model, aug, vocab=None, label_list=["COCO", "ADE", "LVIS"]):

    demo_classes, demo_metadata = build_demo_classes_and_metadata(vocab, label_list)
    with ExitStack() as stack:
        inference_model = OpenPanopticInference(
            model=model,
            labels=demo_classes,
            metadata=demo_metadata,
            semantic_on=False,
            instance_on=False,
            panoptic_on=True,
        )
        stack.enter_context(inference_context(inference_model))
        stack.enter_context(torch.no_grad())

        demo = VisualizationDemo(inference_model, demo_metadata, aug)
        predictions, visualized_output = demo.run_on_image(np.array(image))
        panoptic_seg, segments_info = add_labels_to_segments_info(predictions, demo_metadata)
        predictions["panoptic_seg"] = (panoptic_seg, segments_info)
        return predictions, demo_classes, Image.fromarray(visualized_output.get_image())
    

segment_folder.py:

import numpy as np
import torch
import os
from pathlib import Path
from PIL import Image
from collections import defaultdict

from utils.odise_utils import inference, load_model

def resize_image(img, short_edge=1024):
    w, h = img.size
    if h < w:
        new_h = short_edge
        new_w = int(w * (short_edge / h))
    else:
        new_w = short_edge
        new_h = int(h * (short_edge / w))
    return img.resize((new_w, new_h), Image.LANCZOS)

def save_predictions_checkpoint(pred_src, pred_tgt, src_path, labels, filename='predictions_checkpoint.pth'):
    # Get the directory of the source image
    image_dir = os.path.dirname(src_path)
    
    # Create the full path for the checkpoint file
    checkpoint_path = os.path.join(image_dir, filename)
    
    checkpoint = {
        'pred_src': pred_src,
        'pred_tgt': pred_tgt,
        "classes": labels
    }
    
    torch.save(checkpoint, checkpoint_path)
    print(f"Predictions saved to {checkpoint_path}")

def style_image_seg(data_path, scene):
    data_path = Path(data_path)/scene
    model, aug = load_model() 
    vocab = f"{scene}"
    label_list = ["COCO"]
    for image_name in os.listdir(data_path/"images"):
        rgb_data = Image.open(data_path/'images'/image_name)
        width, height = rgb_data.size
        
        rgb_data = resize_image(rgb_data)
        pred, classes, vis = inference(rgb_data, model, aug, vocab, label_list)
        pred_mask = torch.nn.functional.interpolate(pred['panoptic_seg'][0][None, None].float(), size=(height, width), mode='nearest')
        pred_dict = (pred_mask.squeeze().to(torch.int32), pred['panoptic_seg'][1])
        
        filename = image_name.split('.')[0]
        vis.save(data_path/'visual'/f'{filename}.png')
        checkpoint = {
            'pred_tgt': pred_dict,
            "classes": classes
        }
        
        checkpoint_path = data_path/f'seg_dict/{filename}.pth'
        torch.save(checkpoint, checkpoint_path)
    print()

def rename_file_pairs(folder_path):
    # Group files by their base name (without extension)
    file_groups = defaultdict(list)
    for filename in os.listdir(folder_path):
        base_name = os.path.splitext(filename)[0]
        file_groups[base_name].append(filename)
    
    # Sort the groups to ensure consistent ordering
    sorted_groups = sorted(file_groups.items())
    
    # Rename files with sequential numbers
    for index, (_, files) in enumerate(sorted_groups):
        for file in files:
            # Get the extension
            extension = os.path.splitext(file)[1]
            # Create new filename
            new_filename = f"{index}{extension}"
            
            # Full paths
            old_path = os.path.join(folder_path, file)
            new_path = os.path.join(folder_path, new_filename)
            
            # Rename the file
            os.rename(old_path, new_path)
            print(f"Renamed '{file}' to '{new_filename}'")

def segment_RE10K():
    scenes = ["bedroom", "kitchen", "living_room"]
    data_path = Path("Path to your image folder")
    model, aug = load_model() 
    label_list = ["COCO"]

    for scene in scenes:
        vocab = f"{scene}"
        scene_path = data_path/scene
        for id in os.listdir(scene_path):
            for image_name in os.listdir(scene_path/id/"images"):
                rgb_data = Image.open(scene_path/id/'images'/image_name)
                width, height = rgb_data.size
                
                rgb_data = resize_image(rgb_data)
                pred, classes, vis = inference(rgb_data, model, aug, vocab, label_list)
                pred_mask = torch.nn.functional.interpolate(pred['panoptic_seg'][0][None, None].float(), size=(height, width), mode='nearest')
                pred_dict = (pred_mask.squeeze().to(torch.int32), pred['panoptic_seg'][1])
                
                filename = image_name.split('.')[0]
                vis.save(scene_path/id/'visual'/f'{filename}.png')
                checkpoint = {
                    'pred_tgt': pred_dict,
                    "classes": classes
                }
                
                checkpoint_path = scene_path/id/f'seg_dict/{filename}.pth'
                torch.save(checkpoint, checkpoint_path)


if __name__ == "__main__":
    data_path = "Path to your image folder"
    scenes = ["bedroom", "kitchen", "living_room"]
    segment_RE10K()
    # for scene in scenes:
        # style_image_seg(data_path, scene)

You should be able to get the segmentation pth by running segment_folder.py.