ProgressLabeller for TRansPose dataset generation
bl_info = {
    "name": "ProgressLabeller",
    "description": "Precisely segment object-wise image from sfm results",
    "location": "File/Import and File/Export",
    "version": (1, 0),
    "blender": (2, 92, 0),
}


import bpy
import sys
import os

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import registeration.register 


def register():
    registeration.register.register()


def unregister():
    registeration.register.unregister()


if __name__ == "__main__":
    register()

__init__.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
import json

from kernel.exporter import configuration_export

# ImportHelper is a helper class, defines filename and
# invoke() function which calls the file selector.
from bpy_extras.io_utils import ExportHelper
from bpy.props import StringProperty
from bpy.types import Operator


class ExportConfiguration(Operator, ExportHelper):
    """This appears in the tooltip of the operator and in the generated docs"""
    bl_idname = "export_data.configuration" 
    bl_label = "Save Configuration"

    filename_ext = ".json"

    filter_glob: StringProperty(
        default="*.json",
        options={'HIDDEN'},
        maxlen=255,  
    )

    def execute(self, context):
        assert context.object['type'] == 'setting'
        config_id = context.object['config_id']
        path = context.object['dir']
        config = bpy.context.scene.configuration[config_id]
        configuration_export(config, self.filepath)
        return {'FINISHED'}

    def invoke(self, context, event):
        context.window_manager.fileselect_add(self)
        if "dir" in context.object:
            self.filepath = (context.object["dir"] if context.object["dir"].endswith("/") else context.object["dir"] + "/") + "configuration.json"
        return {'RUNNING_MODAL'}



def register():
    bpy.utils.register_class(ExportConfiguration)

def unregister():
    bpy.utils.unregister_class(ExportConfiguration)
file/fileExporter/configuration_export.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
import json
import os
import numpy as np
import bpy
import json
import os
import yaml
import numpy as np
from PIL import Image
from tqdm import tqdm
import trimesh
import pyrender
import multiprocessing
import subprocess
import sys

from kernel.exporter import data_export
from kernel.logging_utility import log_report

# ImportHelper is a helper class, defines filename and
# invoke() function which calls the file selector.
from bpy_extras.io_utils import ExportHelper
from bpy.props import StringProperty, EnumProperty
from bpy.types import Operator
from kernel.exporter import configuration_export, objectposes_export
from kernel.blender_utility import \
    _get_configuration, _get_reconstruction_insameworkspace, _get_obj_insameworkspace, _get_workspace_name, _apply_trans2obj, \
    _align_reconstruction, _is_progresslabeller_object
import registeration.register 

class PoseOptimization(Operator, ExportHelper):
    """This appears in the tooltip of the operator and in the generated docs"""
    bl_idname = "export_data.poseoptimization" 
    bl_label = "Pose Optimization"


    def execute(self, context):
    
        setting_object = 0
        visible_objects = [obj for obj in bpy.context.view_layer.objects if obj.visible_get()]
        for obj in visible_objects:
            if obj.name == "untitled:Setting":
                setting_object = obj


        config_id = setting_object['config_id']
        config = bpy.context.scene.configuration[config_id]
        
        path = setting_object['dir']
        output_filepath = (setting_object["dir"] if setting_object["dir"].endswith("/") else setting_object["dir"] + "/") + "output"
        obj_pose = np.zeros(7, dtype=np.float32)
        obj_context = 0        

        if context.selected_objects[0]["type"] == "model":

            obj_pose[0] = context.selected_objects[0].location[0]
            obj_pose[1] = context.selected_objects[0].location[1]
            obj_pose[2] = context.selected_objects[0].location[2]
            
            context.selected_objects[0].rotation_mode = 'QUATERNION'

            obj_pose[3] = context.selected_objects[0].rotation_quaternion[0]
            obj_pose[4] = context.selected_objects[0].rotation_quaternion[1]
            obj_pose[5] = context.selected_objects[0].rotation_quaternion[2]
            obj_pose[6] = context.selected_objects[0].rotation_quaternion[3]

            obj_context = context.selected_objects[0]

        else:

            log_report("Error", "error", None)
            return {'FINISHED'}
            
            
        output_dict = {
            'projectname': config.projectname,
            'environment': {
                'modelsrc':config.modelsrc,
                "reconstructionsrc":config.reconstructionsrc,
                "datasrc": config.datasrc,           
            },
            'camera':{
                "resolution": [config.resX, config.resY],
                "intrinsic": [[config.fx, 0, config.cx],
                              [0, config.fy, config.cy],
                              [0, 0, 1]],
            },
            'object':{
                "name":obj_context.name.split(":")[1],
                "path":obj_context["path"],
                "translation":[float(obj_pose[0]), float(obj_pose[1]), float(obj_pose[2])],
                "quaternion":[float(obj_pose[3]), float(obj_pose[4]), float(obj_pose[5]), float(obj_pose[6])],
            },
            'optimization':{
                "stop_criteria":context.selected_objects[0].optimizationparas.stop_criteria,
                "translation_learning_rate":context.selected_objects[0].optimizationparas.translation_learning_rate,
                "rotation_learning_rate":context.selected_objects[0].optimizationparas.rotation_learning_rate
            }
        }            
        
        with open("/tmp/progresslabeller_pose_opt.json", "w") as f:
            json.dump(output_dict, f, indent = True)

        source = os.path.dirname(os.path.dirname(__file__))
        code_path = os.path.join(source, "../", "offline", "mask_opt.py")
        subprocess.call("{} {} {} {}".format(sys.executable, code_path, "/tmp/progresslabeller_pose_opt.json", output_filepath), shell=True)

        os.remove("/tmp/progresslabeller_pose_opt.json")

        f = open("/tmp/progresslabeller_mask_result.json")
        result = json.load(f)
        
        objworkspacename = result["projectname"] + ":" + result["object"]["name"]        
        bpy.data.objects[objworkspacename].location = np.array(result["object"]["translation"])
        bpy.data.objects[objworkspacename].rotation_mode = 'QUATERNION'
        bpy.data.objects[objworkspacename].rotation_quaternion = np.array(result["object"]["quaternion"]) / np.linalg.norm(np.array(result["object"]["quaternion"]))
        
        os.remove("/tmp/progresslabeller_mask_result.json")


        log_report("Info", "Pose Optimization", None)          
        
        
        return {'FINISHED'}



class MaskOutput(Operator, ExportHelper):
    """This appears in the tooltip of the operator and in the generated docs"""
    bl_idname = "export_data.savemask" 
    bl_label = "Save Mask"


    def execute(self, context):

        setting_object = 0
        visible_objects = [obj for obj in bpy.context.view_layer.objects if obj.visible_get()]
        for obj in visible_objects:
            if obj.name == "untitled:Setting":
                setting_object = obj


        config_id = setting_object['config_id']
        config = bpy.context.scene.configuration[config_id]
        # config.keys()
        # ['projectname', 'modelsrc', 'reconstructionsrc', 'datasrc', 'resX', 'resY', 'fx', 'fy', 'cx', 'cy', 'lens', 'inverse_pose', 'reconstructionscale', 'cameradisplayscale', 'recon_trans', 'sample_rate', 'depth_scale', 'depth_ignore']


        path = setting_object['dir']
        output_filepath = (setting_object["dir"] if setting_object["dir"].endswith("/") else setting_object["dir"] + "/") + "output"
        cam_pose = np.zeros(7, dtype=np.float32)
        obj_pose = np.zeros(7, dtype=np.float32)
        obj_context = 0
        cam_context = 0

        if context.selected_objects[0]["type"] == "camera" and context.selected_objects[1]["type"] == "model":

            cam_context = context.selected_objects[0]
            obj_context = context.selected_objects[1]
            	
            cam_pose[0] = context.selected_objects[0].location[0]
            cam_pose[1] = context.selected_objects[0].location[1]
            cam_pose[2] = context.selected_objects[0].location[2]

            cam_pose[3] = context.selected_objects[0].rotation_quaternion[0]
            cam_pose[4] = context.selected_objects[0].rotation_quaternion[1]
            cam_pose[5] = context.selected_objects[0].rotation_quaternion[2]
            cam_pose[6] = context.selected_objects[0].rotation_quaternion[3]

            obj_pose[0] = context.selected_objects[1].location[0]
            obj_pose[1] = context.selected_objects[1].location[1]
            obj_pose[2] = context.selected_objects[1].location[2]
            
            obj_context.rotation_mode = 'QUATERNION'

            obj_pose[3] = context.selected_objects[1].rotation_quaternion[0]
            obj_pose[4] = context.selected_objects[1].rotation_quaternion[1]
            obj_pose[5] = context.selected_objects[1].rotation_quaternion[2]
            obj_pose[6] = context.selected_objects[1].rotation_quaternion[3]



        elif context.selected_objects[1]["type"] == "camera" and context.selected_objects[0]["type"] == "model":

            cam_context = context.selected_objects[1]
            obj_context = context.selected_objects[0]

            cam_pose[0] = context.selected_objects[1].location[0]
            cam_pose[1] = context.selected_objects[1].location[1]
            cam_pose[2] = context.selected_objects[1].location[2]

            cam_pose[3] = context.selected_objects[1].rotation_quaternion[0]
            cam_pose[4] = context.selected_objects[1].rotation_quaternion[1]
            cam_pose[5] = context.selected_objects[1].rotation_quaternion[2]
            cam_pose[6] = context.selected_objects[1].rotation_quaternion[3]

            obj_pose[0] = context.selected_objects[0].location[0]
            obj_pose[1] = context.selected_objects[0].location[1]
            obj_pose[2] = context.selected_objects[0].location[2]
            
            obj_context.rotation_mode = 'QUATERNION'

            obj_pose[3] = context.selected_objects[0].rotation_quaternion[0]
            obj_pose[4] = context.selected_objects[0].rotation_quaternion[1]
            obj_pose[5] = context.selected_objects[0].rotation_quaternion[2]
            obj_pose[6] = context.selected_objects[0].rotation_quaternion[3]

            #obj_context = context.selected_objects[0]
            #cam_context = context.selected_objects[1]

        else:

            log_report("Error", "error", None)
            return {'FINISHED'}

        output_dict = {
            'projectname': config.projectname,
            'environment': {
                'modelsrc':config.modelsrc,
                "reconstructionsrc":config.reconstructionsrc,
                "datasrc": config.datasrc,           
            },
            'camera':{
                "name":cam_context.name.split(":")[1],
                "resolution": [config.resX, config.resY],
                "intrinsic": [[config.fx, 0, config.cx],
                              [0, config.fy, config.cy],
                              [0, 0, 1]],
                "inverse_pose": False,
                "lens": 30.0, 
                "translation":[float(cam_pose[0]), float(cam_pose[1]), float(cam_pose[2])],
                "quaternion":[float(cam_pose[3]), float(cam_pose[4]), float(cam_pose[5]), float(cam_pose[6])],
            },
            'object':{
                "name":obj_context.name.split(":")[1],
                "path":obj_context["path"],
                "translation":[float(obj_pose[0]), float(obj_pose[1]), float(obj_pose[2])],
                "quaternion":[float(obj_pose[3]), float(obj_pose[4]), float(obj_pose[5]), float(obj_pose[6])],
            },
            'reconstruction':{
	        "scale": 1.0,
                "cameradisplayscale": 0.01,
                "recon_trans": "1,0,0,0;0,1,0,0;0,0,1,0;0,0,0,1;"
            },
            'data':{
                "sample_rate": 0.1,
                "depth_scale": 0.001,
                "depth_ignore": 8.0
            }
        }


        with open("/tmp/progresslabeller_mask.json", "w") as f:
            json.dump(output_dict, f, indent = True)

        source = os.path.dirname(os.path.dirname(__file__))
        code_path = os.path.join(source, "../", "offline", "main_mask.py")
        subprocess.call("{} {} {} {} {}".format(sys.executable, code_path, "/tmp/progresslabeller_mask.json", output_filepath, "Yourtype"), shell=True)

        os.remove("/tmp/progresslabeller_mask.json")

        log_report("Info", "Save Mask", None)  
        return {'FINISHED'}


   



class DataOutput(Operator, ExportHelper):
    """This appears in the tooltip of the operator and in the generated docs"""
    bl_idname = "export_data.dataoutput" 
    bl_label = "Output Data"

    filename_ext = "/"

    filter_glob: StringProperty(
        default="/",
        options={'HIDDEN'},
        maxlen=255,  
    )

    dataformatType: EnumProperty(
        name="Output format",
        description="Choose a output format",
        items=(
            ('BOP', "BOP", "BOP challenge format"),
            ('YCBV', "YCBV", "YCBV dataset format"),
            ('ProgressLabeller', "ProgressLabeller", "ProgressLabeller format"),
            ('Yourtype', "Yourtype", "Your own form (please define your own type first)"),
        ),
        default='ProgressLabeller',
    )

    def execute(self, context):

        assert context.object['type'] == 'setting'
        config_id = context.object['config_id']
        path = context.object['dir']
        config = bpy.context.scene.configuration[config_id]

        files = os.listdir(config.reconstructionsrc)
        if "campose_R.txt" not in files:
            log_report(
                "Error", "Please do reconstruction first", None
            )
        elif "label_pose.yaml" not in files:
            log_report(
                "Error", "Please allocate the object in the scene and save object poses", None
            )
        if not os.path.exists(os.path.join(config.modelsrc, "object_label.json")):
            log_report(
                "Error", "Please define your models' label fine in <PATH/TO/MODEL>/object_label.json", None
            )            
        else:
            configuration_export(config, "/tmp/progresslabeller.json")
            log_report(
                "Info", "Export data to" + self.filepath, None
            )
            #print(self.dataformatType)
            data_export("/tmp/progresslabeller.json", self.filepath, self.dataformatType)
            
            os.remove("/tmp/progresslabeller.json")
        return {'FINISHED'}

    def invoke(self, context, event):
        context.window_manager.fileselect_add(self)
        if "dir" in context.object:
            self.filepath = (context.object["dir"] if context.object["dir"].endswith("/") else context.object["dir"] + "/") + "output"
        # Tells Blender to hang on for the slow user input
        return {'RUNNING_MODAL'}



def register():
    bpy.utils.register_class(PoseOptimization)
    bpy.utils.register_class(MaskOutput)
    bpy.utils.register_class(DataOutput)

def unregister():
    bpy.utils.register_class(PoseOptimization)
    bpy.utils.unregister_class(MaskOutput)
    bpy.utils.unregister_class(DataOutput)

file/fileExporter/data_export.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy

from kernel.exporter import objectposes_export

# ImportHelper is a helper class, defines filename and
# invoke() function which calls the file selector.
from bpy_extras.io_utils import ExportHelper
from bpy.props import StringProperty
from bpy.types import Operator


class ExportObjectPoses(Operator, ExportHelper):
    """This appears in the tooltip of the operator and in the generated docs"""
    bl_idname = "export_data.objectposes"  
    bl_label = "Save Object Poses"

    filename_ext = ".yaml"

    filter_glob: StringProperty(
        default="*.yaml",
        options={'HIDDEN'},
        maxlen=255,  
    )

    def execute(self, context):
        assert context.object['type'] == 'setting'
        config_id = context.object['config_id']
        config = bpy.context.scene.configuration[config_id]
        workspace_name = config.projectname
        objectposes_export(workspace_name, self.filepath)
        return {'FINISHED'}

    def invoke(self, context, event):
        context.window_manager.fileselect_add(self)
        assert context.object['type'] == 'setting'
        config_id = context.object['config_id']
        config = bpy.context.scene.configuration[config_id]
        if "dir" in context.object:
            self.filepath = (config.reconstructionsrc if config.reconstructionsrc.endswith("/") else config.reconstructionsrc + "/") + "label_pose.yaml"
        return {'RUNNING_MODAL'}

def register():
    bpy.utils.register_class(ExportObjectPoses)

def unregister():
    bpy.utils.unregister_class(ExportObjectPoses)
file/fileExporter/objectposes_export.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
from kernel.loader import load_configuration

# ImportHelper is a helper class, defines filename and
# invoke() function which calls the file selector.
from bpy_extras.io_utils import ImportHelper
from bpy.props import StringProperty
from bpy.types import Operator


class ImportConfiguration(Operator, ImportHelper):
    """Load model for pose alignment and segmentation"""
    bl_idname = "import_data.configuration"  
    bl_label = "Import Configuration"

    # ImportHelper mixin class uses this
    filename_ext = ".json"

    filter_glob: StringProperty(
        default="*.json",
        options={'HIDDEN'},
        maxlen=255,  
    )

    def execute(self, context):
        load_configuration(self.filepath)
        return {'FINISHED'}
    
    def invoke(self, context, event):
        context.window_manager.fileselect_add(self)
        return {'RUNNING_MODAL'}


def _menu_func_import(self, context):
    self.layout.operator(ImportConfiguration.bl_idname, text="ProgressLabeller Configuration(.json)")


def register():
    bpy.utils.register_class(ImportConfiguration)
    bpy.types.TOPBAR_MT_file_import.append(_menu_func_import)

def unregister():
    bpy.utils.unregister_class(ImportConfiguration)
    bpy.types.TOPBAR_MT_file_import.remove(_menu_func_import)
file/fileImporter/configuration_loader.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
from kernel.loader import load_model, load_model_from_pose

# ImportHelper is a helper class, defines filename and
# invoke() function which calls the file selector.
from bpy_extras.io_utils import ImportHelper
from bpy.props import StringProperty
from bpy.types import Operator


class ImportModel(Operator, ImportHelper):
    """Load model for pose alignment and segmentation"""
    bl_idname = "import_data.model"  
    bl_label = "Import Model"
    filename_ext = ".obj"

    filter_glob: StringProperty(
        default="*.obj",
        options={'HIDDEN'},
        maxlen=255,  
    )

    def execute(self, context):
        load_model(self.filepath, context.object["config_id"])
        return {'FINISHED'}
    
    def invoke(self, context, event):
        context.window_manager.fileselect_add(self)
        if hasattr(bpy.context.scene.configuration[context.object["config_id"]], 'modelsrc'):
            self.filepath = bpy.context.scene.configuration[context.object["config_id"]].modelsrc + "/"
        return {'RUNNING_MODAL'}


class ImportModelfromPoseFile(Operator, ImportHelper):
    """Load model for pose alignment and segmentation from a pose file"""
    bl_idname = "import_data.modelfrompose"  
    bl_label = "Import Model from pose"

    filename_ext = ".yaml"

    filter_glob: StringProperty(
        default="*.yaml",
        options={'HIDDEN'},
        maxlen=255,  
    )


    def execute(self, context):
        print(self.filepath)
        load_model_from_pose(self.filepath, context.object["config_id"])
        return {'FINISHED'}
    
    def invoke(self, context, event):
        context.window_manager.fileselect_add(self)
        if hasattr(bpy.context.scene.configuration[context.object["config_id"]], 'reconstructionsrc'):
            if bpy.context.scene.configuration[context.object["config_id"]].reconstructionsrc[-1] != "/":
                self.filepath = bpy.context.scene.configuration[context.object["config_id"]].reconstructionsrc + "/"
            else:
                self.filepath = bpy.context.scene.configuration[context.object["config_id"]].reconstructionsrc
        return {'RUNNING_MODAL'}

def _menu_func_import(self, context):
    self.layout.operator(ImportModel.bl_idname, text="ProgressLabeller Model(.obj)")
    self.layout.operator(ImportModelfromPoseFile.bl_idname, text="ProgressLabeller Model from pose file(.yaml)")



def register():
    bpy.utils.register_class(ImportModel)
    bpy.utils.register_class(ImportModelfromPoseFile)



def unregister():
    bpy.utils.unregister_class(ImportModel)
    bpy.utils.unregister_class(ImportModelfromPoseFile)
file/fileImporter/model_loader.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
from kernel.loader import load_reconstruction_result

# ImportHelper is a helper class, defines filename and
# invoke() function which calls the file selector.
from bpy_extras.io_utils import ImportHelper
from bpy.props import StringProperty, EnumProperty, FloatProperty
from bpy.types import Operator


class ImportReconstruction(Operator, ImportHelper):
    """Load model for pose alignment and segmentation"""
    bl_idname = "import_data.reconstruction"  
    bl_label = "Import Reconstruction Result"

    filename_ext = "/"

    filter_glob: StringProperty(
        default="/",
        options={'HIDDEN'},
        maxlen=255,  
    )

    reconstruction_method: EnumProperty(
        name="Reconstruction Method",
        description="Choose the method you use for your reconstruction data",
        items=(
            ('COLMAP', "COLMAP", "COLMAP"),
            ('KinectFusion', "KinectFusion", "KinectFusion"),
            ('Meshroom', "Meshroom", "Meshroom"),
        ),
        default='COLMAP',
    )    

    pointcloudscale: FloatProperty(
        name="PC Scale", 
        description="scale for the loading point cloud", 
        default=0.2, min=0.00, 
        max=1.00, step=2, precision=2)

    camerascale: FloatProperty(
        name="Camera Scale", 
        description="scale for displaying the camera", 
        default=0.1, min=0.00, 
        max=1.00, step=2, precision=2)

    def execute(self, context):
        load_reconstruction_result(filepath = self.filepath, 
                                   reconstruction_method = self.reconstruction_method,
                                   pointcloudscale = self.pointcloudscale,
                                   imagesrc = bpy.context.scene.configuration.imagesrc,
                                   camera_display_scale = self.camerascale,)
        return {'FINISHED'}
    
    def invoke(self, context, event):
        context.window_manager.fileselect_add(self)
        if hasattr(bpy.context.scene.configuration, 'reconstructionsrc'):
            self.filepath = bpy.context.scene.configuration.reconstructionsrc
        return {'RUNNING_MODAL'}

def _menu_func_import(self, context):
    self.layout.operator(ImportReconstruction.bl_idname, text="ProgressLabeller Load Reconstruction Result(package)")



def register():
    bpy.utils.register_class(ImportReconstruction)


def unregister():
    bpy.utils.unregister_class(ImportReconstruction)
file/fileImporter/reconstruction_result_loader.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
from bpy_extras.io_utils import ExportHelper
from bpy.props import StringProperty, BoolProperty, EnumProperty
from bpy.types import Operator
from kernel.loader import create_workspace, init_package
from kernel.blender_utility import _get_configuration
import os

class CreateNewWorkspace(Operator, ExportHelper):
    """This appears in the tooltip of the operator and in the generated docs"""
    bl_idname = "new.new_work_space"  
    bl_label = "Create New Work Space"

    filename_ext = ""

    filter_glob: StringProperty(
        default="/",
        options={'HIDDEN'},
        maxlen=255,  
    )

    def execute(self, context):
        path, name = os.path.split(self.filepath)
        create_workspace(path, name)
        _, config = _get_configuration(bpy.data.objects[name + ":Setting"])
        init_package(path, config)
        
        return {'FINISHED'}


  


# Only needed if you want to add into a dynamic menu
def menu_func_export(self, context):
    self.layout.operator(CreateNewWorkspace.bl_idname, text="ProgressLabeller Create New Workspace")


def register():
    bpy.utils.register_class(CreateNewWorkspace)
    bpy.types.TOPBAR_MT_file_new.append(menu_func_export)


def unregister():
    bpy.utils.unregister_class(CreateNewWorkspace)
    bpy.types.TOPBAR_MT_file_new.remove(menu_func_export)

file/fileNew/new_workspace.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
import numpy as np
import open3d as o3d
import os
import tqdm
from PIL import Image

from kernel.logging_utility import log_report
from kernel.geometry import _pose2Rotation, _rotation2Pose
from kernel.scale import _parseImagesFile, _parsePoints3D, _scaleFordepth, _calculateDepth
from kernel.utility import _transstring2trans,  _trans2transstring


def _is_progresslabeller_object(obj):
    if "type" in obj:
        return True
    else:
        return False

def _is_in_blender(name):
    if name in bpy.data.objects:
        return True
    else:
        log_report(
            "Error", "Object {0} is in Blender".format(name), None
        )  
        return False

def _get_workspace_name(obj):
    if _is_progresslabeller_object(obj):
        return obj.name.split(":")[0]
    else:
        return None

def _get_configuration(obj):
    if _is_progresslabeller_object(obj):
        workspaceName = _get_workspace_name(obj)
        config_id = bpy.data.objects[workspaceName + ":Setting"]['config_id']
        return config_id, bpy.context.scene.configuration[config_id]
    else:
        return None, None

def _get_reconstruction_insameworkspace(obj):
    if _is_progresslabeller_object(obj):
        workspaceName = _get_workspace_name(obj)
        if _is_in_blender(workspaceName + ":reconstruction_depthfusion"):
            return bpy.data.objects[workspaceName + ":reconstruction_depthfusion"]
        elif _is_in_blender(workspaceName + ":reconstruction"):
            return bpy.data.objects[workspaceName + ":reconstruction"]
        else:
            return None
    else:
        return None

def _is_obj_type(obj, types):
    if _is_progresslabeller_object(obj) and obj["type"] in types:
        return True
    else:
        return False

def _get_obj_insameworkspace(obj, types):
    if _is_progresslabeller_object(obj):
        obj_list = []
        workspaceName = _get_workspace_name(obj)
        for o in bpy.data.objects:
            if "type" in o and o.name.split(":")[0] == workspaceName and o["type"] in types:
                obj_list.append(o)
        return obj_list
    else:
        return []

def _apply_trans2obj(obj, trans):
    origin_pose = [list(obj.location), list(obj.rotation_quaternion)]
    origin_trans = _pose2Rotation(origin_pose)
    after_align_trans = trans.dot(origin_trans)
    after_align_pose = _rotation2Pose(after_align_trans)
    obj.location = after_align_pose[0]
    obj.rotation_quaternion = after_align_pose[1]/np.linalg.norm(after_align_pose[1])

def _get_allrgb_insameworkspace(config):
    im_list = []
    workspace_name = config.projectname
    for im in bpy.data.images:
        if im.name.split(":")[0] == workspace_name:
            im_list.append(im)
    return im_list

def _clear_allrgbdcam_insameworkspace(config):
    workspace_name = config.projectname
    bpy.ops.object.select_all(action='DESELECT')
    for obj in bpy.data.objects:
        if _is_obj_type(obj, "camera") and _get_workspace_name(obj) == workspace_name:
            obj.select_set(True)
    bpy.ops.object.delete() 
    for im in bpy.data.images:
        if im.name.split(":")[0] == workspace_name:
            bpy.data.images.remove(im)

def clear_initial_object():
    bpy.ops.object.select_all(action='DESELECT')
    for obj in ["Camera", "Cube", "Light"]:
        if obj in bpy.data.objects:
            bpy.data.objects[obj].select_set(True)
    bpy.ops.object.delete()

    if "Collection" in bpy.data.collections:
        collection = bpy.data.collections.get("Collection")
        bpy.data.collections.remove(collection)
def _clear_recon_output(recon_dir):
    not_delete_file = ["label_pose.yaml", "image-list.txt"] 
    files = os.listdir(recon_dir)
    for f in files:
        p = os.path.join(recon_dir, f)
        if os.path.isfile(p) and f not in not_delete_file:
            os.remove(p)
        elif os.path.isdir(p):
            _delete(p)

def _delete(dir):
    files = os.listdir(dir)
    for f in files:
        os.remove(os.path.join(dir, f))
    os.rmdir(dir)

def _initreconpose(config):
    trans = np.identity(4)
    config.recon_trans = _trans2transstring(trans)   

def _align_reconstruction(config, scene, THRESHOLD = 0.01, NUM_THRESHOLD = 5):
    filepath = config.reconstructionsrc

    imagefilepath = os.path.join(filepath, "images.txt")
    points3Dpath = os.path.join(filepath, "points3D.txt")
    datapath = config.datasrc  

    intrinsic = np.array([
        [config.fx, 0, config.cx],
        [0, config.fy, config.cy],
        [0, 0, 1],
    ])

    depth_scale = config.depth_scale

    Camera_dict = {}
    PointsDict = {}
    PointsDepth = {}
    _parseImagesFile(imagefilepath, Camera_dict, PointsDict)
    _parsePoints3D(points3Dpath, PointsDict)
    print("Auto Aligning the scale:....")
    for camera_idx in tqdm.tqdm(Camera_dict):
            rgb_name = Camera_dict[camera_idx]["name"]
            with Image.open(os.path.join(datapath, "depth", rgb_name)) as im:
                depth = np.array(im) * depth_scale
                _scaleFordepth(depth, camera_idx, intrinsic, Camera_dict, PointsDict, PointsDepth, POSE_INVERSE = config.inverse_pose)
        
    scale = _calculateDepth(THRESHOLD = THRESHOLD, NUM_THRESHOLD = NUM_THRESHOLD, PointsDepth = PointsDepth)
    print("The scale is : {0}".format(scale))
    return scale

def _getsameinstance(config, objName): 
    workspace_name = config.projectname
    objlists = bpy.data.objects.keys()
    objNameStart = workspace_name + ":" + objName
    goalobjlist = []
    for objname in objlists:
        if objname.startswith(objNameStart):
            goalobjlist.append(objname)
    return goalobjlist

def _getnextperfixforinstance(config, objName):
    goalobjlist = _getsameinstance(config, objName)
    perfixlist = [int(name[-3:]) for name in goalobjlist]
    if len(perfixlist) == 0:
        return 1
    else:
        return (np.array(perfixlist)).max() + 1


kernel/blender_utility.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
import json
import os
import yaml
import numpy as np
from PIL import Image
from tqdm import tqdm
import trimesh
import pyrender
import multiprocessing
import subprocess
import sys

from registeration.init_configuration import config_json_dict, encode_dict
from kernel.logging_utility import log_report
from kernel.geometry import _loadModel, _pose2Rotation, _render
from kernel.blender_utility import _is_progresslabeller_object


def configuration_export(config, path):
    configuration = encode_dict(config)
    with open(path, "w") as f:
        json.dump(configuration, f, indent = True)


def objectposes_export(name, path):
    pose = {}
    for obj in bpy.data.objects:
        if _is_progresslabeller_object(obj) and obj["type"] == "model" and obj.name.split(":")[0] == name:
            obj.rotation_mode = 'QUATERNION'
            pose[obj.name.split(":")[1]] = {'pose': [list(obj.location), list(obj.rotation_quaternion)], 'type':obj["modeltype"]}
    with open(path, 'w') as file:
        documents = yaml.dump(pose, file)


def data_export(config_path, target_dir, data_format):
    source = os.path.dirname(os.path.dirname(__file__))
    code_path = os.path.join(source, "offline", "main.py")
    subprocess.call("{} {} {} {} {}".format(sys.executable, code_path, config_path, target_dir, data_format), shell=True)



def _createpose(path, perfix, T):
    posefileName = os.path.join(path, perfix + ".txt")
    np.savetxt(posefileName, np.linalg.inv(T), fmt='%f', delimiter=' ')

def _createbgpyrender(image, segement, dirpath, perfix):
    image[segement != 0] = 0
    colorfileName = os.path.join(dirpath, perfix + ".png")
    img = Image.fromarray(image.astype(np.uint8))
    img.save(colorfileName)

def _createrbg(image, model, dirpath, perfix, T, intrinsic):
    segement = _render(image, T, intrinsic, model)
    colorfileName = os.path.join(dirpath, perfix + ".png")
    img = Image.fromarray(segement.astype(np.uint8))
    img.save(colorfileName)

kernel/exporter.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import numpy as np
import open3d as o3d
from PIL import Image
import tqdm
import os
from scipy.spatial.transform import Rotation as R

def _pose2Rotation(pose):
    x, y, z = pose[0]
    q0, q1, q2, q3 = pose[1] / np.linalg.norm(pose[1])

    r00 = 2 * (q0 * q0 + q1 * q1) - 1
    r01 = 2 * (q1 * q2 - q0 * q3)
    r02 = 2 * (q1 * q3 + q0 * q2)

    r10 = 2 * (q1 * q2 + q0 * q3)
    r11 = 2 * (q0 * q0 + q2 * q2) - 1
    r12 = 2 * (q2 * q3 - q0 * q1)

    r20 = 2 * (q1 * q3 - q0 * q2)
    r21 = 2 * (q2 * q3 + q0 * q1)
    r22 = 2 * (q0 * q0 + q3 * q3) - 1

    rot_matrix = np.array([[r00, r01, r02],
                            [r10, r11, r12],
                            [r20, r21, r22]])

    move = np.array([x, y, z])
    move.resize((3, 1))
    T = np.hstack((rot_matrix, move))
    T = np.vstack((T, np.array([0, 0, 0, 1])))
    return T


def _rotation2Pose(rotation):
    m = rotation.copy()
    tr = rotation[0, 0] + rotation[1, 1] + rotation[2, 2]
    if tr > 0:
        S = np.sqrt(tr + 1.) * 2
        qw = S/4
        qx = (m[2, 1] - m[1, 2])/S
        qy = (m[0, 2] - m[2, 0])/S
        qz = (m[1, 0] - m[0, 1])/S
    elif (m[0, 0] > m[1, 1]) and (m[0, 0] > m[2, 2]):
        S = np.sqrt(1. + m[0, 0] - m[1, 1] - m[2, 2]) * 2
        qw = (m[2, 1] - m[1, 2])/S
        qx = S/4
        qy = (m[0, 1] + m[1, 0])/S
        qz = (m[0, 2] + m[2, 0])/S
    elif m[1, 1] > m[2, 2]:
        S = np.sqrt(1. + m[1, 1] - m[2, 2] - m[0, 0]) * 2
        qw = (m[0, 2] - m[2, 0])/S
        qx = (m[0, 1] + m[1, 0])/S
        qy = S/4
        qz = (m[1, 2] + m[2, 1])/S
    else:
        S = np.sqrt(1. + m[2, 2] - m[1, 1] - m[0, 0]) * 2
        qw = (-m[0, 1] + m[1, 0])/S
        qx = (m[0, 2] + m[2, 0])/S
        qy = (m[1, 2] + m[2, 1])/S
        qz = S/4
    location = rotation[:3, 3]
    return [[float(location[0]), float(location[1]), float(location[2])], [float(qw), float(qx), float(qy), float(qz)]]

def _render(image, pose, intrinsic, model):
    homoP = np.dot(pose[0:3, :], model.T)
    homoP = homoP / homoP[2]
    pixel = intrinsic.dot(homoP)
    pixel = (np.around(pixel[0:2])).astype(np.int32)
    pixel[0, pixel[0] < 0] = 0
    pixel[1, pixel[1] < 0] = 0
    pixel[0, pixel[0] >=image.shape[1]] = image.shape[1] - 1
    pixel[1, pixel[1] >=image.shape[0]] = image.shape[0] - 1
    ## remove the repeate pixel
    pixel_sort = pixel[:, np.argsort(pixel[0] * 10000 + pixel[1], axis=0)]
    _unique = np.unique(pixel[0] * 10000 + pixel[1])
    pixel_sort_unique = np.vstack((np.floor(_unique / 10000).astype(np.uint32), (_unique % 10000).astype(np.uint32)))
    ## mask
    maskimage = np.zeros_like(image)
    maskimage[pixel_sort_unique[1], pixel_sort_unique[0], :] = 1
    segimage = maskimage * image
    segimage[maskimage == 0] = 1.
    return segimage

def _loadModel(modelPath):
    ## load vertex from .obj
    pointCloud = []
    f = open(modelPath, "r")
    line = f.readline()
    while line:
        if line.startswith("v "):
            line = line.replace("\n", "")
            point = line.split(" ")[1:]
            pointCloud.append([float(point[i]) for i in range(3)] + [1.])
        else:
            pass
        line = f.readline()
    return np.array(pointCloud)

def plane_alignment(filepath, scale, alignT, threshold, n, iteration):
    pcd = o3d.io.read_point_cloud(filepath)
    scaled_xyz = (alignT[:3, :3].dot(np.asarray(pcd.points).T * scale) + alignT[:3, [3]]).T 
    scaled_pcd = o3d.geometry.PointCloud()
    scaled_pcd.points = o3d.utility.Vector3dVector(scaled_xyz)
    plane_model, inliers = scaled_pcd.segment_plane(distance_threshold=threshold,
                                         ransac_n=n,
                                         num_iterations=iteration)
    [a, b, c, d] = plane_model
    plane_cloud = scaled_pcd.select_by_index(inliers)
    plane_xyz = np.asarray(plane_cloud.points)
    plane_center = np.mean(plane_xyz, axis = 0)
    plane_center[2] = -(a * plane_center[0] + b * plane_center[1] + d)/c
    
    return [a, b, c, d], list(plane_center)

