eval_ginr.py

'''
Evaluates a trained model by predicting on one sample from the specified dataset.
- Bunny and Protein only

Arguments:
    - checkpoint: path to a Pytorch Lightning checkpoint file
    - idx: integer index of a sample from the dataset to predict and plotting
    - mesh: if not None, a path to a mesh file (obj, ply, etc.) on which to plotting the
            predicted signal. If None, the signal is plotted as a scatter plotting.
    - batch_size: number of points per batch when predicting with the trained model
                  (higher is better)

Note: requires the --dataset_dir flag to be specified as well.
'''

import pytorch_lightning as pl

from argparse import ArgumentParser
from plotly import graph_objects as go

from src.data.graph_dataset import GraphDataset
from src.models.graph_inr import GraphINR
from src.models.spatial_graph_inr import SpatialGraphINR
from src.plotting import traces
from src.plotting.figures import PLOT_CONFIGS, draw_mesh
from src.utils.data_generation import load_mesh
from src.utils.get_predictions import get_batched_predictions



pl.seed_everything(1234)

parser = ArgumentParser()
parser.add_argument('checkpoint', type = str)
parser.add_argument('--idx', default = 0, type = int)
parser.add_argument('--mesh', default = None, type = str)
parser.add_argument('--key', type = str, default = 'bunny')
parser.add_argument('--batch_size', default = 1000000, type = int)
parser = pl.Trainer.add_argparse_args(parser)
parser = GraphDataset.add_dataset_specific_args(parser)
parser = GraphINR.add_graph_inr_model_specific_args(parser)
parser = SpatialGraphINR.add_spatial_graph_inr_model_specific_args(parser)

args = parser.parse_args()

assert args.key in [
    'bunny',
    'protein_1AA7_A',
], '--key must be in ["bunny", "protein_1AA7_A"]'

# Model
if args.model == 'GINR':
    model = GraphINR.load_from_checkpoint(args.checkpoint)
elif args.model == 'Spatial_Graph_INR':
    model = SpatialGraphINR.load_from_checkpoint(args.checkpoint)
else:  
    raise ValueError('Inappropriate model type.')

# Data
dataset = GraphDataset(**vars(args))
points = dataset.ginr_npzs[0]['points']
inputs = dataset.get_inputs(0)

# Predict
_, pred = get_batched_predictions(model, inputs, 0, batch_size = args.batch_size)

if args.mesh is not None:
    mesh = load_mesh(args.mesh)
    key = args.key
    rot = PLOT_CONFIGS[key]['rot']
    
    fig = draw_mesh(
        mesh,
        intensity = pred,
        rot = rot,
        colorscale = PLOT_CONFIGS[key]['colorscale'],
        lower_camera = PLOT_CONFIGS[key]['lower_camera'],
    )
    
    fig.show()
    fig.write_html(f'{args.key}.html')
else:
    fig = go.Figure(traces.scatter_3d(points, pred))
    fig.show()