Partial Mesh Completion

[kajalsanklecha]

Thank you for the work and for releasing its code.

Can you please give the steps on how to run the code for Partial Mesh Completion as mentioned in the paper?

[ch3cook-fdu]

Thanks for your interest in our paper!

In our implementation, we truncate 25% - 50% of the tokenized input_ids here. Then, we feed the truncated tokens to the generation function as the prefix by specifying an additional keyword input_ids here.

[kajalsanklecha]

Can you please share a sample code for running the same?

Like the one shared for text to mesh as sample_t2m.py

[ch3cook-fdu]

Here is an example code for shape completion:

@torch.no_grad()
def generate(self, data_dict: dict=None, n_samples: int=8) -> dict:

    data_dict = self.tokenizer.tokenize(data_dict)
    input_ids = data_dict['input_ids']              # 1 x ntoken
    attention_mask = data_dict['attention_mask']    # 1 x ntoken

    # replace padding tokens
    input_ids[:, 0] = self.bos_token_id                                 # <sos> xxx <pad> <pad>
    eos_pos_id = attention_mask.sum(1, keepdim=True) - 1
    input_ids = torch.scatter(
        input_ids, 
        1, 
        eos_pos_id.long(), 
        torch.ones_like(input_ids) * self.eos_token_id
    )


    # conditioned on 1/4 the shape
    input_ids = input_ids[:, attention_mask[0] == 1]    # 1 x [<bos> ... <eos>]
    num_faces = (input_ids.shape[1] - 2) // 9
    kept_length = (num_faces // 4) * 9 + 1
    input_ids = input_ids[:, :kept_length]          # 1 x [<bos> ...]
    
    net_device = next(self.parameters()).device
    max_length = 7202
    outputs = torch.ones(n_samples, max_length).long().to(net_device) * self.eos_token_id
    # batch x ntokens
    results = self.transformer.generate(
        input_ids=input_ids,
        max_new_tokens=max_length-input_ids.shape[1],
        do_sample=True,
        top_k=50,
        top_p=0.95,
        num_return_sequences=n_samples,
        # num_beams=n_samples,
        # no_repeat_ngram_size=9,
        bos_token_id=self.bos_token_id,
        eos_token_id=self.eos_token_id,
        pad_token_id=self.eos_token_id,
    )
    outputs[:, :results.shape[1]] = results
    # batch x ntokens ====> batch x ntokens x D
    outputs = outputs[:, 1: -1]
    outputs[outputs == self.eos_token_id] = self.tokenizer.pad_id
    decoder_output = self.tokenizer.detokenize(outputs)
    
    condition_output = self.tokenizer.detokenize(input_ids[:, 1:])
    decoder_output['partial_mesh'] = condition_output['recon_faces']
    
    return decoder_output     

[DerKleineLi]

Hey @ch3cook-fdu,

thanks for the great work and the script for shape completion. I can run the code to regenerate some results sampled from the pre-trained model. However, I got trouble while completing a custom mesh. I think it's because I didn't provide the input_ids in correct order. Could you share a script to transform a common triangle mesh to input_ids?

Thanks a lot!