orenkatzir / VN-SPD

Shape-Pose Disentanglement using SE (3)-equivariant Vector Neurons
16 stars 2 forks source link

Unable to reproduce VN-SPD results #1

Open erintomlinson opened 1 year ago

erintomlinson commented 1 year ago

Hello, thank you for publishing this interesting paper and codebase! I am trying to reproduce your results on the "plane" category using the "planes_paper" pretrained weights, but I am currently seeing lower-quality reconstructions than those shown in your paper (see below for examples). Since there are two versions of ShapeNet and the dataset links in the AtlasNetV2 repository are partially broken, I am beginning to wonder if the dataset I built using the CanonicalCapsules dataloader is the same as the one you used for training and evaluation. Would it be possible for you to provide me with a link (say, on Google Drive) to the "train/planes.h5" and "valid/planes.h5" files that you used to train your "planes_paper" weights? I have tried building this dataset myself several times and am still seeing lower-quality results. Thank you kindly!

My reconstructions on samples from "valid/planes.h5" using "planes_paper" weights:

image

Reconstructions from Figure 7 in the paper whose quality I am trying to match:

Screenshot 2023-04-20 at 8 59 35 PM

erintomlinson commented 1 year ago

My method of generating the results above is as follows. I wanted to experiment and visualize these results in a notebook, so I wrote a few helper routines and added support to the options/base_options.py module for passing in command-line arguments manually as a list of strings. Please let me know if I am making any obvious mistakes in my use of your codebase, especially with regard to generating the model outputs in the first few lines of my show_model_outputs routine. Thanks again!

from data import create_dataset as _create_dataset
from models import create_model as _create_model
from options.test_options import TestOptions

def create_dataset(model='shape_pose', batch_size=1, shuffle=False):
    args = f'--demo --dataroot data/shapenet --class_choice plane --model {model}'
    opt = TestOptions().parse(args.split())

    opt.split = 'test'
    opt.num_threads = 0
    opt.batch_size = batch_size
    opt.serial_batches = not(shuffle)
    opt.remove_knn = 0
    dataset = _create_dataset(opt)
    return dataset, opt

def create_model(name, opt):
    opt.name = name
    model = _create_model(opt)
    model.setup(opt)
    model.eval()
    return model

def show_model_outputs(data, model, use_rand_trans=False, height=500, width=1200, showaxis=False):

    model.eval()
    with torch.no_grad():
        _, trot = model.set_input(data, use_rand_trans=use_rand_trans)
        model.forward()

    subplots = {
        'Input': model.pc, 
        'Re-posed Input': model.pc_at_inv, 
        'Pose-Invariant Recon.': model.recon_pc_inv, 
        'Input Recon.': model.recon_pc}

    batch_size = len(data[0])
    fig = make_subplots(
        rows=batch_size, cols=len(subplots), 
        subplot_titles=[subplot_name for _ in range(batch_size) for subplot_name in subplots.keys()], 
        specs=[[{'type': 'scene'} for _ in range(len(subplots))] for _ in range(batch_size)], 
        vertical_spacing=0, horizontal_spacing=0)

    for subplot_idx, (subplot_name, subplot_data) in enumerate(subplots.items()):
        for pc_idx in range(batch_size):
            x, y, z = subplot_data[pc_idx].cpu().numpy()
            trace = go.Scatter3d(
                x=x, y=y, z=z, 
                mode='markers', 
                marker=dict(size=2, opacity=0.8, color=colors[subplot_idx]))
            fig.add_trace(trace, row=pc_idx+1, col=subplot_idx+1)

    fig.update_layout(height=height*batch_size, width=width, showlegend=False)
    fig.update_scenes(aspectmode='data', xaxis=dict(visible=showaxis), yaxis=dict(visible=showaxis), zaxis=dict(visible=showaxis))
    fig.show()