Inference script on a single file

[ellkrauze]

Hello! Is it possible to get the OBJ file from a CC file using pretrained model? If so, how do I do it properly?

I've managed to create this script so far, but it creates a TestOnlyDataset which is not a proper way to get the model's inference, I suppose...

import os
import torch
import pickle
import numpy as np
from pathlib import Path

import yaml
import hydra
from torch_geometric.loader import DataLoader
from network import PolyGNN
from dataset import TestOnlyDataset
from utils import Sampler
from abspy import VertexGroup, CellComplex
from omegaconf import DictConfig

from utils import init_device, Sampler, set_seed

def load_model(cfg: DictConfig, device: str='cpu'):
    """Load the PolyGNN model from a checkpoint."""
    model = PolyGNN(cfg)
    state = torch.load(cfg.checkpoint_path, map_location=device)
    state_dict = state['state_dict']
    model.load_state_dict(state_dict, strict=False)
    model.eval()
    return model

def preprocess_data(file_path, config):
    """Prepare the dataset for inference."""
    sampler = Sampler(strategy=config.sample.strategy,
                      length=config.sample.length,
                      ratio=config.sample.ratio,
                      resolutions=config.sample.resolutions,
                      duplicate=config.sample.duplicate,
                      seed=config.seed)
    transform = sampler.sample if config.sample.transform else None
    dataset = TestOnlyDataset(pre_transform=None, transform=transform, root=file_path, split=None)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=False)
    return dataloader

def run_inference(model, dataloader, device, output_dir):
    """Run inference on the given dataset."""
    Path(output_dir).mkdir(exist_ok=True, parents=True)

    for batch in dataloader:
        batch = batch.to(device)
        with torch.no_grad():
            outputs = model(batch)
            predictions = outputs.argmax(dim=1)

        # Save predictions
        for i, name in enumerate(batch.name):
            output_path = Path(output_dir) / f"{name}.npy"
            np.save(output_path, predictions[i].cpu().numpy())

def convert_data(src_path, dst_path=None):
    if dst_path is None:
        base, ext = os.path.splitext(src_path)
        dst_path = os.path.join(base + ".cc")
    vertex_group = VertexGroup(src_path, quiet=True)
    cell_complex = CellComplex(vertex_group.planes, vertex_group.aabbs,
                               vertex_group.points_grouped, build_graph=True, quiet=True)
    cell_complex.prioritise_planes(prioritise_verticals=True)
    cell_complex.construct()
    cell_complex.save(dst_path)
    print(f"New file written to {dst_path}")
    return dst_path

