app.py

import multiprocessing
import objaverse
import random
import numpy
import trimesh
import huggingface_hub
from huggingface_hub import Repository
import os
import pickle
import json
from tqdm import tqdm

DATASET_REPO_URL = "https://huggingface.co/datasets/Gary3410/object_tmp"
DATA_FILENAME = "object_tmp_22_3.pkl"
DATA_FILE = os.path.join("data", DATA_FILENAME)


def random_sampling(points, num_samples):
    # 如果点云数量小于采样数量，允许重复选择点
    if len(points) < num_samples:
        indices = numpy.random.choice(len(points), num_samples, replace=True)
    else:
        indices = numpy.random.choice(len(points), num_samples, replace=False)
    sampled_points = points[indices]
    return sampled_points

def model_to_pc(mesh: trimesh.Trimesh, n_sample_points=10000):
    f32 = numpy.float32
    rad = numpy.sqrt(mesh.area / (3 * n_sample_points))
    for _ in range(24):
        pcd, face_idx = trimesh.sample.sample_surface_even(mesh, n_sample_points, rad)
        rad *= 0.85
        if len(pcd) == n_sample_points:
            break
    else:
        raise ValueError("Bad geometry, cannot finish sampling.", mesh.area)
    if isinstance(mesh.visual, trimesh.visual.ColorVisuals):
        rgba = mesh.visual.face_colors[face_idx]
    elif isinstance(mesh.visual, trimesh.visual.TextureVisuals):
        bc = trimesh.proximity.points_to_barycentric(mesh.triangles[face_idx], pcd)
        if mesh.visual.uv is None or len(mesh.visual.uv) < mesh.faces[face_idx].max():
            uv = numpy.zeros([len(bc), 2])
            # st.warning("Invalid UV, filling with zeroes")
        else:
            uv = numpy.einsum('ntc,nt->nc', mesh.visual.uv[mesh.faces[face_idx]], bc)
        material = mesh.visual.material
        if hasattr(material, 'materials'):
            if len(material.materials) == 0:
                rgba = numpy.ones_like(pcd) * 0.8
                texture = None
                # st.warning("Empty MultiMaterial found, falling back to light grey")
            else:
                material = material.materials[0]
        if hasattr(material, 'image'):
            texture = material.image
            if texture is None:
                rgba = numpy.zeros([len(uv), len(material.main_color)]) + material.main_color
        elif hasattr(material, 'baseColorTexture'):
            texture = material.baseColorTexture
            if texture is None:
                rgba = numpy.zeros([len(uv), len(material.main_color)]) + material.main_color
        else:
            texture = None
            rgba = numpy.ones_like(pcd) * 0.8
            # st.warning("Unknown material, falling back to light grey")
        if texture is not None:
            rgba = trimesh.visual.uv_to_interpolated_color(uv, texture)
    if rgba.max() > 1:
        if rgba.max() > 255:
            rgba = rgba.astype(f32) / rgba.max()
        else:
            rgba = rgba.astype(f32) / 255.0
    return numpy.concatenate([numpy.array(pcd, f32), numpy.array(rgba, f32)[:, :3]], axis=-1)


def trimesh_to_pc(scene_or_mesh):
    if isinstance(scene_or_mesh, trimesh.Scene):
        meshes = []
        for node_name in scene_or_mesh.graph.nodes_geometry:
            # which geometry does this node refer to
            transform, geometry_name = scene_or_mesh.graph[node_name]

            # get the actual potential mesh instance
            geometry = scene_or_mesh.geometry[geometry_name].copy()
            if not hasattr(geometry, 'triangles'):
                continue
            geometry: trimesh.Trimesh
            geometry = geometry.apply_transform(transform)
            meshes.append(geometry)
        total_area = sum(geometry.area for geometry in meshes)
        if total_area < 1e-6:
            raise ValueError("Bad geometry: total area too small (< 1e-6)")
        pcs = []
        for geometry in meshes:
            pcs.append(model_to_pc(geometry, max(1, round(geometry.area / total_area * 10000))))
        if not len(pcs):
            raise ValueError("Unsupported mesh object: no triangles found")
        return numpy.concatenate(pcs)
    else:
        assert isinstance(scene_or_mesh, trimesh.Trimesh)
        return model_to_pc(scene_or_mesh, 10000)



processes = multiprocessing.cpu_count()

# uids = objaverse.load_uids()
# random_object_uids = random.sample(uids, 100)
uids = []
object_id_tmp_dict = {}

# 解析json文件
with open('object_22.json','r') as file:
    str = file.read()
    obj_data = json.loads(str)

for ints_key in obj_data.keys():
    ints_dict_one = obj_data[ints_key]
    # ints_id_dict = ints_dict_one["obj_id"]
    # for ints_one in ints_id_dict.keys():
    #     uid_one = ints_id_dict[ints_one]
    #     uids.append(uid_one)
    for ints_one in ints_dict_one.keys():
        uid_one = ints_dict_one[ints_one]
        uids.append(uid_one)

uids = list(set(uids))
uids = sorted(uids)
# 2195
# uids = uids[:809]
# uids = uids[809:1619]
uids = uids[1619:]

# dict = {
#         "cracker_box": "dbbfa480a2e4476c8b0ca91667087f44",
#         "packet": "82ab1bb22a22407e8ebd86cae46dc9d3",
#         "mug": "a831091dabdb40598ed1681e2aa35058",
#         "chips_can": "6848fd0ec6df45bcb53fc0e47b51e844",
#         "baseball": "4799dc6cb5f046bd88609381020feab1",
#         "paper_towel": "d020130cdc0b4937a062f8d744d707c6",
#         "banana": "1de9ac4b1a4e46279c271747becbf0eb",
#         "bowl": "9fa4ca710fa9458aaeb9554c80583c76",
#         "pear": "fa74a39529f5462e978b4f5dbc0b8add",
#         "sponge": "939ca8d1e0c2424ba28db9afd19e7e53",
#         "tennis_ball": "05eb4a68e56249c5b3db4ea937a57b6e",
#         "plastic_bag": "686c497301da4fa5b6927b79340f473f"
#     }
# new_uids = []
# for dict_key_one in dict.keys():
#     new_uids_one = dict[dict_key_one]
#     if new_uids_one not in uids:
#         new_uids.append(new_uids_one)


objects = objaverse.load_objects(
    uids=uids,
    download_processes=processes
)

repo = Repository(
    local_dir="data", clone_from=DATASET_REPO_URL, use_auth_token=""
)


for objaid in tqdm(uids):
    objamodel = objaverse.load_objects([objaid])[objaid]
    try:
        pc = trimesh_to_pc(trimesh.load(objamodel))
        pc = random_sampling(pc, num_samples=5000)
        object_id_tmp_dict[objaid] = pc
    except:
        continue

with open(DATA_FILE, 'wb') as file:
    pickle.dump(object_id_tmp_dict, file)
commit_url = repo.push_to_hub()

print("Done")