def transform_from_plane(plane, plane_center):
    [a, b, c, d] = plane
    ct = c/np.sqrt(a**2 + b**2 + c**2)
    st = np.sqrt(a**2 + b**2)/np.sqrt(a**2 + b**2 + c**2)
    u1 = b/np.sqrt(a**2 + b**2)
    u2 = -a/np.sqrt(a**2 + b**2)
    move = np.array([[1, 0, 0, -plane_center[0]],
                     [0, 1, 0, -plane_center[1]],
                     [0, 0, 1, -plane_center[2]],
                     [0, 0, 0, 1]])
    rotate = np.array([[ct + u1**2 * (1 - ct), u1 * u2*(1 - ct), u2 * st, 0],
                       [u1 * u2 * (1 - ct), ct + u2**2 * (1 - ct), -u1 * st, 0],
                       [-u2 *st, u1 *st, ct, 0],
                       [0, 0, 0, 1]])
    return rotate.dot(move)

def modelICP(scene_vertices, model_vertices):
    scene = o3d.geometry.PointCloud()
    scene.points = o3d.utility.Vector3dVector(scene_vertices)
    model = o3d.geometry.PointCloud()
    model.points = o3d.utility.Vector3dVector(model_vertices)
    trans_init = np.identity(4)
    threshold = 0.02
    reg_p2p = o3d.pipelines.registration.registration_icp(
        model, scene, threshold, trans_init,
        o3d.pipelines.registration.TransformationEstimationPointToPoint())
    return reg_p2p.transformation


def globalRegisteration(scene_vertices, model_vertices):
    scene = o3d.geometry.PointCloud()
    scene.points = o3d.utility.Vector3dVector(scene_vertices)
    model = o3d.geometry.PointCloud()
    model.points = o3d.utility.Vector3dVector(model_vertices)
    _, modelfpfh = _preprocess_point_cloud(model)
    _, scenefpfh = _preprocess_point_cloud(scene)
    return _register_point_cloud_fpfh(model, scene, modelfpfh, scenefpfh)

def _preprocess_point_cloud(pcd, voxel_size = 0.005):
    pcd_down = pcd.voxel_down_sample(voxel_size)
    pcd_down.estimate_normals(
        o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 2.0,
                                             max_nn=30))
    pcd_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
        pcd_down,
        o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 5.0,
                                             max_nn=100))
    return (pcd_down, pcd_fpfh)

def _register_point_cloud_fpfh(source, target, source_fpfh, target_fpfh, distance_threshold = 0.005 * 1.4):
    result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
        source, target, source_fpfh, target_fpfh, True, distance_threshold,
        o3d.pipelines.registration.TransformationEstimationPointToPoint(
            False), 3,
        [
            o3d.pipelines.registration.
            CorrespondenceCheckerBasedOnEdgeLength(0.9),
            o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(
                distance_threshold)
        ],
        o3d.pipelines.registration.RANSACConvergenceCriteria(
            1000000, 0.999))
    if (result.transformation.trace() == 4.0):
        print("fail")
        return np.identity(4)
    else:
        return result.transformation


def depthfilter(depth, scale, depthignore, ZERODEPTH_IGNORE):
    if not ZERODEPTH_IGNORE:
        return np.array(depth) * scale <= depthignore
    else:
        return (np.array(depth) * scale <= depthignore) * (np.array(depth) > 0)






kernel/geometry.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import numpy as np
import numpy.linalg as la

def get_config(camera='uw'):
    config = dict()
    config['im_w'] = 1280
    config['im_h'] = 720
    if camera == 'uw':
        config['cam_intr'] = np.array([ 
            [1066.778, 0, 312.9869],
            [0, 1067.487, 241.3109],
            [0, 0, 1]
        ])
    elif camera == 'umich':
        config['cam_intr'] = np.array([ 
            [904.572, 0., 635.981],
            [0., 905.295, 353.060],
            [0., 0., 1.]
        ])
    else:
        config['cam_intr'] = np.array([
            [1077.836, 0, 323.7872],
            [0, 1078.189, 279.6921],
            [0, 0, 1]
        ])
    config['tsdf_voxel_size'] = 0.0025  # in meter
    config['tsdf_trunc_margin'] = 0.015 # in meter
    config['pcd_voxel_size'] = 0.005  # in meter

    return config

def set_config(resX, resY, fx, fy, cx, cy, tsdf_voxel_size, tsdf_trunc_margin, pcd_voxel_size):
    config = dict()
    config['im_w'] = resX
    config['im_h'] = resY
    config['cam_intr'] = np.array([ 
        [fx, 0., cx],
        [0., fy, cy],
        [0., 0., 1.]
    ])  
    config['tsdf_voxel_size'] = tsdf_voxel_size  # in meter
    config['tsdf_trunc_margin'] = tsdf_trunc_margin # in meter
    config['pcd_voxel_size'] = pcd_voxel_size  # in meter

    return config          

def print_config(cfg):
    for key, value in cfg.items():
        if isinstance(value, np.ndarray):
            print(f"{key}:\n{value}")
        else:
            print(f"{key}: {value}")


if __name__ == '__main__':
    cfg = get_config("test")
    print_config(cfg)

kernel/kf_pycuda/config.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import numpy as np
# import matplotlib.pyplot as plt
# import matplotlib.animation as animation

import pycuda
import pycuda.autoinit
import pycuda.driver as cuda
from pycuda import gpuarray
from pycuda.compiler import SourceModule

source_module = SourceModule(
    """
    #define _X (threadIdx.x + blockDim.x * blockIdx.x)
    #define _Y (threadIdx.y + blockDim.y * blockIdx.y)
    #define _Z (threadIdx.z + blockDim.z * blockIdx.z)
    #define _ColorConst 65536

    __device__
    size_t getVoxelIndex(int x, int y, int z, int *dim)
    {
        return z*dim[0]*dim[1] + y*dim[0] + x;
    }

    __device__
    float serializeColor(float r, float g, float b)
    {
        return b*_ColorConst + g*256 + r;
    }

    __device__
    void deserializeColor(float color, float *r, float *g, float *b)
    {
        *b = floorf(color / _ColorConst);
        *g = floorf((color - (*b)*_ColorConst) / 256);
        *r = color - (*b)*_ColorConst - (*g)*256;
    }

    __device__
    void matmul(float *A, float *B, float *out, int m, int n, int l)
    {
        /*
            A is a matrix of shape m * n
            B is a matrix of shape n * l
            out is a matrix of shape m * l
            All these three matrix are stored in a row-wise 1D array
        */
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < l; ++j) {
                int sum = 0;
                for (int k = 0; k < n; ++k) {
                    sum += A[i*n + k] * B[k*l + j];
                }
                out[i*l + j] = sum;
            }
        }
    }

    __device__
    void inv_T(float *A, float *out)
    {
        /*
           inverse a transformation matrix
           [[R, t]   -->  [[R.T, -R.T @ t]
            [0, 1]]        [0,       1]]
        */
        out[0] = A[0];
        out[1] = A[4];
        out[2] = A[8];
        out[3] = -(out[0]*A[3] + out[1]*A[7] + out[2]*A[11]);
        out[4] = A[1];
        out[5] = A[5];
        out[6] = A[9];
        out[7] = -(out[4]*A[3] + out[5]*A[7] + out[6]*A[11]);
        out[8] = A[2];
        out[9] = A[6];
        out[10] = A[10];
        out[11] = -(out[8]*A[3] + out[9]*A[7] + out[10]*A[11]);
        out[12] = 0;
        out[13] = 0;
        out[14] = 0;
        out[15] = 1;
    }

    __global__
    void integrate(float *tsdf_vol, float *weight_vol, float *color_vol, int *dim, float *vol_origin,
                   float voxel_size, float *cam_intr, float *cam_pose, int im_h, int im_w,
                   float *color_im, float *depth_im, float trunc_margin, float weight)
    {
        const int x = _X;
        const int y = _Y;
        const int z = _Z;

        if (x >= dim[0] || y >= dim[1] || z >= dim[2]) {
            return;
        }

        size_t voxel_idx = getVoxelIndex(x, y, z, dim);

        // voxel grid coordinates to world coordinates
        float pt_x = vol_origin[0] + x * voxel_size;
        float pt_y = vol_origin[1] + y * voxel_size;
        float pt_z = vol_origin[2] + z * voxel_size;

        // world coordinates to camera coordinates
        float tmp_pt_x = pt_x - cam_pose[0*4+3];
        float tmp_pt_y = pt_y - cam_pose[1*4+3];
        float tmp_pt_z = pt_z - cam_pose[2*4+3];
        float cam_pt_x = cam_pose[0*4+0]*tmp_pt_x + cam_pose[1*4+0]*tmp_pt_y + cam_pose[2*4+0]*tmp_pt_z;
        float cam_pt_y = cam_pose[0*4+1]*tmp_pt_x + cam_pose[1*4+1]*tmp_pt_y + cam_pose[2*4+1]*tmp_pt_z;
        float cam_pt_z = cam_pose[0*4+2]*tmp_pt_x + cam_pose[1*4+2]*tmp_pt_y + cam_pose[2*4+2]*tmp_pt_z;

        // camera coordinates to image pixels
        // cam_intr = [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]
        float fx = cam_intr[0];
        float fy = cam_intr[4];
        float cx = cam_intr[2];
        float cy = cam_intr[5];

        int pixel_x = __float2int_rn(fx*(cam_pt_x/cam_pt_z) + cx);
        int pixel_y = __float2int_rn(fy*(cam_pt_y/cam_pt_z) + cy);
        int pixel_idx = pixel_y*im_w + pixel_x;

        // skip if outside view frustum
        if (pixel_x < 0 || pixel_x >= im_w || pixel_y < 0 || pixel_y >= im_h) return;
        if (cam_pt_z < 0) return;

        float depth_value = depth_im[pixel_idx];
        if (depth_value == 0) return;

        // skip voxels that are beyond truncated value
        float depth_diff = depth_value - cam_pt_z;
        if (depth_diff < -trunc_margin) return;

        // update weights
        float w_old = weight_vol[voxel_idx];
        float w_new = w_old + weight;
        weight_vol[voxel_idx] = w_new;

        // integrate TSDF
        float dist = fminf(1.0f, depth_diff / trunc_margin);
        tsdf_vol[voxel_idx] = (tsdf_vol[voxel_idx]*w_old + weight*dist) / w_new;

        // integrate color
        float old_color = color_vol[voxel_idx];
        float old_r, old_g, old_b;
        deserializeColor(old_color, &old_r, &old_g, &old_b);

        float new_color = color_im[pixel_idx];
        float new_r, new_g, new_b;
        deserializeColor(new_color, &new_r, &new_g, &new_b);

        new_b = fminf(roundf((old_b*w_old + weight*new_b) / w_new), 255.0f);
        new_g = fminf(roundf((old_g*w_old + weight*new_g) / w_new), 255.0f);
        new_r = fminf(roundf((old_r*w_old + weight*new_r) / w_new), 255.0f);
        color_vol[voxel_idx] = serializeColor(new_r, new_g, new_b);
    }

    __device__
    float getMinTime(float3 &min_range, float3 &max_range, float3 &cam_pos, float3 &ray)
    {
        float txmin = ((ray.x > 0 ? min_range.x : max_range.x) - cam_pos.x) / ray.x;
        float tymin = ((ray.y > 0 ? min_range.y : max_range.y) - cam_pos.y) / ray.y;
        float tzmin = ((ray.z > 0 ? min_range.z : max_range.z) - cam_pos.z) / ray.z;
        return fmaxf(fmaxf(txmin, tymin), tzmin);
    }

    __device__
    float getMaxTime(float3 &min_range, float3 &max_range, float3 &cam_pos, float3 &ray)
    {
        float txmax = ((ray.x > 0 ? max_range.x : min_range.x) - cam_pos.x) / ray.x;
        float tymax = ((ray.y > 0 ? max_range.y : min_range.y) - cam_pos.y) / ray.y;
        float tzmax = ((ray.z > 0 ? max_range.z : min_range.z) - cam_pos.z) / ray.z;
        return fminf(fminf(txmax, tymax), tzmax);
    }

    __global__
    void batchInlierMetric(int im_h, int im_w, int batch_size, float *batch_depth_im, float *depth_im,
                           float *weights, float inlier_thresh)
    {
        const int x = _X;
        if (x >= batch_size) return;

        int num_pixels = im_h * im_w;
        size_t start_idx = x * num_pixels;

        float num_inliers = 0;
        float num_total = 0;
        float num_obs = 0;
        for (int i = 0; i < num_pixels; ++i) {
            if (depth_im[i] != 0) num_obs += 1;

            float depth_render = batch_depth_im[i + start_idx];
            if (depth_render != 0) num_total += 1;
            else continue;

            if (fabsf(depth_render - depth_im[i]) < inlier_thresh) {
                num_inliers += 1;
            }
        }
        float ratio = 0.8;
        if (num_total == 0 || num_obs == 0) {
            weights[x] = 0;
        } else {
            // weights[x] = num_inliers;
            // weights[x] = 2 * num_inliers  / (num_total + num_obs);
            // weights[x] = num_inliers * num_inliers / (num_total * num_obs);
            weights[x] = ratio * (num_inliers / num_total) + (1 - ratio) * (num_inliers / num_obs);
        }
        return;
    }

    __global__
    void batchRayCasting(float *tsdf_vol, float *color_vol, float *weight_vol, int *dim, float* vol_origin, float voxel_size,
                         float *cam_intr, float *cam_pose, float *inv_cam_pose,
                         int start_row, int start_col, int im_h, int im_w,
                         char *batch_color_im, float *batch_depth_im, int batch_size)
    {
        const int x = _X;
        const int y = _Y;
        const int z = _Z;

        if (x >= im_w || y >= im_h || z > batch_size) {
            return;
        }

        int cam_pose_const = 16;  // 4x4 matrix
        cam_pose = cam_pose + z*cam_pose_const;
        inv_cam_pose = inv_cam_pose + z*cam_pose_const;

        int im_size_const = im_h * im_w;
        size_t pixel_idx = y*im_w + x + z*im_size_const;
        batch_depth_im[pixel_idx] = 0;

        size_t r_pixel_idx = y*im_w + x + 3*z*im_size_const;
        size_t g_pixel_idx = r_pixel_idx + im_size_const;
        size_t b_pixel_idx = g_pixel_idx + im_size_const;
        batch_color_im[r_pixel_idx] = 0;
        batch_color_im[g_pixel_idx] = 0;
        batch_color_im[b_pixel_idx] = 0;

        float3 min_range = make_float3(vol_origin[0], vol_origin[1], vol_origin[2]);
        float3 max_range = make_float3(vol_origin[0] + voxel_size * dim[0],
                                       vol_origin[1] + voxel_size * dim[1],
                                       vol_origin[2] + voxel_size * dim[2]);

        // camera coordinates to image pixels
        // cam_intr = [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]
        float fx = cam_intr[0];
        float fy = cam_intr[4];
        float cx = cam_intr[2];
        float cy = cam_intr[5];

        float3 pixel_pos = make_float3((start_col + x - cx) / fx, (start_row + y - cy) / fy, 1);
        float3 ray = make_float3(
            cam_pose[0]*pixel_pos.x + cam_pose[1]*pixel_pos.y + cam_pose[2]*pixel_pos.z,
            cam_pose[4]*pixel_pos.x + cam_pose[5]*pixel_pos.y + cam_pose[6]*pixel_pos.z,
            cam_pose[8]*pixel_pos.x + cam_pose[9]*pixel_pos.y + cam_pose[10]*pixel_pos.z
        );
        float normalize_const = sqrtf(ray.x*ray.x + ray.y*ray.y + ray.z*ray.z);
        ray.x = voxel_size * ray.x / normalize_const;
        ray.y = voxel_size * ray.y / normalize_const;
        ray.z = voxel_size * ray.z / normalize_const;

        float3 cam_pos = make_float3(cam_pose[3], cam_pose[7], cam_pose[11]);
        int t_min = __float2int_rd(fmaxf(getMinTime(min_range, max_range, cam_pos, ray), 0));
        int t_max = __float2int_ru(getMaxTime(min_range, max_range, cam_pos, ray));
        if (t_min >= t_max) {
            return;
        }

        float dist_x = cam_pos.x - vol_origin[0];
        float dist_y = cam_pos.y - vol_origin[1];
        float dist_z = cam_pos.z - vol_origin[2];
        
        int3 prev_voxel = make_int3(
            __float2int_rn((dist_x + ray.x * t_min) / voxel_size),
            __float2int_rn((dist_y + ray.y * t_min) / voxel_size),
            __float2int_rn((dist_z + ray.z * t_min) / voxel_size)
        );

        bool ray_inside_box = false;
        int t_step = 5;  // 5 is relatively safe. 10 is much faster but may miss the crossing point.
        for (int t = t_min + 1; t < t_max - 1; t += t_step) {
            int3 curr_voxel = make_int3(
                __float2int_rn((dist_x + ray.x * t) / voxel_size),
                __float2int_rn((dist_y + ray.y * t) / voxel_size),
                __float2int_rn((dist_z + ray.z * t) / voxel_size)
            );

            if (!ray_inside_box) {
                if (prev_voxel.x >= 0 && prev_voxel.x < dim[0] && 
                    prev_voxel.y >= 0 && prev_voxel.y < dim[1] &&
                    prev_voxel.z >= 0 && prev_voxel.z < dim[2])
                {
                    ray_inside_box = true;
                } else {
                    prev_voxel = curr_voxel;
                    continue;
                }
            }

            // Box has a convex shape. A ray cannot re-enter a box after leaving it.
            if (ray_inside_box &&
                (curr_voxel.x < 0 || curr_voxel.x >= dim[0] ||
                 curr_voxel.y < 0 || curr_voxel.y >= dim[1] ||
                 curr_voxel.z < 0 || curr_voxel.z >= dim[2]))
            {
                return;
            }

            size_t prev_voxel_idx = getVoxelIndex(prev_voxel.x, prev_voxel.y, prev_voxel.z, dim);
            size_t curr_voxel_idx = getVoxelIndex(curr_voxel.x, curr_voxel.y, curr_voxel.z, dim);

            float prev_tsdf = tsdf_vol[prev_voxel_idx];
            float curr_tsdf = tsdf_vol[curr_voxel_idx];

            // zero crossing
            if (prev_tsdf > 0 && curr_tsdf < 0) {

                // find the exact zero crossing time
                for (int i = max(t_min + 1, t - t_step); i <= t; ++i) {
                    int3 voxel = make_int3(
                        __float2int_rn((dist_x + ray.x * i) / voxel_size),
                        __float2int_rn((dist_y + ray.y * i) / voxel_size),
                        __float2int_rn((dist_z + ray.z * i) / voxel_size)
                    );
                    size_t voxel_index = getVoxelIndex(voxel.x, voxel.y, voxel.z, dim);

                    if (tsdf_vol[voxel_index] < 0) {
                        if (tsdf_vol[voxel_index] > -1) {
                            float3 p = make_float3(
                                vol_origin[0] + voxel.x * voxel_size,
                                vol_origin[1] + voxel.y * voxel_size,
                                vol_origin[2] + voxel.z * voxel_size
                            );

                            float depth = inv_cam_pose[8]*p.x + inv_cam_pose[9]*p.y + inv_cam_pose[10]*p.z + inv_cam_pose[11];
                            batch_depth_im[pixel_idx] = depth;

                            float r, g, b;
                            deserializeColor(color_vol[voxel_index], &r, &g, &b);
                            batch_color_im[r_pixel_idx] = (char)r;
                            batch_color_im[g_pixel_idx] = (char)g;
                            batch_color_im[b_pixel_idx] = (char)b;
                            return;
                        } else {
                            break;  // false surface
                        }
                    }
                }
            }
            prev_voxel = curr_voxel;
        }
    }

    __global__
    void rayCasting(float *tsdf_vol, float *color_vol, float *weight_vol, int *dim, float* vol_origin, float voxel_size,
                    float *cam_intr, float *cam_pose, float *inv_cam_pose,
                    int start_row, int start_col, int im_h, int im_w,
                    char *color_im, float *depth_im)
    {
        const int x = _X;
        const int y = _Y;

        if (x >= im_w || y >= im_h) {
            return;
        }
        
        size_t pixel_idx = y*im_w + x;
        depth_im[pixel_idx] = 0;

        size_t r_pixel_idx = pixel_idx;
        size_t g_pixel_idx = im_w * im_h + pixel_idx;
        size_t b_pixel_idx = 2 * im_w * im_h + pixel_idx;
        color_im[r_pixel_idx] = 0;
        color_im[g_pixel_idx] = 0;
        color_im[b_pixel_idx] = 0;

        float3 min_range = make_float3(vol_origin[0], vol_origin[1], vol_origin[2]);
        float3 max_range = make_float3(vol_origin[0] + voxel_size * dim[0],
                                       vol_origin[1] + voxel_size * dim[1],
                                       vol_origin[2] + voxel_size * dim[2]);

        // camera coordinates to image pixels
        // cam_intr = [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]
        float fx = cam_intr[0];
        float fy = cam_intr[4];
        float cx = cam_intr[2];
        float cy = cam_intr[5];

        float3 pixel_pos = make_float3((start_col + x - cx) / fx, (start_row + y - cy) / fy, 1);
        float3 ray = make_float3(
            cam_pose[0]*pixel_pos.x + cam_pose[1]*pixel_pos.y + cam_pose[2]*pixel_pos.z,
            cam_pose[4]*pixel_pos.x + cam_pose[5]*pixel_pos.y + cam_pose[6]*pixel_pos.z,
            cam_pose[8]*pixel_pos.x + cam_pose[9]*pixel_pos.y + cam_pose[10]*pixel_pos.z
        );
        float normalize_const = sqrtf(ray.x*ray.x + ray.y*ray.y + ray.z*ray.z);

        // ray contains information of both direction and magnitute
        ray.x = voxel_size * ray.x / normalize_const;
        ray.y = voxel_size * ray.y / normalize_const;
        ray.z = voxel_size * ray.z / normalize_const;

        float3 cam_pos = make_float3(cam_pose[3], cam_pose[7], cam_pose[11]);
        int t_min = __float2int_ru(fmaxf(getMinTime(min_range, max_range, cam_pos, ray), 0));
        int t_max = __float2int_rd(getMaxTime(min_range, max_range, cam_pos, ray));
        if (t_min >= t_max) {
            return;
        }

        float dist_x = cam_pos.x - vol_origin[0];
        float dist_y = cam_pos.y - vol_origin[1];
        float dist_z = cam_pos.z - vol_origin[2];

        int3 prev_voxel = make_int3(
            __float2int_rn((dist_x + ray.x * t_min) / voxel_size),
            __float2int_rn((dist_y + ray.y * t_min) / voxel_size),
            __float2int_rn((dist_z + ray.z * t_min) / voxel_size)
        );

        bool ray_inside_box = false;
        for (int t = t_min + 1; t < t_max; t += 1) {
            int3 curr_voxel = make_int3(
                __float2int_rn((dist_x + ray.x * t) / voxel_size),
                __float2int_rn((dist_y + ray.y * t) / voxel_size),
                __float2int_rn((dist_z + ray.z * t) / voxel_size)
            );

            if (!ray_inside_box) {
                if (prev_voxel.x >= 0 && prev_voxel.x < dim[0] && 
                    prev_voxel.y >= 0 && prev_voxel.y < dim[1] &&
                    prev_voxel.z >= 0 && prev_voxel.z < dim[2])
                {
                    ray_inside_box = true;
                } else {
                    prev_voxel = curr_voxel;
                    continue;
                }
            }

            // Box has a convex shape. A ray cannot re-enter a box after leaving it.
            if (ray_inside_box &&
                (curr_voxel.x < 0 || curr_voxel.x >= dim[0] ||
                 curr_voxel.y < 0 || curr_voxel.y >= dim[1] ||
                 curr_voxel.z < 0 || curr_voxel.z >= dim[2]))
            {
                return;
            }

            size_t prev_voxel_idx = getVoxelIndex(prev_voxel.x, prev_voxel.y, prev_voxel.z, dim);
            size_t curr_voxel_idx = getVoxelIndex(curr_voxel.x, curr_voxel.y, curr_voxel.z, dim);

            float prev_tsdf = tsdf_vol[prev_voxel_idx];
            float curr_tsdf = tsdf_vol[curr_voxel_idx];

            // zero crossing from front
            if (prev_tsdf > 0 && curr_tsdf < 0 && curr_tsdf > -1) {
                float3 p = make_float3(
                    vol_origin[0] + curr_voxel.x * voxel_size,
                    vol_origin[1] + curr_voxel.y * voxel_size,
                    vol_origin[2] + curr_voxel.z * voxel_size
                );
                float depth = inv_cam_pose[8]*p.x + inv_cam_pose[9]*p.y + inv_cam_pose[10]*p.z + inv_cam_pose[11];
                depth_im[pixel_idx] = depth;

                float r, g, b;
                deserializeColor(color_vol[curr_voxel_idx], &r, &g, &b);
                color_im[r_pixel_idx] = (char)r;
                color_im[g_pixel_idx] = (char)g;
                color_im[b_pixel_idx] = (char)b;
                return;
            }

            prev_voxel = curr_voxel;
        }
    }

    __device__
    bool checkInBox(float3 p, float *origin, float *vectors, float *dims)
    {
        /*
            bbox is represented using an origin and three normalized vectors from this origin, and dims
            origin: [x, y, z]
            vectors: [x1, y1, z1, x2, y2, z2, x3, y3, z3]
            dims: [dim_x, dim_y, dim_z]
        */

        float3 v = make_float3(p.x - origin[0], p.y - origin[1], p.z - origin[2]);

        // calculate dot products of v with v1, v2, v3
        float prod1 = v.x*vectors[0] + v.y*vectors[1] + v.z*vectors[2];
        float prod2 = v.x*vectors[3] + v.y*vectors[4] + v.z*vectors[5];
        float prod3 = v.x*vectors[6] + v.y*vectors[7] + v.z*vectors[8];

        // point is outside the oriented bbox
        if (prod1 < 0 || prod1 > dims[0]*dims[0] || prod2 < 0 || prod2 > dims[1]*dims[1] || prod3 < 0 || prod3 > dims[2]*dims[2]) {
            return false;
        }
        return true;
    }

    __global__
    void renderBoxes(float *depth_im, int im_h, int im_w, float *cam_intr,
                     int num_boxes, float *cam_poses, float *box_dims)
    {
        const int x = _X;
        const int y = _Y;

        if (x >= im_w || y >= im_h) {
            return;
        }

        size_t pixel_idx = y*im_w + x;
        depth_im[pixel_idx] = 0;

        for (int i = 0; i < num_boxes; ++i) {
            float dim_x = box_dims[3*i];
            float dim_y = box_dims[3*i + 1];
            float dim_z = box_dims[3*i + 2];
            float3 min_range = make_float3(-dim_x / 2, -dim_y / 2, -dim_z / 2);
            float3 max_range = make_float3(dim_x / 2, dim_y / 2, dim_z / 2);

            // camera coordinates to image pixels
            // cam_intr = [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]
            float fx = cam_intr[0];
            float fy = cam_intr[4];
            float cx = cam_intr[2];
            float cy = cam_intr[5];

            float *cam_pose = cam_poses + 16*i;
            float3 pixel_pos = make_float3((x - cx) / fx, (y - cy) / fy, 1);
            float3 cam_pos = make_float3(cam_pose[3], cam_pose[7], cam_pose[11]);
            float3 ray = make_float3(
                cam_pose[0]*pixel_pos.x + cam_pose[1]*pixel_pos.y + cam_pose[2]*pixel_pos.z,
                cam_pose[4]*pixel_pos.x + cam_pose[5]*pixel_pos.y + cam_pose[6]*pixel_pos.z,
                cam_pose[8]*pixel_pos.x + cam_pose[9]*pixel_pos.y + cam_pose[10]*pixel_pos.z
            );

            float t_min = fmaxf(getMinTime(min_range, max_range, cam_pos, ray), 0);
            float t_max = getMaxTime(min_range, max_range, cam_pos, ray);
            if (t_min >= t_max) {
                continue;
            }

            if (depth_im[pixel_idx] == 0 || depth_im[pixel_idx] > t_min) {
                depth_im[pixel_idx] = t_min;
            }
        }
    }

    __global__
    void filterOcclusion(float *depth_im, float *occ_im, int im_h, int im_w)
    {
        const int x = _X;
        const int y = _Y;

        if (x >= im_w || y >= im_h) {
            return;
        }

        size_t pixel_idx = y*im_w + x;
        if (occ_im[pixel_idx] != 0 && depth_im[pixel_idx] > occ_im[pixel_idx]) {
            depth_im[pixel_idx] = 0;
        }
    }
    
    __global__
    void filterPointsInBoxes(float *depth_im, int im_w, int im_h, float *cam_intr,
                             int num_boxes, float *origins, float *vectors, float *dims, int *inverse)
    {
        /*
            cam_intr = [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]

            bbox is represented using an origin and three normalized vectors from this origin, and dims
            origins: [[x1, y1, z1], ...]
            vectors: [[x1, y1, z1, x2, y2, z2, x3, y3, z3], ...]
            dims: [[dim_x, dim_y, dim_z], ...]
            inverse: a list of in integer, 0 means keep pts in box, 1 means filter pts in box
        */

        const int x = _X;
        const int y = _Y;

        if (x >= im_w || y >= im_h) {
            return;
        }

        size_t pixel_idx = y*im_w + x;
        float depth = depth_im[pixel_idx];
        if (depth == 0) return;

        float fx = cam_intr[0];
        float fy = cam_intr[4];
        float cx = cam_intr[2];
        float cy = cam_intr[5];

        float3 p = make_float3((x - cx)*depth/fx, (y - cy)*depth/fy, depth);

        for (int i = 0; i < num_boxes; ++i) {
            bool inBox = checkInBox(p, origins + 3*i, vectors + 9*i, dims + 3*i);
            if ((inBox && inverse[i] == 1) || (!inBox && inverse[i] == 0)) {
                depth_im[pixel_idx] = 0;
                return;
            }
        }
        return;
    }

    __global__
    void colorDistanceL2(char *rendered_color, char *obs_color, float* rendered_depth, float *diff, int im_w, int im_h)
    {
        const int x = _X;
        const int y = _Y;

        if (x >= im_w || y >= im_h) {
            return;
        }

        size_t pixel_idx = y*im_w + x;
        diff[pixel_idx] = 0;

        if (rendered_depth[pixel_idx] == 0) {
            return;
        }

        size_t r_pixel_idx = pixel_idx;
        size_t g_pixel_idx = im_w * im_h + pixel_idx;
        size_t b_pixel_idx = 2 * im_w * im_h + pixel_idx;

        float r_diff = (float)obs_color[r_pixel_idx] - (float)rendered_color[r_pixel_idx];
        float g_diff = (float)obs_color[g_pixel_idx] - (float)rendered_color[g_pixel_idx];
        float b_diff = (float)obs_color[b_pixel_idx] - (float)rendered_color[b_pixel_idx];
        diff[pixel_idx] = sqrt(r_diff*r_diff + g_diff*g_diff + b_diff*b_diff);
    }

    __global__
    void indexTest(float *output, int *dim)
    {
        const int x = _X;
        const int y = _Y;
        const int z = _Z;

        if (x >= dim[0] || y >= dim[1] || z >= dim[2]) {
            return;
        }

        size_t global_idx = getVoxelIndex(x, y, z, dim);

        output[global_idx] = z * dim[0] * dim[1] + y * dim[0] + x + 1;
        // output[global_idx] = x*x + y*y;
    }

    """
)

kernel/kf_pycuda/cuda_kernels.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import os
import pickle
import copy
import json
import open3d as o3d
import numpy as np
# import cv2
# import matplotlib.pyplot as plt
import scipy.linalg as la
from scipy.io import loadmat
from tqdm import tqdm
from tqdm.contrib import tzip

import pycuda.autoinit
import pycuda.driver as cuda
from pycuda import gpuarray

from kernel.kf_pycuda.tsdf_lib import TSDFVolume
from kernel.kf_pycuda.cuda_kernels import source_module
from kernel.kf_pycuda import utils
from kernel.kf_pycuda.config import get_config, print_config

from kernel.geometry import _rotation2Pose

# open3d version should be at least 0.11.0
reg = o3d.pipelines.registration

