infer.py

import os
import torch
from contextlib import nullcontext
from tqdm import tqdm
from da2 import (
    prepare_to_run,
    load_model,
    load_infer_data,
    distance2pointcloud,
    colorize_distance,
    concatenate_images
)


def infer(model, config, accelerator, output_dir):
    model.eval()
    if accelerator.is_main_process:
        if torch.backends.mps.is_available():
            autocast_ctx = nullcontext()
        else:
            autocast_ctx = torch.autocast(accelerator.device.type)
        with autocast_ctx, torch.no_grad():
            infer_data = load_infer_data(config, accelerator.device)
            pred_distances = []
            for i in tqdm(range(infer_data['size']), desc='Predicting 360° depth'):
                distances = model(infer_data['images']['torch'][i])
                pred_distances.append(distances.cpu().numpy())
            pred_distance_images = []
            pred_normal_images = []
            for i in tqdm(range(infer_data['size']), desc='Visualizing 360° depth'):
                pred_distance_images.append(colorize_distance(pred_distances[i], infer_data['masks'][i]))
            for i in tqdm(range(infer_data['size']), desc='Saving 3D points'):
                normal_image = distance2pointcloud(pred_distances[i], 
                    infer_data['images']['cv2'][i], infer_data['masks'][i], 
                    save_path=os.path.join(output_dir, f'3dpc/{infer_data['filenames'][i]}.ply'), return_normal=True, save_distance=True)
                pred_normal_images.append(normal_image)
            concatenate_images(infer_data['images']['PIL'], pred_distance_images, pred_normal_images).save(os.path.join(output_dir, 'vis_all.png'))

if __name__ == '__main__':
    config, accelerator, output_dir = prepare_to_run()
    model = load_model(config, accelerator)
    infer(model, config, accelerator, output_dir)