@hydra.main(config_path='./conf', config_name='infer_config', version_base='1.2')
def main(cfg: DictConfig):
    set_seed(cfg.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    model = load_model(cfg).to(device)

    cc_file = convert_data(cfg.data_file)
    print(f"Converted {cfg.data_file} to cc")

    dataloader = preprocess_data(cc_file, cfg)
    print("Created dataloader")

    run_inference(model, dataloader, device, cfg.output_dir)
    print(f"Inference completed. Results saved to {cfg.output_dir}")


if __name__ == "__main__":
    main()

[chenzhaiyu]

Hi, yes, it’s possible. The model uses a PolyGraph as input. If you have the CC file, you can provide it as shown in this line. You can also modify the TestOnlyDataset implementation to handle multiple samples if needed.

[ellkrauze]

Thanks! I'll try this one :)

[ellkrauze]

I’ve successfully run inference on a custom .vg file, but the model’s output consists only of 0s without any errors. Could this indicate an issue with my preprocessing or inference setup?

Here’s my inference script:

import os
import torch
import pickle
import numpy as np
from pathlib import Path
from tqdm import tqdm
import hydra
from torch_geometric.loader import DataLoader
from network import PolyGNN
from dataset import PolyGraph
from abspy import VertexGroup, CellComplex
from omegaconf import DictConfig
from utils import set_seed
from torchmetrics.classification import BinaryAccuracy
import logging
import multiprocessing

logger = logging.getLogger(__name__)

BACKEND = 'nccl' if torch.cuda.is_available() else 'gloo'

class PredictionSaver:
    """
    Asynchronous prediction saving.
    """
    def __init__(self, processes):
        self.pool = multiprocessing.Pool(processes=processes)

    @staticmethod
    def save(args):
        pred, name, cfg = args
        indices_cells = np.where(pred.cpu().numpy())[0]
        if len(indices_cells) > 0:
            complex_path = f'{cfg.complex_dir}/{name}.cc'
            with open(complex_path, 'rb') as handle:
                cell_complex = pickle.load(handle)
            output_path = f'{cfg.output_dir}/{name}.npy'
            output = np.zeros([cell_complex.num_cells], dtype=int)
            output[indices_cells] = 1

            if cfg.evaluate.seal:
                cells_boundary = cell_complex.cells_boundary()
                output[cells_boundary] = 0
            logger.info(f"Saving to {output_path}")
            np.save(output_path, output)

def load_model(cfg: DictConfig, rank):
    model = PolyGNN(cfg)
    metric = BinaryAccuracy()
    model.metric = metric
    model = model.to(rank)

    torch.compile(model, dynamic=True, fullgraph=True)

    map_location = f'cuda:{rank}'
    if rank == 0:
        logger.info(f'Resuming from {cfg.checkpoint_path}')

    state = torch.load(cfg.checkpoint_path, map_location=map_location)
    state_dict = state['state_dict']
    model.load_state_dict(state_dict, strict=False)
    return model

def run_inference(model, dataloader, rank, output_dir, cfg):
    """Run inference on the given dataset."""
    Path(output_dir).mkdir(exist_ok=True, parents=True)

    if cfg.sample.strategy == 'grid':
        points_suffix = f'_{cfg.sample.resolution}'
    elif cfg.sample.strategy == 'random':
        points_suffix = f'_{cfg.sample.length}'
    else:
        points_suffix = ''

    pbar = tqdm(dataloader, desc=f'eval', disable = rank!=0)
    world_size = len(cfg.gpu_ids)
    prediction_saver = PredictionSaver(processes=cfg.num_workers // world_size)

    model.eval()

    with torch.no_grad():
        for batch in pbar:
            batch = batch.to(rank, f'points{points_suffix}', f'batch_points{points_suffix}', 'queries',
                             'edge_index', 'batch', 'y')
            batch = batch.to(rank)
            outs = model(batch)
            outs = outs.argmax(dim=1)
            targets = batch.y

            Path(cfg.output_dir).mkdir(exist_ok=True)
            _, boundary_indices = torch.unique(batch.batch, return_counts=True)

            preds = torch.split(outs, split_size_or_sections=boundary_indices.tolist(), dim=0)

            names = ["test"]

            for pred, name in zip(preds, names):
                logger.info(f'Evaluating {name}')
                logger.info(f'Prediction: {pred}')
                prediction_saver.save((pred, name, cfg))

def bvg2cc(src_path, dst_path=None):
    if dst_path is None:
        base, ext = os.path.splitext(src_path)
        dst_path = os.path.join(base + ".cc")
    vertex_group = VertexGroup(src_path, quiet=True)
    cell_complex = CellComplex(vertex_group.planes, vertex_group.aabbs,
                               vertex_group.points_grouped, build_graph=True, quiet=True)
    cell_complex.prioritise_planes(prioritise_verticals=True)
    cell_complex.construct()
    cell_complex.save(dst_path)
    print(f"New file written to {dst_path}")
    return dst_path

@hydra.main(config_path='./conf', config_name='infer_config', version_base='1.2')
def main(cfg: DictConfig):
    set_seed(cfg.seed)
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    rank = 0
    print(f"Using device: {device}")

    model = load_model(cfg, rank)
    complex_path = "data/CHE2_sampled.cc"
    vertex_group_path = "data/CHE2_sampled.vg"

    # bvg2cc(vertex_group_path, complex_path)

    pg = PolyGraph(num_queries=cfg.num_queries)
    data = pg.data_loader(
        cloud_path=None,
        # complex_path=complex_path,  # returns "points = self.vertex_group.points AttributeError: 'NoneType' object has no attribute 'points'"
        vertex_group_path=vertex_group_path
    )
    print(f"Data: {data}")
    dataloader = DataLoader([data], batch_size=32, shuffle=True)
    run_inference(model, dataloader, rank, cfg.output_dir, cfg)


if __name__ == "__main__":
    main()

Model prediction is

tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0], device='cuda:0')

Test .vg file (inside zip archive): CHE2_sampled.vg.zip

Any insights would be appreciated!

[chenzhaiyu]

Just a quick view at your data I see the coordinate range isn't right: if you did normalize the data properly you wouldn't have a Z range (0.3, 1.7).