eval_data.py

import os
from pathlib import Path
import trimesh
from tqdm import tqdm
import shutil
from eval_utils.utils import get_image_description, get_object_mask

class DataItem:
    def __init__(self, path,output_path,use_cached=True):
        self.valid = True
        self.id = os.path.basename(path)
        self.path = path
        self.output_path = os.path.join(output_path,self.id)
        self.use_cached = use_cached
        self.singapo_obj_path = None
        self.easitex_obj_path = None
        self.cosine_similarity = None
        self.cosine_similarity_no_easitex = None
        self.naive_cosine_similarity = None
        self.singapo_dict = None
        self.naive_texturing_path = os.path.join(self.output_path,"naive_texturing","object.obj")
        os.makedirs(os.path.dirname(self.naive_texturing_path), exist_ok=True)

        self.TEXTure_path = None

        

        if not os.path.exists(os.path.join(path,"imgs","00.png")):
            self.valid = False
            print(f"Invalid data item: {path}")
            return

        self.img_path = os.path.join(self.output_path,"image.png")

        os.makedirs(self.output_path, exist_ok=True)
        
        shutil.copy(os.path.join(path,"imgs","00.png"), self.img_path)

        self.parts_path = os.path.join(path,"objs")

        if not os.path.exists(self.parts_path):
            self.valid = False
            print(f"Invalid data item: {path}")
            return

        self.scene_path = self._construct_full_obj()
        self.gt_dict = os.path.join(path,"object.json")
        

        if use_cached and os.path.exists(os.path.join(self.output_path,"description.txt")):
            with open(os.path.join(self.output_path,"description.txt"), "r") as f:
                self.description = f.read()
        else:
            self.description = get_image_description(self.img_path)
            with open(os.path.join(self.output_path,"description.txt"), "w") as f:
                f.write(self.description)
        
        self.mask_path = os.path.join(self.output_path,"image_mask.png")
        if not (use_cached and os.path.exists(self.mask_path)):
            get_object_mask(self.img_path,self.mask_path)
    
    def _construct_full_obj(self):
        """
        Constructs the full object mesh from the parts.

        Returns:
            str: Path to the constructed full object mesh.
        """
        scene_path = os.path.abspath(os.path.join(self.output_path, "full_obj", f"{self.id}.glb"))

        if self.use_cached and os.path.exists(scene_path):
            return scene_path

        os.makedirs(os.path.join(self.output_path, "full_obj"), exist_ok=True)

        scene = trimesh.Scene()
        for part in os.listdir(self.parts_path):
            part_path = os.path.join(self.parts_path, part)
            if part.endswith('.obj'):
                mesh = trimesh.load(part_path, force='mesh')
                if not isinstance(mesh, trimesh.Trimesh):
                    print(f"{part_path} did not load as a Trimesh object")
                    continue
                scene.add_geometry(mesh,geom_name=part.replace('.obj','.ply'))

        scene.export(scene_path)

        return scene_path

    
    def set_singapo_obj_path(self, path):
        """
        Set the path to the Singapo generated object.

        Args:
            path (str): Path to the Singapo generated object.
        """
        self.singapo_obj_path = path

    def set_easitex_obj_path(self, path):
        """
        Set the path to the Easi-Tex generated object.

        Args:
            path (str): Path to the Easi-Tex generated object.
        """
        self.easitex_obj_path = path
    
    def set_cosine_similarity(self, similarity):
        """
        Set the cosine similarity between the generated object and the original object.

        Args:
            similarity (float): Cosine similarity value.
        """
        self.cosine_similarity = similarity

    def set_cosine_similarity_no_easitex(self, similarity):
        """
        Set the cosine similarity between the generated object and the original object without Easi-Tex.

        Args:
            similarity (float): Cosine similarity value.
        """
        self.cosine_similarity_no_easitex = similarity

    def set_naive_cosine_similarity(self, similarity):
        """
        Set the naive cosine similarity between the generated object and the original object.

        Args:
            similarity (float): Cosine similarity value.
        """
        self.naive_cosine_similarity = similarity

    def set_singapo_dict(self, path):
        """
        Set the path to the Singapo generated object dictionary.

        Args:
            path (str): Path to the Singapo generated object dictionary.
        """
        self.singapo_dict = path
    
    def set_TEXTure_path(self, path):
        """
        Set the path to the TEXTure generated object.

        Args:
            path (str): Path to the TEXTure generated object.
        """
        self.TEXTure_path = path
    



class EvaluationData:
    def __init__(self, data_path,output_path,use_cached=True):
        self.data_path = data_path
        self.output_path = output_path
        self.use_cached = use_cached
        self.items = []

        self._load_items()

    def _load_items(self):
        """
        Load all items in the dataset directory.
        """

        print(f"Loading data from {self.data_path}...")

        root = Path(self.data_path)

        item_paths = []

        for dataset in root.iterdir():
            if dataset.is_dir():
                for class_dir in dataset.iterdir():
                    if class_dir.is_dir():
                        for item in class_dir.iterdir():
                            if item.is_dir():
                                item_paths.append(item)
        
        for item_path in tqdm(item_paths):
            item = DataItem(item_path,self.output_path,self.use_cached)
            if item.valid:
                self.items.append(item)
            else:
                print(f"Invalid item: {item_path}, skipping...")


    def get_data_items(self):
        return self.items