class KinectFusion:

    def __init__(self, cfg: dict = None):
        self.cfg = cfg
        self.tsdf_volume = None
        self.init_transformation = None
        self.transformation = None
        self.prev_pcd = None
        self.cam_poses = []  
    
    def initialize_tsdf_volume(self, color_im, depth_im, pose=None, visualize=False):
        if pose is not None:
            pcd = utils.create_pcd(depth_im, self.cfg['cam_intr'], color_im, depth_trunc = self.cfg['tsdf_trunc_margin']*10000)
            
            transformed_pcd = copy.deepcopy(pcd).transform(pose)
            transformed_pts = np.asarray(transformed_pcd.points)
            
            vol_bnds = np.zeros((3, 2), dtype=np.float32)
            vol_bnds[:, 0] = np.mean(transformed_pts, axis = 0) - np.abs(transformed_pts - np.mean(transformed_pts, axis = 0)).max(0)*1.2
            vol_bnds[:, 1] = np.mean(transformed_pts, axis = 0) + np.abs(transformed_pts - np.mean(transformed_pts, axis = 0)).max(0)*1.2
            self.init_transformation = la.inv(pose).copy()
            self.transformation = la.inv(pose).copy()
            self.tsdf_volume = TSDFVolume(vol_bnds=vol_bnds,
                                        voxel_size=self.cfg['tsdf_voxel_size'],
                                        trunc_margin=self.cfg['tsdf_trunc_margin'])
            self.tsdf_volume.integrate(color_im, depth_im, self.cfg['cam_intr'], pose)
            self.prev_pcd = pcd
            self.cam_poses.append(pose)
        else:
            pcd = utils.create_pcd(depth_im, self.cfg['cam_intr'], color_im)
            plane_frame, inlier_ratio = utils.plane_detection_ransac(pcd, inlier_thresh=0.005, early_stop_thresh=0.6,
                                                                    visualize=visualize)
            
            cam_pose = la.inv(plane_frame)
            transformed_pcd = copy.deepcopy(pcd).transform(la.inv(plane_frame))
            transformed_pts = np.asarray(transformed_pcd.points)

            vol_bnds = np.zeros((3, 2), dtype=np.float32)
            vol_bnds[:, 0] = np.mean(transformed_pts, axis = 0) - np.abs(transformed_pts - np.mean(transformed_pts, axis = 0)).max(0)*1.2
            vol_bnds[:, 1] = np.mean(transformed_pts, axis = 0) + np.abs(transformed_pts - np.mean(transformed_pts, axis = 0)).max(0)*1.2

            if visualize:
                vol_box = o3d.geometry.OrientedBoundingBox()
                vol_box.center = vol_bnds.mean(1)
                vol_box.extent = vol_bnds[:, 1] - vol_bnds[:, 0]
                o3d.visualization.draw_geometries([vol_box, transformed_pcd])

            self.init_transformation = plane_frame.copy()
            self.transformation = plane_frame.copy()
            self.tsdf_volume = TSDFVolume(vol_bnds=vol_bnds,
                                        voxel_size=self.cfg['tsdf_voxel_size'],
                                        trunc_margin=self.cfg['tsdf_trunc_margin'])
            self.tsdf_volume.integrate(color_im, depth_im, self.cfg['cam_intr'], cam_pose)
            self.prev_pcd = pcd
            self.cam_poses.append(cam_pose)            
    

    @staticmethod
    def multiscale_icp(src: o3d.geometry.PointCloud,
                       tgt: o3d.geometry.PointCloud,
                       voxel_size_list: list, 
                       max_iter_list: list,
                       init: np.ndarray = np.eye(4),
                       inverse: bool = False):

        if len(src.points) > len(tgt.points):
            return KinectFusion.multiscale_icp(tgt, src, voxel_size_list, max_iter_list,
                                               init=la.inv(init), inverse=True)

        transformation = init
        result_icp = None
        for i, (voxel_size, max_iter) in enumerate(zip(voxel_size_list, max_iter_list)):
            src_down = src.voxel_down_sample(voxel_size)
            tgt_down = tgt.voxel_down_sample(voxel_size)

            src_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size*3, max_nn=30))
            tgt_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size*3, max_nn=30))

            estimation_method = reg.TransformationEstimationPointToPlane()
            result_icp = reg.registration_icp(
                src_down, tgt_down, max_correspondence_distance=voxel_size*3,
                init=transformation,
                estimation_method=estimation_method,
                criteria=reg.ICPConvergenceCriteria(max_iteration=max_iter))
            transformation = result_icp.transformation

        if inverse and result_icp is not None:
            result_icp.transformation = la.inv(result_icp.transformation)

        return result_icp

    def update_pose_using_icp(self, depth_im):
        curr_pcd = utils.create_pcd(depth_im, self.cfg['cam_intr'])

        cam_pose = la.inv(self.transformation)
        rendered_depth, _ = self.tsdf_volume.ray_casting(self.cfg['im_w'], self.cfg['im_h'], self.cfg['cam_intr'], 
                                                         cam_pose, to_host=True)
        rendered_pcd = utils.create_pcd(rendered_depth, self.cfg['cam_intr'])

        result_icp = self.multiscale_icp(rendered_pcd,
                                         curr_pcd,
                                         voxel_size_list=[0.025, 0.01],
                                         max_iter_list=[5, 10])
        if result_icp is None:
            return False

        self.transformation = result_icp.transformation @ self.transformation
        self.prev_observation = curr_pcd
        return True

    
    
    def update(self, color_im, depth_im, pose = None):
        if pose is not None:
            ## update with pose
            if self.tsdf_volume is None:
                self.initialize_tsdf_volume(color_im, depth_im, pose, visualize=False)
                return True

            self.tsdf_volume.integrate(color_im, depth_im, self.cfg['cam_intr'], pose, weight=1)
        else:
            if self.tsdf_volume is None:
                self.initialize_tsdf_volume(color_im, depth_im, visualize=False)
                return True

            success = self.update_pose_using_icp(depth_im)
            if success:
                cam_pose = la.inv(self.transformation)
                self.cam_poses.append(cam_pose)
                self.tsdf_volume.integrate(color_im, depth_im, self.cfg['cam_intr'], cam_pose, weight=1)
            else:
                self.cam_poses.append(np.eye(4))            
     

    def interpolation_update(self, color_im, depth_im, accurate_pose = None):
        if accurate_pose is None:
            self.update(color_im, depth_im)
        else:
            if self.tsdf_volume is None:
                assert(accurate_pose is not None)
                self.initialize_tsdf_volume_withaccuratepose(color_im, depth_im, accurate_pose, None)
                return True
            else:
                cam_pose = accurate_pose
                curr_pcd = utils.create_pcd(depth_im, self.cfg['cam_intr'])
                self.transformation = la.inv(cam_pose)
                self.prev_observation = curr_pcd                
                self.cam_poses.append(cam_pose)
                self.tsdf_volume.integrate(color_im, depth_im, self.cfg['cam_intr'], cam_pose, weight=1)            
    

    def initialize_tsdf_volume_withaccuratepose(self, color_im, depth_im, accurate_pose = None, visualize=False):
        pcd = utils.create_pcd(depth_im, self.cfg['cam_intr'], color_im)
        
        cam_pose = accurate_pose
        transformed_pcd = copy.deepcopy(pcd).transform(cam_pose)
        transformed_pts = np.asarray(transformed_pcd.points)

        vol_bnds = np.zeros((3, 2), dtype=np.float32)
        vol_bnds[:, 0] = transformed_pts.min(0)
        vol_bnds[:, 1] = transformed_pts.max(0)


        if visualize:
            vol_box = o3d.geometry.OrientedBoundingBox()
            vol_box.center = vol_bnds.mean(1)
            vol_box.extent = vol_bnds[:, 1] - vol_bnds[:, 0]
            o3d.visualization.draw_geometries([vol_box, transformed_pcd])

        self.init_transformation = la.inv(cam_pose.copy())
        self.transformation = la.inv(cam_pose.copy())
        self.tsdf_volume = TSDFVolume(vol_bnds=vol_bnds,
                                      voxel_size=self.cfg['tsdf_voxel_size'],
                                      trunc_margin=self.cfg['tsdf_trunc_margin'])
        self.tsdf_volume.integrate(color_im, depth_im, self.cfg['cam_intr'], cam_pose)
        self.prev_pcd = pcd
        self.cam_poses.append(cam_pose)

    def F2FPoseEstimation(self, color_im, depth_im, accurate_pose = None):
        if accurate_pose is None:
            curr_pcd = utils.create_pcd(depth_im, self.cfg['cam_intr'])
            # #------------------------------ frame to frame ICP (open loop) ------------------------------
            result_icp = self.multiscale_icp(self.prev_pcd, curr_pcd,
                                             voxel_size_list=[0.025, 0.01, 0.005],
                                             max_iter_list=[10, 10, 10], init=np.eye(4))
            if result_icp is not None:
                self.transformation = result_icp.transformation @ self.transformation
            self.prev_pcd = curr_pcd
            cam_pose = la.inv(self.transformation)
            return cam_pose
        else:
            cam_pose = accurate_pose
            curr_pcd = utils.create_pcd(depth_im, self.cfg['cam_intr'])
            self.transformation = la.inv(cam_pose)
            self.prev_pcd = curr_pcd   
            return True


    def save(self, output_folder, prefix_list):

        cam_poses = np.stack(self.cam_poses)
        camout = open(os.path.join(output_folder, "campose.txt"), "w+")
        camout.writelines("# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n")
        camout.writelines("\n")
        for index, rot, prefix in zip(range(1, cam_poses.shape[0] + 1), cam_poses, prefix_list):
            pose = _rotation2Pose(rot)
            line = str(index) + " " + str(pose[1][0]) + " " + str(pose[1][1])\
                    + " " + str(pose[1][2]) + " " + str(pose[1][3])\
                    + " " + str(pose[0][0]) + " " + str(pose[0][1]) + " " + str(pose[0][2])\
                    + " " + "1" + " " + prefix + ".png\n"
            camout.writelines(line)
        camout.close()
        surface = self.tsdf_volume.get_surface_cloud_marching_cubes(voxel_size=0.005)
        o3d.io.write_point_cloud(os.path.join(output_folder, 'fused.ply'), surface)
        print(f"Results have been saved to {output_folder}.")


kernel/kf_pycuda/kinect_fusion.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import open3d as o3d
import numpy as np
import scipy.linalg as la
import time
import os
import json
from skimage import measure

import pycuda.autoinit
import pycuda.driver as cuda
from pycuda import gpuarray, cumath

from kernel.kf_pycuda.cuda_kernels import source_module


class TSDFVolume:

    def __init__(self, vol_bnds, voxel_size, trunc_margin=0.015):
        """
        Args:
            vol_bnds (ndarray): An ndarray of shape (3, 2). Specifies the xyz bounds (min/max) in meters.
            voxel_size (float): The volume discretization in meters.
        """
        vol_bnds = np.asarray(vol_bnds)
        assert vol_bnds.shape == (3, 2), "[!] `vol_bnds` should be of shape (3, 2)."

        # Define voxel volume parameters
        self._vol_bnds = vol_bnds
        self._voxel_size = voxel_size
        self._trunc_margin = trunc_margin
        self._color_const = np.float32(256 * 256)

        # Adjust volume bounds and ensure C-order contiguous
        self._vol_dim = np.ceil((self._vol_bnds[:,1] - self._vol_bnds[:,0])/self._voxel_size).copy(order='C').astype(int)
        self._vol_bnds[:,1] = self._vol_bnds[:,0] + self._vol_dim*self._voxel_size
        self._vol_origin = self._vol_bnds[:,0].copy(order='C').astype(np.float32)

        # print("Voxel volume size: {} x {} x {} - # points: {:,}".format(
        #     self._vol_dim[0], self._vol_dim[1], self._vol_dim[2],
        #     self._vol_dim[0]*self._vol_dim[1]*self._vol_dim[2])
        # )

        # Copy voxel volumes to GPU
        self.block_x, self.block_y, self.block_z = 8, 8, 16  # block_x * block_y * block_z must equal to 1024

        x_dim, y_dim, z_dim = int(self._vol_dim[0]), int(self._vol_dim[1]), int(self._vol_dim[2])
        xyz = x_dim * y_dim * z_dim
        self.grid_x = int(np.ceil(x_dim / self.block_x))
        self.grid_y = int(np.ceil(y_dim / self.block_y))
        self.grid_z = int(np.ceil(z_dim / self.block_z))

        # initialize tsdf values to be -1
        self._tsdf_vol_gpu = gpuarray.zeros(shape=(xyz), dtype=np.float32) - 1
        self._weight_vol_gpu = gpuarray.zeros(shape=(xyz), dtype=np.float32)
        self._color_vol_gpu= gpuarray.zeros(shape=(xyz), dtype=np.float32)

        # integrate function using PyCuda
        self._cuda_integrate = source_module.get_function("integrate")
        self._cuda_ray_casting = source_module.get_function("rayCasting")
        self._cuda_batch_ray_casting = source_module.get_function("batchRayCasting")

    def integrate(self, color_im, depth_im, cam_intr, cam_pose, weight=1.0):
        """ Integrate an RGB-D frame into the TSDF volume.

        Args:
            color_im (np.ndarray): input RGB image of shape (H, W, 3)
            depth_im (np.ndarray): input depth image of shape (H, W)
            cam_intr (np.ndarray): Camera intrinsics matrix of shape (3, 3)
            cam_pose (np.ndarray): Camera pose of shape (4, 4)
            weight (float, optional): weight to be assigned for the current observation. Defaults to 1.0.
        """
        im_h, im_w = depth_im.shape

        # color image is always from host
        color_im = color_im.astype(np.float32)
        color_im = np.floor(color_im[...,2]*self._color_const + color_im[...,1]*256 + color_im[...,0])

        if isinstance(depth_im, np.ndarray):
            depth_im_gpu = gpuarray.to_gpu(depth_im.astype(np.float32))

        self._cuda_integrate(
            self._tsdf_vol_gpu,
            self._weight_vol_gpu,
            self._color_vol_gpu,
            cuda.InOut(self._vol_dim.astype(np.int32)), 
            cuda.InOut(self._vol_origin.astype(np.float32)),
            np.float32(self._voxel_size),
            cuda.InOut(cam_intr.reshape(-1).astype(np.float32)),
            cuda.InOut(cam_pose.reshape(-1).astype(np.float32)),
            np.int32(im_h),
            np.int32(im_w),
            cuda.InOut(color_im.astype(np.float32)),
            depth_im_gpu,
            np.float32(self._trunc_margin),
            np.float32(weight),
            block=(self.block_x, self.block_y, self.block_z),
            grid=(self.grid_x, self.grid_y, self.grid_z)
        )

    def batch_ray_casting(self, im_w, im_h, cam_intr, cam_poses, inv_cam_poses, start_row=0, start_col=0, batch_size=1, to_host=True):
        batch_color_im_gpu = gpuarray.zeros(shape=(batch_size, 3, im_h, im_w), dtype=np.uint8)
        batch_depth_im_gpu = gpuarray.zeros(shape=(batch_size, im_h, im_w), dtype=np.float32)

        self._cuda_batch_ray_casting(
            self._tsdf_vol_gpu,
            self._color_vol_gpu,
            self._weight_vol_gpu,
            cuda.InOut(self._vol_dim.astype(np.int32)), 
            cuda.InOut(self._vol_origin.astype(np.float32)),
            np.float32(self._voxel_size),
            cuda.InOut(cam_intr.reshape(-1).astype(np.float32)),
            cuda.InOut(cam_poses.reshape(-1).astype(np.float32)),
            cuda.InOut(inv_cam_poses.reshape(-1).astype(np.float32)),
            np.int32(start_row),
            np.int32(start_col),
            np.int32(im_h),
            np.int32(im_w),
            batch_color_im_gpu,
            batch_depth_im_gpu,
            np.int32(batch_size),
            block=(8, 8, 16),
            grid=(int(np.ceil(im_w / 8)), int(np.ceil(im_h / 8)), int(np.ceil(batch_size / 16)))
        )
        if not to_host:
            return batch_depth_im_gpu, batch_color_im_gpu

        batch_depth_im = batch_depth_im_gpu.get()
        batch_color_im = batch_color_im_gpu.get().transpose(0, 2, 3, 1)  # b, h, w, c
        return batch_depth_im, batch_color_im

    def ray_casting(self, im_w, im_h, cam_intr, cam_pose, start_row=0, start_col=0, to_host=True):
        """
        Render an image patch
        """
        depth_im_gpu = gpuarray.zeros(shape=(im_h, im_w), dtype=np.float32)
        color_im_gpu = gpuarray.zeros(shape=(3, im_h, im_w), dtype=np.uint8)

        self._cuda_ray_casting(
            self._tsdf_vol_gpu,
            self._color_vol_gpu,
            self._weight_vol_gpu,
            cuda.InOut(self._vol_dim.astype(np.int32)), 
            cuda.InOut(self._vol_origin.astype(np.float32)),
            np.float32(self._voxel_size),
            cuda.InOut(cam_intr.reshape(-1).astype(np.float32)),
            cuda.InOut(cam_pose.reshape(-1).astype(np.float32)),
            cuda.InOut(la.inv(cam_pose).reshape(-1).astype(np.float32)),
            np.int32(start_row),
            np.int32(start_col),
            np.int32(im_h),
            np.int32(im_w),
            color_im_gpu,
            depth_im_gpu,
            block=(32, 32, 1),
            grid=(int(np.ceil(im_w / 32)), int(np.ceil(im_h / 32)), 1)
        )
        if not to_host:
            return depth_im_gpu, color_im_gpu

        depth_im = depth_im_gpu.get()
        color_im = color_im_gpu.get().transpose(1, 2, 0)
        return depth_im, color_im

    def get_volume(self):
        x_dim, y_dim, z_dim = self._vol_dim
        tsdf_vol_cpu = self._tsdf_vol_gpu.get().reshape((z_dim, y_dim, x_dim)).transpose(2, 1, 0)
        color_vol_cpu = self._color_vol_gpu.get().reshape((z_dim, y_dim, x_dim)).transpose(2, 1, 0)
        weight_vol_cpu = self._weight_vol_gpu.get().reshape((z_dim, y_dim, x_dim)).transpose(2, 1, 0)
        return tsdf_vol_cpu, color_vol_cpu, weight_vol_cpu

    def get_surface_cloud_marching_cubes(self, voxel_size=0.005):
        tsdf_vol, color_vol, weight_vol = self.get_volume()

        # Marching cubes
        verts = measure.marching_cubes_classic(tsdf_vol, level=0)[0]
        verts_ind = np.round(verts).astype(int)

        # remove false surface
        verts_weight = weight_vol[verts_ind[:, 0], verts_ind[:, 1], verts_ind[:, 2]]
        verts_val = tsdf_vol[verts_ind[:, 0], verts_ind[:, 1], verts_ind[:, 2]]
        valid_idx = (verts_weight > 0) & (np.abs(verts_val) < 0.2)
        verts_ind = verts_ind[valid_idx]
        verts = verts[valid_idx]
        verts = verts*self._voxel_size + self._vol_origin

        # Get vertex colors
        rgb_vals = color_vol[verts_ind[:, 0], verts_ind[:, 1], verts_ind[:, 2]]
        colors_b = np.floor(rgb_vals / self._color_const)
        colors_g = np.floor((rgb_vals - colors_b*self._color_const) / 256)
        colors_r = rgb_vals - colors_b*self._color_const - colors_g*256
        colors = np.floor(np.asarray([colors_r, colors_g, colors_b])).T

        surface_cloud = o3d.geometry.PointCloud()
        surface_cloud.points = o3d.utility.Vector3dVector(verts)
        surface_cloud.colors = o3d.utility.Vector3dVector(colors / 255)
        surface_cloud = surface_cloud.voxel_down_sample(voxel_size=voxel_size)
        return surface_cloud

    def get_conservative_volume(self, voxel_size=0.01):
        tsdf_vol, _, _ = self.get_volume()
        verts = np.vstack(np.where(tsdf_vol < -0.2)).T
        verts = verts*self._voxel_size + self._vol_origin

        conservative_volume = o3d.geometry.PointCloud()
        conservative_volume.points = o3d.utility.Vector3dVector(verts)
        conservative_volume = conservative_volume.voxel_down_sample(voxel_size=voxel_size)
        # conservative_volume.paint_uniform_color(np.array([0.3, 0.3, 0.3]))
        return conservative_volume

    def save(self, output_path):
        np.savez_compressed(output_path,
            vol_bounds=self._vol_bnds,
            voxel_size=self._voxel_size,
            trunc_margin=self._trunc_margin,
            tsdf_vol=self._tsdf_vol_gpu.get(),
            weight_vol=self._weight_vol_gpu.get(),
            color_vol=self._color_vol_gpu.get()
        )        
        print(f"tsdf volume has been saved to: {output_path}")

    @classmethod
    def load(cls, input_path):
        loaded = np.load(input_path)
        print('loaded vol bounds:', loaded['vol_bounds'])
        print('loaded voxel_size:', loaded['voxel_size'])
        print('loaded trunc_margin:', loaded['trunc_margin'])
        print("shape", loaded['tsdf_vol'].shape)
        obj = cls(loaded['vol_bounds'], loaded['voxel_size'], loaded['trunc_margin'])
        obj._tsdf_vol_gpu = gpuarray.to_gpu(loaded['tsdf_vol'])
        obj._weight_vol_gpu = gpuarray.to_gpu(loaded['weight_vol'])
        obj._color_vol_gpu = gpuarray.to_gpu(loaded['color_vol'])
        print(f"tsdf volume has been loaded from: {input_path}")
        return obj
        
kernel/kf_pycuda/tsdf_lib.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import time
import copy
from collections import deque
import open3d as o3d
import numpy as np
from numba import njit, prange
import scipy.linalg as la

def timeit(f, n=1, need_compile=False):
    def wrapper(*args, **kwargs):
        if need_compile:  # ignore the first run if needs compile
            result = f(*args, **kwargs)
        print("------------------------------------------------------------------------")
        tic = time.time()
        for i in range(n):
            result = f(*args, **kwargs)
        total_time = time.time() - tic
        print(f"time elapsed: {total_time}s. Average running time: {total_time / n}s")
        return result
    return wrapper


def create_pcd(depth_im: np.ndarray,
               cam_intr: np.ndarray,
               color_im: np.ndarray = None,
               depth_scale: float = 1,
               depth_trunc: float = 1.5,
               cam_extr: np.ndarray = np.eye(4)):

    intrinsic_o3d = o3d.camera.PinholeCameraIntrinsic()
    intrinsic_o3d.intrinsic_matrix = cam_intr
    depth_im_o3d = o3d.geometry.Image(depth_im)
    if color_im is not None:
        color_im_o3d = o3d.geometry.Image(color_im)
        rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(color_im_o3d, depth_im_o3d, 
            depth_scale=depth_scale, depth_trunc=depth_trunc, convert_rgb_to_intensity=False)
        pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd, intrinsic_o3d, extrinsic=cam_extr)
    else:
        pcd = o3d.geometry.PointCloud.create_from_depth_image(depth_im_o3d, intrinsic_o3d, extrinsic=cam_extr,
                                                              depth_scale=depth_scale, depth_trunc=depth_trunc)
    return pcd


@njit   # @njit(parallel=True) will be slower for whatever reason
def batch_compute_iou(roi: np.ndarray, proposals: np.ndarray, iou_threshold=0.8):
    """
    compute IoU between one roi and N proposals

    roi: np.ndarray of shape (4,)
    proposals: np.ndarray of shape (N, 4)
    """
    x1, y1, w, h = roi
    boxA = np.array([x1, y1, x1 + w, y1 + h])

    ious = np.zeros(len(proposals))
    for i in range(len(proposals)):
        boxB = proposals[i]
        xA = max(boxA[0], boxB[0])
        yA = max(boxA[1], boxB[1])
        xB = min(boxA[2], boxB[2])
        yB = min(boxA[3], boxB[3])

        # compute the area of intersection rectangle
        interArea = abs(max((xB - xA, 0)) * max((yB - yA), 0))

        # compute the area of both the prediction and ground-truth
        # rectangles
        boxAArea = abs((boxA[2] - boxA[0]) * (boxA[3] - boxA[1]))
        boxBArea = abs((boxB[2] - boxB[0]) * (boxB[3] - boxB[1]))

        # compute the intersection over union by taking the intersection
        # area and dividing it by the sum of prediction + ground-truth
        # areas - the interesection area
        iou = interArea / float(boxAArea + boxBArea - interArea)
        ious[i] = iou
    return ious


def plane_detection_ransac(pcd: o3d.geometry.PointCloud,
                           inlier_thresh: float,
                           max_iterations: int = 1000,
                           early_stop_thresh: float = 0.3,
                           visualize: bool = False):
    """detect plane (table) in a pointcloud for background removal

    Args:
        pcd (o3d.geometry.PointCloud): [input point cloud, assuming in the camera frame]
        inlier_thresh (float): [inlier distance threshold between a point to a plain]
        max_iterations (int): [max number of iteration to perform for RANSAC]
        early_stop_thresh (float): [inlier ratio to stop RANSAC early]
    
    Return:
        frame (np.ndarray): [z_dir is the estimated plane normal towards the camera, x_dir and y_dir randomly sampled]
        inlier ratio (float): [ratio of inliers in the estimated plane]
    """
    num_pts = len(pcd.points)
    pts = np.asarray(pcd.points)

    origin = None
    plane_normal = None
    max_inlier_ratio = 0
    max_num_inliers = 0

    for _ in range(max_iterations):
        # sample 3 points from the point cloud
        selected_indices = np.random.choice(num_pts, size=3, replace=False)
        selected_pts = pts[selected_indices]
        p1 = selected_pts[0]
        v1 = selected_pts[1] - p1
        v2 = selected_pts[2] - p1
        normal = np.cross(v1, v2)
        normal /= la.norm(normal)

        dist = np.abs((pts - p1) @ normal)
        num_inliers = np.sum(dist < inlier_thresh)
        inlier_ratio = num_inliers / num_pts

        if num_inliers > max_num_inliers:
            max_num_inliers = num_inliers
            origin = p1
            plane_normal = normal
            max_inlier_ratio = inlier_ratio

        if inlier_ratio > early_stop_thresh:
            break
    
    if plane_normal[2] < 0:
        plane_normal *= -1
    # if plane_normal @ origin > 0:
    #     plane_normal *= -1

    # randomly sample x_dir and y_dir given plane normal as z_dir
    x_dir = np.array([-plane_normal[2], 0, plane_normal[0]])
    x_dir /= la.norm(x_dir)
    y_dir = np.cross(plane_normal, x_dir)
    plane_frame = np.eye(4)
    plane_frame[:3, 0] = x_dir
    plane_frame[:3, 1] = y_dir
    plane_frame[:3, 2] = plane_normal
    plane_frame[:3, 3] = origin

    if visualize:
        dist = (pts - origin) @ plane_normal
        pcd_vis = copy.deepcopy(pcd)
        colors = np.asarray(pcd_vis.colors)
        colors[np.abs(dist) < inlier_thresh] = np.array([0, 0, 1])
        colors[dist < -inlier_thresh] = np.array([1, 0, 0])

        plane_frame_vis = generate_coordinate_frame(plane_frame, scale=0.05)
        cam_frame_vis = generate_coordinate_frame(np.eye(4), scale=0.05)
        o3d.visualization.draw_geometries([cam_frame_vis, plane_frame_vis, pcd_vis])

    return plane_frame, max_inlier_ratio


def extract_euclidean_clusters(pcd: o3d.geometry.PointCloud,
                               search_radius: 0.03,
                               min_pts_per_cluster=100,
                               max_pts_per_cluster=np.inf):
    """extract euclidean clusters from a point cloud
    https://pcl.readthedocs.io/en/latest/cluster_extraction.html
    https://github.com/PointCloudLibrary/pcl/blob/master/segmentation/include/pcl/segmentation/impl/extract_clusters.hpp

    Args:
        pcd (o3d.geometry.PointCloud): [input point cloud]
        search_radius (0.03): [max radius during clustering]
        min_pts_per_cluster (int, optional): [min number of points per cluster]. Defaults to 100.
        max_pts_per_cluster ([type], optional): [max number of points per cluster]. Defaults to np.inf.

    Returns:
        [list]: [a list of clusters, each cluster is a list of point indices]
    """
    pts = np.asarray(pcd.points)
    N = pts.shape[0]
    kd_tree = o3d.geometry.KDTreeFlann(pcd)

    clustered_indices = []  # a list of clusters, each cluster is a list of point indices

    processed = np.zeros(N, dtype=bool)

    for i in range(N):
        if processed[i]:
            continue

        processed[i] = True
        cluster = []
        queue = deque()
        queue.append(i)

        while queue:
            pt_idx = queue.popleft()
            pt = pts[pt_idx]
            cluster.append(pt_idx)
            
            # search_hybrid is much faster than search_radius
            _, neighbor_indices, _ = kd_tree.search_hybrid_vector_3d(pt, search_radius, max_nn=30)
            # _, neighbor_indices, _ = kd_tree.search_radius_vector_3d(pt, search_radius)

            for neighbor_idx in neighbor_indices:
                if not processed[neighbor_idx]:
                    queue.append(neighbor_idx)
                    processed[neighbor_idx] = True
        
        if min_pts_per_cluster < len(cluster) < max_pts_per_cluster:
            clustered_indices.append(cluster)

    return clustered_indices


def generate_coordinate_frame(T, scale=0.05):
    mesh = o3d.geometry.TriangleMesh.create_coordinate_frame()
    mesh.scale(scale, center=np.array([0, 0, 0]))
    return mesh.transform(T)


def get_view_frustum(depth_im, cam_intr, cam_pose):
    """
    Get corners of 3D camera view frustum of depth image
    """

    def rigid_transform(xyz, transform):
        """Applies a rigid transform to an (N, 3) pointcloud. """
        xyz_h = np.hstack([xyz, np.ones((len(xyz), 1), dtype=np.float32)])
        xyz_t_h = np.dot(transform, xyz_h.T).T
        return xyz_t_h[:, :3]

    im_h = depth_im.shape[0]
    im_w = depth_im.shape[1]
    max_depth = np.max(depth_im)
    view_frust_pts = np.array([
        (np.array([0,0,0,im_w,im_w])-cam_intr[0,2])*np.array([0,max_depth,max_depth,max_depth,max_depth])/cam_intr[0,0],
        (np.array([0,0,im_h,0,im_h])-cam_intr[1,2])*np.array([0,max_depth,max_depth,max_depth,max_depth])/cam_intr[1,1],
        np.array([0,max_depth,max_depth,max_depth,max_depth])
    ])
    view_frust_pts = rigid_transform(view_frust_pts.T, cam_pose).T
    return view_frust_pts


if __name__ == '__main__':

    pcd = o3d.io.read_point_cloud('demo.pcd')
    clusters = timeit(extract_euclidean_clusters)(pcd, search_radius=0.05)
    print(len(clusters))
    pcds = []
    for cluster in clusters:
        cluster_pcd = pcd.select_by_index(cluster)
        color = np.random.random((3,))
        cluster_pcd.paint_uniform_color(color)
        pcds.append(cluster_pcd)
    o3d.visualization.draw_geometries(pcds)

kernel/kf_pycuda/utils.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
from PIL import Image
import bpy
import os
import yaml
import numpy as np
from mathutils import Vector
import json
from tqdm import tqdm

from kernel.geometry import _pose2Rotation, _rotation2Pose
from kernel.ply_importer.point_data_file_handler import(
    PointDataFileHandler
)
from kernel.ply_importer.utility import(
    draw_points
)
from kernel.utility import _transstring2trans, _parse_camfile

from kernel.logging_utility import log_report
from registeration.init_configuration import config_json_dict, decode_dict
from kernel.blender_utility import \
    _get_configuration, _get_obj_insameworkspace, _apply_trans2obj, \
    _clear_allrgbdcam_insameworkspace, _getsameinstance, _getnextperfixforinstance
from kernel.utility import _select_sample_files, _generate_image_list
import time


def load_configuration(filepath):
    if os.path.exists(filepath):
        with open(filepath) as file:
            configuration = json.load(file)

        if configuration['projectname'] not in bpy.data.collections:
            config = bpy.context.scene.configuration.add()
            config_from_file(config, configuration)

        else:
            for config in bpy.context.scene.configuration:
                if config.projectname == configuration['projectname']:
                    config_from_file(config, configuration)
                    break
        
        create_workspace(os.path.dirname(filepath), 
                         configuration['projectname'], 
                         config = config)


def load_model(filepath, config_id):
    config = bpy.context.scene.configuration[config_id]
    workspace_name = config.projectname
    objFilename = filepath.split("/")[-1]
    objname = objFilename.split(".")[0]
    obj_instancename = workspace_name + ":" + objname + ".instance{0:03d}".format(_getnextperfixforinstance(config, objname))
    # if objName in bpy.data.objects:
    #     print("Unsupported for same object loaded several times")
    bpy.ops.import_scene.obj(filepath=filepath)
    bpy.context.selected_objects[0].name = obj_instancename
    bpy.ops.object.select_all(action='DESELECT')
    bpy.data.objects[obj_instancename].rotation_mode = 'QUATERNION'
    bpy.data.objects[obj_instancename].rotation_quaternion = [0.707, 0.707, 0., 0.]
    bpy.data.objects[obj_instancename]["type"] = "model"
    bpy.data.objects[obj_instancename]["path"] = filepath
    ## first unlink all collection, then link to Model collection
    for collection in bpy.data.objects[obj_instancename].users_collection:
        collection.objects.unlink(bpy.data.objects[obj_instancename])

    create_collection(workspace_name + ":Model", parent_collection = workspace_name)

    bpy.data.collections[workspace_name + ":Model"].objects.link(bpy.data.objects[obj_instancename])

    ### check whether the object is normal, split or URDF
    files = os.listdir(os.path.dirname(filepath))
    if 'split' in files:
        bpy.data.objects[obj_instancename]["modeltype"] = "split"
    else:
        bpy.data.objects[obj_instancename]["modeltype"] = "normal"
        
    bpy.data.objects[obj_instancename].optimizationparas.stop_criteria = 100
    bpy.data.objects[obj_instancename].optimizationparas.translation_learning_rate = 0.0005 
    bpy.data.objects[obj_instancename].optimizationparas.rotation_learning_rate = 0.005 

    

def load_model_from_pose(filepath, config_id):
    config = bpy.context.scene.configuration[config_id]
    workspace_name = config.projectname
    with open(filepath, 'r') as file:
        poses = yaml.load(file, Loader=yaml.FullLoader)
    if config.modelsrc == "":
        log_report(
            "INFO", "You should initialize the modelsrc before using this function", None
        )       
    else:
        model_dir = os.listdir(config.modelsrc)
        print(poses)
        for obj_instancename in poses:
            objname = obj_instancename.split(".")[0]
            if objname in model_dir:
                objworkspacename = workspace_name + ":" + obj_instancename
                if objworkspacename in bpy.data.objects:
                    bpy.data.objects[objworkspacename].location = poses[obj_instancename]['pose'][0]
                    bpy.data.objects[objworkspacename].rotation_quaternion = poses[obj_instancename]['pose'][1]/np.linalg.norm(poses[obj_instancename]['pose'][1])
                else:
                    load_model(os.path.join(config.modelsrc, objname, objname + ".obj" ), config_id)
                    bpy.data.objects[objworkspacename].location = poses[obj_instancename]['pose'][0]
                    bpy.data.objects[objworkspacename].rotation_quaternion = poses[obj_instancename]['pose'][1]/np.linalg.norm(poses[obj_instancename]['pose'][1])


def load_pc(filepath, pointcloudscale, config_id):
    workspace_name = bpy.context.scene.configuration[config_id].projectname

    if workspace_name + ":" + 'reconstruction' in bpy.data.objects:
        bpy.ops.object.select_all(action='DESELECT')
        bpy.data.objects[workspace_name + ":" + 'reconstruction'].select_set(True)
        bpy.ops.object.delete() 
    points = PointDataFileHandler.parse_point_data_file(filepath)
    draw_points(points = points, 
                point_size = 3, 
                add_points_to_point_cloud_handle = True, 
                reconstruction_collection = bpy.data.collections[workspace_name + ":" + "Pointcloud"], 
                object_anchor_handle_name=workspace_name + ":" + "reconstruction", op=None)
    bpy.data.objects[workspace_name + ":" + 'reconstruction']["type"] = "reconstruction"
    bpy.data.objects[workspace_name + ":" + 'reconstruction'].scale = Vector((pointcloudscale, 
                                                        pointcloudscale, 
                                                        pointcloudscale))

    bpy.data.objects[workspace_name + ":" + 'reconstruction']["path"] = filepath
    bpy.data.objects[workspace_name + ":" + 'reconstruction']["scale"] = pointcloudscale
    bpy.data.objects[workspace_name + ":" + 'reconstruction'].rotation_mode = 'QUATERNION'
    bpy.data.objects[workspace_name + ":" + 'reconstruction']["alignT"] = [[1., 0., 0., 0.],
                                                                            [0., 1., 0., 0.], 
                                                                            [0., 0., 1., 0.], 
                                                                            [0., 0., 0., 1.]]

def load_pc(filepath, pointcloudscale, config_id, name = 'reconstruction'):
    if os.path.exists(filepath):
        workspace_name = bpy.context.scene.configuration[config_id].projectname

        if workspace_name + ":" + name in bpy.data.objects:
            bpy.ops.object.select_all(action='DESELECT')
            bpy.data.objects[workspace_name + ":" + name].select_set(True)
            bpy.ops.object.delete() 
        points = PointDataFileHandler.parse_point_data_file(filepath)
        draw_points(points = points, 
                    point_size = 3, 
                    add_points_to_point_cloud_handle = True, 
                    reconstruction_collection = bpy.data.collections[workspace_name + ":" + "Pointcloud"], 
                    object_anchor_handle_name=workspace_name + ":" + name, op=None)
        bpy.data.objects[workspace_name + ":" + name]["type"] = "reconstruction"
        bpy.data.objects[workspace_name + ":" + name].scale = Vector((pointcloudscale, 
                                                                    pointcloudscale, 
                                                                    pointcloudscale))

        bpy.data.objects[workspace_name + ":" + name]["path"] = filepath
        bpy.data.objects[workspace_name + ":" + name]["scale"] = pointcloudscale
        bpy.data.objects[workspace_name + ":" + name].rotation_mode = 'QUATERNION'
        bpy.data.objects[workspace_name + ":" + name]["alignT"] = [[1., 0., 0., 0.],
                                                                    [0., 1., 0., 0.], 
                                                                    [0., 0., 1., 0.], 
                                                                    [0., 0., 0., 1.]]

