sam_3d.py

import os
import copy
import random
from PIL import Image

import torch
import matplotlib.pyplot as plt
import numpy as np
from plyfile import PlyData

from segment_anything import SamPredictor, sam_model_registry

def get_image_ids(path):
    files = os.listdir(path)
    files = [f.split('.')[0] for f in files if os.path.isfile(path+'/'+f)] #Filtering only the files.
    return sorted(files)

def load_align_matrix_from_txt(path):
    lines = open(path).readlines()
    # test set data doesn't have align_matrix
    axis_align_matrix = np.eye(4)
    for line in lines:
        if 'axisAlignment' in line:
            axis_align_matrix = [
                float(x)
                for x in line.rstrip().strip('axisAlignment = ').split(' ')
            ]
            break
    axis_align_matrix = np.array(axis_align_matrix).reshape((4, 4))
    return axis_align_matrix

def load_matrix_from_txt(path, shape=(4, 4)):
    with open(path) as f:
        txt = f.readlines()
    txt = ''.join(txt).replace('\n', ' ')
    matrix = [float(v) for v in txt.split()]
    return np.array(matrix).reshape(shape)

def load_image(path):
    image = Image.open(path)
    return np.array(image)

def convert_from_uvd(u, v, d, intr, pose, align):
    extr = np.linalg.inv(pose)
    if d == 0:
        return None, None, None
    
    fx = intr[0, 0]
    fy = intr[1, 1]
    cx = intr[0, 2]
    cy = intr[1, 2]
    depth_scale = 1000
    
    z = d / depth_scale
    x = (u - cx) * z / fx
    y = (v - cy) * z / fy
    
    world = (align @ pose @ np.array([x, y, z, 1]))
    return world[:3] / world[3]

# Find the cloest point in the cloud with select
def find_closest_point(point, point_cloud, num=1):
    # calculate the Euclidean distances between the input vector and each row of the matrix
    distances = np.linalg.norm(point_cloud - point, axis=1)

    # find the index of the row with the minimum distance
    closest_index = np.argsort(distances)[:num]

    # get the closest vector from the matrix
    closest_vector = point_cloud[closest_index]
    
    return closest_index, closest_vector
    
def plot_3d(xdata, ydata, zdata, color=None, b_min=2, b_max=8, view=(45, 45)):
        fig, ax = plt.subplots(subplot_kw={"projection": "3d"}, dpi=200)
        ax.view_init(view[0], view[1])
        ax.set_xlim(b_min, b_max)
        ax.set_ylim(b_min, b_max)
        ax.set_zlim(b_min, b_max)
        ax.scatter3D(xdata, ydata, zdata, c=color, cmap='rgb', s=0.1)

