|
|
import os |
|
|
import copy |
|
|
import sys |
|
|
import json |
|
|
import importlib |
|
|
import argparse |
|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
import numpy as np |
|
|
import pandas as pd |
|
|
import utils3d |
|
|
from tqdm import tqdm |
|
|
from easydict import EasyDict as edict |
|
|
from concurrent.futures import ThreadPoolExecutor |
|
|
from queue import Queue |
|
|
from torchvision import transforms |
|
|
from PIL import Image |
|
|
|
|
|
|
|
|
torch.set_grad_enabled(False) |
|
|
|
|
|
|
|
|
def get_data(frames, sha256): |
|
|
with ThreadPoolExecutor(max_workers=16) as executor: |
|
|
def worker(view): |
|
|
image_path = os.path.join(opt.output_dir, 'renders', sha256, view['file_path']) |
|
|
try: |
|
|
image = Image.open(image_path) |
|
|
except: |
|
|
print(f"Error loading image {image_path}") |
|
|
return None |
|
|
image = image.resize((518, 518), Image.Resampling.LANCZOS) |
|
|
image = np.array(image).astype(np.float32) / 255 |
|
|
image = image[:, :, :3] * image[:, :, 3:] |
|
|
image = torch.from_numpy(image).permute(2, 0, 1).float() |
|
|
|
|
|
c2w = torch.tensor(view['transform_matrix']) |
|
|
c2w[:3, 1:3] *= -1 |
|
|
extrinsics = torch.inverse(c2w) |
|
|
fov = view['camera_angle_x'] |
|
|
intrinsics = utils3d.torch.intrinsics_from_fov_xy(torch.tensor(fov), torch.tensor(fov)) |
|
|
|
|
|
return { |
|
|
'image': image, |
|
|
'extrinsics': extrinsics, |
|
|
'intrinsics': intrinsics |
|
|
} |
|
|
|
|
|
datas = executor.map(worker, frames) |
|
|
for data in datas: |
|
|
if data is not None: |
|
|
yield data |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument('--output_dir', type=str, required=True, |
|
|
help='Directory to save the metadata') |
|
|
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None, |
|
|
help='Filter objects with aesthetic score lower than this value') |
|
|
parser.add_argument('--model', type=str, default='dinov2_vitl14_reg', |
|
|
help='Feature extraction model') |
|
|
parser.add_argument('--instances', type=str, default=None, |
|
|
help='Instances to process') |
|
|
parser.add_argument('--batch_size', type=int, default=16) |
|
|
parser.add_argument('--rank', type=int, default=0) |
|
|
parser.add_argument('--world_size', type=int, default=1) |
|
|
opt = parser.parse_args() |
|
|
opt = edict(vars(opt)) |
|
|
|
|
|
feature_name = opt.model |
|
|
os.makedirs(os.path.join(opt.output_dir, 'features', feature_name), exist_ok=True) |
|
|
|
|
|
|
|
|
dinov2_model = torch.hub.load('facebookresearch/dinov2', opt.model) |
|
|
dinov2_model.eval().cuda() |
|
|
transform = transforms.Compose([ |
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), |
|
|
]) |
|
|
n_patch = 518 // 14 |
|
|
|
|
|
|
|
|
if os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')): |
|
|
metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv')) |
|
|
else: |
|
|
raise ValueError('metadata.csv not found') |
|
|
if opt.instances is not None: |
|
|
with open(opt.instances, 'r') as f: |
|
|
instances = f.read().splitlines() |
|
|
metadata = metadata[metadata['sha256'].isin(instances)] |
|
|
else: |
|
|
if opt.filter_low_aesthetic_score is not None: |
|
|
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score] |
|
|
if f'feature_{feature_name}' in metadata.columns: |
|
|
metadata = metadata[metadata[f'feature_{feature_name}'] == False] |
|
|
metadata = metadata[metadata['voxelized'] == True] |
|
|
metadata = metadata[metadata['rendered'] == True] |
|
|
|
|
|
start = len(metadata) * opt.rank // opt.world_size |
|
|
end = len(metadata) * (opt.rank + 1) // opt.world_size |
|
|
metadata = metadata[start:end] |
|
|
records = [] |
|
|
|
|
|
|
|
|
sha256s = list(metadata['sha256'].values) |
|
|
for sha256 in copy.copy(sha256s): |
|
|
if os.path.exists(os.path.join(opt.output_dir, 'features', feature_name, f'{sha256}.npz')): |
|
|
records.append({'sha256': sha256, f'feature_{feature_name}' : True}) |
|
|
sha256s.remove(sha256) |
|
|
|
|
|
|
|
|
load_queue = Queue(maxsize=4) |
|
|
try: |
|
|
with ThreadPoolExecutor(max_workers=8) as loader_executor, \ |
|
|
ThreadPoolExecutor(max_workers=8) as saver_executor: |
|
|
def loader(sha256): |
|
|
try: |
|
|
with open(os.path.join(opt.output_dir, 'renders', sha256, 'transforms.json'), 'r') as f: |
|
|
metadata = json.load(f) |
|
|
frames = metadata['frames'] |
|
|
data = [] |
|
|
for datum in get_data(frames, sha256): |
|
|
datum['image'] = transform(datum['image']) |
|
|
data.append(datum) |
|
|
positions = utils3d.io.read_ply(os.path.join(opt.output_dir, 'voxels', f'{sha256}.ply'))[0] |
|
|
load_queue.put((sha256, data, positions)) |
|
|
except Exception as e: |
|
|
print(f"Error loading data for {sha256}: {e}") |
|
|
|
|
|
loader_executor.map(loader, sha256s) |
|
|
|
|
|
def saver(sha256, pack, patchtokens, uv): |
|
|
pack['patchtokens'] = F.grid_sample( |
|
|
patchtokens, |
|
|
uv.unsqueeze(1), |
|
|
mode='bilinear', |
|
|
align_corners=False, |
|
|
).squeeze(2).permute(0, 2, 1).cpu().numpy() |
|
|
pack['patchtokens'] = np.mean(pack['patchtokens'], axis=0).astype(np.float16) |
|
|
save_path = os.path.join(opt.output_dir, 'features', feature_name, f'{sha256}.npz') |
|
|
np.savez_compressed(save_path, **pack) |
|
|
records.append({'sha256': sha256, f'feature_{feature_name}' : True}) |
|
|
|
|
|
for _ in tqdm(range(len(sha256s)), desc="Extracting features"): |
|
|
sha256, data, positions = load_queue.get() |
|
|
positions = torch.from_numpy(positions).float().cuda() |
|
|
indices = ((positions + 0.5) * 64).long() |
|
|
assert torch.all(indices >= 0) and torch.all(indices < 64), "Some vertices are out of bounds" |
|
|
n_views = len(data) |
|
|
N = positions.shape[0] |
|
|
pack = { |
|
|
'indices': indices.cpu().numpy().astype(np.uint8), |
|
|
} |
|
|
patchtokens_lst = [] |
|
|
uv_lst = [] |
|
|
for i in range(0, n_views, opt.batch_size): |
|
|
batch_data = data[i:i+opt.batch_size] |
|
|
bs = len(batch_data) |
|
|
batch_images = torch.stack([d['image'] for d in batch_data]).cuda() |
|
|
batch_extrinsics = torch.stack([d['extrinsics'] for d in batch_data]).cuda() |
|
|
batch_intrinsics = torch.stack([d['intrinsics'] for d in batch_data]).cuda() |
|
|
features = dinov2_model(batch_images, is_training=True) |
|
|
uv = utils3d.torch.project_cv(positions, batch_extrinsics, batch_intrinsics)[0] * 2 - 1 |
|
|
patchtokens = features['x_prenorm'][:, dinov2_model.num_register_tokens + 1:].permute(0, 2, 1).reshape(bs, 1024, n_patch, n_patch) |
|
|
patchtokens_lst.append(patchtokens) |
|
|
uv_lst.append(uv) |
|
|
patchtokens = torch.cat(patchtokens_lst, dim=0) |
|
|
uv = torch.cat(uv_lst, dim=0) |
|
|
|
|
|
|
|
|
saver_executor.submit(saver, sha256, pack, patchtokens, uv) |
|
|
|
|
|
saver_executor.shutdown(wait=True) |
|
|
except: |
|
|
print("Error happened during processing.") |
|
|
|
|
|
records = pd.DataFrame.from_records(records) |
|
|
records.to_csv(os.path.join(opt.output_dir, f'feature_{feature_name}_{opt.rank}.csv'), index=False) |
|
|
|