def load_cam_img_depth(packagepath, config_id, camera_display_scale, sample_rate):

    workspace_name = bpy.context.scene.configuration[config_id].projectname

    work_space_collection = create_collection(workspace_name, parent_collection = None)
    recon_collection = create_collection(workspace_name + ":Reconstruction", 
                                         parent_collection = work_space_collection)
    cam_collection = create_collection(workspace_name + ":Camera", 
                                       parent_collection = recon_collection)
    thermal_collection = create_collection(workspace_name + ":Thermal", 
                                       parent_collection = recon_collection)                                       
    

    rgb_path = os.path.join(packagepath, "R/rgb_undistort")
    depth_path = os.path.join(packagepath, "R/depth")
    thermal_path = os.path.join(packagepath, "thermal/undistort")    
    
    rgb_files = os.listdir(rgb_path)
    depth_files = os.listdir(depth_path)
    thermal_files = os.listdir(thermal_path)

    rgb_files.sort()
    rgb_sample_files = _select_sample_files(rgb_files, sample_rate)
    
    thermal_files.sort()
    thermal_sample_files = _select_sample_files(thermal_files, sample_rate)
    
    os.system("mkdir -p " + bpy.context.scene.configuration[config_id].reconstructionsrc)
    _generate_image_list(bpy.context.scene.configuration[config_id].reconstructionsrc, rgb_sample_files)
    _clear_allrgbdcam_insameworkspace(bpy.context.scene.configuration[config_id])
    log_report(
        "INFO", "Loading camera, rgb and depth images", None
    )
    for rgb in tqdm(rgb_sample_files):
        perfix = rgb.split(".")[0]
        if perfix + ".png":# in depth_files:
            cam_name = workspace_name + ":view" + perfix
            if cam_name not in bpy.data.objects:
                cam_data = bpy.data.cameras.new(cam_name)
                cam_data.lens = bpy.context.scene.configuration[config_id].lens
                f = (bpy.context.scene.configuration[config_id].fx + bpy.context.scene.configuration[config_id].fy)/2
                cam_data.sensor_width = cam_data.lens * bpy.context.scene.configuration[config_id].resX/f
                cam_data.display_size = camera_display_scale
                
                cam_data.shift_x = (bpy.context.scene.configuration[config_id].resX/2 - bpy.context.scene.configuration[config_id].cx)/bpy.context.scene.configuration[config_id].resX
                cam_data.shift_y = (bpy.context.scene.configuration[config_id].cy - bpy.context.scene.configuration[config_id].resY/2)/bpy.context.scene.configuration[config_id].resX
                cam_data.background_images.new()
                
                cam_object = bpy.data.objects.new(cam_name, cam_data)
                cam_collection.objects.link(cam_object)
                cam_object.rotation_mode = 'QUATERNION'
                cam_object.location = [0, 0, 0]
                cam_object.rotation_quaternion = [1., 0, 0, 0]
            else:
                cam_object = bpy.data.objects[cam_name]

            ## load rgb
            rgb_name = workspace_name + ":rgb_undistort" + perfix
            if rgb_name not in bpy.data.images:
                bpy.ops.image.open(filepath=os.path.join(rgb_path, perfix + ".png"), show_multiview=False)
                bpy.data.images[perfix + ".png"].name = rgb_name
            bpy.data.images[rgb_name]["UPDATEALPHA"] = True
            bpy.data.images[rgb_name]["alpha"] = [0.5]
            ## load depth
            #depth_name = workspace_name + ":depth" + perfix
            #if depth_name not in bpy.data.images:
                #bpy.ops.image.open(filepath=os.path.join(depth_path, perfix + ".png"), 
                #                    directory=depth_path, 
                #                    files=[{"name":perfix + ".png"}], 
                #                    relative_path=True, show_multiview=False)
                #bpy.data.images[perfix + ".png"].name = depth_name
                
                #depth = np.array(Image.open(os.path.join(depth_path, perfix + ".png")))
                #depth = depth[::-1, ::]
                #bpy.data.images[depth_name]["depth"] = depth.flatten().astype(np.float32)
                ## change transparency
            #bpy.data.images[depth_name]["UPDATEALPHA"] = True
            #bpy.data.images[depth_name]["alpha"] = [0.5]
            #cam_object["depth"] = bpy.data.images[depth_name]
            cam_object["rgb"] = bpy.data.images[rgb_name]
            cam_object["type"] = "camera"  
            
            
def load_reconstruction_result(filepath, 
                               pointcloudscale, 
                               datasrc,
                               config_id,
                               camera_display_scale = 0.1,
                               IMPORT_RATIO = 1.0,
                               CAMPOSE_INVERSE = False
                               ):
                            
    packagepath = bpy.context.scene.configuration[config_id].datasrc

    workspace_name = bpy.context.scene.configuration[config_id].projectname
    work_space_collection = create_collection(workspace_name, parent_collection = None)
    recon_collection = create_collection(workspace_name + ":Reconstruction", 
                                         parent_collection = work_space_collection)
    cam_R_collection = create_collection(workspace_name + ":R", 
                                       parent_collection = recon_collection)
    thermal_collection = create_collection(workspace_name + ":T", 
                                       parent_collection = recon_collection)                                                                              
    cam_L_collection = create_collection(workspace_name + ":L", 
                                       parent_collection = recon_collection)                                       
    #pc_collection = create_collection(workspace_name + ":Pointcloud", 
    #                                  parent_collection = recon_collection)    
    ## load reconstruction result
    
    camera_R_file = os.path.join(filepath, "campose_R.txt")
    camera_thermal_file = os.path.join(filepath, "campose_T.txt")
    camera_L_file = os.path.join(filepath, "campose_L.txt")
    #reconstruction_path = os.path.join(filepath, "fused.ply")
    #load_pc(reconstruction_path, pointcloudscale, config_id)
    bpy.ops.object.select_all(action='DESELECT')

    R_rgb_path = os.path.join(packagepath, "R/undistort_centering")
    depth_path = os.path.join(packagepath, "R/depth")
    L_rgb_path = os.path.join(packagepath, "L/undistort_centering")
    thermal_path = os.path.join(packagepath, "T/undistort_centering")     
    
    #rgb_files = os.listdir(R_rgb_path)
    #depth_files = os.listdir(depth_path)    
    #thermal_files = os.listdir(thermal_path)         
    
    ## load camera and image result
    camera_lines = _parse_camfile(camera_R_file)
    camera_selected_lines = _select_sample_files(camera_lines, IMPORT_RATIO)
    _clear_allrgbdcam_insameworkspace(bpy.context.scene.configuration[config_id])
    for l in tqdm(camera_selected_lines):
        data = l.split(" ")
        if data[0].isnumeric():
            pose = [[float(data[5]) * pointcloudscale, float(data[6]) * pointcloudscale, float(data[7]) * pointcloudscale], 
                    [float(data[1]), float(data[2]), float(data[3]), float(data[4])]]
            Axis_align = np.array([[1, 0, 0, 0],
                                    [0, -1, 0, 0],
                                    [0, 0, -1, 0],
                                    [0, 0, 0, 1],]
            )
            Trans = _pose2Rotation(pose).dot(Axis_align) if not CAMPOSE_INVERSE else np.linalg.inv(_pose2Rotation(pose)).dot(Axis_align)
            pose = _rotation2Pose(Trans)
            framename = data[-1]
            perfix = framename.split(".")[0]

            cam_name = workspace_name + ":R" + perfix
            if cam_name in bpy.data.objects:
                cam_object = bpy.data.objects[cam_name]
                cam_object.location = pose[0]
                cam_object.rotation_quaternion = pose[1]
            else:
                fx = 579.2955827
                fy = 580.11605653
                cx = 320
                cy = 240
                resX = 640
                resY = 480
            
                cam_data = bpy.data.cameras.new(cam_name)
                cam_data.lens = bpy.context.scene.configuration[config_id].lens
                #f = (bpy.context.scene.configuration[config_id].fx + bpy.context.scene.configuration[config_id].fy)/2
                f = (fx + fy) / 2
                #cam_data.sensor_width = cam_data.lens * bpy.context.scene.configuration[config_id].resX/f
                cam_data.sensor_width = cam_data.lens * resX / f
                #cam_data.shift_x = (bpy.context.scene.configuration[config_id].resX/2 - bpy.context.scene.configuration[config_id].cx)/bpy.context.scene.configuration[config_id].resX
                #cam_data.shift_x = (resX/2 - cx)/resX
                cam_data.shift_x = (resX/2 - cx)/resX
                ### divide resX not resY
                #cam_data.shift_y = (bpy.context.scene.configuration[config_id].cy - bpy.context.scene.configuration[config_id].resY/2)/bpy.context.scene.configuration[config_id].resX
                cam_data.shift_y = (cy - resY/2)/resX
                #cam_data.shift_y = (resY/2-cy)/resX
                cam_data.display_size = camera_display_scale
                ## allow background display
                cam_data.background_images.new()

                cam_object = bpy.data.objects.new(cam_name, cam_data)
                cam_R_collection.objects.link(cam_object)
                cam_object.rotation_mode = 'QUATERNION'
                cam_object.location = pose[0]
                cam_object.rotation_quaternion = pose[1]
                ## load rgb
                rgb_name = workspace_name + ":R_undistort" + perfix
                if rgb_name not in bpy.data.images:
                    bpy.ops.image.open(filepath=os.path.join(R_rgb_path, perfix + ".png"), 
                                        directory=R_rgb_path, 
                                        files=[{"name":perfix + ".png"}], 
                                        relative_path=True, show_multiview=False)
                    bpy.data.images[perfix + ".png"].name = rgb_name
                bpy.data.images[rgb_name]["UPDATEALPHA"] = True
                bpy.data.images[rgb_name]["alpha"] = [0.5]
                cam_object["rgb"] = bpy.data.images[rgb_name]
                cam_object["type"] = "camera" 
                
                
                
                
    #################################################################################            
    ## load camera and image result
    camera_lines = _parse_camfile(camera_thermal_file)
    camera_selected_lines = _select_sample_files(camera_lines, IMPORT_RATIO)
    #_clear_allrgbdcam_insameworkspace(bpy.context.scene.configuration[config_id])
    for l in tqdm(camera_selected_lines):
        data = l.split(" ")
        if data[0].isnumeric():
            pose = [[float(data[5]) * pointcloudscale, float(data[6]) * pointcloudscale, float(data[7]) * pointcloudscale], 
                    [float(data[1]), float(data[2]), float(data[3]), float(data[4])]]
            Axis_align = np.array([[1, 0, 0, 0],
                                    [0, -1, 0, 0],
                                    [0, 0, -1, 0],
                                    [0, 0, 0, 1],]
            )
            Trans = _pose2Rotation(pose).dot(Axis_align) if not CAMPOSE_INVERSE else np.linalg.inv(_pose2Rotation(pose)).dot(Axis_align)
            pose = _rotation2Pose(Trans)
            framename = data[-1]
            perfix = framename.split(".")[0]

            cam_name = workspace_name + ":T" + perfix
            if cam_name in bpy.data.objects:
                cam_object = bpy.data.objects[cam_name]
                cam_object.location = pose[0]
                cam_object.rotation_quaternion = pose[1]
            else:
            
            
                fx = 275.00512361
                fy = 274.90509712
                cx = 320
                cy = 240
                resX = 640
                resY = 480
                
                #[[296.26782227   0.         317.36209511]
                # [  0.         297.97677612 260.14224065]
                # [  0.           0.           1.        ]]

            
                cam_data = bpy.data.cameras.new(cam_name)
                cam_data.lens = bpy.context.scene.configuration[config_id].lens
                #f = (bpy.context.scene.configuration[config_id].fx + bpy.context.scene.configuration[config_id].fy)/2
                f = (fx + fy) / 2
                #cam_data.sensor_width = cam_data.lens * bpy.context.scene.configuration[config_id].resX/f
                cam_data.sensor_width = cam_data.lens * resX / f
                #cam_data.shift_x = (bpy.context.scene.configuration[config_id].resX/2 - bpy.context.scene.configuration[config_id].cx)/bpy.context.scene.configuration[config_id].resX
                cam_data.shift_x = (resX/2 - cx)/resX
                ### divide resX not resY
                #cam_data.shift_y = (bpy.context.scene.configuration[config_id].cy - bpy.context.scene.configuration[config_id].resY/2)/bpy.context.scene.configuration[config_id].resX
                cam_data.shift_y = (cy - resY/2)/resX
                #cam_data.shift_y = (resY/2-cy)/resX
                cam_data.display_size = camera_display_scale
                ## allow background display
                cam_data.background_images.new()

                cam_object = bpy.data.objects.new(cam_name, cam_data)
                thermal_collection.objects.link(cam_object)
                cam_object.rotation_mode = 'QUATERNION'
                cam_object.location = pose[0]
                cam_object.rotation_quaternion = pose[1]
                ## load thermal
                thermal_name = workspace_name + ":T_undistort" + perfix
                if thermal_name not in bpy.data.images:
                    bpy.ops.image.open(filepath=os.path.join(thermal_path, perfix + ".png"), 
                                        directory=thermal_path, 
                                        files=[{"name":perfix + ".png"}], 
                                        relative_path=True, show_multiview=False)
                    bpy.data.images[perfix + ".png"].name = thermal_name
                bpy.data.images[thermal_name]["UPDATEALPHA"] = True
                bpy.data.images[thermal_name]["alpha"] = [0.5]
                cam_object["rgb"] = bpy.data.images[thermal_name]
                cam_object["type"] = "camera"                 
                
                
                
         
         
                
    #################################################################################            
    ## load camera and image result
    ## load camera and image result
    camera_lines = _parse_camfile(camera_L_file)
    camera_selected_lines = _select_sample_files(camera_lines, IMPORT_RATIO)
    for l in tqdm(camera_selected_lines):
        data = l.split(" ")
        if data[0].isnumeric():
            pose = [[float(data[5]) * pointcloudscale, float(data[6]) * pointcloudscale, float(data[7]) * pointcloudscale], 
                    [float(data[1]), float(data[2]), float(data[3]), float(data[4])]]
            Axis_align = np.array([[1, 0, 0, 0],
                                    [0, -1, 0, 0],
                                    [0, 0, -1, 0],
                                    [0, 0, 0, 1],]
            )
            Trans = _pose2Rotation(pose).dot(Axis_align) if not CAMPOSE_INVERSE else np.linalg.inv(_pose2Rotation(pose)).dot(Axis_align)
            pose = _rotation2Pose(Trans)
            framename = data[-1]
            perfix = framename.split(".")[0]

            cam_name = workspace_name + ":L" + perfix
            if cam_name in bpy.data.objects:
                cam_object = bpy.data.objects[cam_name]
                cam_object.location = pose[0]
                cam_object.rotation_quaternion = pose[1]
            else:
                fx = 599.77885008
                fy = 600.2724334
                cx = 320
                cy = 240
                resX = 640
                resY = 480
            
                cam_data = bpy.data.cameras.new(cam_name)
                cam_data.lens = bpy.context.scene.configuration[config_id].lens
                #f = (bpy.context.scene.configuration[config_id].fx + bpy.context.scene.configuration[config_id].fy)/2
                f = (fx + fy) / 2
                #cam_data.sensor_width = cam_data.lens * bpy.context.scene.configuration[config_id].resX/f
                cam_data.sensor_width = cam_data.lens * resX / f
                #cam_data.shift_x = (bpy.context.scene.configuration[config_id].resX/2 - bpy.context.scene.configuration[config_id].cx)/bpy.context.scene.configuration[config_id].resX
                cam_data.shift_x = (resX/2 - cx)/resX
                ### divide resX not resY
                #cam_data.shift_y = (bpy.context.scene.configuration[config_id].cy - bpy.context.scene.configuration[config_id].resY/2)/bpy.context.scene.configuration[config_id].resX
                #cam_data.shift_y = (cy - resY/2)/resX
                cam_data.shift_y = (resY/2-cy)/resX
                cam_data.display_size = camera_display_scale
                ## allow background display
                cam_data.background_images.new()

                cam_object = bpy.data.objects.new(cam_name, cam_data)
                cam_L_collection.objects.link(cam_object)
                cam_object.rotation_mode = 'QUATERNION'
                cam_object.location = pose[0]
                cam_object.rotation_quaternion = pose[1]
                ## load rgb
                rgb_name = workspace_name + ":L_undistort" + perfix
                if rgb_name not in bpy.data.images:
                    bpy.ops.image.open(filepath=os.path.join(L_rgb_path, perfix + ".png"), 
                                        directory=L_rgb_path, 
                                        files=[{"name":perfix + ".png"}], 
                                        relative_path=True, show_multiview=False)
                    bpy.data.images[perfix + ".png"].name = rgb_name
                bpy.data.images[rgb_name]["UPDATEALPHA"] = True
                bpy.data.images[rgb_name]["alpha"] = [0.5]
                cam_object["rgb"] = bpy.data.images[rgb_name]
                cam_object["type"] = "camera"     
    
    
    
    
    
    
    
    
    
    
    
    
                    
    #################################################################################            
         
         
         
                
    
    obj_lists = _get_obj_insameworkspace(cam_object, ["reconstruction", "R", "L", "T"])
    _, config = _get_configuration(cam_object)

    trans = _transstring2trans(config.recon_trans)
    for obj in obj_lists:
        _apply_trans2obj(obj, trans)  
        if obj['type'] == 'reconstruction':
            obj["alignT"] = trans.tolist()      
    


def create_workspace(path, name, 
                     config = None):
    if not config:
        config = bpy.context.scene.configuration.add()
        config_id = len(bpy.context.scene.configuration) - 1
        config.projectname = name
    else:
        for idx, c in enumerate(bpy.context.scene.configuration):
            if c.projectname == name:
                config_id = idx
    
    work_space_collection = create_collection(name, parent_collection = None)
    setting = setting_init(name + ":Setting", work_space_collection, path)
    setting["config_id"] = config_id
    model_collection = create_collection(name + ":Model", 
                                         parent_collection = work_space_collection)
    recon_collection = create_collection(name + ":Reconstruction", 
                                         parent_collection = work_space_collection)
    #pc_collection = create_collection(name + ":Pointcloud", 
    #                                  parent_collection = recon_collection)
    cam_R_collection = create_collection(name + ":R", 
                                       parent_collection = recon_collection)
    thermal_collection = create_collection(name + ":T", 
                                       parent_collection = recon_collection)                                       
    cam_L_collection = create_collection(name + ":L", 
                                       parent_collection = recon_collection)                                       

def init_package(path, config):
    create_packages(path, ["model", "recon", "data"])
    config.modelsrc = os.path.join(path, "model")
    config.datasrc = os.path.join(path, "camera/")
    config.reconstructionsrc = os.path.join(path, "recon")

def create_packages(path, packages):
    for package in packages:
        dir = os.path.join(path, package)
        if not os.path.exists(dir):
            os.mkdir(dir)


def create_collection(new_name, parent_collection = None):
    if new_name not in bpy.data.collections:
        new_collection = bpy.data.collections.new(new_name)
        if not parent_collection:
            bpy.context.scene.collection.children.link(new_collection)
        elif parent_collection.name in bpy.data.collections:
            parent_collection.children.link(new_collection)
        else:
            raise Exception("parent collection not exists")
        return new_collection
    else: 
        return bpy.data.collections[new_name]

def setting_init(name, collection, path):
    if name not in bpy.data.objects:
        setting = bpy.data.objects.new(name, None)
        setting["type"] = "setting"
        setting["dir"] = path
        collection.objects.link(setting)
    else:
        setting = bpy.data.objects[name]
        print(name.split(":")[0] + " is currently in the Blender")
    return setting


def config_from_file(config, configuration):
    for item in config_json_dict:
        value = decode_dict(configuration, config_json_dict[item])
        setattr(config, item, value)



def updateprojectname():
    current_object = bpy.context.object.name
    
    old_name = current_object.split(":")[0]
    config_id = bpy.data.objects[old_name + ":Setting"]['config_id']

    new_name = bpy.context.scene.configuration[config_id].projectname
    for config in bpy.context.scene.configuration:
        if config.projectname == new_name and config != bpy.context.scene.configuration[config_id]:
            log_report(
                "INFO", "Same name is existing in Blender", None
            )                   
            bpy.context.scene.configuration[config_id].projectname = old_name
            return 
    for collection in bpy.data.collections:
        if collection.name.split(":")[0] == old_name:
            collection.name = collection.name.replace(old_name, new_name)
    
    for obj in bpy.data.objects:
        if obj.name.split(":")[0] == old_name:
            obj.name = obj.name.replace(old_name, new_name)    

    for img in bpy.data.images:
        if img.name.split(":")[0] == old_name:
            img.name = img.name.replace(old_name, new_name)          
    return

def removeworkspace(name):
    for obj in bpy.data.objects:
        if obj.name.split(":")[0] == name:
            bpy.data.objects.remove(obj)    
    for collection in bpy.data.collections:
        if collection.name.split(":")[0] == name:
            bpy.data.collections.remove(collection)
    for img in bpy.data.images:
        if img.name.split(":")[0] == name:
            bpy.data.images.remove(img)         
    return 




kernel/loader.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import logging
import bpy

logging.basicConfig(level=logging.INFO)
_logger = logging.getLogger()

def ShowMessageBox(message = "", title = "An ProgressLabeller error", icon = 'ERROR'):

    def draw(self, context):
        self.layout.label(text = message)

    bpy.context.window_manager.popup_menu(draw, title = title, icon = icon)


def log_report(output_type, some_str, op=None):
    """ Write a string to the console and to Blender's info area."""
    # output_type is one of: 'INFO', 'WARNING' or 'ERROR'
    if output_type in ['ERROR', 'Error']:
        ShowMessageBox(message=some_str)
    _logger.info(output_type + ": " + some_str)
    if op is not None:
        op.report({output_type}, some_str)





kernel/logging_utility.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
if __name__ == "__main__":
    import sys
    import os
    sys.path.append(os.path.dirname(os.path.abspath(__file__)))
    from build.orb_extension import orb_slam_recon
    orb_slam_recon(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], sys.argv[5], float(sys.argv[6]), float(sys.argv[7]))

kernel/orb_slam/orb_slam.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import os

def orbslam_yaml(file_path, fx, fy, cx,cy, width, height, depthscale, frequency):
    f = open(file_path, "w")
    f.write("%YAML:1.0\n")
    f.write("\n")
    f.write("#--------------------------------------------------------------------------------------------\n")
    f.write("# Camera Parameters. Adjust them!\n")
    f.write("#--------------------------------------------------------------------------------------------\n")
    f.write("\n")
    f.write("# Camera calibration and distortion parameters (OpenCV)\n")
    f.write("Camera.fx: {0:.3f}\n".format(fx))
    f.write("Camera.fy: {0:.3f}\n".format(fy))
    f.write("Camera.cx: {0:.3f}\n".format(cx))
    f.write("Camera.cy: {0:.3f}\n\n".format(cy))

    f.write("Camera.k1: 0\n")
    f.write("Camera.k2: 0\n")
    f.write("Camera.p1: 0\n")
    f.write("Camera.p2: 0\n")
    f.write("Camera.k3: 0\n\n")

    f.write("Camera.width: {0}\n".format(width))
    f.write("Camera.height: {0}\n\n".format(height))

    f.write("# Camera frames per second\n") 
    f.write("Camera.fps: {0:.1f}\n\n".format(frequency))

    f.write("# IR projector baseline times fx (aprox.)\n")
    f.write("Camera.bf: 40.0\n\n")

    f.write("# Color order of the images (0: BGR, 1: RGB. It is ignored if images are grayscale)\n")
    f.write("Camera.RGB: 1\n\n")

    f.write("# Close/Far threshold. Baseline times.\n")
    f.write("ThDepth: 40.0\n\n")

    f.write("# Deptmap values factor\n") 
    f.write("DepthMapFactor: {0:.1f}\n".format(1/depthscale))

    f.write("#--------------------------------------------------------------------------------------------\n")
    f.write("# ORB Parameters\n")
    f.write("#--------------------------------------------------------------------------------------------\n\n")

    f.write("# ORB Extractor: Number of features per image\n")
    f.write("ORBextractor.nFeatures: 2000\n\n")

    f.write("# ORB Extractor: Scale factor between levels in the scale pyramid\n") 	
    f.write("ORBextractor.scaleFactor: 1.2\n\n")

    f.write("# ORB Extractor: Number of levels in the scale pyramid\n")	
    f.write("ORBextractor.nLevels: 8\n\n")

    f.write("# ORB Extractor: Fast threshold\n")
    f.write("# Image is divided in a grid. At each cell FAST are extracted imposing a minimum response.\n")
    f.write("# Firstly we impose iniThFAST. If no corners are detected we impose a lower value minThFAST\n")
    f.write("# You can lower these values if your images have low contrast\n")			
    f.write("ORBextractor.iniThFAST: 20\n")
    f.write("ORBextractor.minThFAST: 7\n\n")

    f.write("#--------------------------------------------------------------------------------------------\n")
    f.write("# Viewer Parameters\n")
    f.write("#--------------------------------------------------------------------------------------------\n")
    f.write("Viewer.KeyFrameSize: 0.05\n")
    f.write("Viewer.KeyFrameLineWidth: 1\n")
    f.write("Viewer.GraphLineWidth: 0.9\n")
    f.write("Viewer.PointSize: 2\n")
    f.write("Viewer.CameraSize: 0.08\n")
    f.write("Viewer.CameraLineWidth: 3\n")
    f.write("Viewer.ViewpointX: 0\n")
    f.write("Viewer.ViewpointY: -0.7\n")
    f.write("Viewer.ViewpointZ: -1.8\n")
    f.write("Viewer.ViewpointF: 500\n")
    f.close()

def orbslam_associatefile(file_path, dataset_path, frequency):
    rgb_files = os.listdir(os.path.join(dataset_path, "rgb"))
    depth_files = os.listdir(os.path.join(dataset_path, "depth"))
    rgb_files.sort()
    depth_files.sort()
    index = 1/frequency
    with open(file_path, "w") as f:
        for perfix in rgb_files:
            if perfix in depth_files:
                f.write(f'{index} ' + 'rgb/{0} '.format(perfix) + f'{index} ' + 'depth/{0}\n'.format(perfix))
                index += 1/frequency
kernel/orb_slam/orbslam_utility.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
if __name__ == "__main__":
    import sys
    import os
    sys.path.append(os.path.dirname(os.path.abspath(__file__)))
    from build.orb3_extension import orb3_slam_recon
    orb3_slam_recon(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], sys.argv[5], float(sys.argv[6]), float(sys.argv[7]))

kernel/orb_slam3/orb_slam3.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import os

def orbslam3_yaml(file_path, fx, fy, cx,cy, width, height, depthscale, frequency):
    f = open(file_path, "w")

    f.write("%YAML:1.0\n")
    f.write("\n")
    f.write("#--------------------------------------------------------------------------------------------\n")
    f.write("# Camera Parameters. Adjust them!\n")
    f.write("#--------------------------------------------------------------------------------------------\n")
    f.write("File.version: \"1.0\"\n")
    f.write("\n")
    f.write("Camera.type: \"PinHole\"\n")
    f.write("\n")
    f.write("# Right Camera calibration and distortion parameters (OpenCV)\n")
    f.write("Camera1.fx: {0:.3f}\n".format(fx))
    f.write("Camera1.fy: {0:.3f}\n".format(fy))
    f.write("Camera1.cx: {0:.3f}\n".format(cx))
    f.write("Camera1.cy: {0:.3f}\n".format(cy))
    f.write("\n")
    f.write("\n")
    f.write("# distortion parameters\n")
    f.write("Camera1.k1: 0.0\n")
    f.write("Camera1.k2: 0.0\n")
    f.write("Camera1.p1: 0.0\n")
    f.write("Camera1.p2: 0.0\n")
    f.write("\n")
    f.write("# Camera resolution\n")
    f.write("Camera.width: {0}\n".format(width))
    f.write("Camera.height: {0}\n".format(height))
    f.write("\n")
    f.write("# Camera frames per second\n") 
    f.write("Camera.fps: {0:d}\n".format(int(frequency)))
    f.write("\n")
    f.write("# Color order of the images (0: BGR, 1: RGB. It is ignored if images are grayscale)\n")
    f.write("Camera.RGB: 1\n")
    f.write("\n")
    f.write("# ******************* TODO: may need change *******************\n")
    f.write("# Close/Far threshold. Baseline times.\n")
    f.write("Stereo.ThDepth: 40.0\n")
    f.write("Stereo.b: 0.0745\n")
    f.write("\n")
    f.write("# Depth map values factor\n")
    f.write("RGBD.DepthMapFactor: {0:.1f}\n".format(1/depthscale))
    f.write("\n")
    f.write("#--------------------------------------------------------------------------------------------\n")
    f.write("# ORB Parameters\n")
    f.write("#--------------------------------------------------------------------------------------------\n")
    f.write("# ORB Extractor: Number of features per image\n")
    f.write("ORBextractor.nFeatures: 2500\n")
    f.write("\n")
    f.write("# ORB Extractor: Scale factor between levels in the scale pyramid\n") 	
    f.write("ORBextractor.scaleFactor: 1.2\n")
    f.write("\n")
    f.write("# ORB Extractor: Number of levels in the scale pyramid\n")	
    f.write("ORBextractor.nLevels: 8\n")
    f.write("\n")
    f.write("# ORB Extractor: Fast threshold\n")
    f.write("# Image is divided in a grid. At each cell FAST are extracted imposing a minimum response.\n")
    f.write("# Firstly we impose iniThFAST. If no corners are detected we impose a lower value minThFAST\n")
    f.write("# You can lower these values if your images have low contrast\n")			
    f.write("ORBextractor.iniThFAST: 20\n")
    f.write("ORBextractor.minThFAST: 7\n")
    f.write("\n")
    f.write("#--------------------------------------------------------------------------------------------\n")
    f.write("# Viewer Parameters\n")
    f.write("#--------------------------------------------------------------------------------------------\n")
    f.write("Viewer.KeyFrameSize: 0.05\n")
    f.write("Viewer.KeyFrameLineWidth: 1.0\n")
    f.write("Viewer.GraphLineWidth: 0.9\n")
    f.write("Viewer.PointSize: 2.0\n")
    f.write("Viewer.CameraSize: 0.08\n")
    f.write("Viewer.CameraLineWidth: 3.0\n")
    f.write("Viewer.ViewpointX: 0.0\n")
    f.write("Viewer.ViewpointY: -0.7\n")
    f.write("Viewer.ViewpointZ: -3.5\n")
    f.write("Viewer.ViewpointF: 500.0\n")
    f.close()


def orbslam3_associatefile(file_path, dataset_path, frequency):
    rgb_files = os.listdir(os.path.join(dataset_path, "rgb"))
    depth_files = os.listdir(os.path.join(dataset_path, "depth"))
    rgb_files.sort()
    depth_files.sort()
    index = 1/frequency
    with open(file_path, "w") as f:
        for perfix in rgb_files:
            if perfix in depth_files:
                f.write(f'{index} ' + 'rgb/{0} '.format(perfix) + f'{index} ' + 'depth/{0}\n'.format(perfix))
                index += 1/frequency
kernel/orb_slam3/orbslam3_utility.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import numpy as np
import atexit
import bpy
import bgl
import gpu
from gpu_extras.batch import batch_for_shader
from kernel.logging_utility import log_report


def _compute_transformed_coords(object_anchor_matrix_world, positions):

    if len(positions) == 0:
        return []

    pos_arr = np.asarray(positions)
    ones_arr = np.ones((pos_arr.shape[0], 1))
    pos_hom_arr = np.hstack((pos_arr, ones_arr))

    # Transpose the matrix to transform the coordinates
    # with a single matrix multiplication
    pos_hom_arr_transposed = np.transpose(pos_hom_arr)
    transf_pos_hom_transposed_arr = np.matmul(
        object_anchor_matrix_world, pos_hom_arr_transposed
    )
    transf_pos_arr_hom = transf_pos_hom_transposed_arr.T

    # Delete the homogeneous entries
    transf_pos_arr = np.delete(transf_pos_arr_hom, -1, axis=1)
    transf_pos_list = transf_pos_arr.tolist()
    return transf_pos_list


class DrawManager:
    """Class that allows to represent point clouds with OpenGL in Blender."""

    def __init__(self):
        self._anchor_to_draw_callback_handler = {}
        self._anchor_to_point_coords = {}
        self._anchor_to_point_colors = {}

    @classmethod
    def get_singleton(cls):
        """Return a singleton of this class."""
        if hasattr(bpy.types.Object, "current_draw_manager"):
            draw_manger = bpy.types.Object.current_draw_manager
        else:
            draw_manger = cls()
            bpy.types.Object.current_draw_manager = draw_manger
        return draw_manger

    def register_points_draw_callback(
        self, object_anchor, coords, colors, point_size
    ):
        """Register a callback to draw a point cloud."""
        draw_callback_handler = _DrawCallBackHandler()
        draw_callback_handler.register_points_draw_callback(
            self, object_anchor, coords, colors, point_size
        )
        self._anchor_to_draw_callback_handler[
            object_anchor
        ] = draw_callback_handler
        self._anchor_to_point_coords[object_anchor] = coords
        self._anchor_to_point_colors[object_anchor] = colors

    def get_coords_and_colors(self, visible_only=False):
        """Return the coordinates and the colors of the maintained points."""
        transf_coord_list = []
        color_list = []
        for object_anchor in self._anchor_to_point_coords:

            if visible_only and not object_anchor.visible_get():
                continue

            coords = self._anchor_to_point_coords[object_anchor]
            transf_coord_list = (
                transf_coord_list
                + _compute_transformed_coords(
                    object_anchor.matrix_world, coords
                )
            )

            colors = self._anchor_to_point_colors[object_anchor]
            color_list = color_list + colors

        return transf_coord_list, color_list

    def delete_anchor(self, object_anchor):
        """Delete the anchor used to control the pose of the point cloud."""
        del self._anchor_to_point_coords[object_anchor]
        del self._anchor_to_point_colors[object_anchor]
        # del self._anchor_to_draw_callback_handler[object_anchor]

    def get_draw_callback_handler(self, object_anchor):
        """Get the draw callback handler corresponding to the object anchor."""
        return self._anchor_to_draw_callback_handler[object_anchor]


class _DrawCallBackHandler:
    """Class that allows to handle point drawing callbacks."""

    def __init__(self):
        self._shader = gpu.shader.from_builtin("3D_FLAT_COLOR")

        # Handle to the function
        self._draw_handler_handle = None

        # Handle to the object
        self._object_anchor_pose_previous = np.array([])
        self._batch_cached = None
        self._point_size = 5

        # If Blender is closed and self._batch_cached is not properly deleted,
        # this causes something like the following:
        # "Error: Not freed memory blocks: 2, total unfreed memory 0.001358 MB"
        atexit.register(self._clean_batch_cached)

    def _clean_batch_cached(self):
        """Clean the cached batch used to draw the points."""
        self._batch_cached = None

    def set_point_size(self, point_size):
        """Set the point size used to draw the points in the 3D point cloud."""
        self._point_size = point_size

    def _draw_points_callback(
        self, draw_manager, object_anchor, positions, colors
    ):
        """A callback function to draw a point cloud in Blender's 3D view."""
        handle_is_valid = True
        try:
            # Check if object still exists
            object_anchor_name = object_anchor.name
        except:
            handle_is_valid = False

        if handle_is_valid:
            if object_anchor_name in bpy.data.objects:

                # Use the visibility of the object to enable /
                # disable the drawing of the point cloud
                if bpy.data.objects[object_anchor_name].visible_get():

                    # Update the batch depending on the anchor pose (only if
                    # necessary)
                    object_anchor_has_changed = not np.array_equal(
                        self._object_anchor_pose_previous,
                        object_anchor.matrix_world,
                    )
                    if self._batch_cached is None or object_anchor_has_changed:

                        self._object_anchor_pose_previous = np.copy(
                            object_anchor.matrix_world
                        )
                        transf_pos_list = _compute_transformed_coords(
                            object_anchor.matrix_world, positions
                        )

                        self._batch_cached = batch_for_shader(
                            self._shader,
                            "POINTS",
                            {"pos": transf_pos_list, "color": colors},
                        )

                    self._shader.bind()
                    bgl.glPointSize(self._point_size)
                    bgl.glEnable(bgl.GL_DEPTH_TEST)
                    bgl.glDepthMask(bgl.GL_TRUE)

                    self._batch_cached.draw(self._shader)

        else:
            if self._draw_handler_handle is not None:
                # log_report(
                #     "INFO",
                #     "Removing draw handler of deleted point cloud handle",
                # )
                # Using bpy.types.SpaceView3D.draw_handler_remove causes
                # Blender to crash (frequently).
                # bpy.types.SpaceView3D.draw_handler_remove(
                #     self._draw_handler_handle, "WINDOW"
                # )
                # Note:
                # Not removing the _draw_handler_handle is not that bad, since
                # a) after removing the point cloud anchor the
                # "_draw_points_callback" becomes very cheap
                # b) closing Blender frees the _draw_handler_handle and it
                # is not restored even if the Blend file is reloaded from disk
                self._draw_handler_handle = None
                self._batch_cached = None
                draw_manager.delete_anchor(object_anchor)

    def register_points_draw_callback(
        self, draw_manager, object_anchor, positions, colors, point_size
    ):
        """Register a callback to draw a point cloud."""
        self.set_point_size(point_size)
        args = (draw_manager, object_anchor, positions, colors)
        self._draw_handler_handle = bpy.types.SpaceView3D.draw_handler_add(
            self._draw_points_callback, args, "WINDOW", "POST_VIEW"
        )

kernel/ply_importer/draw_manager.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import os
import numpy as np
import importlib

from kernel.ply_importer.point import Point
from kernel.logging_utility import log_report
from kernel.ply_importer.type_utility import is_float, is_int