class SAM3DDemo(object):
    def __init__(self, sam_model, sam_ckpt, scene_name):
        sam = sam_model_registry[sam_model](checkpoint=sam_ckpt)
        self.predictor = SamPredictor(sam)
        self.scene_name = scene_name
        scene_path = os.path.join('./scannet_data', scene_name)
        self.color_path = os.path.join(scene_path, 'color')
        self.depth_path = os.path.join(scene_path, 'depth')
        self.pose_path = os.path.join(scene_path, 'pose')
        self.intrinsic_path = os.path.join(scene_path, 'intrinsic')
        self.align_matirx_path = f'{scene_path}/{scene_name}.txt'
        self.img_ids = get_image_ids(self.color_path)
        self.align_matrix = load_align_matrix_from_txt(self.align_matirx_path)
        self.intrinsic_depth = load_matrix_from_txt(os.path.join(self.intrinsic_path, 'intrinsic_depth.txt'))
        self.poses = [load_matrix_from_txt(os.path.join(self.pose_path, f'{i}.txt')) for i in self.img_ids]
        self.rgb_images = [load_image(os.path.join(self.color_path, f'{i}.jpg')) for i in self.img_ids]
        self.depth_images = [load_image(os.path.join(self.depth_path, f'{i}.png'))for i in self.img_ids]

    def project_3D_to_images(self, select_points, valid_margin=20):
        valid_img_ids = []
        valid_points = {}
        for img_i in range(len(self.img_ids)):
            rgb_img = self.rgb_images[img_i]
            depth_img = self.depth_images[img_i]
            extrinsics = self.poses[img_i]
            projection_matrix = self.intrinsic_depth @ np.linalg.inv(self.align_matrix @ extrinsics)
            raw_points = np.vstack((select_points.T, np.ones((1, select_points.T.shape[1]))))
            raw_points = np.dot(projection_matrix, raw_points)
            # bounding simplest 
            points = raw_points[:2, :] / raw_points[2, :]
            points = np.round(points).astype(np.int32)
            valid = (points[0] >= valid_margin).all() & (points[1] >= valid_margin).all() \
                & (points[0] < (rgb_img.shape[1] - valid_margin)).all() & (points[1] < (rgb_img.shape[0] - valid_margin)).all() \
                & (raw_points[2, :] > 0).all()
            if valid:
                depth_margin = 0.4
                gt_depths = depth_img[points[1], points[0]] / 1000
                proj_depths = raw_points[2, :]
                if (proj_depths[0] > (1 - depth_margin / 2.0) * gt_depths[0]) & (proj_depths[0] < (1 + depth_margin / 2.0) * gt_depths[0]):
                    valid_img_ids.append(img_i)
                    valid_points[img_i] = points

        show_id = valid_img_ids[-1]
        show_points = valid_points[show_id]
        rgb_img = self.rgb_images[show_id]
        fig, ax = plt.subplots()
        ax.imshow(rgb_img)
        for x, y  in zip(show_points[0], show_points[1]):
            ax.plot(x, y, 'ro')
        canvas = fig.canvas
        canvas.draw()
        w, h = canvas.get_width_height()
        rgb_img_w_points = np.frombuffer(canvas.tostring_rgb(), dtype='uint8').reshape(h, w, 3)
        print("projecting 3D point to images successfully...")
        return valid_img_ids, valid_points, rgb_img_w_points

    def process_img_w_sam(self, valid_img_ids, valid_points, granularity):
        mask_colors = []
        for img_i in range(len(self.img_ids)):
            rgb_img = self.rgb_images[img_i]
            msk_color = np.full(rgb_img.shape, 0.5)
            if img_i in valid_img_ids:
                self.predictor.set_image(rgb_img)
                point_coor = valid_points[img_i].T[0][None]
                masks, _, _ = self.predictor.predict(point_coords=point_coor, point_labels=np.array([1]))
                # fig, axs = plt.subplots(nrows=1, ncols=3, figsize=(10, 5))
                # for i in range(3):
                #     mask_img = masks[i][:,:,None] * rgb_img
                #     axs[i].set_title(f'granularity {i}')
                #     axs[i].imshow(mask_img)
                m = masks[granularity]
                msk_color[m] = [0, 0, 1.0]
            mask_colors.append(msk_color)
        show_id = valid_img_ids[-1]
        rgb_img = self.rgb_images[show_id]
        fig, axs = plt.subplots(nrows=1, ncols=3, figsize=(24, 8))
        for i in range(3):
            mask_img = masks[i][:,:,None] * rgb_img
            axs[i].set_title(f'granularity {i}')
            axs[i].imshow(mask_img)
        canvas = fig.canvas
        canvas.draw()
        w, h = canvas.get_width_height()
        rgb_img_w_masks = np.frombuffer(canvas.tostring_rgb(), dtype='uint8').reshape(h, w, 3)
        print("processing images with SAM successfully...")
        return mask_colors, rgb_img_w_masks

    def project_mask_to_3d(self, mask_colors, sample_ratio=0.002):
        x_data, y_data, z_data, c_data = [], [], [], []
        for img_i in range(len(self.img_ids)):
            id = self.img_ids[img_i]
            # RGBD 
            d = self.depth_images[img_i]
            c = self.rgb_images[img_i]
            p = self.poses[img_i]
            msk_color = mask_colors[img_i]
            # Projecting RGB features into the point space
            for i in range(d.shape[0]):
                for j in range(d.shape[1]):
                    if random.random() < sample_ratio:
                        x, y, z = convert_from_uvd(j, i, d[i, j], self.intrinsic_depth, p, self.align_matrix)
                        if x is None:
                            continue
                        x_data.append(x)
                        y_data.append(y)
                        z_data.append(z)
                        ci = int(i * c.shape[0] / d.shape[0])
                        cj = int(j * c.shape[1] / d.shape[1])
                        c_data.append([msk_color[ci, cj]])
        print("reprojecting images to 3D points successfully...")
        return x_data, y_data, z_data, c_data

    def match_projected_point_to_gt_point(self, x_data, y_data, z_data, c_data, gt_coords):

        c_data = torch.tensor(np.concatenate(c_data, axis=0))
        img_coords = np.array([x_data, y_data, z_data], dtype=np.float32).T
        gt_quant_coords = np.floor_divide(gt_coords, 0.2)
        img_quant_coords = np.floor_divide(img_coords, 0.2)

        # Remove the reduandant coords
        unique_gt_coords, gt_inverse_indices = np.unique(gt_quant_coords, axis=0, return_inverse=True)
        unique_img_coords, img_inverse_indices = np.unique(img_quant_coords, axis=0, return_inverse=True)

        # Match the coords in gt_coords to img_corrds
        def find_loc(vec):
            obj = np.empty((), dtype=object)
            out = np.where((unique_img_coords == vec).all(1))[0]
            obj[()] = out
            return obj
        
        gt_2_img_map = np.apply_along_axis(find_loc, 1, unique_gt_coords)
        # Since some places are empty, using the simple round interplation
        gt_2_img_map_filled = []
        start_id = np.array([0])
        for loc in gt_2_img_map:
            if not np.any(loc):
                loc = start_id
            else:
                start_id = loc
            gt_2_img_map_filled.append(int(loc))
        
        mean_colors = []
        for i in range(unique_img_coords.shape[0]):
            valid_locs = np.where(img_inverse_indices == i)
            mean_f = torch.mean(c_data[valid_locs], axis=0)
            # mean_f, _ = torch.mode(c_data[valid_locs], dim=0)
            mean_colors.append(mean_f.unsqueeze(0))
        mean_colors = torch.cat(mean_colors)
        # Project the averaged features back to groundtruth point clouds
        img_2_gt_colors = mean_colors[gt_2_img_map_filled]
        projected_gt_colors = img_2_gt_colors[gt_inverse_indices]
        print("convert projected points to GT points successfully...")

        return projected_gt_colors
    
    def render_point_cloud(self, data, color):
        data_copy = copy.copy(data)
        uint_color = torch.round(torch.tensor(color) * 255).to(torch.uint8)
        data_copy['red'] = uint_color[:, 0]
        data_copy['green'] = uint_color[:, 1]
        data_copy['blue'] = uint_color[:, 2]
        return data_copy
    
    def run_with_coord(self, point, granularity):
        x_data, y_data, z_data, c_data = [], [], [], []

        plydata = PlyData.read(f"./scannet_data/{self.scene_name}/{self.scene_name}.ply")
        data = plydata.elements[0].data

        # gt_coords stand for the groudtruth point clouds coordinates 
        gt_coords = np.array([data['x'], data['y'], data['z']], dtype=np.float32).T
        gt_color = np.array([data['red'], data['green'], data['blue']], dtype=np.float32).T
        blank_color = np.full(gt_color.shape, 0.5)

        select_index, select_points = find_closest_point(point, gt_coords, num=10)
        point_select_color = blank_color.copy()
        point_select_color[select_index] = [1.0, 0, 0]
        data_point_select = self.render_point_cloud(data, point_select_color)

        valid_img_ids, valid_points, rgb_img_w_points = self.project_3D_to_images(select_points)
        mask_colors, rgb_img_w_masks = self.process_img_w_sam(valid_img_ids, valid_points, granularity)
        x_data, y_data, z_data, c_data = self.project_mask_to_3d(mask_colors)
        projected_gt_colors = self.match_projected_point_to_gt_point(x_data, y_data, z_data, c_data, gt_coords)
        
        data_final = self.render_point_cloud(data, projected_gt_colors)

        return data_point_select, rgb_img_w_points, rgb_img_w_masks, data_final