Diffusion-3D-Features/extract_tosca.py at main · niladridutt/Diffusion-3D-Features · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import torch
import os
import glob
from diff3f import get_features_per_vertex
from time import time
from utils import convert_mesh_container_to_torch_mesh
import configparser
from datetime import datetime
from dataloaders.mesh_container import MeshContainer
from diffusion import init_pipe
from dino import init_dino


config = configparser.ConfigParser()
now = datetime.now()
date_time = now.strftime("%d-%m-%Y_%H-%M")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
min_points = 10000
num_samples = "all"
num_views = 100
H = 512
W = 512

make_prompt_null = False
num_images_per_prompt = 1
ensemble_size = 1
divide_by_sum = False
tolerance = 0.01
random_seed = 42
rotated = False
use_normal_map = True
dataset = "TOSCA"
prompt = "animal extracted from file name"
target_path = (
    "datasets/tosca/Meshes/*.off"
)
save_path = (
    f"output/tosca/{date_time}"
)
add_path = dataset

if rotated:
    add_path += "-rotated"
save_path += add_path
if not os.path.isdir(save_path):
    os.makedirs(save_path)
target_files = sorted(glob.glob(target_path))


def write_config_file():
    config["Parameters"] = {
        "dataset": dataset,
        "method": "dino",
        "rotated": str(rotated),
        "prompt": prompt,
        "num_samples": str(num_samples),
        "num_views": str(num_views),
        "H": str(H),
        "W": str(W),
        "tolerance": str(tolerance),
        "make_prompt_null": str(make_prompt_null),
        "num_images_per_prompt": str(num_images_per_prompt),
        "ensemble_size": str(ensemble_size),
        "divide_by_sum": str(divide_by_sum),
        "target_path": target_path,
        "save_path": save_path,
        "random_seed": str(random_seed),
        "min_points": str(min_points),
        "use_normal_map": str(use_normal_map),
    }
    print(config)
    # Write the config file to disk
    with open(f"configs/config{add_path}-{date_time}.ini", "w") as f:
        config.write(f)
    with open(f"{save_path}/config{add_path}-{date_time}.ini", "w") as f:
        config.write(f)


def compute_features():
    t1 = time()
    pipe = init_pipe(device)
    dino_model = init_dino(device)
    for file in target_files:
        try:
            print(f"Processing {file}")
            filename = file.split("/")[-1].split(".")[0]
            m = MeshContainer().load_from_file(str(file))
            mesh = convert_mesh_container_to_torch_mesh(m, device=device, is_tosca=True)
            mesh_vertices = mesh.verts_list()[0]
            prompt = "".join(filter(str.isalpha, filename))
            print("Prompt: ", prompt)
            features = get_features_per_vertex(
                device=device,
                pipe=pipe,
                dino_model=dino_model,
                mesh=mesh,
                prompt=prompt,
                mesh_vertices=mesh_vertices,
                num_views=num_views,
                H=H,
                W=W,
                tolerance=tolerance,
                num_images_per_prompt=num_images_per_prompt,
                use_normal_map=use_normal_map,
            )
            save_name = f"{save_path}/{filename}.pt"
            torch.save(features, save_name)
        except Exception as e:
            print(e)
    time_taken = (time() - t1) / 60
    print("Time taken to complete in minutes", time_taken)


if __name__ == "__main__":
    write_config_file()
    compute_features()