class _DataSemantics:
    def __init__(self):
        self.x_idx = None
        self.y_idx = None
        self.z_idx = None
        self.r_idx = None
        self.g_idx = None
        self.b_idx = None
        self.num_data_entries = None
        self.pseudo_color = None

    def is_initialized(self):
        return not None in [
            self.x_idx,
            self.y_idx,
            self.z_idx,
            self.r_idx,
            self.g_idx,
            self.b_idx,
            self.num_data_entries,
            self.pseudo_color,
        ]


class PointDataFileHandler:
    """Class to read and write common point data files."""

    @staticmethod
    def _read_lines_as_tuples(ifc, delimiter):
        lines_as_tup = []
        for line in ifc.readlines():
            elements = line.split(delimiter)
            lines_as_tup.append(elements)
        return lines_as_tup

    @staticmethod
    def _guess_data_semantics_from_tuple(data_tuple):
        data_semantics = _DataSemantics()
        data_semantics.num_data_entries = len(data_tuple)
        # Data must start with subsequent float values
        # representing the three-dimensional position
        for idx in [0, 1, 2]:
            assert is_float(data_tuple[idx])
        data_semantics.x_idx = 0
        data_semantics.y_idx = 1
        data_semantics.z_idx = 2

        num_data_entries = len(data_tuple)

        # Search for three subsequent int values between 0 and 255
        # (which indicate RGB color information)
        for idx in [3, num_data_entries - 3]:
            if not is_int(data_tuple[idx]):
                continue
            if not is_int(data_tuple[idx + 1]):
                continue
            if not is_int(data_tuple[idx + 2]):
                continue
            if not 0 <= int(data_tuple[idx]) <= 255:
                continue
            if not 0 <= int(data_tuple[idx + 1]) <= 255:
                continue
            if not 0 <= int(data_tuple[idx + 2]) <= 255:
                continue
            data_semantics.r_idx = idx
            data_semantics.g_idx = idx + 1
            data_semantics.b_idx = idx + 2
            data_semantics.pseudo_color = False
            break

        if data_semantics.is_initialized():
            return data_semantics

        # If not int values are found, we assume that the color information
        # is stored as pseudo colors, i.e. we are looking for 3 subsequent
        # floats between (0,1).
        for idx in [3, num_data_entries - 3]:
            assert is_float(data_tuple[idx])
            if not 0 <= float(data_tuple[idx]) <= 1:
                continue
            if not 0 <= float(data_tuple[idx + 1]) <= 1:
                continue
            if not 0 <= float(data_tuple[idx + 2]) <= 1:
                continue
            data_semantics.r_idx = idx
            data_semantics.g_idx = idx + 1
            data_semantics.b_idx = idx + 2
            data_semantics.pseudo_color = True
            break
        assert data_semantics.is_initialized()
        return data_semantics

    @staticmethod
    def _get_data_semantics_from_header(line):
        data_semantics = _DataSemantics()
        data_tuple = line.lstrip("//").rstrip().split(" ")
        data_semantics.num_data_entries = len(data_tuple)

        for idx, val in enumerate(data_tuple):
            val = val.lower()
            if val == "x":
                data_semantics.x_idx = idx
            elif val == "y":
                data_semantics.y_idx = idx
            elif val == "z":
                data_semantics.z_idx = idx
            elif val in ["r", "red"]:
                data_semantics.r_idx = idx
                data_semantics.pseudo_color = False
            elif val in ["g", "green"]:
                data_semantics.g_idx = idx
                data_semantics.pseudo_color = False
            elif val in ["b", "blue"]:
                data_semantics.b_idx = idx
                data_semantics.pseudo_color = False
            elif val == "rf":
                data_semantics.r_idx = idx
                data_semantics.pseudo_color = True
            elif val == "gf":
                data_semantics.g_idx = idx
                data_semantics.pseudo_color = True
            elif val == "bf":
                data_semantics.b_idx = idx
                data_semantics.pseudo_color = True
        assert data_semantics.is_initialized()
        return data_semantics

    @staticmethod
    def _get_data_semantics_from_ascii(ifp, delimiter, has_header, op=None):
        with open(ifp, "r") as ifc:
            data_semantics = None
            if has_header:
                line = ifc.readline()
                if line.startswith("//"):
                    log_report(
                        "INFO", "Reading data semantics from header", op
                    )
                    data_semantics = (
                        PointDataFileHandler._get_data_semantics_from_header(
                            line
                        )
                    )
                    line = ifc.readline()

            if data_semantics is None:
                log_report(
                    "INFO", "No header available, guessing data semantics", op
                )
                lines_as_tuples = PointDataFileHandler._read_lines_as_tuples(
                    ifc, delimiter=delimiter
                )
                data_semantics = (
                    PointDataFileHandler._guess_data_semantics_from_tuple(
                        lines_as_tuples[0]
                    )
                )
            return data_semantics

    @staticmethod
    def _convert_data_semantics_to_list(data_semantics):

        named_list = [
            "s" + str(idx) for idx in range(data_semantics.num_data_entries)
        ]
        named_list[data_semantics.x_idx] = "x"
        named_list[data_semantics.y_idx] = "y"
        named_list[data_semantics.z_idx] = "z"
        named_list[data_semantics.r_idx] = "red"
        named_list[data_semantics.g_idx] = "green"
        named_list[data_semantics.b_idx] = "blue"
        return named_list

    @staticmethod
    def parse_point_data_file(ifp, op=None):
        """Parse a point data file.

        Supported file formats are: :code:`.ply`, :code:`.pcd`, :code:`.las`,
        :code:`.laz`, :code:`.asc`, :code:`.pts` and :code:`.csv`.

        Relies on the :code:`pyntcloud`, the :code:`pylas` and/or the
        :code:`lazrs` library to parse the different file formats.
        """

        log_report("INFO", "Parse Point Data File: ...", op)
        # https://pyntcloud.readthedocs.io/en/latest/io.html
        # https://www.cloudcompare.org/doc/wiki/index.php?title=FILE_I/O
        module_spec = importlib.util.find_spec("pyntcloud")
        if module_spec is None:
            log_report(
                "ERROR",
                "Importing this file type requires the pyntcloud library.",
                op,
            )
            assert False
        from pyntcloud import PyntCloud

        assert os.path.isfile(ifp)
        ext = os.path.splitext(ifp)[1].lower()
        if ext in [".asc", ".pts"]:
            sep = " "
            data_semantics = (
                PointDataFileHandler._get_data_semantics_from_ascii(
                    ifp, sep, has_header=True, op=op
                )
            )
            names = PointDataFileHandler._convert_data_semantics_to_list(
                data_semantics
            )
            point_cloud = PyntCloud.from_file(
                ifp, sep=sep, header=0, names=names
            )
            pseudo_color = data_semantics.pseudo_color
        elif ext == ".csv":
            sep = ","
            data_semantics = (
                PointDataFileHandler._get_data_semantics_from_ascii(
                    ifp, sep, has_header=False, op=op
                )
            )
            names = PointDataFileHandler._convert_data_semantics_to_list(
                data_semantics
            )
            point_cloud = PyntCloud.from_file(
                ifp, sep=sep, header=0, names=names
            )
            pseudo_color = data_semantics.pseudo_color
        else:
            pseudo_color = False
            point_cloud = PyntCloud.from_file(ifp)
        xyz_arr = point_cloud.points.loc[:, ["x", "y", "z"]].to_numpy()
        if set(["red", "green", "blue"]).issubset(point_cloud.points.columns):
            color_arr = point_cloud.points.loc[
                :, ["red", "green", "blue"]
            ].to_numpy()
            if pseudo_color:
                color_arr *= 255
        else:
            color_arr = np.ones_like(xyz_arr) * 255
        num_points = xyz_arr.shape[0]
        points = []
        for idx in range(num_points):
            point = Point(
                coord=xyz_arr[idx].astype("float64"),
                color=color_arr[idx].astype("int"),
                id=idx,
                scalars=dict(),
            )
            points.append(point)
        log_report("INFO", f"Number Points {len(points)}", op)
        log_report("INFO", "Parse Point Data File: Done", op)
        return points
kernel/ply_importer/point_data_file_handler.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
import bgl
import gpu
from bpy.app.handlers import persistent
from mathutils import Matrix
from gpu.types import GPUOffScreen
from gpu_extras.batch import batch_for_shader

from kernel.ply_importer.point import Point
from kernel.ply_importer.draw_manager import DrawManager
from kernel.ply_importer.object_utility import add_empty
from kernel.logging_utility import log_report


def _draw_coords_with_color(
    coords,
    colors,
    point_size,
    add_points_to_point_cloud_handle,
    reconstruction_collection=None,
    object_anchor_handle_name="OpenGL Point Cloud",
    op=None,
):

    object_anchor_handle = add_empty(
        object_anchor_handle_name, reconstruction_collection
    )
    if add_points_to_point_cloud_handle:
        object_anchor_handle["particle_coords"] = coords
        object_anchor_handle["particle_colors"] = colors
        object_anchor_handle["point_size"] = point_size
        bpy.context.scene["contains_opengl_point_clouds"] = True

    draw_manager = DrawManager.get_singleton()
    draw_manager.register_points_draw_callback(
        object_anchor_handle, coords, colors, point_size
    )
    return object_anchor_handle


def draw_points(
    points,
    point_size,
    add_points_to_point_cloud_handle,
    reconstruction_collection=None,
    object_anchor_handle_name="OpenGL Point Cloud",
    op=None,
):
    """Draw points using OpenGL."""
    log_report("INFO", "Add particle draw handlers", op)

    coords, colors = Point.split_points(points, normalize_colors=True)
    object_anchor_handle = _draw_coords_with_color(
        coords,
        colors,
        point_size,
        add_points_to_point_cloud_handle,
        reconstruction_collection,
        object_anchor_handle_name,
        op=op,
    )
    return object_anchor_handle


def draw_coords(
    coords,
    color=(0, 0, 255, 1.0),
    point_size=1,
    add_points_to_point_cloud_handle=True,
    reconstruction_collection=None,
    object_anchor_handle_name="OpenGL Coord Point Cloud",
    op=None,
):
    """Draw coordinates using OpenGL."""
    if len(color) == 3:
        color = (color[0], color[1], color[2], 1)
    assert len(color) == 4
    colors = [color for coord in coords]
    object_anchor_handle = _draw_coords_with_color(
        coords,
        colors,
        point_size,
        add_points_to_point_cloud_handle,
        reconstruction_collection,
        object_anchor_handle_name,
        op=op,
    )
    return object_anchor_handle


@persistent
def redraw_points(dummy):
    """Redraw points of the previous Blender session."""

    # This test is very cheap, so it will not cause
    # huge overheads for scenes without point clouds
    if "contains_opengl_point_clouds" in bpy.context.scene:

        log_report(
            "INFO",
            "Checking scene for missing point cloud draw handlers",
            op=None,
        )
        for obj in bpy.data.objects:
            if (
                "particle_coords" in obj
                and "particle_colors" in obj
                and "point_size" in obj
            ):
                coords = obj["particle_coords"]
                colors = obj["particle_colors"]
                point_size = obj["point_size"]

                draw_manager = DrawManager.get_singleton()
                draw_manager.register_points_draw_callback(
                    obj, coords, colors, point_size
                )

        for area in bpy.context.screen.areas:
            if area.type == "VIEW_3D":
                area.tag_redraw()
                break


def render_opengl_image(image_name, cam, coords, colors, point_size):
    """Render the given coordinates with :code:`OpenGL`."""

    render = bpy.context.scene.render

    width = render.resolution_x
    height = render.resolution_y
    # TODO Provide an option to render from the 3D view perspective
    # width = bpy.context.region.width
    # height = bpy.context.region.height

    offscreen = gpu.types.GPUOffScreen(width, height)
    with offscreen.bind():

        bgl.glPointSize(point_size)
        bgl.glEnable(bgl.GL_DEPTH_TEST)
        # bgl.glClear(bgl.GL_COLOR_BUFFER_BIT)
        # bgl.glClear(bgl.GL_DEPTH_BUFFER_BIT)

        view_matrix = cam.matrix_world.inverted()
        projection_matrix = cam.calc_matrix_camera(
            bpy.context.evaluated_depsgraph_get(), x=width, y=height
        )
        perspective_matrix = projection_matrix @ view_matrix

        gpu.matrix.load_matrix(perspective_matrix)
        gpu.matrix.load_projection_matrix(Matrix.Identity(4))

        shader = gpu.shader.from_builtin("3D_FLAT_COLOR")
        shader.bind()
        batch = batch_for_shader(
            shader, "POINTS", {"pos": coords, "color": colors}
        )
        batch.draw(shader)

        buffer = bgl.Buffer(bgl.GL_BYTE, width * height * 4)
        bgl.glReadPixels(
            0, 0, width, height, bgl.GL_RGBA, bgl.GL_UNSIGNED_BYTE, buffer
        )

    offscreen.free()

    image = _create_image_lazy(image_name, width, height)
    _copy_buffer_to_pixel(buffer, image)


def _create_image_lazy(image_name, width, height):
    if image_name not in bpy.data.images:
        image = bpy.data.images.new(image_name, width, height)
    else:
        image = bpy.data.images[image_name]
        if image.size[0] != width or image.size[1] != height:
            image.scale(width, height)
    return image


def _copy_buffer_to_pixel(buffer, image):

    # According to
    #   https://developer.blender.org/D2734
    #   https://docs.blender.org/api/current/gpu.html#copy-offscreen-rendering-result-back-to-ram
    # the buffer protocol is currently not implemented for
    # bgl.Buffer and bpy.types.Image.pixels
    # (this makes the extraction very slow)

    # # from ply_importer.utility.timing_utility import StopWatch
    # Option 1 (faster)
    # sw = StopWatch()
    image.pixels = [v / 255 for v in buffer]
    # log_report('INFO', 'sw.get_elapsed_time(): ' + str(sw.get_elapsed_time()))

    # Option 2 (slower)
    # sw = StopWatch()
    # image.pixels = (np.asarray(buffer, dtype=np.uint8) / 255).tolist()
    # log_report('INFO', 'sw.get_elapsed_time(): ' + str(sw.get_elapsed_time()))

kernel/ply_importer/utility.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
from kernel.kf_pycuda.config import set_config, print_config
from kernel.kf_pycuda.kinect_fusion import KinectFusion
from offline.parse import offlineParam
from kernel.geometry import _pose2Rotation
from tqdm import tqdm
from PIL import Image
import os
import numpy as np
import open3d as o3d


def KinectfusionRecon(
    data_folder, save_folder, prefix_list,
    resX, resY, fx, fy, cx, cy, 
    tsdf_voxel_size, tsdf_trunc_margin, pcd_voxel_size, depth_scale, depth_ignore, 
    DISPLAY, frame_per_display,
    ):
    depth_path = os.path.join(data_folder, "depth")
    color_path = os.path.join(data_folder, "rgb")

    config = set_config(resX, resY, fx, fy, cx, cy, tsdf_voxel_size, tsdf_trunc_margin, pcd_voxel_size)
    kf = KinectFusion(cfg=config)
    print_config(config)


    for idx, prefix in enumerate(tqdm(prefix_list)):
        color_im_path = os.path.join(color_path, prefix + '.png')
        depth_im_path = os.path.join(depth_path, prefix + '.png')
        
        color_im = np.asarray(Image.open(color_im_path))
        depth_im = np.asarray(Image.open(depth_im_path)).astype(np.float32) * depth_scale
        depth_im[depth_im > depth_ignore] = 0

        kf.update(color_im, depth_im)
        # if (idx + 1) % 500 == 0:
        #     kf.save(save_folder)
    kf.save(save_folder, prefix_list)

def poseFusion(
    param_path,
    tsdf_voxel_size, tsdf_trunc_margin, pcd_voxel_size, depth_ignore, MAXFRAME = 200
    ):
    def parsecamfile(param):
        param.camposes = {}
        f = open(os.path.join(param.reconstructionsrc, "campose.txt"))
        lines = f.readlines()
        for l in lines:
            datas = l.split(" ")
            if datas[0].isnumeric():
                param.camposes[datas[-1].split("\n")[0]] = _pose2Rotation([[float(datas[5]), float(datas[6]), float(datas[7])],\
                                                        [float(datas[1]), float(datas[2]), float(datas[3]), float(datas[4])]])

    def applytrans2cam(param):
        scale = param.recon["scale"]
        trans = param.recon["trans"]
        for cam in param.camposes:
            origin_pose = param.camposes[cam]
            origin_pose[:3, 3] = origin_pose[:3, 3] * scale
            # origin_pose = np.linalg.inv(origin_pose).dot(Axis_align)
            if param.camera["inverse_pose"]:
                origin_pose = np.linalg.inv(origin_pose)
            param.camposes[cam] = trans.dot(origin_pose)  
    
    param = offlineParam(param_path)
    parsecamfile(param)
    applytrans2cam(param)
    depth_path = os.path.join(param.datasrc, "depth")
    rgb_path = os.path.join(param.datasrc, "rgb")

    config = set_config(param.camera["resolution"][0], param.camera["resolution"][1], 
                        param.camera["intrinsic"][0, 0], param.camera["intrinsic"][1, 1],
                        param.camera["intrinsic"][0, 2], param.camera["intrinsic"][1, 2], 
                        tsdf_voxel_size = tsdf_voxel_size, tsdf_trunc_margin = tsdf_trunc_margin, pcd_voxel_size = pcd_voxel_size)
    depth_ignore = depth_ignore
    kf = KinectFusion(cfg=config)
    frame_keys = list(param.camposes.keys())
    ratio = len(frame_keys)/MAXFRAME
    
    # for frame in tqdm(param.camposes):
    #     rgb = np.array(Image.open(os.path.join(rgb_path, frame)))
    #     depth = np.array(Image.open(os.path.join(depth_path, frame))).astype(np.float32) * param.data['depth_scale']

    #     depth[depth > depth_ignore] = 0
    #     kf.update(rgb, depth, param.camposes[frame])
    for frame_id in tqdm(range(min(len(param.camposes), MAXFRAME))):
        frame = frame_keys[int(frame_id)] if ratio <= 1 else frame_keys[int(frame_id * ratio)]
        rgb = np.array(Image.open(os.path.join(rgb_path, frame)))
        depth = np.array(Image.open(os.path.join(depth_path, frame))).astype(np.float32) * param.data['depth_scale']

        depth[depth > depth_ignore] = 0
        kf.update(rgb, depth, param.camposes[frame])
    surface = kf.tsdf_volume.get_surface_cloud_marching_cubes(voxel_size=pcd_voxel_size)
    o3d.io.write_point_cloud(os.path.join(param.reconstructionsrc, 'depthfused.ply'), surface)
    



kernel/reconstruction.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
import numpy as np
from PIL import Image
import os
import itertools
from kernel.logging_utility import log_report

def render_fromblender():
    pass

def save_img(img_render, img_origin, img_segment, img_segment_inverse, name):
    renderim = Image.fromarray(img_render)
    renderim.save(os.path.join("tmp", name + "_render.png"))
    originim = Image.fromarray(img_origin)
    originim.save(os.path.join("tmp", name + "_origin.png"))
    segmentim = Image.fromarray(img_segment)
    segmentim.save(os.path.join("tmp", name + "_segment.png"))
    segmentim_inverse = Image.fromarray(img_segment_inverse)
    segmentim_inverse.save(os.path.join("tmp", name + "_segmentinverse.png"))


def metric_test():
    dx_list = [-0.002, -0.001, 0, 0.001, 0.002]
    dy_list = [-0.002, -0.001, 0, 0.001, 0.002]
    dz_list = [-0.002, -0.001, 0, 0.001, 0.002]
    dt_list = [-0.02, -0.01, 0, 0.01, 0.02]
    x, y, z = bpy.data.objects["mustard_bottle"].location
    _, tx, ty, _ = bpy.data.objects["mustard_bottle"].rotation_quaternion
    t = np.arctan2(tx, ty)

    current_object = bpy.context.object.name
    pair_image_name = bpy.data.objects[current_object]["image"].name
    img_rgb = np.array(bpy.data.images[pair_image_name].pixels).reshape(bpy.context.scene.configuration.resY, bpy.context.scene.configuration.resX, 4)
    
    for dx, dy, dz, dt in list(itertools.product(dx_list, dy_list, dz_list, dt_list)):
        log_report("INFO", f"On view {pair_image_name} : dx = {dx:.3f}, dy = {dy:.3f}, dz = {dz:.3f}, dt = {dt:.3f}", None)        
        bpy.data.objects["mustard_bottle"].location = [x + dx, y + dy, z + dz]
        t_test = t + dt
        bpy.data.objects["mustard_bottle"].rotation_quaternion = [0, np.sin(t_test), np.cos(t_test), 0]
        
        tmpimgname_render = './tmp/tmp.png'
        bpy.context.scene.render.filepath = tmpimgname_render
        bpy.context.scene.render.image_settings.file_format='PNG' 
        bpy.context.scene.render.image_settings.color_mode ='RGBA'
        bpy.context.scene.render.film_transparent = True
        bpy.context.scene.render.resolution_x = bpy.context.scene.configuration.resX
        bpy.context.scene.render.resolution_y = bpy.context.scene.configuration.resY
        res = bpy.ops.render.render(write_still = True)
        img_render = np.array(Image.open(tmpimgname_render))
        os.system('rm ' + tmpimgname_render)
        foreground_mask = img_render[:, :, 3] == 0
        img_origin = (img_rgb[::-1, :, :3]*255).astype(np.uint8)
        img_segment = img_origin.copy()
        img_segment[foreground_mask, :3] = 0
        img_segment_inverse = img_origin.copy()
        img_segment_inverse[~foreground_mask, :3] = 0
        save_img(img_render[:, :, :3], img_origin, img_segment, img_segment_inverse,pair_image_name + f"dx{dx:.3f}dy{dy:.3f}dz{dz:.3f}dt{dt:.3f}")
kernel/render.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
from scipy.spatial.transform import Rotation as R
import numpy as np
from PIL import Image
import os
from kernel.logging_utility import log_report

def _parseImagesFile(filename, Camera_dict, PointsDict):
    with open(filename) as f:
        lines = f.readlines()
        for line in lines:
            if line.startswith("#"):
                pass
            else:
                words = line.split(" ")
                if '\n' in words:
                    words = words[:-1]
                if words[0].isnumeric() and words[-1].endswith(".png\n"): 
                    camera_idx = int(words[0])
                    Camera_dict[camera_idx] = {}
                    Camera_dict[camera_idx]["name"] = words[-1].split("\n")[0]
                    r = R.from_quat([float(words[2]), float(words[3]), float(words[4]), float(words[1])])
                    T = np.hstack((r.as_matrix(), np.array([[float(words[5])], [float(words[6])], [float(words[7])]])))
                    T = np.vstack((T, np.array([[0, 0, 0, 1]])))
                    Camera_dict[camera_idx]["pose"] = T
                    if "points" not in Camera_dict[camera_idx]: Camera_dict[camera_idx]["points"] = []

                else:
                    for i in range(0, len(words), 3):
                        if int(words[i + 2]) != -1:
                            pointidx = int(words[i + 2])
                            if pointidx not in PointsDict:
                                PointsDict[pointidx] = {}
                                PointsDict[pointidx]["camera"] = {}
                            PointsDict[pointidx]["camera"][camera_idx] = {}
                            PointsDict[pointidx]["camera"][camera_idx]["px"] = float(words[i])
                            PointsDict[pointidx]["camera"][camera_idx]["py"] = float(words[i + 1])
                            Camera_dict[camera_idx]["points"].append(pointidx)


def _parsePoints3D(filename, PointsDict):
    with open(filename) as f:
        lines = f.readlines()
        for line in lines:
            words = line.split(" ")
            if not words[0] == '#':
                PointsDict[int(words[0])]["location"] = np.array([float(words[1]), float(words[2]), float(words[3])])
                PointsDict[int(words[0])]["error"] = float(words[7])

def _depthInterpolation(depth, px, py):
    px_0 = int(np.floor(px))
    py_0 = int(np.floor(py))
    depth_inter = depth[px_0, py_0]
    return depth_inter

def _scaleFordepth(depth, index, intrinsic, Camera_dict, PointsDict, PointsDepth, POSE_INVERSE = True):
    if POSE_INVERSE:
        cam_pose = np.linalg.inv(Camera_dict[index]["pose"])
    else:
        cam_pose = Camera_dict[index]["pose"]
    for point_idx in Camera_dict[index]["points"]:
        point = PointsDict[point_idx]
        location = point['location']
        recon_error = point['error']
        pixel = point['camera'][index]
        depth_real = _depthInterpolation(depth, pixel['py'], pixel['px'])
        depth_recon = np.linalg.norm(location - cam_pose[:3, 3])
        if not point_idx in PointsDepth:
            PointsDepth[point_idx] = {}
        PointsDepth[point_idx][index] = {}
        PointsDepth[point_idx][index]['recon'] = depth_recon
        PointsDepth[point_idx][index]['real'] = depth_real
        PointsDepth[point_idx][index]['scale'] = depth_real/depth_recon

def _calculateDepth(THRESHOLD, NUM_THRESHOLD, PointsDepth):
    number = 0
    total_scale = 0
    for point in PointsDepth:
        scale_point = []
        number_point = 0
        for frame in PointsDepth[point]:
            scale = PointsDepth[point][frame]['scale']
            if scale == 0:
                pass
            else:
                scale_point.append(scale)
                number_point += 1
        if scale_point != []:
            std_point = np.std(np.array(scale_point))
            si = std_point/np.mean(np.array(scale_point))
            if si < THRESHOLD and si > 0 and number_point > NUM_THRESHOLD:
                number += number_point
                total_scale += np.sum(np.array(scale_point))
    print("remain {0:.2f}% vaild points".format(number * 100/len(PointsDepth)))
    if number == 0:
        return 0
    else:
        return total_scale/number



if __name__ == "__main__":
    Camera_dict = {}
    PointsDict = {}
    PointsDepth = {}
    imagefilepath = "/home/huijie/Desktop/testdataset/recon/images.txt"
    points3Dpath = "/home/huijie/Desktop/testdataset/recon/points3D.txt"
    datapath = "/home/huijie/Desktop/testdataset/data"
    intrinsic = np.array([
        [904.572, 0, 635.981 ],
        [0, 905.295, 353.06],
        [0, 0, 1]
    ])
    _parseImagesFile(imagefilepath, Camera_dict, PointsDict)
    _parsePoints3D(points3Dpath, PointsDict)

    for camera_idx in Camera_dict:
        rgb_name = Camera_dict[camera_idx]["name"]
        with Image.open(os.path.join(datapath, "depth", rgb_name)) as im:
            depth = np.array(im) * 0.00025
            _scaleFordepth(depth, camera_idx, intrinsic, Camera_dict, PointsDict, PointsDepth)
    
    scale = _calculateDepth(THRESHOLD = 0.005, NUM_THRESHOLD = 3, PointsDepth = PointsDepth)
    pass
kernel/scale.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import numpy as np
import math
import os

def _transstring2trans(transstring):
    vectorstring = transstring.split(';')[:-1]
    vectorstring = [v.split(',') for v in vectorstring]
    return np.array(vectorstring).astype(np.float32)

def _trans2transstring(trans):
    trans_u1 = trans.astype(dtype='<U8')
    transstring = ''
    for i in range(trans_u1.shape[0]):
        for j in range(trans_u1.shape[1]):
            transstring+=trans_u1[i, j]
            if j == trans_u1.shape[1] - 1:
                transstring+=';'
            else:
                transstring+=','
    return transstring

def _select_sample_files(files, sample_rate):
    ii = 0
    selected_files = []
    while ii < len(files):
        selected_files.append(files[math.floor(ii)])
        ii += 1/sample_rate
    if files[-1] not in selected_files:
        selected_files.append(files[-1])
    return selected_files

def _generate_image_list(path, files):
    f= open(os.path.join(path, "image-list.txt"),"w+")
    for file in files:
        f.write(file + "\n")

def _parse_camfile(camera_rgb_file):
    camera_lines = []
    file = open(camera_rgb_file, "r")
    lines = file.read().split("\n")
    for l in lines:
        data = l.split(" ")
        if data[0].isnumeric():
            camera_lines.append(l)
    return camera_lines

kernel/utility.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import sys
import os
import json
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from parse import offlineParam
#from render import offlineRender
from render_custom import offlineRender
from offlineRecon import offlineRecon

import sys
import os
import json
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import pyrender
import trimesh
import numpy as np
from kernel.geometry import _pose2Rotation, _rotation2Pose
import cv2
from PIL import Image
import time
os.environ['PYOPENGL_PLATFORM'] = 'egl'

if __name__ == "__main__":
    config_path = sys.argv[1]
    output_dir = sys.argv[2]
    data_format = sys.argv[3]
    param = offlineParam(config_path)
    interpolation_type = "all"
    #offlineRecon(param, interpolation_type)
    #offlineRender(param, output_dir, interpolation_type, pkg_type=data_format)  
    offlineRender(param, output_dir)  
    

offline/main.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import sys
import os
import json
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import pyrender
import trimesh
import numpy as np
from kernel.geometry import _pose2Rotation, _rotation2Pose
import cv2
from PIL import Image
import time
os.environ['PYOPENGL_PLATFORM'] = 'egl'

if __name__ == "__main__":
    config_path = sys.argv[1]
    output_dir = sys.argv[2]

    f = open(config_path)
    config = json.load(f)

    #print(config)
    
    if str(config["camera"]["name"][0]) == "R":
    
        r = pyrender.OffscreenRenderer(640, 480)
        scene = pyrender.Scene()
        cam_intrinsic = np.array([
                    [603.77647837,   0.,         329.25940456],
                    [  0.,         604.63162524, 246.48479807],
                    [  0. ,          0. ,          1.        ]], dtype=np.float64)      

        cam = pyrender.camera.IntrinsicsCamera(cam_intrinsic[0, 0],
	                               cam_intrinsic[1, 1], 
	                               cam_intrinsic[0, 2], 
	                               cam_intrinsic[1, 2], 
	                               znear=0.05, zfar=100.0, name=None)
    
    if str(config["camera"]["name"][0]) == "L":
    
        r = pyrender.OffscreenRenderer(640, 480)
        scene = pyrender.Scene()
        cam_intrinsic = np.array([
                    [609.0447980649861, 0.0, 326.84169183471016],
                    [0.0, 609.546006727566, 241.56572826467618],
                    [0, 0, 1]], dtype=np.float64)        

        cam = pyrender.camera.IntrinsicsCamera(cam_intrinsic[0, 0],
	                               cam_intrinsic[1, 1], 
	                               cam_intrinsic[0, 2], 
	                               cam_intrinsic[1, 2], 
	                               znear=0.05, zfar=100.0, name=None)            

    if str(config["camera"]["name"][0]) == "T":
    
        r = pyrender.OffscreenRenderer(640, 512)
        scene = pyrender.Scene()
        cam_intrinsic = np.array([
                    [441.62619164519106, 0.0, 318.41319462358956],
                    [0.0, 441.46556074040666, 258.6163562473],
                    [0, 0, 1]], dtype=np.float64)      

        cam = pyrender.camera.IntrinsicsCamera(cam_intrinsic[0, 0],
	                               cam_intrinsic[1, 1], 
	                               cam_intrinsic[0, 2], 
	                               cam_intrinsic[1, 2], 
	                               znear=0.05, zfar=100.0, name=None)   


    #r = pyrender.OffscreenRenderer(config["camera"]["resolution"][0], config["camera"]["resolution"][1])
    #scene = pyrender.Scene()
    #cam_intrinsic = np.array(config["camera"]["intrinsic"])

    #cam = pyrender.camera.IntrinsicsCamera(cam_intrinsic[0, 0],
    #                               cam_intrinsic[1, 1], 
    #                               cam_intrinsic[0, 2], 
    #                               cam_intrinsic[1, 2], 
    #                               znear=0.05, zfar=100.0, name=None)

    axis_align = np.array([[1, 0, 0, 0],
                           [0, -1, 0, 0],
                           [0, 0, -1, 0],
                           [0, 0, 0, 1]]
                                )
       
    cam_pose = [list(config["camera"]["translation"]), list(config["camera"]["quaternion"])]
    cam_T = _pose2Rotation(cam_pose)
    #cam_T = cam_T.dot(axis_align)

    mesh_pose = [list(config["object"]["translation"]), list(config["object"]["quaternion"])]
    mesh_T = _pose2Rotation(mesh_pose)

    node_cam = pyrender.Node(camera=cam, matrix=cam_T)
    scene.add_node(node_cam)


    obj_mesh = trimesh.load(config["object"]["path"], force = 'mesh')
    mesh = pyrender.Mesh.from_trimesh(obj_mesh, smooth = False)
    mesh_node = pyrender.Node(mesh=mesh, matrix=mesh_T)
    scene.add_node(mesh_node)
    mesh_node.mesh.is_visible = True


    flags = pyrender.constants.RenderFlags.DEPTH_ONLY
    depth = r.render(scene, flags = flags)
    segment = depth == 0    
    
    #print(config["camera"]["name"])
    #print(config["camera"]["name"][0])
    #print(config["camera"]["name"][1:])
    rgb_img_path = config["environment"]["datasrc"] + str(config["camera"]["name"][0]) + "/undistort/" + str(config["camera"]["name"][1:]) + ".png"
    rgb_img = np.array(Image.open(rgb_img_path))
    rgb_img[segment] = 0

    output_path = output_dir + "/" + config["object"]["name"] + "/mask/"
    os.makedirs(output_path, exist_ok=True)  

    img = Image.fromarray(rgb_img)
    img.save(output_path + str(config["camera"]["name"]) + ".png")

    






offline/main_mask.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import os
import sys
import pyglet
os.environ["PYOPENGL_PLATFORM"] = "egl"
pyglet.options['shadow_window'] = False

import numpy as np
import time
import glob
import random
import sys
import cv2
import json

import torch 
import torch.nn as nn
from torchvision.ops import RoIPool
import pytorch3d
from pytorch3d.io import load_objs_as_meshes, load_obj, load_ply
from pytorch3d.structures import Meshes
from pytorch3d.vis.plotly_vis import AxisArgs, plot_batch_individually, plot_scene
from pytorch3d.vis.texture_vis import texturesuv_image_matplotlib
from pytorch3d.renderer import (
    look_at_view_transform,
    FoVPerspectiveCameras, 
    PointLights, 
    DirectionalLights, 
    Materials, 
    RasterizationSettings, 
    MeshRenderer, 
    MeshRasterizer,  
    SoftPhongShader,
    Textures,
    TexturesUV,
    TexturesVertex,
    PerspectiveCameras,
    MeshRendererWithFragments,
    SoftSilhouetteShader,
    TexturesAtlas,
    TexturesUV,
    BlendParams
)

class Model(nn.Module):
    def __init__(self, mesh, AA_World2Object, t_World2Object):  
        super().__init__()
        self.mesh = mesh
        self.device = mesh.device
                
        self.AA_World2Object = nn.Parameter(AA_World2Object.to(self.device))
        #self.AA_World2Object = AA_World2Object.to(self.device)
        self.t_World2Object = nn.Parameter(t_World2Object.to(self.device))
        #self.t_World2Object = t_World2Object.to(self.device)
               
        self.keys = ["R", "L", "T"]       
               
    def forward(self, renderers, gt_masks, R_World2Cams, t_World2Cams):
        
        loss_all = 0
        for k in renderers.keys():
            for i in range(len(gt_masks[k])):
              
                R_World2Object = pytorch3d.transforms.axis_angle_to_matrix(self.AA_World2Object)  
                R_Cam2Object = torch.bmm(R_World2Cams[k][i].inverse().view(1, 3, 3), R_World2Object.view(1, 3, 3))
        
                R_Cam2Object_render = R_Cam2Object.permute(0, 2, 1)
                R_Cam2Object_render[:, :, :2] *= -1    
                
                r1t2 = torch.bmm(R_World2Cams[k][i].inverse().view(1, 3, 3), self.t_World2Object.view(1, 3, 1))
                r1t1 = torch.bmm(R_World2Cams[k][i].inverse().view(1, 3, 3), t_World2Cams[k][i].view(1, 3, 1))
                t_Cam2Object_render = r1t2-r1t1
                t_Cam2Object_render[:, :2] *= -1   
                
                obj_mask = renderers[k](meshes_world=self.mesh, R=R_Cam2Object_render.view(1, 3, 3), T=t_Cam2Object_render.view(1, 3)*torch.tensor(1000.))                     
                
                pred_flat = torch.flatten(obj_mask[:, :, :, 3]/(obj_mask[:, :, :, 3] + 1e-4), start_dim=1)  
                gt_flat = torch.flatten(torch.from_numpy(gt_masks[k][i]).to(self.device).view(1, gt_masks[k][i].shape[0], gt_masks[k][i].shape[1]), start_dim=1)   
                loss = nn.MSELoss(reduction='none')(pred_flat, gt_flat)
                values, indices = torch.topk(loss, int(loss.shape[1]), dim=1)       
                loss = torch.sum(values)        
                loss_all += loss       
   
        return loss_all        



