import torch
import numpy as np
from Inpainters.LBAM.models.LBAMModel import LBAMModel

def inpaint_struct_rgb(tensor,mask,model):
    torch.cuda.empty_cache()
    result_inp = torch.zeros_like(tensor)
    tensor_in = tensor.permute([2,0,1]).clone()
    mask_in = mask[:,:,0].unsqueeze(0).clone()
    input_tensor = torch.cat((tensor_in, mask_in),0).unsqueeze(0)
    mask_out = mask.permute(2,0,1).unsqueeze(0).clone()
    with torch.no_grad():
        output = model(input_tensor, mask_out)
        result_inp = output.data.squeeze(0).permute([1,2,0]).float().clamp_(0, 1)

    result_inp = result_inp*(1-mask)+tensor*mask

    return result_inp