if __name__ == "__main__":

    config_path = sys.argv[1]
    output_dir = sys.argv[2]

    f = open(config_path)
    config = json.load(f)
    
    #print(config)

    device = "cuda:0"
    Deg2Rad = torch.tensor(np.pi / 180.).to(device)
     
    verts, faces, aux = load_obj(config["object"]["path"])
    verts = verts * 1000.                         
               
    tex_maps = aux.texture_images
    if tex_maps is not None and len(tex_maps) > 0:
        verts_uvs = aux.verts_uvs.to(device) 
        faces_uvs = faces.textures_idx.to(device)
        image = list(tex_maps.values())[0].to(device)[None]
        tex = TexturesUV(verts_uvs=[verts_uvs], faces_uvs=[faces_uvs], maps=image)
        obj_mesh = Meshes(verts=[verts.to(device)], faces=[faces.verts_idx.to(device)], textures=tex)
    else:
        obj_mesh = Meshes(verts=[verts.to(device)], faces=[faces.verts_idx.to(device)])    

    ####################################################################################################################################################################################
        
    blend_params = BlendParams(sigma=1e-4, gamma=1e-4)

    raster_settings_R = RasterizationSettings(image_size=[480, 640], blur_radius=0.000, faces_per_pixel=1, bin_size=None)   
    cameras_R = PerspectiveCameras(
                            focal_length=np.array([603.77647837, 604.63162524]).reshape((1, 2)), 
                            principal_point=np.array([329.25940456, 246.48479807]).reshape((1, 2)), 
                            image_size=np.array([480, 640]).reshape((1, 2)), 
                            device=device,
                            in_ndc=False                            
                            )                        
    renderer_R = MeshRenderer(
        rasterizer=MeshRasterizer(
            cameras=cameras_R,
            raster_settings=raster_settings_R
        ),
        shader=SoftSilhouetteShader(blend_params=blend_params)
    )
    cameras_R.zfar = 10000
    cameras_R.znear = 1  
        
    raster_settings_L = RasterizationSettings(image_size=[480, 640], blur_radius=0.000, faces_per_pixel=1, bin_size=None)       
    cameras_L = PerspectiveCameras(
                            focal_length=np.array([609.0447980649861, 609.546006727566]).reshape((1, 2)), 
                            principal_point=np.array([326.84169183471016, 241.56572826467618]).reshape((1, 2)), 
                            image_size=np.array([480, 640]).reshape((1, 2)), 
                            device=device,
                            in_ndc=False                            
                            )                        
    renderer_L = MeshRenderer(
        rasterizer=MeshRasterizer(
            cameras=cameras_L,
            raster_settings=raster_settings_L
        ),
        shader=SoftSilhouetteShader(blend_params=blend_params)
    )
    cameras_L.zfar = 10000
    cameras_L.znear = 1      
        

    raster_settings_T = RasterizationSettings(image_size=[512, 640], blur_radius=0.000, faces_per_pixel=1, bin_size=None)   
    cameras_T = PerspectiveCameras(
                            focal_length=np.array([441.62619164519106, 441.46556074040666]).reshape((1, 2)), 
                            principal_point=np.array([318.41319462358956, 258.6163562473]).reshape((1, 2)), 
                            image_size=np.array([512, 640]).reshape((1, 2)), 
                            device=device,
                            in_ndc=False                            
                            )                        
    renderer_T = MeshRenderer(
        rasterizer=MeshRasterizer(
            cameras=cameras_T,
            raster_settings=raster_settings_T
        ),
        shader=SoftSilhouetteShader(blend_params=blend_params)
    )
    cameras_T.zfar = 10000
    cameras_T.znear = 1          
      
    renderers = {"R":renderer_R, "L":renderer_L, "T":renderer_T}
    
    ####################################################################################################################################################################################    	   
    
    file_list_idx = {"R":list(), "L":list(), "T":list()}
    file_list = os.listdir(output_dir + "/" + config["object"]["name"] + "/mask/") 
    n_pose = len(file_list)
    for i in range(len(file_list)):
    
        cam_idx = file_list[i].split(".")[0][0]    
        img_idx = file_list[i].split(".")[0][1:]        
        file_list_idx[str(cam_idx)].append(int(img_idx))


    ####################################################################################################################################################################################    	   

    Q_World2Cams = {"R":torch.zeros([len(file_list_idx["R"]), 4], dtype=torch.float32).to(device),
                    "L":torch.zeros([len(file_list_idx["L"]), 4], dtype=torch.float32).to(device),
                    "T":torch.zeros([len(file_list_idx["T"]), 4], dtype=torch.float32).to(device)}
    AA_World2Cams = {"R":torch.zeros([len(file_list_idx["R"]), 3], dtype=torch.float32).to(device),
                    "L":torch.zeros([len(file_list_idx["L"]), 3], dtype=torch.float32).to(device),
                    "T":torch.zeros([len(file_list_idx["T"]), 3], dtype=torch.float32).to(device)} 
    R_World2Cams = {"R":torch.zeros([len(file_list_idx["R"]), 3, 3], dtype=torch.float32).to(device),
                    "L":torch.zeros([len(file_list_idx["L"]), 3, 3], dtype=torch.float32).to(device),
                    "T":torch.zeros([len(file_list_idx["T"]), 3, 3], dtype=torch.float32).to(device)}                          
    t_World2Cams = {"R":torch.zeros([len(file_list_idx["R"]), 3], dtype=torch.float32).to(device),
                    "L":torch.zeros([len(file_list_idx["L"]), 3], dtype=torch.float32).to(device),
                    "T":torch.zeros([len(file_list_idx["T"]), 3], dtype=torch.float32).to(device)}      
    T_World2Cams = {"R":torch.zeros([len(file_list_idx["R"]), 4, 4], dtype=torch.float32).to(device),
                    "L":torch.zeros([len(file_list_idx["L"]), 4, 4], dtype=torch.float32).to(device),
                    "T":torch.zeros([len(file_list_idx["T"]), 4, 4], dtype=torch.float32).to(device)}                                                       

    for k in file_list_idx.keys():
    
        f = open(config["environment"]["reconstructionsrc"] + "/campose_" + k + ".txt", 'r')
        lines = f.readlines() 
        f.close()    
    
        for i in range(len(file_list_idx[k])):
        
            token = lines[file_list_idx[k][i] + 1].split(" ")
        
            Q_World2Cams[k][i, 0] = float(token[1])
            Q_World2Cams[k][i, 1] = float(token[2])
            Q_World2Cams[k][i, 2] = float(token[3])
            Q_World2Cams[k][i, 3] = float(token[4])        
            R_World2Cams[k][i]  = pytorch3d.transforms.quaternion_to_matrix(Q_World2Cams[k][i]).to(device)
            AA_World2Cams[k][i] = pytorch3d.transforms.quaternion_to_axis_angle(Q_World2Cams[k][i]).to(device)
        
            t_World2Cams[k][i, 0] = float(token[5])# * 1000.
            t_World2Cams[k][i, 1] = float(token[6])# * 1000.
            t_World2Cams[k][i, 2] = float(token[7])# * 1000.
        
            T_World2Cams[k][i, :3, :4]  = torch.cat((R_World2Cams[k][i].reshape(-1, 3, 3), t_World2Cams[k][i].reshape(-1, 3, 1)), 2)
            T_World2Cams[k][i, 3, 3] = 1.
                     
    ####################################################################################################################################################################################    	   

    gt_masks = {"R":np.zeros([len(file_list_idx["R"]), 480, 640], dtype=np.float32),
                "L":np.zeros([len(file_list_idx["L"]), 480, 640], dtype=np.float32),
                "T":np.zeros([len(file_list_idx["T"]), 512, 640], dtype=np.float32)
               }
               
    for k in file_list_idx.keys():               
        for i in range(len(file_list_idx[k])):
            
            gt_mask = cv2.imread(output_dir + "/" + config["object"]["name"] + "/mask/" + k + "%06d.png"%file_list_idx[k][i])
            gt_mask = cv2.cvtColor(gt_mask, cv2.COLOR_BGR2GRAY)  
            
            #cv2.imshow(k + "%06d.png"%file_list_idx[k][i], gt_mask)
            #cv2.waitKey(0)
            
            gt_masks[k][i] = (gt_mask != 0)# * 0.5
            
    ####################################################################################################################################################################################        
        
    Q_World2Object = torch.tensor([config["object"]["quaternion"][0],  config["object"]["quaternion"][1], config["object"]["quaternion"][2], config["object"]["quaternion"][3]]).reshape(-1, 4).to(device)  # w x y z
    R_World2Object = pytorch3d.transforms.quaternion_to_matrix(Q_World2Object).to(device)
    AA_World2Object = pytorch3d.transforms.quaternion_to_axis_angle(Q_World2Object).to(device)
    t_World2Object = torch.tensor([config["object"]["translation"][0],  config["object"]["translation"][1], config["object"]["translation"][2]]).reshape(-1, 3).to(device) # x y z
    T_World2Object = torch.cat((R_World2Object, t_World2Object.reshape(-1, 3, 1).to(device)), 2)
    T_World2Object = torch.cat((T_World2Object, torch.tensor([0., 0., 0., 1.]).reshape(-1, 1, 4).to(device)), 1)        
    
    #for k in file_list_idx.keys(): 
    #    for i in range(len(file_list_idx[k])):
    #        print(file_list_idx[k][i])
        
    #    print()    
    #    print(Q_World2Cams[k])
        
    #    print()    
    #    print(R_World2Cams[k])
        
    #    print()    
    #    print(AA_World2Cams[k])
        
    #    print()    
    #    print(t_World2Cams[k])
        
    #    print()    
    #    print(T_World2Cams[k])        
                             

    epoch = 800
    model = Model(obj_mesh, AA_World2Object, t_World2Object).to(device)    
    optimizer = torch.optim.RMSprop(model.parameters(), lr=float(config["optimization"]["translation_learning_rate"]))
    optimizer_r = torch.optim.RMSprop(model.parameters(), lr=float(config["optimization"]["rotation_learning_rate"]))
    #optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["optimization"]["translation_learning_rate"]))
    #optimizer_r = torch.optim.AdamW(model.parameters(), lr=float(config["optimization"]["rotation_learning_rate"]))
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[500, 700], gamma=0.5)
    scheduler_r = torch.optim.lr_scheduler.MultiStepLR(optimizer_r, milestones=[500, 700], gamma=0.5)
    
    smallest_loss = 1000000000
    optimal_t_World2Object = 0
    optimal_AA_World2Object = 0
    
    for i in range(epoch):

        model.AA_World2Object.requires_grad = False
        model.t_World2Object.requires_grad = True        
        for kkk in range(1):
            optimizer.zero_grad()        
            loss = model(renderers, gt_masks, R_World2Cams, t_World2Cams)
            loss.backward()
            optimizer.step()   
            scheduler.step()                               
        
            
        
        model.AA_World2Object.requires_grad = True
        model.t_World2Object.requires_grad = False
        for kkk in range(1):
            optimizer.zero_grad()        
            loss = model(renderers, gt_masks, R_World2Cams, t_World2Cams)
            loss.backward()
            optimizer_r.step()   
            scheduler_r.step()

        
        loss = model(renderers, gt_masks, R_World2Cams, t_World2Cams)
        if loss < smallest_loss:
            smallest_loss = loss
            optimal_t_World2Object = model.t_World2Object.clone().detach()
            optimal_AA_World2Object = model.AA_World2Object.clone().detach()
            
            for k in renderers.keys():
                for j in range(len(gt_masks[k])):
              
                    R_World2Object = pytorch3d.transforms.axis_angle_to_matrix(model.AA_World2Object)  
                    R_Cam2Object = torch.bmm(R_World2Cams[k][j].inverse().view(1, 3, 3), R_World2Object.view(1, 3, 3))
        
                    R_Cam2Object_render = R_Cam2Object.permute(0, 2, 1)
                    R_Cam2Object_render[:, :, :2] *= -1    
                
                    r1t2 = torch.bmm(R_World2Cams[k][j].inverse().view(1, 3, 3), model.t_World2Object.view(1, 3, 1))
                    r1t1 = torch.bmm(R_World2Cams[k][j].inverse().view(1, 3, 3), t_World2Cams[k][j].view(1, 3, 1))
                    t_Cam2Object_render = r1t2-r1t1
                    t_Cam2Object_render[:, :2] *= -1   
                
                    obj_mask = renderers[k](meshes_world=obj_mesh, R=R_Cam2Object_render.view(1, 3, 3), T=t_Cam2Object_render.view(1, 3)*torch.tensor(1000.))
                                        
                    cv2.imshow(k + "%06d"%file_list_idx[k][j], cv2.addWeighted(gt_masks[k][j], 0.5, (obj_mask[0, :, :, 3]/(obj_mask[0, :, :, 3] + 1e-4)).data.cpu().numpy(), 0.5, 0.0))
                    cv2.waitKey(30)
            
        
        if loss < float(config["optimization"]["stop_criteria"] * n_pose):
            break        


         
        
        
        print()
        print("t_lr: ", optimizer.param_groups[0]['lr'], " ", "r_lr: ", optimizer_r.param_groups[0]['lr'])
        print(i, "loss:", loss.item())    
        print("t:", model.t_World2Object.data.cpu().numpy()*1000.)
        print("R:", model.AA_World2Object.data.cpu().numpy()*(180/np.pi))
    
        #if i%10 == 0:
            
        #    for k in renderers.keys():
        #        for j in range(len(gt_masks[k])):
              
        #            R_World2Object = pytorch3d.transforms.axis_angle_to_matrix(model.AA_World2Object)  
        #            R_Cam2Object = torch.bmm(R_World2Cams[k][j].inverse().view(1, 3, 3), R_World2Object.view(1, 3, 3))
        
        #            R_Cam2Object_render = R_Cam2Object.permute(0, 2, 1)
        #            R_Cam2Object_render[:, :, :2] *= -1    
                
        #            r1t2 = torch.bmm(R_World2Cams[k][j].inverse().view(1, 3, 3), model.t_World2Object.view(1, 3, 1))
        #            r1t1 = torch.bmm(R_World2Cams[k][j].inverse().view(1, 3, 3), t_World2Cams[k][j].view(1, 3, 1))
        #            t_Cam2Object_render = r1t2-r1t1
        #            t_Cam2Object_render[:, :2] *= -1   
                
        #            obj_mask = renderers[k](meshes_world=obj_mesh, R=R_Cam2Object_render.view(1, 3, 3), T=t_Cam2Object_render.view(1, 3)*torch.tensor(1000.))
                                        
        #            cv2.imshow(k + "%06d"%file_list_idx[k][j], cv2.addWeighted(gt_masks[k][j], 0.5, (obj_mask[0, :, :, 3]/(obj_mask[0, :, :, 3] + 1e-4)).data.cpu().numpy(), 0.5, 0.0))
        #            cv2.waitKey(1)
    
    
    print("============================================================")
    print("Optimization End")
    print("loss:", smallest_loss.item())   
    print("t:", optimal_t_World2Object.data.cpu().numpy()*1000.)
    print("R:", optimal_AA_World2Object.data.cpu().numpy()*(180/np.pi))
    print("============================================================")
    
    result_t = optimal_t_World2Object.data.cpu().numpy().reshape(3)  
    result_q = pytorch3d.transforms.axis_angle_to_quaternion(optimal_AA_World2Object).data.cpu().numpy().reshape(4)                      
    output_dict = {
            'projectname': config["projectname"],
            'object':{
                "name":config['object']['name'],
                "path":config['object']['path'],
                "translation":[float(result_t[0]), float(result_t[1]), float(result_t[2])],
                "quaternion":[float(result_q[0]), float(result_q[1]), float(result_q[2]), float(result_q[3])],
            }
        }

    with open("/tmp/progresslabeller_mask_result.json", "w") as f:
        json.dump(output_dict, f, indent = True)    
	                         
    
    
    


offline/mask_opt.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import os
import numpy as np
from tqdm import tqdm
from PIL import Image
from kernel.geometry import _pose2Rotation, _rotation2Pose
from kernel.utility import _select_sample_files

class offlineRecon:
    def __init__(self, param, interpolation_type = "KF_forward") -> None:
        print("Start offline reconstruction interpolation")
        self.param = param
        self.datasrc = self.param.datasrc
        self.reconstructionsrc = self.param.reconstructionsrc
        self.CAM_INVERSE = self.param.camera["inverse_pose"]
        self.wholemap = {} ## mapping the keyframe pair to frame under estimation
        self.wholecam = {} ## mapping all frames to their poses
        self.keyposes = {} ## mapping all key frames to their poses
        self._parsecamfile()
        self._applytrans2cam()
        self._parsewholeimg()
        self.rgb_path = os.path.join(self.datasrc, "rgb_undistort")
        self.depth_path = os.path.join(self.datasrc, "depth")
        self.depth_scale = self.param.data['depth_scale']
        self._interpolation(type = interpolation_type)
        self._savecampose("campose_all_{0}.txt".format(interpolation_type))
    
    def _parsewholeimg(self):
        imagefiles = os.listdir(os.path.join(self.datasrc, "rgb_undistort"))
        imagefiles.sort()
        keyframe = sorted(self.keyposes.keys())
        index_whole = 0
        index_key = 0
        key_pair = None
        while(index_whole < len(imagefiles)):
            current_image_name = imagefiles[index_whole]
            if (current_image_name in keyframe) and (index_key + 1 < len(keyframe)):
                key_pair = (keyframe[index_key], keyframe[index_key+1])
                self.wholemap[key_pair] = []
                index_key += 1
            elif (current_image_name in keyframe) and (index_key + 1 >= len(keyframe)):
                pass
            elif (current_image_name not in keyframe) and key_pair is not None:
                self.wholemap[key_pair].append(current_image_name)
            else:
                pass
            index_whole += 1
        for img in imagefiles:
            self.wholecam[img] = np.empty((0, 0))
            if img in self.keyposes:
                self.wholecam[img] = self.keyposes[img]

    def _parsecamfile(self):
        f = open(os.path.join(self.reconstructionsrc, "campose.txt"))
        lines = f.readlines()
        for l in lines:
            datas = l.split(" ")
            if datas[0].isnumeric():
                self.keyposes[datas[-1].split("\n")[0]] = _pose2Rotation([[float(datas[5]), float(datas[6]), float(datas[7])],\
                                                        [float(datas[1]), float(datas[2]), float(datas[3]), float(datas[4])]])
                pass

    def _applytrans2cam(self):
        # Axis_align = np.array([[1, 0, 0, 0],
        #                       [0, -1, 0, 0],
        #                       [0, 0, -1, 0],
        #                       [0, 0, 0, 1],]
        #     )
        scale = self.param.recon["scale"]
        trans = self.param.recon["trans"]
        for cam in self.keyposes:
            origin_pose = self.keyposes[cam]
            origin_pose[:3, 3] = origin_pose[:3, 3] * scale
            # origin_pose = np.linalg.inv(origin_pose).dot(Axis_align)
            if self.CAM_INVERSE:
                origin_pose = np.linalg.inv(origin_pose)
            self.keyposes[cam] = trans.dot(origin_pose)
    
    def _interpolation(self, type):
        assert type in ["KF_forward_m2f", "KF_forward_f2f", "all"]
        if type == "all":
            ## don't need do anything
            pass
        if type == "KF_forward_m2f":
            try: 
                from kernel.kf_pycuda.kinect_fusion import KinectFusion
                from kernel.kf_pycuda.config import set_config
            except:
                print(
                    "Error", "Please successfully install pycuda", None
                )        
            else:  
                
                config = set_config(resX = self.param.camera["resolution"][0], 
                    resY = self.param.camera["resolution"][0], 
                    fx = self.param.camera["intrinsic"][0, 0], 
                    fy = self.param.camera["intrinsic"][1, 1], 
                    cx = self.param.camera["intrinsic"][0, 2], 
                    cy = self.param.camera["intrinsic"][1, 2], 
                    tsdf_voxel_size = 0.0025, 
                    tsdf_trunc_margin = 0.015, 
                    pcd_voxel_size = 0.005) 
                self.kf = KinectFusion(cfg=config)
                for keypair in tqdm(self.wholemap):
                    self._interpolationKFForwardM2F(keypair, self.wholemap[keypair])
                for idx, prefix in enumerate(self.wholecam.keys()):
                    if idx < len(self.kf.cam_poses):
                        self.wholecam[prefix] = self.kf.cam_poses[idx]
        if type == "KF_forward_f2f":
            try: 
                from kernel.kf_pycuda.kinect_fusion import KinectFusion
                from kernel.kf_pycuda.config import set_config
            except:
                print(
                    "Error", "Please successfully install pycuda", None
                )        
            else:  
                
                config = set_config(resX = self.param.camera["resolution"][0], 
                    resY = self.param.camera["resolution"][0], 
                    fx = self.param.camera["intrinsic"][0, 0], 
                    fy = self.param.camera["intrinsic"][1, 1], 
                    cx = self.param.camera["intrinsic"][0, 2], 
                    cy = self.param.camera["intrinsic"][1, 2], 
                    tsdf_voxel_size = 0.0025, 
                    tsdf_trunc_margin = 0.015, 
                    pcd_voxel_size = 0.005) 
                self.kf = KinectFusion(cfg=config)
                for keypair in tqdm(self.wholemap):
                    self._interpolationKFForwardF2F(keypair, self.wholemap[keypair])
                for idx, prefix in enumerate(self.wholecam.keys()):
                    if idx < len(self.kf.cam_poses):
                        self.wholecam[prefix] = self.kf.cam_poses[idx]

    def _interpolationKFForwardM2F(self, keypair, imglists):
        for imgname in [keypair[0]] + imglists:
            color_im_path = os.path.join(self.rgb_path, imgname)
            depth_im_path = os.path.join(self.depth_path, imgname)
            color_im = np.array(Image.open(color_im_path))
            depth_im = np.array(Image.open(depth_im_path)).astype(np.float32) * self.depth_scale
            depth_im[depth_im > 1.5] = 0
            if imgname == keypair[0]:
                self.kf.interpolation_update(color_im, depth_im, self.keyposes[keypair[0]])
            else:
                self.kf.interpolation_update(color_im, depth_im, None)

    def _interpolationKFForwardF2F(self, keypair, imglists):
        for imgname in [keypair[0]] + imglists[:int(len(imglists)/2)] + (imglists[int(len(imglists)/2):] + [keypair[1]])[::-1]:
            color_im_path = os.path.join(self.rgb_path, imgname)
            depth_im_path = os.path.join(self.depth_path, imgname)
            color_im = np.array(Image.open(color_im_path))
            depth_im = np.array(Image.open(depth_im_path)).astype(np.float32) * self.depth_scale
            depth_im[depth_im > 1.5] = 0
            if imgname == keypair[0]:
                self.kf.F2FPoseEstimation(color_im, depth_im, self.keyposes[keypair[0]])
            else:
                pose = self.kf.F2FPoseEstimation(color_im, depth_im, None)
                self.wholecam[imgname] = pose
              
            
    
    def _savecampose(self, campose_filename):
        with open(os.path.join(self.reconstructionsrc, campose_filename), 'w') as f:
            f.write('# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n')
            f.write('\n')
            for idx, cam in enumerate(self.wholecam.keys()):
                if self.wholecam[cam].any():
                    pose = _rotation2Pose(self.wholecam[cam])
                    f.write(str(idx) + " {0} {1} {2} {3} {4} {5} {6} 1 ".format(*pose[1], *pose[0])+ cam + "\n")


offline/offlineRecon.py
MyungHwanJeon/ProgressLabeller_for_TRansPose

import json
import yaml
import numpy as np
import os


from kernel.utility import _transstring2trans
from kernel.geometry import _pose2Rotation, _rotation2Pose

class offlineParam:
    def __init__(self, config_path) -> None:
        f = open(config_path)
        configuration = json.load(f)
        self.config = configuration
        self.dir = os.path.dirname(config_path)
        self.parsecamera()
        self.parseenv()
        self.parsereconpara()
        self.parsedatapara()
        self.parseobj()
        with open(os.path.join(self.config["environment"]["modelsrc"], "object_label.json")) as objlabelf:
            object_label = json.load(objlabelf)
        for obj in object_label:
            object_label[obj] = int(object_label[obj])  
        self.object_label = object_label
        


    def parsecamera(self):
        self.camera = {}
        self.camera["resolution"] = self.config["camera"]["resolution"]
        self.camera["intrinsic"] = np.array(self.config["camera"]["intrinsic"])
        self.camera["inverse_pose"] = self.config["camera"]["inverse_pose"]

    def parseenv(self):
        self.modelsrc = self.config["environment"]["modelsrc"]
        self.modelposesrc = self.config["environment"]["reconstructionsrc"]
        self.reconstructionsrc = self.config["environment"]["reconstructionsrc"]
        self.datasrc = self.config["environment"]["datasrc"]


    def parsereconpara(self):
        self.recon = {}
        self.recon["scale"] = self.config['reconstruction']['scale']
        self.recon["trans"] = _transstring2trans(self.config['reconstruction']['recon_trans'])

    def parsedatapara(self):
        self.data = {}
        self.data['sample_rate'] = self.config['data']['sample_rate']
        self.data['depth_scale'] = self.config['data']['depth_scale']
    
    def parseobj(self):
        self.objs = {}
        objnames_list = []
        if os.path.exists(os.path.join(self.modelposesrc, "label_pose.yaml")):
            f = open(os.path.join(self.modelposesrc, "label_pose.yaml"))
            poses = yaml.safe_load(f)
            model_dir = os.listdir(self.modelsrc)
            for obj_instancename in poses:
                objname = obj_instancename.split(".")[0]
                if objname not in objnames_list:
                    objnames_list.append(objname)
                if objname in model_dir:
                    self.objs[obj_instancename] = {}
                    self.objs[obj_instancename]['type'] = poses[obj_instancename]['type']
                    self.objs[obj_instancename]['trans'] = _pose2Rotation(poses[obj_instancename]['pose'])
                    files = os.listdir(os.path.join(self.modelsrc, objname))

        



offline/parse.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
from cv2 import INTER_NEAREST
import pyrender
import numpy as np
import os
import json
import trimesh
import pyrender
from kernel.geometry import _pose2Rotation
from PIL import Image
from tqdm import tqdm
from scipy.io import savemat
import copy
import open3d as o3d
import cv2
import time
os.environ['PYOPENGL_PLATFORM'] = 'egl'

class offlineRender:
    def __init__(self, param, outputdir, interpolation_type, pkg_type = "BOP") -> None:
        assert(pkg_type in ["ProgressLabeller", "BOP", "YCBV", "Yourtype"])
        print("Start offline rendering")
        self.param = param
        self.interpolation_type = interpolation_type
        self.outputpath = outputdir
        self.modelsrc = self.param.modelsrc
        self.reconstructionsrc = self.param.reconstructionsrc
        self.datasrc = self.param.datasrc
        self.intrinsic = self.param.camera["intrinsic"]
        self.objects = self.param.objs
        
        self._parsecamfile()
        if pkg_type == "ProgressLabeller":
            self._prepare_scene()
            self._createallpkgs()
            self.render = pyrender.OffscreenRenderer(self.param.camera["resolution"][0], self.param.camera["resolution"][1])
            self.renderAll()
        elif pkg_type == "BOP":
            self.object_label = param.object_label
            self._prepare_scene_BOP()
            self.render = pyrender.OffscreenRenderer(self.param.camera["resolution"][0], self.param.camera["resolution"][1])
            self.renderBOP()
        elif pkg_type == "YCBV":    
            self.object_label = param.object_label
            self._prepare_scene()
            self.render = pyrender.OffscreenRenderer(self.param.camera["resolution"][0], self.param.camera["resolution"][1])
            self.renderYCBV()
        elif pkg_type == "Yourtype":
            self.object_label = param.object_label
            self._prepare_scene()
            self.render = pyrender.OffscreenRenderer(self.param.camera["resolution"][0], self.param.camera["resolution"][1])            
            self.renderYourtype()
    
    def data_export(self, target_dir):
        if not os.path.exists(target_dir):
            os.mkdir(target_dir)
        
    def _prepare_scene(self):
        self.objectmap = {}
        object_index = 0
        self.scene = pyrender.Scene()

        cam = pyrender.camera.IntrinsicsCamera(self.param.camera["intrinsic"][0, 0],
                                               self.param.camera["intrinsic"][1, 1], 
                                               self.param.camera["intrinsic"][0, 2], 
                                               self.param.camera["intrinsic"][1, 2], 
                                               znear=0.05, zfar=100.0, name=None)
        self.nc = pyrender.Node(camera=cam, matrix=np.eye(4))
        self.scene.add_node(self.nc)
        for obj_instancename in self.objects:
            obj = obj_instancename.split(".")[0]
            ## for full model
            if self.objects[obj_instancename]['type'] == 'normal':
                tm = trimesh.load(os.path.join(self.modelsrc, obj, obj+".obj"))
                mesh = pyrender.Mesh.from_trimesh(tm, smooth = False)
                node = pyrender.Node(mesh=mesh, matrix=self.objects[obj_instancename]['trans'])
                self.objectmap[node] = {"index":object_index, "name":obj_instancename, "trans":self.objects[obj_instancename]['trans']}
                self.scene.add_node(node)
                object_index += 1

    def _parsecamfile(self):
        self.camposes = {}
        f = open(os.path.join(self.reconstructionsrc, "campose_all_{0}.txt".format(self.interpolation_type)))
        # f = open(os.path.join(self.reconstructionsrc, "campose.txt"))
        lines = f.readlines()
        for l in lines:
            datas = l.split(" ")
            if datas[0].isnumeric():
                self.camposes[datas[-1].split("\n")[0]] = _pose2Rotation([[float(datas[5]), float(datas[6]), float(datas[7])],\
                                                           [float(datas[1]), float(datas[2]), float(datas[3]), float(datas[4])]])
    
    def _applytrans2cam(self):
        Axis_align = np.array([[1, 0, 0, 0],
                               [0, -1, 0, 0],
                               [0, 0, -1, 0],
                               [0, 0, 0, 1],]
            )
        scale = self.param.recon["scale"]
        trans = self.param.recon["trans"]
        for cam in self.camposes:
            origin_pose = self.camposes[cam]
            origin_pose[:3, 3] = origin_pose[:3, 3] * scale
            if self.CAM_INVERSE:
                origin_pose = np.linalg.inv(origin_pose).dot(Axis_align)
            else:
                origin_pose = origin_pose.dot(Axis_align)
            self.camposes[cam] = trans.dot(origin_pose)

    def _render(self, cam_pose, scene):
        ##segimg is the instance segmentation for each part(normal or each part for the split)
        scene.set_pose(self.nc, pose=cam_pose)
        flags = pyrender.constants.RenderFlags.DEPTH_ONLY
        segimg = np.zeros((self.param.camera["resolution"][1], self.param.camera["resolution"][0]), dtype=np.uint8)

        full_depth = self.render.render(scene, flags = flags)
        for node in self.objectmap:
            node.mesh.is_visible = False
        
        for node in self.objectmap:
            node.mesh.is_visible = True
            depth = self.render.render(scene, flags = flags)
            mask = np.logical_and(
                (np.abs(depth - full_depth) < 1e-6), np.abs(full_depth) > 0
            )
            segimg[mask] = self.objectmap[node]['index'] + 1
            node.mesh.is_visible = False
        
        for node in self.objectmap:
            node.mesh.is_visible = True
        return segimg

    def _createpkg(self, dir):
        if os.path.exists(dir):
            return True
        else:
            if self._createpkg(os.path.dirname(dir)):
                os.mkdir(dir)
                return self._createpkg(os.path.dirname(dir))

    def _createallpkgs(self):
        for node in self.objectmap:
            self._createpkg(os.path.join(self.outputpath, self.objectmap[node]["name"], "pose"))
            self._createpkg(os.path.join(self.outputpath, self.objectmap[node]["name"], "rgb"))
    
    def renderAll(self):
        ## generate whole output dataset
        Axis_align = np.array([[1, 0, 0, 0],
                               [0, -1, 0, 0],
                               [0, 0, -1, 0],
                               [0, 0, 0, 1],]
                                )
       
        for cam in tqdm(self.camposes):
            camT = self.camposes[cam].dot(Axis_align)
            segment = self._render(camT, self.scene)
            perfix = cam.split(".")[0]
            inputrgb = np.array(Image.open(os.path.join(self.datasrc, "rgb_undistort", cam)))

            for node in self.objectmap:
                posepath = os.path.join(self.outputpath, self.objectmap[node]["name"], "pose")
                rgbpath = os.path.join(self.outputpath, self.objectmap[node]["name"], "rgb")
                modelT = self.objectmap[node]["trans"]
                model_camT = np.linalg.inv(modelT).dot(self.camposes[cam])
                self._createpose(posepath, perfix, model_camT)
                self._createrbg(inputrgb, segment, os.path.join(rgbpath, cam), self.objectmap[node]["index"] + 1)

    def _createpose(self, path, perfix, T):
        posefileName = os.path.join(path, perfix + ".txt")
        # np.savetxt(posefileName, np.linalg.inv(T), fmt='%f', delimiter=' ')
        np.savetxt(posefileName, T, fmt='%f', delimiter=' ')

    def _createrbg(self, inputrgb, segment, outputpath, segment_index):
        rgb = inputrgb.copy()
        mask = np.repeat((segment != segment_index)[:, :, np.newaxis], 3, axis=2)
        rgb[mask] = 0
        img = Image.fromarray(rgb)
        img.save(outputpath)
    

    def renderBOP(self):
        ## should be change by user, 
        self._createpkg(os.path.join(self.outputpath, "depth"))
        self._createpkg(os.path.join(self.outputpath, "mask"))
        self._createpkg(os.path.join(self.outputpath, "mask_visib"))
        self._createpkg(os.path.join(self.outputpath, "rgb"))
        scene_camera = {}
        scene_gt = {}
        scene_gt_info = {}
        for idx, cam_name in tqdm(enumerate(self.camposes)):
            os.system('cp ' + os.path.join(self.datasrc, "rgb", cam_name) + ' ' + os.path.join(self.outputpath, "rgb", "{0:06d}.png".format(idx)))
            os.system('cp ' + os.path.join(self.datasrc, "depth", cam_name) + ' ' + os.path.join(self.outputpath, "depth", "{0:06d}.png".format(idx)))
            inputdepth = Image.open(os.path.join(self.datasrc, "depth", cam_name))
            ### 
            scene_camera[idx] = {
                "cam_K": self.intrinsic.flatten().tolist(),
                "cam_R_w2c": (self.camposes[cam_name][:3, :3]).flatten().tolist(),
                "cam_t_w2c": (self.camposes[cam_name][:3, 3]).flatten().tolist(),
                "depth_scale": np.round(self.param.data['depth_scale'], 5),
                "mode": 0
            }
            scene_gt[idx] = list()
            scene_gt_info[idx] = list()
            ## render
            Axis_align = np.array([[1, 0, 0, 0],
                               [0, -1, 0, 0],
                               [0, 0, -1, 0],
                               [0, 0, 0, 1],]
            )
            camT = self.camposes[cam_name].dot(Axis_align)
            self.scene.set_pose(self.nc, pose=camT)
            flags = pyrender.constants.RenderFlags.DEPTH_ONLY
            for node in self.objectmap:
                node.mesh.is_visible = True
            full_depth = self.render.render(self.scene, flags = flags)

            for node in self.objectmap:
                node.mesh.is_visible = False
            
            for obj_idx, node in enumerate(self.objectmap):
                node.mesh.is_visible = True
                depth = self.render.render(self.scene, flags = flags)
                mask = np.logical_and(
                    (np.abs(depth - full_depth) < 1e-6), np.abs(full_depth) > 0
                )

                ## calculate ketpoints location
                instance_name = self.objectmap[node]['name']
                cam_world_T = np.linalg.inv(self.camposes[cam_name])

                mask_trim = (np.abs(depth) > 0)
                mask_visiable_trim = mask
                mask_pillow = Image.fromarray((mask_trim * 255).astype('uint8'))
                mask_pillow.save(os.path.join(self.outputpath, "mask", "{0:06d}_{1:06d}.png".format(idx ,obj_idx)))
                mask_visiable_pillow = Image.fromarray((mask_visiable_trim * 255).astype('uint8'))
                mask_visiable_pillow.save(os.path.join(self.outputpath, "mask_visib", "{0:06d}_{1:06d}.png".format(idx ,obj_idx)))
                node.mesh.is_visible = False
                if not self._getbbx(mask_trim)[0]:
                    continue
                scene_gt_info[idx].append({
                    "bbox_obj": self._getbbx(mask_trim)[1], 
                    "bbox_visib": self._getbbx(mask_visiable_trim)[1],
                    "px_count_all": int(np.sum(depth > 0)),
                    "px_count_valid": int(np.sum(np.array(inputdepth)[mask_trim] != 0)),
                    "px_count_visib": int(np.sum(mask_visiable_trim)),
                    "visib_fract": float(np.sum(mask_visiable_trim)/np.sum(depth > 0)),
                })
                modelT = self.objectmap[node]["trans"]
                model_camT = np.linalg.inv(modelT).dot(self.camposes[cam_name])
                scene_gt[idx].append({
                    "cam_R_m2c": (model_camT[:3, :3]).flatten().tolist(),
                    "cam_t_m2c":(model_camT[:3, 3]).flatten().tolist(),
                    "obj_id": self.objectmap[node]['index']
                })
            # if idx == 1: # TODO: remove later, rerun entire scene
            #     break
        with open(os.path.join(self.outputpath, 'scene_camera.json'), 'w', encoding='utf-8') as f:
            json.dump(scene_camera, f, ensure_ascii=False, indent=1)
        with open(os.path.join(self.outputpath, 'scene_gt.json'), 'w', encoding='utf-8') as f:
            json.dump(scene_gt, f, ensure_ascii=False, indent=1)
        with open(os.path.join(self.outputpath, 'scene_gt_info.json'), 'w', encoding='utf-8') as f:
            json.dump(scene_gt_info, f, ensure_ascii=False, indent=1)

    def _prepare_scene_BOP(self):
        self.objectmap = {}
        self.scene = pyrender.Scene()
        cam = pyrender.camera.IntrinsicsCamera(self.param.camera["intrinsic"][0, 0],
                                            self.param.camera["intrinsic"][1, 1], 
                                            self.param.camera["intrinsic"][0, 2], 
                                            self.param.camera["intrinsic"][1, 2], 
                                            znear=0.05, zfar=100.0, name=None)
        self.nc = pyrender.Node(camera=cam, matrix=np.eye(4))
        self.scene.add_node(self.nc)
        for obj_instancename in self.objects:
            ## for full model
            obj = obj_instancename.split(".")[0]
            if self.objects[obj_instancename]['type'] == 'normal':
                tm = trimesh.load(os.path.join(self.modelsrc, obj, obj+".obj"))
                mesh = pyrender.Mesh.from_trimesh(tm)
                node = pyrender.Node(mesh=mesh, matrix=self.objects[obj_instancename]['trans'])
                self.objectmap[node] = {"index":self.object_label[obj], "name":obj_instancename , "trans":self.objects[obj_instancename ]['trans']}
                self.scene.add_node(node)

    
    def _getbbx(self, mask):
        pixel_list = np.where(mask)
        if np.any(pixel_list):
            top = pixel_list[0].min()
            bottom = pixel_list[0].max()
            left = pixel_list[1].min()
            right = pixel_list[1].max()
            return True, [int(left), int(top), int(right - left), int(bottom - top)]
        else:
            return False, []
    
    def renderYCBV(self):
        self._createpkg(self.outputpath)
        for idx, cam_name in tqdm(enumerate(self.camposes)):
            os.system('cp ' + os.path.join(self.datasrc, "rgb", cam_name) + ' ' + os.path.join(self.outputpath, "{0:06d}-color.png".format(idx)))
            os.system('cp ' + os.path.join(self.datasrc, "depth", cam_name) + ' ' + os.path.join(self.outputpath, "{0:06d}-depth.png".format(idx)))
            ## render
            Axis_align = np.array([[1, 0, 0, 0],
                               [0, -1, 0, 0],
                               [0, 0, -1, 0],
                               [0, 0, 0, 1],]
            )
            camT = self.camposes[cam_name].dot(Axis_align)
            self.scene.set_pose(self.nc, pose=camT)
            flags = pyrender.constants.RenderFlags.DEPTH_ONLY

            segimg = np.zeros((self.param.camera["resolution"][1], self.param.camera["resolution"][0]), dtype=np.uint8)
            vertmap = np.zeros((self.param.camera["resolution"][1], self.param.camera["resolution"][0]), dtype=np.float32)
            for node in self.objectmap:
                node.mesh.is_visible = True
            full_depth = self.render.render(self.scene, flags = flags)

            for node in self.objectmap:
                node.mesh.is_visible = False
            
            ## create -label.txt
            txtfile = open(os.path.join(self.outputpath, "{0:06d}-box.txt".format(idx)),"w+")
            ## create -meta.mat
            mat = {}
            mat['cls_indexes'] = np.empty((0, 1), dtype = np.uint8)
            mat['center'] = np.empty((0, 2))
            mat['factor_depth'] = np.array([[np.around(1/self.param.data['depth_scale'])]], dtype = np.uint16)
            mat['intrinsic_matrix'] = self.intrinsic
            mat['poses'] = np.empty((3, 4, 0))
            mat['rotation_translation_matrix'] = self.camposes[cam_name][:3, :]
            for obj_idx, node in enumerate(self.objectmap):
                node.mesh.is_visible = True
                depth = self.render.render(self.scene, flags = flags)
                mask = np.logical_and(
                    (np.abs(depth - full_depth) < 1e-6), np.abs(full_depth) > 0
                )
                mask_visiable = (mask * 255).astype('uint8')
                segimg[mask] = self.object_label[self.objectmap[node]["name"].split(".")[0]]
                node.mesh.is_visible = False
                if not self._getbbxycb(mask_visiable)[0]:
                    continue
                else:
                    bbx = self._getbbxycb(mask_visiable)[1]
                    txtfile.write(self.objectmap[node]["name"].split(".")[0] + f' {bbx[0]} {bbx[1]} {bbx[2]} {bbx[3]}\n')
                    mat['cls_indexes'] = np.vstack((mat['cls_indexes'], np.array([[self.object_label[self.objectmap[node]["name"].split(".")[0]]]], dtype = np.uint8)))
                    
                    modelT = self.objectmap[node]["trans"]
                    model_camT = np.linalg.inv(self.camposes[cam_name]).dot(modelT)
                    center_homo = self.intrinsic @ model_camT[:3, 3]
                    center = center_homo[:2]/center_homo[2]
                    mat['center'] = np.vstack((mat['center'], center))
                    mat['poses'] = np.concatenate((mat['poses'], model_camT[:3, :, np.newaxis]), axis = 2)
                    pass

            txtfile.close()
            savemat(os.path.join(self.outputpath, "{0:06d}-meta.mat".format(idx)), mat)
            segimg_pillow = Image.fromarray(segimg)
            segimg_pillow.save(os.path.join(self.outputpath, "{0:06d}-label.png".format(idx)))

    def _getbbxycb(self, mask):
        pixel_list = np.where(mask)
        if np.any(pixel_list):
            top = pixel_list[0].min()
            bottom = pixel_list[0].max()
            left = pixel_list[1].min()
            right = pixel_list[1].max()
            return True, [int(left), int(top), int(right), int(bottom)]
        else:
            return False, []
    
    def renderYourtype(self):
        '''
        Define your own type
        Parameters you might needed
            self.camposes: dictionary with camera name(image_name) as key and camera
            self.objectmap: dictionary with object_node (node for pyrender) as key
                            you could access the:
                                 name: self.objectmap[node]["name"]
                                 Transformation (objec poses in the world frame): self.objectmap[node]["trans"]
            self.intrinsic: camera intrinsic
            os.path.join(self.datasrc, "rgb", cam_name): path to the original rgb
            os.path.join(self.datasrc, "depth", cam_name): path to the original depth

            Remember to use the following because the camera pose we use and the camera pose for pyrender have different coordinate system:
            Axis_align = np.array([[1, 0, 0, 0],
                                [0, -1, 0, 0],
                                [0, 0, -1, 0],
                                [0, 0, 0, 1],]
            )
            camT = self.camposes[cam_name].dot(Axis_align)

        '''
        


        raise NotImplementedError


        


offline/render.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
from cv2 import INTER_NEAREST
import pyrender
import numpy as np
import os
import json
import trimesh
import pyrender
from kernel.geometry import _pose2Rotation
from PIL import Image
from tqdm import tqdm
from scipy.io import savemat
import copy
import open3d as o3d
import cv2
import time
os.environ['PYOPENGL_PLATFORM'] = 'egl'


class offlineRender:
    def __init__(self, param, outputdir) -> None:
        print("Start offline rendering")
        self.param = param
        self.outputpath = outputdir
        self.modelsrc = self.param.modelsrc
        self.reconstructionsrc = self.param.reconstructionsrc
        self.datasrc = self.param.datasrc
        self.objects = self.param.objs
        
        self._parsecamfile()
        
        ###############################################
        
        self.renderer_R = pyrender.OffscreenRenderer(640, 480)
        self.objectmap_R = {}
        object_index = 0
        self.scene_R = pyrender.Scene()
        cam = pyrender.camera.IntrinsicsCamera(603.77647837,
                                               604.63162524, 
                                               329.25940456, 
                                               246.48479807, 
                                               znear=0.01, zfar=100.0, name=None)
        self.cam_node_R = pyrender.Node(camera=cam, matrix=np.eye(4))
        self.scene_R.add_node(self.cam_node_R)
        for obj_instancename in self.objects:
            obj = obj_instancename.split(".")[0]
            ## for full model
            if self.objects[obj_instancename]['type'] == 'normal':
                tm = trimesh.load(os.path.join(self.modelsrc, obj, obj+".obj"), force="mesh")
                mesh = pyrender.Mesh.from_trimesh(tm, smooth = False)
                node = pyrender.Node(mesh=mesh, matrix=self.objects[obj_instancename]['trans'])
                self.objectmap_R[node] = {"index":object_index, "name":obj_instancename, "trans":self.objects[obj_instancename]['trans']}
                self.scene_R.add_node(node)
                object_index += 1                                                       
        
        ###############################################
        
        self.renderer_L = pyrender.OffscreenRenderer(640, 480)
        self.objectmap_L = {}
        object_index = 0
        self.scene_L = pyrender.Scene()
        cam = pyrender.camera.IntrinsicsCamera(609.04479806,
                                               609.54600673, 
                                               326.84169183, 
                                               241.56572826, 
                                               znear=0.01, zfar=100.0, name=None)
        self.cam_node_L = pyrender.Node(camera=cam, matrix=np.eye(4))
        self.scene_L.add_node(self.cam_node_L)
        for obj_instancename in self.objects:
            obj = obj_instancename.split(".")[0]
            ## for full model
            if self.objects[obj_instancename]['type'] == 'normal':
                tm = trimesh.load(os.path.join(self.modelsrc, obj, obj+".obj"), force="mesh")
                mesh = pyrender.Mesh.from_trimesh(tm, smooth = False)
                node = pyrender.Node(mesh=mesh, matrix=self.objects[obj_instancename]['trans'])
                self.objectmap_L[node] = {"index":object_index, "name":obj_instancename, "trans":self.objects[obj_instancename]['trans']}
                self.scene_L.add_node(node)
                object_index += 1             
                
        ###############################################                
        
        self.renderer_T = pyrender.OffscreenRenderer(640, 512)
        self.objectmap_T = {}
        object_index = 0
        self.scene_T = pyrender.Scene()
        cam = pyrender.camera.IntrinsicsCamera(441.62619165,
                                               441.46556074, 
                                               318.41319462, 
                                               258.61635625, 
                                               znear=0.01, zfar=100.0, name=None)
        self.cam_node_T = pyrender.Node(camera=cam, matrix=np.eye(4))
        self.scene_T.add_node(self.cam_node_T)
        for obj_instancename in self.objects:
            obj = obj_instancename.split(".")[0]
            ## for full model
            if self.objects[obj_instancename]['type'] == 'normal':
                tm = trimesh.load(os.path.join(self.modelsrc, obj, obj+".obj"), force="mesh")
                mesh = pyrender.Mesh.from_trimesh(tm, smooth = False)
                node = pyrender.Node(mesh=mesh, matrix=self.objects[obj_instancename]['trans'])
                self.objectmap_T[node] = {"index":object_index, "name":obj_instancename, "trans":self.objects[obj_instancename]['trans']}
                self.scene_T.add_node(node)
                object_index += 1           
        
        self._createallpkgs()
        self.renderAll()

    def data_export(self, target_dir):
        if not os.path.exists(target_dir):
            os.mkdir(target_dir)       

    def _parsecamfile(self):
        self.camposes_R = {}
        #f = open(os.path.join(self.reconstructionsrc, "campose_all_{0}.txt".format(self.interpolation_type)))
        f = open(os.path.join(self.reconstructionsrc, "campose_R.txt"))
        lines = f.readlines()
        for l in lines:
            datas = l.split(" ")
            if datas[0].isnumeric():
                self.camposes_R[datas[-1].split("\n")[0]] = _pose2Rotation([[float(datas[5]), float(datas[6]), float(datas[7])],\
                                                           [float(datas[1]), float(datas[2]), float(datas[3]), float(datas[4])]])    
        self.camposes_L = {}                                                   
        f = open(os.path.join(self.reconstructionsrc, "campose_L.txt"))
        lines = f.readlines()
        for l in lines:
            datas = l.split(" ")
            if datas[0].isnumeric():
                self.camposes_L[datas[-1].split("\n")[0]] = _pose2Rotation([[float(datas[5]), float(datas[6]), float(datas[7])],\
                                                           [float(datas[1]), float(datas[2]), float(datas[3]), float(datas[4])]]) 
        
        self.camposes_T = {}                                                      
        f = open(os.path.join(self.reconstructionsrc, "campose_T.txt"))
        lines = f.readlines()
        for l in lines:
            datas = l.split(" ")
            if datas[0].isnumeric():
                self.camposes_T[datas[-1].split("\n")[0]] = _pose2Rotation([[float(datas[5]), float(datas[6]), float(datas[7])],\
                                                           [float(datas[1]), float(datas[2]), float(datas[3]), float(datas[4])]]) 
                                                           
                                                                                                                                                                                 
    def _applytrans2cam(self):
        Axis_align = np.array([[1, 0, 0, 0],
                               [0, -1, 0, 0],
                               [0, 0, -1, 0],
                               [0, 0, 0, 1],]
            )
        scale = self.param.recon["scale"]
        trans = self.param.recon["trans"]
        for cam in self.camposes:
            origin_pose = self.camposes[cam]
            origin_pose[:3, 3] = origin_pose[:3, 3] * scale
            if self.CAM_INVERSE:
                origin_pose = np.linalg.inv(origin_pose).dot(Axis_align)
            else:
                origin_pose = origin_pose.dot(Axis_align)
            self.camposes[cam] = trans.dot(origin_pose)

    def _render(self, renderer, cam_node, cam_pose, object_map, scene, resolution):
        ##segimg is the instance segmentation for each part(normal or each part for the split)
        scene.set_pose(cam_node, pose=cam_pose)
        flags = pyrender.constants.RenderFlags.DEPTH_ONLY
        segimg = np.zeros((resolution[1], resolution[0]), dtype=np.uint8)

        full_depth = renderer.render(scene, flags = flags)
        for node in object_map:
            node.mesh.is_visible = False
        
        for node in object_map:
            node.mesh.is_visible = True
            depth = renderer.render(scene, flags = flags)
            mask = np.logical_and(
                (np.abs(depth - full_depth) < 1e-6), np.abs(full_depth) > 0
            )
            segimg[mask] = object_map[node]['index'] + 1
            node.mesh.is_visible = False
        
        for node in object_map:
            node.mesh.is_visible = True
        return segimg

    def _createpkg(self, dir):
        if os.path.exists(dir):
            return True
        else:
            if self._createpkg(os.path.dirname(dir)):
                os.mkdir(dir)
                return self._createpkg(os.path.dirname(dir))

    def _createallpkgs(self):
        for node in self.objectmap_R:            
            self._createpkg(os.path.join(self.outputpath, self.objectmap_R[node]["name"], "R"))
            self._createpkg(os.path.join(self.outputpath, self.objectmap_R[node]["name"], "R/pose"))
            self._createpkg(os.path.join(self.outputpath, self.objectmap_R[node]["name"], "R/mask"))
            self._createpkg(os.path.join(self.outputpath, self.objectmap_R[node]["name"], "R/depth"))   
            
        for node in self.objectmap_L:    
            self._createpkg(os.path.join(self.outputpath, self.objectmap_L[node]["name"], "L"))
            self._createpkg(os.path.join(self.outputpath, self.objectmap_L[node]["name"], "L/pose"))
            self._createpkg(os.path.join(self.outputpath, self.objectmap_L[node]["name"], "L/mask"))
            
        for node in self.objectmap_T:      
            self._createpkg(os.path.join(self.outputpath, self.objectmap_T[node]["name"], "T"))
            self._createpkg(os.path.join(self.outputpath, self.objectmap_T[node]["name"], "T/pose"))
            self._createpkg(os.path.join(self.outputpath, self.objectmap_T[node]["name"], "T/mask"))
            
    
    def renderAll(self):
        ## generate whole output dataset
        Axis_align = np.array([[1, 0, 0, 0],
                               [0, -1, 0, 0],
                               [0, 0, -1, 0],
                               [0, 0, 0, 1],]
                                )
        
        print("Start offline rendering - Right Cam")
        for cam in tqdm(self.camposes_R):
            camT = self.camposes_R[cam].dot(Axis_align)
            #renderer, cam_node, cam_pose, object_map, scene, resolution
            segment_R = self._render(self.renderer_R, self.cam_node_R, camT, self.objectmap_R, self.scene_R, [640, 480])
            perfix = cam.split(".")[0]
            input_R = np.array(Image.open(os.path.join(self.datasrc, "R/undistort", cam)))
            #input_R = np.array(Image.open(os.path.join(self.datasrc, "R/rgb", cam)))
            
            self.scene_R.set_pose(self.cam_node_R, pose=camT)
            full_depth = self.renderer_R.render(self.scene_R, flags = pyrender.constants.RenderFlags.DEPTH_ONLY) *1000
            full_depth = full_depth.astype(np.uint16)       
            input_depth_R = cv2.imread(os.path.join(self.datasrc, "R/depth_undistort", cam), cv2.IMREAD_UNCHANGED)    
            #input_depth_R = cv2.imread(os.path.join(self.datasrc, "R/depth", cam), cv2.IMREAD_UNCHANGED)    
            
            input_depth_R[full_depth != 0] = full_depth[full_depth != 0]       
            	
            	

            for node in self.objectmap_R:
                posepath = os.path.join(self.outputpath, self.objectmap_R[node]["name"], "R/pose")
                R_path = os.path.join(self.outputpath, self.objectmap_R[node]["name"], "R/mask")
                modelT = self.objectmap_R[node]["trans"]
                model_camT = np.linalg.inv(modelT).dot(self.camposes_R[cam])
                self._createpose(posepath, perfix, model_camT)
                self._createrbg(input_R, segment_R, os.path.join(R_path, cam), self.objectmap_R[node]["index"] + 1)
                
                cv2.imwrite(os.path.join(self.outputpath, self.objectmap_R[node]["name"], "R/depth", cam), input_depth_R)
                
                
        print("Start offline rendering - Left Cam")        
        for cam in tqdm(self.camposes_L):
            camT = self.camposes_L[cam].dot(Axis_align)
            #renderer, cam_node, cam_pose, object_map, scene, resolution
            segment_L = self._render(self.renderer_L, self.cam_node_L, camT, self.objectmap_L, self.scene_L, [640, 480])
            perfix = cam.split(".")[0]
            input_L = np.array(Image.open(os.path.join(self.datasrc, "L/undistort", cam)))

            for node in self.objectmap_L:
                posepath = os.path.join(self.outputpath, self.objectmap_L[node]["name"], "L/pose")
                R_path = os.path.join(self.outputpath, self.objectmap_L[node]["name"], "L/mask")
                modelT = self.objectmap_L[node]["trans"]
                model_camT = np.linalg.inv(modelT).dot(self.camposes_L[cam])
                self._createpose(posepath, perfix, model_camT)
                self._createrbg(input_L, segment_L, os.path.join(R_path, cam), self.objectmap_L[node]["index"] + 1)
                
        

        print("Start offline rendering - Thermal Cam")
        for cam in tqdm(self.camposes_T):
            camT = self.camposes_T[cam].dot(Axis_align)
            #renderer, cam_node, cam_pose, object_map, scene, resolution
            segment_T = self._render(self.renderer_T, self.cam_node_T, camT, self.objectmap_T, self.scene_T, [640, 512])
            perfix = cam.split(".")[0]
            input_T = np.array(Image.open(os.path.join(self.datasrc, "T/undistort", cam)))

            for node in self.objectmap_T:
                posepath = os.path.join(self.outputpath, self.objectmap_T[node]["name"], "T/pose")
                R_path = os.path.join(self.outputpath, self.objectmap_T[node]["name"], "T/mask")
                modelT = self.objectmap_T[node]["trans"]
                model_camT = np.linalg.inv(modelT).dot(self.camposes_T[cam])
                self._createpose(posepath, perfix, model_camT)
                self._createrbg_thermal(input_T, segment_T, os.path.join(R_path, cam), self.objectmap_T[node]["index"] + 1)        
                
                
                
                
                
                                                          

    def _createpose(self, path, perfix, T):
        posefileName = os.path.join(path, perfix + ".txt")
        # np.savetxt(posefileName, np.linalg.inv(T), fmt='%f', delimiter=' ')
        np.savetxt(posefileName, T, fmt='%f', delimiter=' ')

    def _createrbg(self, inputrgb, segment, outputpath, segment_index):
        rgb = inputrgb.copy()
        mask = np.repeat((segment != segment_index)[:, :, np.newaxis], 3, axis=2)
        rgb[mask] = 0
        img = Image.fromarray(rgb)
        img.save(outputpath)    
        
        
    def _createrbg_thermal(self, inputrgb, segment, outputpath, segment_index):
        rgb = inputrgb.copy()
        
        h, w = rgb.shape
        min_thresh = np.min(rgb)
        max_thresh = np.max(rgb)

        if (min_thresh != None) and (max_thresh != None):
            for i in range(h):
                for j in range(w):
                    if (rgb[i][j]) > max_thresh:
                        rgb[i][j] = max_thresh

                    elif (rgb[i][j]) < min_thresh:
                        rgb[i][j] = min_thresh

        normalized_array = (rgb - min_thresh) / (max_thresh - min_thresh)

        norm_img = (normalized_array * 255).astype(np.uint8)
        
        mask = np.repeat((segment != segment_index)[:, :, np.newaxis], 3, axis=2)
        norm_img[mask[:, :, 0]] = 0
        img = Image.fromarray(norm_img)
        img.save(outputpath)          
        
        
    
    def _getbbx(self, mask):
        pixel_list = np.where(mask)
        if np.any(pixel_list):
            top = pixel_list[0].min()
            bottom = pixel_list[0].max()
            left = pixel_list[1].min()
            right = pixel_list[1].max()
            return True, [int(left), int(top), int(right - left), int(bottom - top)]
        else:
            return False, []

    def _getbbxycb(self, mask):
        pixel_list = np.where(mask)
        if np.any(pixel_list):
            top = pixel_list[0].min()
            bottom = pixel_list[0].max()
            left = pixel_list[1].min()
            right = pixel_list[1].max()
            return True, [int(left), int(top), int(right), int(bottom)]
        else:
            return False, []           

offline/render_custom.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import os
import sys
import time

def split_obj(filepath):
    with open(filepath) as f:
        lines = f.readlines()
    obj_files = []
    o_index_list = [] 
    index_l = 0
    while index_l < len(lines):
        l = lines[index_l]
        if l[:2] == "o ":
            o_index_list.append(index_l)
        index_l += 1

    o_index = 1
    while o_index < len(o_index_list):
        obj_files.append(lines[o_index_list[o_index - 1]:o_index_list[o_index]])
        o_index += 1

    obj_files.append(lines[o_index_list[-1]:])    
    obj_start = lines[:o_index_list[0]]
    total_v = 0
    total_from_last = 0
    for file in obj_files:
        file_name = file[0][2:-1] + ".obj"
        f = open(os.path.join(os.path.dirname(filepath), file_name), "w")
        all_file = obj_start+file
        for l in all_file:

            if l[:2] == "v ":
                total_v += 1
                l_new = l
            elif l[:2] == "f ":
                point_pairs = l.split(" ")[1:]
                l_new = "f "
                for point_p in point_pairs:
                    if "//" in point_p:
                        p_num_pair = point_p.split("//")
                        for p_num in p_num_pair:
                            l_new += str(int(p_num) - total_from_last)
                            l_new += "//"
                        l_new = l_new[:-2]
                        l_new += " "
                    elif "/" in point_p:
                        p_num_pair = point_p.split("/")
                        for p_num in p_num_pair:
                            l_new += str(int(p_num) - total_from_last)
                            l_new += "/"
                        l_new = l_new[:-2]
                        l_new += " "                            

                l_new = l_new[:-1]
                l_new += "\n" 
            else:
                l_new = l
            f.write(l_new)
        print("Successfully split {0}".format(file_name))
        f.close()
        total_from_last = total_v

offline/splitobj.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
from bpy.props import StringProperty, EnumProperty, FloatProperty, IntProperty
from bpy_extras.view3d_utils import location_3d_to_region_2d
from bpy.types import Operator
import numpy as np
from PIL import Image
import os
import time
import open3d as o3d
import tqdm
import trimesh
import pyrender
import multiprocessing
import subprocess
import sys
import json



from kernel.render import save_img
from kernel.geometry import plane_alignment, transform_from_plane, _pose2Rotation, _rotation2Pose, modelICP, globalRegisteration
from kernel.logging_utility import log_report
from kernel.loader import load_cam_img_depth, load_reconstruction_result, updateprojectname, removeworkspace, load_pc
from kernel.blender_utility import \
    _get_configuration, _get_reconstruction_insameworkspace, _get_obj_insameworkspace, _get_workspace_name, _apply_trans2obj, \
    _align_reconstruction
from registeration.init_configuration import config
from kernel.utility import _trans2transstring,  _parse_camfile, _select_sample_files
from kernel.blender_utility import _is_progresslabeller_object, _initreconpose, _get_obj_insameworkspace
from panel.FloatScreenPanel import draw_for_area
import registeration.register 


class MaskOptimization(Operator):
    """Using RANSAC to detect the plane and align it"""
    bl_idname = "object_property.maskoptimation"  
    bl_label = "Mask Optimization"

    def execute(self, context):
        
        obj = context.object
        assert obj['type'] == 'model'
        ws_name = _get_workspace_name(obj)
        if len(bpy.context.selected_objects) > 0:
            obj_name = bpy.context.selected_objects[0].name
        
        
            setting_object = 0
            visible_objects = [obj for obj in bpy.context.view_layer.objects if obj.visible_get()]

            for obj_ in visible_objects:
                if obj_.name == "untitled:Setting":
                    setting_object = obj_
        
            config_id = setting_object['config_id']
            config = bpy.context.scene.configuration[config_id]
        

            path = setting_object['dir']
            output_filepath = (setting_object["dir"] if setting_object["dir"].endswith("/") else setting_object["dir"] + "/") + "output"
            obj_pose = np.zeros(7, dtype=np.float32)
            obj_context = 0             
        
        
            if context.selected_objects[0]["type"] == "model":

                obj_pose[0] = context.selected_objects[0].location[0]
                obj_pose[1] = context.selected_objects[0].location[1]
                obj_pose[2] = context.selected_objects[0].location[2]
            
                context.selected_objects[0].rotation_mode = 'QUATERNION'

                obj_pose[3] = context.selected_objects[0].rotation_quaternion[0]
                obj_pose[4] = context.selected_objects[0].rotation_quaternion[1]
                obj_pose[5] = context.selected_objects[0].rotation_quaternion[2]
                obj_pose[6] = context.selected_objects[0].rotation_quaternion[3]

                obj_context = context.selected_objects[0]

            else:

                log_report("Error", "error", None)
                return {'FINISHED'}        
        
        
            output_dict = {
                'projectname': config.projectname,
                'environment': {
                    'modelsrc':config.modelsrc,
                    "reconstructionsrc":config.reconstructionsrc,
                    "datasrc": config.datasrc,           
                },
                'camera':{
                    "resolution": [config.resX, config.resY],
                    "intrinsic": [[config.fx, 0, config.cx],
                                  [0, config.fy, config.cy],
                                  [0, 0, 1]],
                },
                'object':{
                    "name":obj_context.name.split(":")[1],
                    "path":obj_context["path"],
                    "translation":[float(obj_pose[0]), float(obj_pose[1]), float(obj_pose[2])],
                    "quaternion":[float(obj_pose[3]), float(obj_pose[4]), float(obj_pose[5]), float(obj_pose[6])],
                },
                'optimization':{
                    "stop_criteria":context.selected_objects[0].optimizationparas.stop_criteria,
                    "translation_learning_rate":context.selected_objects[0].optimizationparas.translation_learning_rate,
                    "rotation_learning_rate":context.selected_objects[0].optimizationparas.rotation_learning_rate
                }
            }           
            
            with open("/tmp/progresslabeller_pose_opt.json", "w") as f:
                json.dump(output_dict, f, indent = True)

            source = os.path.dirname(os.path.dirname(__file__))
            code_path = os.path.join(source, "./", "offline", "mask_opt.py")
            subprocess.call("{} {} {} {}".format(sys.executable, code_path, "/tmp/progresslabeller_pose_opt.json", output_filepath), shell=True)

            os.remove("/tmp/progresslabeller_pose_opt.json")

            f = open("/tmp/progresslabeller_mask_result.json")
            result = json.load(f)
        
            objworkspacename = result["projectname"] + ":" + result["object"]["name"]        
            bpy.data.objects[objworkspacename].location = np.array(result["object"]["translation"])
            bpy.data.objects[objworkspacename].rotation_mode = 'QUATERNION'
            bpy.data.objects[objworkspacename].rotation_quaternion = np.array(result["object"]["quaternion"]) / np.linalg.norm(np.array(result["object"]["quaternion"]))
        
            os.remove("/tmp/progresslabeller_mask_result.json")


            log_report("Info", "Pose Optimization", None)                           
        
        
        return {'FINISHED'}
        
class MaskOptimizationConfig(bpy.types.PropertyGroup):
    stop_criteria: bpy.props.FloatProperty(name="stop_criteria", 
                                        description="stop_criteria", 
                                        default=100.0, 
                                        min=0., 
                                        max=1000., 
                                        step=1, 
                                        precision=3)

    translation_learning_rate: bpy.props.FloatProperty(name="translation_learning_rate", 
                                description="translation_learning_rate", 
                                default=0.0005, 
                                min=0, 
                                max=1, 
                                step=0.0001,
                                precision=5)
    rotation_learning_rate: bpy.props.FloatProperty(name="rotation_learning_rate", 
                                        description="rotation_learning_rate", 
                                        default=0.005, 
                                        min=0, 
                                        max=1, 
                                        step=0.0001,
                                        precision=5)


class PlaneAlignment(Operator):
    """Using RANSAC to detect the plane and align it"""
    bl_idname = "object_property.planealignment"  
    bl_label = "Plane Alignment"

    def execute(self, context):
        recon = context.object
        log_report(
            "INFO", "Starting calculate the plane function", None
        )      
        
        [a, b, c, d], plane_center = plane_alignment(recon["path"], 
                                                     recon["scale"],
                                                     np.array(recon["alignT"]),
                                                     bpy.context.scene.planalignmentparas.threshold,
                                                     bpy.context.scene.planalignmentparas.n,
                                                     bpy.context.scene.planalignmentparas.iteration)
        log_report(
            "INFO", f"Plane equation: {a:.2f}x + {b:.2f}y + {c:.2f}z + {d:.2f} = 0", None
        )        
        if d < 0:
            trans  = transform_from_plane([-a, -b, -c, -d], plane_center)
        else:
            trans  = transform_from_plane([a, b, c, d], plane_center)
        log_report(
            "INFO", "Starting Transform the scene", None
        )   

        obj_lists = _get_obj_insameworkspace(context.object, ["reconstruction", "camera"])
        
        for obj in obj_lists:
            _apply_trans2obj(obj, trans)


        recons = _get_obj_insameworkspace(recon, ["reconstruction"])

        for obj in recons:
            print(obj.name)
            obj["alignT"] = (trans.dot(np.array(obj["alignT"]))).tolist()

        _, config = _get_configuration(context.object)
        config.recon_trans = _trans2transstring(np.array(obj["alignT"]))

        return {'FINISHED'}

    def invoke(self, context, event):
        return context.window_manager.invoke_props_dialog(self, width = 400)

    def draw(self, context):
        layout = self.layout
        scene = context.scene
        layout.label(text="Set Plane Alignment (ICP) Parameters:")
        box = layout.box() 
        row = box.row()
        row.prop(scene.planalignmentparas, "threshold") 
        row = box.row()
        row.prop(scene.planalignmentparas, "n") 
        row = box.row()
        row.prop(scene.planalignmentparas, "iteration") 

class PlaneAlignmentConfig(bpy.types.PropertyGroup):
    threshold: bpy.props.FloatProperty(name="Inlier Threshold", 
                                        description="Inlier threshold for points aligned to plane", 
                                        default=0.01, 
                                        min=0.001, 
                                        max=1.000, 
                                        step=0.001, 
                                        precision=3)

    n: bpy.props.IntProperty(name="Fit Number", 
                                description="Numbers for fitting a plane", 
                                default=3, 
                                min=3, 
                                max=10, 
                                step=1)
    iteration: bpy.props.IntProperty(name="Fit Iteration", 
                                        description="ICP Iteration", 
                                        default=1000, 
                                        min=10, 
                                        max=10000, 
                                        step=100)


class ImportCamRGBDepth(Operator):
    """This appears in the tooltip of the operator and in the generated docs"""
    bl_idname = "object_property.importcamrgbdepth"  
    bl_label = "Import RGB & Depth"


    def execute(self, context): 
        config_id, config = _get_configuration(context.object)
        packagepath = config.datasrc
        files = os.listdir(packagepath)
        #if "rgb_undistort" not in files or "depth" not in files:
        #    log_report(
        #        "Error", "either rgb or depth package is not in the datasrc", None
        #    )       
        #elif config.resX == 0 or config.resY == 0\
        #    or config.fx == 0 or config.fy == 0\
        #    or config.cx == 0 or config.cy == 0:
        #    log_report(
        #        "ERROR", "Please set the camera parameters before loading the RGB and Camera", None
        #    )   
        #    return {'FINISHED'}      
        #else:   
        #    load_cam_img_depth(packagepath, config_id, camera_display_scale = 0.1, sample_rate=config.sample_rate)
        #return {'FINISHED'}
    
    def invoke(self, context, event):
        config_id, config = _get_configuration(context.object)
        #files = os.listdir(os.path.join(config.datasrc, "rgb_undistort"))
        #self.total_files = len(files)
        return context.window_manager.invoke_props_dialog(self, width = 400)

    def draw(self, context):
        layout = self.layout
        scene = context.scene
        config_id, config = _get_configuration(context.object)
        #layout.label(text="Set Sample Rate for RGB to reconstruct:")
        #box = layout.box() 
        #row = box.row()
        #row.prop(config, "sample_rate") 
        #row = box.row()
        #row.label(text="You would load around {0} cameras".format(int(self.total_files * config.sample_rate)))


class ImportReconResult(Operator):
    """This appears in the tooltip of the operator and in the generated docs"""
    bl_idname = "object_property.importreconresult"  # important since its how bpy.ops.import_test.some_data is constructed
    bl_label = "Import Reconstruction Result"

    def execute(self, context): 
        scene = context.scene
        config_id, config = _get_configuration(context.object)
        datasrc = config.datasrc
        filepath = config.reconstructionsrc
        scene.loadreconparas.pointcloud_scale = 1.0#config.reconstructionscale
        # config.cameradisplayscale = scene.loadreconparas.camera_display_scale

        if not scene.loadreconparas.AUTOALIGN:
            load_pc(os.path.join(config.reconstructionsrc, "depthfused.ply"), scene.loadreconparas.pointcloud_scale, config_id, "reconstruction_depthfusion")
            load_reconstruction_result(filepath = filepath, 
                                        pointcloudscale = scene.loadreconparas.pointcloud_scale, 
                                        datasrc = datasrc,
                                        config_id = config_id,
                                        camera_display_scale = config.cameradisplayscale,
                                        IMPORT_RATIO = scene.loadreconparas.Import_ratio,
                                        CAMPOSE_INVERSE = config.inverse_pose
                                        )
        else:
            log_report(
            "INFO", "Starting aligning the point cloud", None
            )     

            scale =  _align_reconstruction(config, scene, scene.scalealign.THRESHOLD, scene.scalealign.NUM_THRESHOLD)
            config.reconstructionscale = 1.0#scale
            load_pc(os.path.join(config.reconstructionsrc, "depthfused.ply"), scale, config_id, "reconstruction_depthfusion")
            load_reconstruction_result(filepath = filepath, 
                                        pointcloudscale = scale, 
                                        datasrc = datasrc,
                                        config_id = config_id,
                                        camera_display_scale = config.cameradisplayscale,
                                        IMPORT_RATIO = scene.loadreconparas.Import_ratio,
                                        CAMPOSE_INVERSE = config.inverse_pose
                                        )
            

        return {'FINISHED'}
        
    def invoke(self, context, event):
    
        return context.window_manager.invoke_props_dialog(self, width = 400)
        
    def draw(self, context):
        layout = self.layout
        scene = context.scene
        _, config = _get_configuration(context.object)
        layout.label(text="Load Image & Pose data")
        
        box = layout.box() 
        row = box.row()
        row.prop(scene.loadreconparas, "Import_ratio") 




class LoadRecon(bpy.types.PropertyGroup):
    # The properties for this class which is referenced as an 'entry' below.
    pointcloud_scale: bpy.props.FloatProperty(name="Point Cloud Display Scale", 
                                            description="Scale for display reconstruction point cloud", 
                                            default=1.00, 
                                            min=0.00, 
                                            max=1.00, 
                                            step=2, 
                                            precision=2)  
    AUTOALIGN: bpy.props.BoolProperty(name="Auto Align Point Cloud Scale", 
                                      description="Algin the Point Clound from Depth Information", 
                                      default=False)      

    Import_ratio: bpy.props.FloatProperty(name="Import Ratio", 
                                          description="Ratio to import images from campose.txt", 
                                          default=0.1, 
                                          min=0.00, 
                                          max=1.00, 
                                          step=2, 
                                          precision=2)   
                                                     
  


class WorkspaceRename(Operator):
    """This appears in the tooltip of the operator and in the generated docs"""
    bl_idname = "object_property.workspacerename"  # important since its how bpy.ops.import_test.some_data is constructed
    bl_label = "Rename"

    def execute(self, context): 
        updateprojectname()
        return {'FINISHED'}

class RemoveWorkspace(Operator):
    """This appears in the tooltip of the operator and in the generated docs"""
    bl_idname = "object_property.removeworkspace"  # important since its how bpy.ops.import_test.some_data is constructed
    bl_label = "Remove Current Workspace"

    def execute(self, context):
        assert context.object['type'] == 'setting'
        name  = _get_workspace_name(context.object)
        removeworkspace(name)
        for area in registeration.register.area_image_pair:
            bpy.types.SpaceView3D.draw_handler_remove(registeration.register.area_image_pair[area]["handler"], 'WINDOW')
        registeration.register.area_image_pair = {}
        return {'FINISHED'}  

class ModelICP(Operator):
    """This appears in the tooltip of the operator and in the generated docs"""
    bl_idname = "object_property.modelicp"  # important since its how bpy.ops.import_test.some_data is constructed
    bl_label = "Model Alignment"

    def execute(self, context):
        obj = context.object
        assert obj['type'] == 'model'
        name  = _get_workspace_name(obj)
        model_vertices = np.array([list(obj.matrix_world @ v.co) for v in obj.data.vertices])
        if name + ":reconstruction" not in bpy.data.objects:
            log_report(
                "Error", "You should upload the reconstruction result first", None
            )     
            return {'FINISHED'}
        else:
            recon = _get_reconstruction_insameworkspace(obj)
            print(recon.name)
            recon_vertices = np.array(recon["particle_coords"])
            rot = np.array(recon.matrix_world)
            recon_vertices_rotated = (rot[:3, :3].dot(recon_vertices.T) + rot[:3, [3]]).T
            trans_obj_icp = modelICP(recon_vertices_rotated, model_vertices)
            _apply_trans2obj(obj, trans_obj_icp)
        return {'FINISHED'} 

class AllModelsICP(Operator):
    """This appears in the tooltip of the operator and in the generated docs"""
    bl_idname = "object_property.allmodelsicp"  # important since its how bpy.ops.import_test.some_data is constructed
    bl_label = "All Models Alignment"

    def execute(self, context):
        recon = context.object
        assert recon['type'] == 'reconstruction'
        objects = _get_obj_insameworkspace(recon, 'model')

        for obj in objects:
            model_vertices = np.array([list(obj.matrix_world @ v.co) for v in obj.data.vertices])
            recon_vertices = np.array(recon["particle_coords"])
            rot = np.array(recon.matrix_world)
            recon_vertices_rotated = (rot[:3, :3].dot(recon_vertices.T) + rot[:3, [3]]).T
            trans_obj_icp = modelICP(recon_vertices_rotated, model_vertices)
            _apply_trans2obj(obj, trans_obj_icp)
        return {'FINISHED'}    

class CurrentDepthOperator(bpy.types.Operator):
    """Move an object with the mouse, example"""
    bl_idname = "object_property.current_depth_operator"
    bl_label = "Current depth operator"

    current_depth: FloatProperty()

    def modal(self, context, event):
        if _is_progresslabeller_object(context.object) and context.object["type"] == "camera" and event.type == 'MOUSEMOVE':
            config_id, config = _get_configuration(context.object)
            for area in context.screen.areas:
                if area.type == 'VIEW_3D':
                    area.spaces.active.region_3d.view_perspective = 'CAMERA'
                    bpy.context.scene.render.resolution_x = config.resX
                    bpy.context.scene.render.resolution_y = config.resY
                    bpy.context.scene.camera = context.object
                    cam = bpy.context.scene.camera
                    frame = cam.data.view_frame(scene = bpy.context.scene)
                    frame = [cam.matrix_world @ corner for corner in frame]
                    region = bpy.context.region
                    rv3d = bpy.context.region_data
                    frame_px = [location_3d_to_region_2d(region, rv3d, corner) for corner in frame]           
                    bias_X = min([v[0] for v in frame_px])
                    bias_Y = min([v[1] for v in frame_px])
                    res_X = max([v[0] for v in frame_px]) - min([v[0] for v in frame_px])
                    res_Y = max([v[1] for v in frame_px]) - min([v[1] for v in frame_px])
                    print(res_X, res_Y)
            if event.mouse_x > bias_X \
                and event.mouse_x < bias_X + res_X\
                and event.mouse_y > bias_Y\
                and event.mouse_x < bias_Y + res_Y:
                print(event.mouse_x, event.mouse_y) 
        return {'PASS_THROUGH'}
        

    def invoke(self, context, event):
        self.current_depth = 0
        context.window_manager.modal_handler_add(self)
        return {'FINISHED'}

class Lockcurrent3DArea(bpy.types.Operator):
    """Move an object with the mouse, example"""
    bl_idname = "object_property.lockcurrent3darea"
    bl_label = "Lock the view (progresslabelelr)"

    def execute(self, context):
        context.window_manager.modal_handler_add(self)
        return {'RUNNING_MODAL'}  
    
    def modal(self, context, event):
        for area in bpy.context.screen.areas:
            if area.type == 'VIEW_3D':
                area_bottom = area.y
                area_left = area.x
                area_top= area.y + area.height
                area_right = area.x + area.width  
                if event.mouse_x >= area_left and event.mouse_x < area_right\
                    and event.mouse_y >= area_bottom and event.mouse_y < area_top:
                    if context.object is not None and _is_progresslabeller_object(context.object) and context.object["type"] == "camera":
                        #context.object["depth"]["UPDATEALPHA"] = True
                        context.object["rgb"]["UPDATEALPHA"] = True
                        floatscreen_handler = area.spaces[0].draw_handler_add(draw_for_area, (area, context.object), 'WINDOW', 'POST_PIXEL')
                        registeration.register.area_image_pair[area] = {"camera" : context.object,
                                                                        "handler" : floatscreen_handler}
                        print(context.object.keys())
                        log_report(
                            "Info", "Lock current 3D view area", None
                        )
                        return {'FINISHED'}
                    else:
                        log_report(
                            "Info", "Doesn't Lock anything, please select a Progresslabeller camera object", None
                        )
                        return {'FINISHED'}
        log_report(
            "Info", "Doesn't Lock anything, please select a 3D view area", None
        )
        return {'FINISHED'}

class Unlockcurrent3DArea(bpy.types.Operator):
    """Move an object with the mouse, example"""
    bl_idname = "object_property.unlockcurrent3darea"
    bl_label = "Unlock the view (progresslabelelr)"

    def execute(self, context):
        context.window_manager.modal_handler_add(self)
        return {'RUNNING_MODAL'}  
    
    def modal(self, context, event):
        for area in bpy.context.screen.areas:
            if area.type == 'VIEW_3D':
                area_bottom = area.y
                area_left = area.x
                area_top= area.y + area.height
                area_right = area.x + area.width  
                if event.mouse_x >= area_left and event.mouse_x < area_right\
                    and event.mouse_y >= area_bottom and event.mouse_y < area_top:
                    if area in registeration.register.area_image_pair:
                        # area.spaces[0].draw_handler_remove(registeration.register.area_image_pair[area]["handler"], 'WINDOW')
                        bpy.types.SpaceView3D.draw_handler_remove(registeration.register.area_image_pair[area]["handler"], 'WINDOW')
                        registeration.register.area_image_pair.pop(area)
                        log_report(
                            "Info", "Unlock current 3D view area", None
                        )
                        return {'FINISHED'}
                    else:
                        log_report(
                            "Info", "Current 3D view area is not lock", None
                        )
                        return {'FINISHED'}                        
        log_report(
            "Info", "Doesn't Unlock anything, please select a 3D view area", None
        )
        return {'FINISHED'}

            

def register():
    bpy.utils.register_class(PlaneAlignment)
    bpy.utils.register_class(PlaneAlignmentConfig)
    bpy.utils.register_class(MaskOptimizationConfig)
    bpy.types.Object.optimizationparas = bpy.props.PointerProperty(type=MaskOptimizationConfig) 
    bpy.types.Scene.planalignmentparas = bpy.props.PointerProperty(type=PlaneAlignmentConfig) 
    bpy.utils.register_class(ImportCamRGBDepth)
    bpy.utils.register_class(ImportReconResult)
    bpy.utils.register_class(LoadRecon)
    bpy.types.Scene.loadreconparas = bpy.props.PointerProperty(type=LoadRecon)  
    
    bpy.utils.register_class(ModelICP)
    bpy.utils.register_class(AllModelsICP)

    bpy.utils.register_class(WorkspaceRename)
    bpy.utils.register_class(RemoveWorkspace)

    bpy.utils.register_class(CurrentDepthOperator)
    bpy.utils.register_class(Lockcurrent3DArea)
    bpy.utils.register_class(Unlockcurrent3DArea)
    
    bpy.utils.register_class(MaskOptimization)
    

def unregister():
    bpy.utils.unregister_class(PlaneAlignment)
    bpy.utils.unregister_class(PlaneAlignmentConfig)
    bpy.utils.unregister_class(MaskOptimizationConfig)
    bpy.utils.unregister_class(ImportCamRGBDepth)
    bpy.utils.unregister_class(ImportReconResult)
    bpy.utils.unregister_class(LoadRecon)
    bpy.utils.unregister_class(ModelICP)
    bpy.utils.unregister_class(AllModelsICP)
    
    bpy.utils.unregister_class(WorkspaceRename)
    bpy.utils.unregister_class(RemoveWorkspace)
    bpy.utils.unregister_class(CurrentDepthOperator)
    bpy.utils.unregister_class(Lockcurrent3DArea)
    bpy.utils.unregister_class(Unlockcurrent3DArea)
    
    bpy.utils.unregister_class(MaskOptimization)
    
    

operators/ObjectPropertyOperator.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
from typing import DefaultDict
import bpy
from bpy.props import StringProperty, EnumProperty, FloatProperty
from bpy.types import Operator
import os
import multiprocessing
import subprocess
import sys

from kernel.exporter import configuration_export

from kernel.logging_utility import log_report
from kernel.loader import load_reconstruction_result, load_pc
from kernel.blender_utility import _get_configuration, _align_reconstruction, _clear_recon_output, _initreconpose


operators/ReconstructionOperator.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
import bgl
import gpu
from gpu_extras.presets import draw_texture_2d
from bpy_extras.view3d_utils import location_3d_to_region_2d
from bpy.props import StringProperty, EnumProperty, FloatProperty, BoolProperty
from bpy.types import Operator
from kernel.geometry import depthfilter
from kernel.blender_utility import _get_configuration
import registeration.register 
import numpy as np
from mathutils import Matrix

def draw():
    context = bpy.context
    scene = context.scene
    obj = context.object
    if obj and "type" in obj:
        ## the obj is progresslabeller's object
        name = obj.name.split(":")[0]
        config_id =  bpy.data.objects[name + ":Setting"]['config_id']
        config = scene.configuration[config_id]
        if obj['type'] == "camera":
            if scene.floatscreenproperty.viewimage_mode == "RGB Origin":
                show_frame = obj["rgb"]
                if show_frame.bindcode == 0:
                    show_frame.gl_load()
            elif scene.floatscreenproperty.viewimage_mode == "Depth Origin":
                show_frame = obj["depth"]
                if show_frame.bindcode == 0:
                    show_frame.gl_load()
            
            for area in bpy.context.screen.areas:
                if area.type == 'IMAGE_EDITOR': 
                    area.spaces.active.image = show_frame

            if show_frame:
                if show_frame["UPDATEALPHA"] and scene.floatscreenproperty.UPDATE_DEPTHFILTER:
                    #alpha = depthfilter(obj["depth"]["depth"], config.depth_scale, config.depth_ignore, scene.floatscreenproperty.IGNORE_ZERODEPTH)
                    pixels = list(show_frame.pixels) 
                    #for i in range(0, int(len(pixels)/4)):
                    #    pixels[4 * i + 3] = float(alpha[i]) * 1
                    show_frame.pixels[:] = pixels                   
                obj.data.show_background_images = scene.floatscreenproperty.BACKGROUND
                obj.data.background_images[0].image = show_frame
                obj.data.background_images[0].alpha = scene.floatscreenproperty.background_alpha

                if scene.floatscreenproperty.TRACK:
                    for area in bpy.context.screen.areas:
                        ## change camera view
                        if area.type == 'VIEW_3D' and area not in registeration.register.area_image_pair:
                            area.spaces.active.region_3d.view_perspective = 'CAMERA'
                            bpy.context.scene.render.resolution_x = bpy.context.scene.configuration[config_id].resX
                            bpy.context.scene.render.resolution_y = bpy.context.scene.configuration[config_id].resY
                            # bpy.context.scene.camera = obj
                            area.spaces[0].use_local_camera = True
                            area.spaces[0].camera = obj
                show_frame["UPDATEALPHA"] = False

def draw_for_area(area, camera_obj):
    context = bpy.context
    scene = context.scene
    config_id, config = _get_configuration(camera_obj)
    if scene.floatscreenproperty.viewimage_mode == "RGB Origin":
        show_frame = camera_obj["rgb"]
        if show_frame.bindcode == 0:
            show_frame.gl_load()
    elif scene.floatscreenproperty.viewimage_mode == "Depth Origin":
        show_frame = camera_obj["depth"]
        if show_frame.bindcode == 0:
            show_frame.gl_load()
    if show_frame:

        if show_frame["UPDATEALPHA"]:
            pixels = list(show_frame.pixels) 
            for i in range(0, int(len(pixels)/4)):
                pixels[4 * i + 3] = 1
            show_frame.pixels[:] = pixels      
        camera_obj.data.show_background_images = scene.floatscreenproperty.BACKGROUND
        camera_obj.data.background_images[0].image = show_frame
        camera_obj.data.background_images[0].alpha = scene.floatscreenproperty.background_alpha   
        bpy.context.scene.render.resolution_x = bpy.context.scene.configuration[config_id].resX
        bpy.context.scene.render.resolution_y = bpy.context.scene.configuration[config_id].resY
        if len(area.spaces) > 0:
            area.spaces[0].use_local_camera = True
            area.spaces[0].region_3d.view_perspective = 'CAMERA'
            area.spaces[0].camera = camera_obj
            # 
            # area.spaces[0].overlay.show_wireframes = True
            # area.spaces[0].overlay.wireframe_threshold = 06
        show_frame["UPDATEALPHA"] = False        


class FloatScreenProperty(bpy.types.PropertyGroup):
    # The properties for this class which is referenced as an 'entry' below.
    viewimage_mode: EnumProperty(
        name="View Mode",
        description="Choose the mode for viewing the frame",
        items=(
            ('RGB Origin', "RGB Origin", "Display Origin RGB Image"),
            ('Depth Origin', "Depth Origin", "Display Origin Depth Image"),
        ),
        default='RGB Origin',
    )
    DISPLAY:  BoolProperty(
                name="Show Float Screen",
                description="Display the float screen to show rgb and depth image",
                default=True,
            )
    TRACK:  BoolProperty(
                name="Track the camera",
                description="Automaticly change to active camera view",
                default=True,
            )
    BACKGROUND:  BoolProperty(
                name="Show pairwise image in background",
                description="Display the image(either depth or rgb) in the camera view background",
                default=True,
            )
    background_alpha: FloatProperty(name="BackgoundAlhpa", 
                            description="Alhpa value for background", 
                            default=1.00, 
                            min=0.00, 
                            max=1.00, 
                            step=0.1, 
                            precision=2)

    ALIGN:  BoolProperty(
                name="Align the float screen and camera view",
                description="Align the float screen and camera view to verify the model pose and segmentation",
                default=False,
            )
    empty_alpha: FloatProperty(name="EmptyAlpha", 
                                description="Alhpa value for empty regin", 
                                default=0.2, 
                                min=0.00, 
                                max=1.00, 
                                step=0.1, 
                                precision=2)

    segment_alpha: FloatProperty(name="SegmentAlpha", 
                                description="Alhpa value for segmented regin", 
                                default=1.00, 
                                min=0.00, 
                                max=1.00, 
                                step=0.1, 
                                precision=2)

    display_scale: FloatProperty(name="Display Scale", 
                                description="Display Scale for the floating screen", 
                                default=0.50, 
                                min=0.00, 
                                max=4.00, 
                                step=0.1, 
                                precision=2)
    display_X: FloatProperty(name="X position", 
                            description="X bias position for the floating screen", 
                            default=0, 
                            min=0, 
                            max=2000, 
                            step=100, 
                            precision=0)
    display_Y: FloatProperty(name="Y position", 
                            description="Y bias position for the floating screen", 
                            default=0, 
                            min=0, 
                            max=2000, 
                            step=100, 
                            precision=0)
    
    UPDATE_DEPTHFILTER: BoolProperty(
                name="Update depth filter on image",
                description="Display filtered depth on screen (would be slow)",
                default=True,
            )
    
    def ignorezerodepthUpdate(self, context):
        context.object['depth']["UPDATEALPHA"] = True
        context.object['rgb']["UPDATEALPHA"] = True
    IGNORE_ZERODEPTH: BoolProperty(
                name="Ignore zero depth",
                description="If true, depth with zero value would be filtered out",
                default=False,
                update=ignorezerodepthUpdate
            )
    
def right_click_menu_func(self, context):
    layout = self.layout
    layout.separator()
    layout.operator("export_data.poseoptimization")
    layout.operator("object_property.lockcurrent3darea")
    layout.operator("object_property.unlockcurrent3darea")
    layout.operator("export_data.savemask")

def register():
    global floatscreen_handler
    bpy.utils.register_class(FloatScreenProperty)
    bpy.types.Scene.floatscreenproperty = bpy.props.PointerProperty(type=FloatScreenProperty)
    floatscreen_handler = bpy.types.SpaceView3D.draw_handler_add(draw, (), 'WINDOW', 'POST_PIXEL')
    bpy.types.VIEW3D_MT_object_context_menu.append(right_click_menu_func)

def unregister():
    bpy.utils.unregister_class(FloatScreenProperty)
    bpy.types.SpaceView3D.draw_handler_remove(floatscreen_handler, 'WINDOW')
    bpy.types.VIEW3D_MT_object_context_menu.remove(right_click_menu_func)

panel/FloatScreenPanel.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
from kernel.blender_utility import _get_configuration

import numpy as np

class ObjectPropertyPanel(bpy.types.Panel):
    """Creates a Panel in the scene context of the properties editor"""
    bl_label = "Progress Labeller"
    bl_idname = "OBJECT_PROPERTY_layout"
    bl_space_type = 'PROPERTIES'
    bl_region_type = 'WINDOW'
    bl_context = "object"

    def draw(self, context):
        current_object = bpy.context.object.name
        scene = context.scene
        if "type" in bpy.data.objects[current_object]:
            name = current_object.split(":")[0]
            object_type = bpy.data.objects[current_object]["type"]
            config_id, config = _get_configuration(context.object)
            if object_type == "model":
                layout = self.layout
                scene = context.scene
                #obj_name = bpy.context.selected_objects[0].name
                if len(bpy.context.selected_objects) > 0:
                    obj_name = bpy.context.selected_objects[0].name

                    layout.label(text="Mask Optimization")
                    box = layout.box() 
                    row = box.row()
                    row.prop(bpy.data.objects[obj_name].optimizationparas, "stop_criteria")
                    row = box.row()
                    row.prop(bpy.data.objects[obj_name].optimizationparas, "translation_learning_rate")
                    row = box.row()
                    row.prop(bpy.data.objects[obj_name].optimizationparas, "rotation_learning_rate")
                    row = layout.row()
                    row.operator("object_property.maskoptimation")
                
                
            elif object_type == "reconstruction":
                workspace_name = current_object.split(":")[0]   
                config_id = bpy.data.objects[workspace_name + ":Setting"]['config_id']

                layout = self.layout
                scene = context.scene
                layout.label(text="Set scale:")
                row = layout.row()
                row.prop(scene.configuration[config_id], 'reconstructionscale')
                row.prop(scene.configuration[config_id], 'cameradisplayscale')
                layout.label(text="Align plane in reconstruction to X-Y:")
                row = layout.row()
                row.operator("object_property.planealignment")

                layout.label(text="Align all models:")
                row = layout.row()
                row.operator("object_property.allmodelsicp")
                

            elif object_type == "camera":
                layout = self.layout
                scene = context.scene
                layout.label(text="Set FloatingScreen Property:")
                box = layout.box() 
                row = box.row(align=True)
                row.prop(scene.floatscreenproperty, "viewimage_mode")
                row = box.row()
                row.prop(scene.floatscreenproperty, "DISPLAY")
                row = box.row()
                row.prop(scene.floatscreenproperty, "TRACK")
                if scene.floatscreenproperty.TRACK:
                    row = box.row()
                    row.prop(scene.floatscreenproperty, "ALIGN")
                    # row.operator("object_property.current3darea")
                    
                row = box.row(align=True)
                row.prop(scene.floatscreenproperty, "BACKGROUND")
                row.prop(scene.floatscreenproperty, "background_alpha")

                if not scene.floatscreenproperty.ALIGN:
                    row = box.row(align=True)
                    row.prop(scene.floatscreenproperty, "display_scale")
                    row = box.row(align=True)
                    row.prop(scene.floatscreenproperty, "display_X")
                    row.prop(scene.floatscreenproperty, "display_Y")
                
                layout.label(text="Set depth filter parameter:")
                
                box = layout.box() 
                row = box.row(align=True)
                row = box.row()
                row.prop(scene.floatscreenproperty, "UPDATE_DEPTHFILTER")
                row = box.row()
                row.prop(scene.floatscreenproperty, "IGNORE_ZERODEPTH")
                row = box.row()
                row.prop(config, "depth_scale")
                row = box.row()
                row.prop(config, "depth_ignore")
                #meandepth = np.mean(np.array(context.object["depth"]["depth"])) * config.depth_scale
                #box.label(text="Mean depth for current image is : {0:.3f}m".format(meandepth))
                

            elif object_type == "setting":
                object = bpy.data.objects[current_object]
                scene = context.scene
                layout = self.layout
                layout.label(text="Set Workspace Name:")
                row = layout.row()
                row.prop(config, 'projectname')
                row.operator("object_property.workspacerename")
                layout.label(text="Set Environment:")

                row = layout.row()
                row.prop(config, 'modelsrc')
                row.operator("import_data.model")

                row = layout.row()
                # row.prop(config, 'modelposesrc')
                row.operator("import_data.modelfrompose")
                row.operator("export_data.objectposes")

                row = layout.row()
                row.prop(config, 'datasrc')
                #row.operator("object_property.importcamrgbdepth")

                row = layout.row()
                row.prop(config, 'reconstructionsrc')
                
                row = layout.row()
                row.operator("object_property.importreconresult")


                


                layout.label(text="Set Camera Parameters:")
                box = layout.box() 
                row = box.row(align=True)
                row.prop(scene.configuration[object["config_id"]], "fx")
                row.prop(scene.configuration[object["config_id"]], "fy")
                row = box.row(align=True)
                row.prop(scene.configuration[object["config_id"]], "cx")
                row.prop(scene.configuration[object["config_id"]], "cy")
                row = box.row(align=True)
                row.prop(scene.configuration[object["config_id"]], "resX")
                row.prop(scene.configuration[object["config_id"]], "resY")
                


                layout.label(text="Export Configuration:")
                box = layout.box() 
                row = box.row(align=True)
                row.operator("export_data.configuration")

                layout.label(text="Remove Current Workspace:")
                box = layout.box() 
                row = box.row(align=True)
                row.operator("object_property.removeworkspace")


                layout.label(text="Data Output:")
                box = layout.box() 
                row = box.row(align=True)
                row.operator("export_data.dataoutput")  

def register():
    bpy.utils.register_class(ObjectPropertyPanel)


def unregister():
    bpy.utils.unregister_class(ObjectPropertyPanel)

panel/ObjectPropertyPanel.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
import bpy
import json
import os
import numpy as np
from mathutils import Vector
from kernel.geometry import _pose2Rotation, _rotation2Pose
from kernel.blender_utility import _is_progresslabeller_object, _get_allrgb_insameworkspace, _get_configuration, _get_obj_insameworkspace

config_json_dict = {
    'projectname': [['projectname']],
    'modelsrc': [['environment', 'modelsrc']],
    'reconstructionsrc':[['environment', 'reconstructionsrc']] ,
    'datasrc': [['environment', 'datasrc']],
    'resX':[['camera', 'resolution'], ['0']],  
    'resY':[['camera', 'resolution'], ['1']], 
    'fx':[['camera', 'intrinsic'], ['0', '0']],  
    'fy':[['camera', 'intrinsic'], ['1', '1']],   
    'cx':[['camera', 'intrinsic'], ['0', '2']],    
    'cy':[['camera', 'intrinsic'], ['1', '2']],    
    'lens':[['camera', 'lens']],
    'inverse_pose':[['camera', 'inverse_pose']],
    'reconstructionscale': [['reconstruction', 'scale']],
    'cameradisplayscale': [['reconstruction', 'cameradisplayscale']],
    'recon_trans': [['reconstruction', 'recon_trans']],
    'sample_rate': [['data', 'sample_rate']],
    'depth_scale': [['data', 'depth_scale']],
    'depth_ignore': [['data', 'depth_ignore']],

}

def decode_dict(configuration, code):
    value = configuration
    for item in code[0]:
        value = value[item]
    if len(code) == 2:
        for item in code[1]:
            value = value[int(item)]
    return value

def encode_dict(configuration):
    output_dict = {
        'projectname': configuration.projectname,
        'environment': {
            'modelsrc':configuration.modelsrc,
            "reconstructionsrc":configuration.reconstructionsrc,
            "datasrc": configuration.datasrc,           
        },
        'camera':{
        "resolution": [configuration.resX, configuration.resY],
        "intrinsic": [[configuration.fx, 0, configuration.cx],
                      [0, configuration.fy, configuration.cy],
                      [0, 0, 1]],
        "inverse_pose": configuration.inverse_pose,
        "lens": configuration.lens, 
        },
        'reconstruction':{
        "scale": configuration.reconstructionscale,
        "cameradisplayscale": configuration.cameradisplayscale,
        "recon_trans": configuration.recon_trans
        },
        'data':{
        "sample_rate": configuration.sample_rate,
        "depth_scale": configuration.depth_scale,
        "depth_ignore": configuration.depth_ignore,
        }
    }
    return output_dict


class config(bpy.types.PropertyGroup):
    # The properties for this class which is referenced as an 'entry' below.

    projectname: bpy.props.StringProperty(name = "projectname")
    modelsrc: bpy.props.StringProperty(name = "modelsrc", 
                subtype = "DIR_PATH")
    reconstructionsrc: bpy.props.StringProperty(name = "reconstructionsrc", 
                subtype = "DIR_PATH") 
    datasrc: bpy.props.StringProperty(name = "datasrc", 
                subtype = "DIR_PATH") 


    resX: bpy.props.IntProperty(name = "resX", default= 640)
    resY: bpy.props.IntProperty(name = "resY", default= 480)
    fx: bpy.props.FloatProperty(name="fx", description="camera intrinsic fx.", 
        min=0.00, max=1500.00, step=10, precision=2, default= 601.17364502)    
    fy: bpy.props.FloatProperty(name="fy", description="camera intrinsic fy.", 
        min=0.00, max=1500.00, step=10, precision=2, default= 600.0300293)   
    cx: bpy.props.FloatProperty(name="cx", description="camera intrinsic cx.", 
        min=0.00, max=1000.00, step=10, precision=2, default= 332.08534315)    
    cy: bpy.props.FloatProperty(name="cy", description="camera intrinsic cy.", 
        min=0.00, max=1000.00, step=10, precision=2, default= 243.56048719)
    lens: bpy.props.FloatProperty(name="lens", description="camera lens length", 
        min=0.000, max=1000.000, step=3, precision=3, default = 30)

    inverse_pose: bpy.props.BoolProperty(
        name="Inverse Camera Pose",
        description="Need when given poses are from world to camera",
        default=False,
    )       
    
    recon_trans: bpy.props.StringProperty(name = "recon_trans", default = "1,0,0,0;0,1,0,0;0,0,1,0;0,0,0,1;")
    
    sample_rate: bpy.props.FloatProperty(name="Sample Rate", 
                                        description="Sample rate for loading in blender for visualize", 
                                        default=0.10, 
                                        min=0.00, 
                                        max=1.00, 
                                        step=0.1, 
                                        precision=2)
    
    def depthInfoUpdate(self, context):
        if context.object != None and _is_progresslabeller_object(context.object):
            _, config = _get_configuration(context.object)
            im_list = _get_allrgb_insameworkspace(config)
            for im in im_list:
                im["UPDATEALPHA"] = True
                        
    depth_scale: bpy.props.FloatProperty(name="Depth Scale", 
                                        description="Scale for depth image", 
                                        default=0.001, 
                                        min=0.000000, 
                                        max=10.000000, 
                                        step=6, 
                                        precision=6,
                                        update=depthInfoUpdate)  

    depth_ignore: bpy.props.FloatProperty(name="Depth Ignore (m)", 
                                        description="Use depth as a filter for rgb", 
                                        default=10, 
                                        min=0.0, 
                                        max=10000.0, 
                                        step=3, 
                                        precision=3,
                                        update=depthInfoUpdate)  

    cameradisplayscale: bpy.props.FloatProperty(name="reconstruction scale", 
                                        description="reconstruction scale for the fused.ply", 
                                        default=0.05, 
                                        min=0.00, 
                                        max=1000.00, 
                                        step=0.01, 
                                        precision=2)   



    def scale_update(self, context):
        
        if hasattr(context, "object") \
            and context.object \
            and "type" in context.object \
            and context.object["type"] == "reconstruction":

            recon = context.object
            workspacename = recon.name.split(":")[0]
            old_scale = recon["scale"]
            new_scale = self.reconstructionscale
            alignT = np.asarray(recon["alignT"])

            self.cameradisplayscale = self.cameradisplayscale * new_scale/old_scale

            for obj in bpy.data.objects:
                if obj.name.startswith(workspacename):
                    if obj["type"] == "reconstruction":
                        obj.scale = Vector((new_scale, new_scale, new_scale))
                    if obj["type"] == "camera":
                        current_location = np.asarray(obj.location).reshape((3, 1))
                        origin_location = np.linalg.inv(alignT[:3, :3]).dot(current_location - alignT[:3, [3]])/old_scale
                        scaled_location = origin_location * new_scale
                        current_scaled_location = (alignT[:3, :3].dot(scaled_location) + alignT[:3, [3]]).reshape((3, ))
                        obj.location = Vector(current_scaled_location)

            recons = _get_obj_insameworkspace(recon, ["reconstruction"])

            for obj in recons:
                obj["scale"] = self.reconstructionscale   

    
    reconstructionscale: bpy.props.FloatProperty(name="reconstruction scale", description="scale of the reconstruction model", 
        min=0.01, max=1000.000, step=3, precision=3, default = 1.000, 
        update=scale_update)

    def cameradisplayscale_update(self, context):
        
        if hasattr(context, "object") \
            and context.object \
            and "type" in context.object \
            and context.object["type"] == "reconstruction":
            recon = context.object
            workspacename = recon.name.split(":")[0]
            for obj in bpy.data.objects:
                if obj.name.startswith(workspacename):
                    if obj["type"] == "camera":
                        obj.data.display_size = self.cameradisplayscale
                    

    cameradisplayscale: bpy.props.FloatProperty(name="camera display scale", description="scale of the camera display size", 
        min=0.01, max=1.000, step=0.1, precision=3, default = 0.1, 
        update=cameradisplayscale_update)



def register():
    bpy.utils.register_class(config)
    bpy.types.Scene.configuration = bpy.props.CollectionProperty(type=config)  

def unregister():
    bpy.utils.unregister_class(config)


registeration/init_configuration.py
MyungHwanJeon/ProgressLabeller_for_TRansPose
from typing import Dict
import bpy
from file.fileImporter import model_loader, reconstruction_result_loader, configuration_loader
from file.fileNew import new_workspace
from file.fileExporter import configuration_export, objectposes_export, data_export
from registeration import init_configuration
from panel import ObjectPropertyPanel, FloatScreenPanel
from operators import ObjectPropertyOperator, ReconstructionOperator
from kernel.blender_utility import clear_initial_object

area_image_pair = {
}

def register():
    init_configuration.register()
    configuration_loader.register()
    model_loader.register()
    reconstruction_result_loader.register()
    new_workspace.register()
    ObjectPropertyOperator.register()
    #ReconstructionOperator.register()
    ObjectPropertyPanel.register()
    FloatScreenPanel.register()

    configuration_export.register()
    objectposes_export.register()
    data_export.register()

def unregister():
    init_configuration.unregister()
    configuration_loader.unregister()
    model_loader.unregister()
    reconstruction_result_loader.unregister()
    new_workspace.unregister()
    ObjectPropertyPanel.unregister()
    ObjectPropertyOperator.unregister()
    #ReconstructionOperator.unregister()
    configuration_export.unregister()
    objectposes_export.unregister()
    FloatScreenPanel.unregister()
    data_export.unregister()

registeration/register.py
MyungHwanJeon/ProgressLabeller_for_TRansPose