inference.py

import os, sys
sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
import json
import torch
import argparse
import numpy as np
from PIL import Image, ImageOps
import imageio
# from my_utils.plot import viz_graph
from my_utils.misc import load_config
import torchvision.transforms as T
from diffusers import DDPMScheduler
from models.denoiser import Denoiser
from scripts.json2urdf import create_urdf_from_json, pybullet_render
from dataset.utils import make_white_background, load_input_from, convert_data_range, parse_tree
import models
import torch.nn.functional as F
from io import BytesIO
import base64
from scripts.graph_pred.api import predict_graph_twomode, gpt_infer_image_category
import subprocess
import spaces
import time


cat_ref = {
    "Table": 0,
    "Dishwasher": 1,
    "StorageFurniture": 2,
    "Refrigerator": 3,
    "WashingMachine": 4,
    "Microwave": 5,
    "Oven": 6,
}

def run_retrieve(src_dir, json_name, data_root):
    fn_call = ['python', 'scripts/mesh_retrieval/retrieve.py', '--src_dir', src_dir, '--json_name', json_name, '--gt_data_root', data_root]
    try:
        subprocess.run(fn_call, check=True,  stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError as e:
        print(f'Error from run_retrieve: {src_dir}')
        print(f'Error: {e}')

def make_white_background(src_img):
    '''Make the white background for the input RGBA image.'''
    src_img.load() 
    background = Image.new("RGB", src_img.size, (255, 255, 255))
    background.paste(src_img, mask=src_img.split()[3]) # 3 is the alpha channel
    return background

def pad_to_square(img, fill=0):
    """Pad image to square with given fill value (default: 0 = black)."""
    width, height = img.size
    if width == height:
        return img
    max_side = max(width, height)
    delta_w = max_side - width
    delta_h = max_side - height
    padding = (delta_w // 2, delta_h // 2, delta_w - delta_w // 2, delta_h - delta_h // 2)
    return ImageOps.expand(img, padding, fill=fill)

def load_img(img_path):
    transform = T.Compose([
        T.Resize((224, 224), interpolation=T.InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ])
    with Image.open(img_path) as img:
        if img.mode == 'RGBA':
            img = make_white_background(img)
        img = img.convert('RGB')  # Ensure it's 3-channel for normalization
        img = pad_to_square(img, fill=0)
        img = transform(img)
    img_batch = img.unsqueeze(0).cuda()

    return img_batch


def load_frame_with_imageio(frame):
    """
    将单帧图像处理为符合 DINO 模型输入的格式。
    """
    transform = T.Compose([
        T.Resize((224, 224), interpolation=T.InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ])
    img = Image.fromarray(frame)  # 转为 PIL 图像
    if img.mode == 'RGBA':
        img = make_white_background(img)
    img = transform(img)  # 应用预处理
    return img.unsqueeze(0).cuda()  # 增加 batch 维度

def read_video_as_batch_with_imageio(video_path):
    """
    使用 imageio 读取视频并将所有帧处理为 batch 格式 (B, C, H, W)。
    """
    reader = imageio.get_reader(video_path)
    batch_frames = []

    try:
        for frame in reader:
            # 加载帧并处理为 (1, C, H, W)
            processed_frame = load_frame_with_imageio(frame)
            batch_frames.append(processed_frame)

        reader.close()
        if batch_frames:
            return torch.cat(batch_frames, dim=0).cuda()  # 在 batch 维度堆叠，并转移到 GPU
        else:
            print("视频没有有效帧")
            return None
    except Exception as e:
        print(f"处理视频时出错: {e}")
        return None

def extract_dino_feature(img_path_1, img_path_2):
    print('Extracting DINO feature...')
    feat_1 = load_img(img_path_1)
    feat_2 = load_img(img_path_2)
    frames = torch.cat([feat_1, feat_2], dim=0)
    dinov2_vitb14_reg = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14_reg', pretrained=True).cuda()
    print('step4')
    with torch.no_grad():
        feat = dinov2_vitb14_reg.forward_features(frames)["x_norm_patchtokens"]
    # release the GPU memory of the model
    feat_input = torch.cat([feat[0], feat[-1]], dim=0).unsqueeze(0)
    print('Extracting DINO feature over')
    torch.cuda.empty_cache()
    return feat_input

def set_scheduler(n_steps=100):
    scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule='linear', prediction_type='epsilon')
    scheduler.set_timesteps(n_steps)
    return scheduler

def prepare_model_input(data, cond, feat, n_samples):
    # attention masks
    attr_mask = torch.from_numpy(cond['attr_mask']).unsqueeze(0).repeat(n_samples, 1, 1)
    key_pad_mask = torch.from_numpy(cond['key_pad_mask'])
    graph_mask = torch.from_numpy(cond['adj_mask'])
    # input image feature
    f = feat.repeat(n_samples, 1, 1)
    # input noise
    B, C = data.shape
    noise = torch.randn([n_samples, B, C], dtype=torch.float32)
    # dummy image feature (used for guided diffusion)
    dummy_feat = torch.from_numpy(np.load('systems/dino_dummy.npy').astype(np.float32))
    dummy_feat = dummy_feat.unsqueeze(0).repeat(n_samples, 1, 1)
    # dummy object category
    cat = torch.zeros(1, dtype=torch.long).repeat(n_samples)
    return {
        "noise": noise.cuda(),
        "attr_mask": attr_mask.cuda(),
        "key_pad_mask": key_pad_mask.cuda(),
        "graph_mask": graph_mask.cuda(),
        "dummy_f": dummy_feat.cuda(),
        'cat': cat.cuda(),
        'f': f.cuda(),  
    }

def prepare_model_input_nocond(feat, n_samples):
    # attention masks
    cond_example = np.zeros((32*5, 32*5), dtype=bool)
    attr_mask = np.eye(32, 32, dtype=bool)
    attr_mask = attr_mask.repeat(5, axis=0).repeat(5, axis=1)
    attr_mask = torch.from_numpy(attr_mask).unsqueeze(0).repeat(n_samples, 1, 1)
    key_pad_mask = torch.from_numpy(cond_example).unsqueeze(0).repeat(n_samples, 1, 1)
    graph_mask = torch.from_numpy(cond_example).unsqueeze(0).repeat(n_samples, 1, 1)
    # input image feature
    f = feat.repeat(n_samples, 1, 1)
    # input noise
    data = np.zeros((32*5, 6), dtype=bool)
    noise = torch.randn(data.shape, dtype=torch.float32).repeat(n_samples, 1, 1)
    # dummy image feature (used for guided diffusion)
    dummy_feat = torch.from_numpy(np.load('systems/dino_dummy.npy').astype(np.float32))
    dummy_feat = dummy_feat.unsqueeze(0).repeat(n_samples, 1, 1)
    # dummy object category
    cat = torch.zeros(1, dtype=torch.long).repeat(n_samples)
    return {
        "noise": noise.cuda(),
        "attr_mask": attr_mask.cuda(),
        "key_pad_mask": key_pad_mask.cuda(),
        "graph_mask": graph_mask.cuda(),
        "dummy_f": dummy_feat.cuda(),
        'cat': cat.cuda(),
        'f': f.cuda(),  
    }

def save_graph(pred_graph, save_dir):
    print(f'Saving the predicted graph to {save_dir}/pred_graph.json')
    # save the response
    with open(os.path.join(save_dir, "pred_graph.json"), "w") as f:
        json.dump(pred_graph, f, indent=4)
    # Visualize the graph
    # img_graph = Image.fromarray(viz_graph(pred_graph))
    # img_graph.save(os.path.join(save_dir, "pred_graph.png"))

def forward(model, scheduler, inputs, omega=0.5):
    print('Running inference...')
    noisy_x = inputs['noise']
    for t in scheduler.timesteps:
        timesteps = torch.tensor([t], device=inputs['noise'].device)
        outputs_cond = model(
            x=noisy_x,
            cat=inputs['cat'],
            timesteps=timesteps,
            feat=inputs['f'], 
            key_pad_mask=inputs['key_pad_mask'],
            graph_mask=inputs['graph_mask'],
            attr_mask=inputs['attr_mask'],
            label_free=True,
        ) # take condtional image as input
        if omega != 0:
            outputs_free = model(
                x=noisy_x,
                cat=inputs['cat'],
                timesteps=timesteps,
                feat=inputs['dummy_f'], 
                key_pad_mask=inputs['key_pad_mask'],
                graph_mask=inputs['graph_mask'],
                attr_mask=inputs['attr_mask'],
                label_free=True,
            ) # take the dummy DINO features for the condition-free mode
            noise_pred = (1 + omega) * outputs_cond['noise_pred'] - omega * outputs_free['noise_pred']
        else:
            noise_pred = outputs_cond['noise_pred']
        noisy_x = scheduler.step(noise_pred, t, noisy_x).prev_sample
    return noisy_x

def _convert_json(x, c):
    out = {"meta": {}, "diffuse_tree": []}
    n_nodes = c["n_nodes"]
    par = c["parents"].tolist()
    adj = c["adj"]
    np.fill_diagonal(adj, 0) # remove self-loop for the root node
    if "obj_cat" in c:
        out["meta"]["obj_cat"] = c["obj_cat"]

    # convert the data to original range
    data = convert_data_range(x)
    # parse the tree
    out["diffuse_tree"] = parse_tree(data, n_nodes, par, adj)
    return out

def post_process(output, cond, save_root, gt_data_root, visualize=False):
    print('Post-processing...')
    # N = output.shape[0]
    N = 1
    for i in range(N):
        cond_n = {}
        cond_n['n_nodes'] = cond['n_nodes'][i] 
        cond_n['parents'] = cond['parents'][i]
        cond_n['adj'] = cond['adj'][i]
        cond_n['obj_cat'] = cond['cat']
        # convert the raw model output to the json format
        out_json = _convert_json(output, cond_n)
        save_dir = os.path.join(save_root, str(i))
        os.makedirs(save_dir, exist_ok=True)
        with open(os.path.join(save_dir, "object.json"), "w") as f:
            json.dump(out_json, f, indent=4)
        

        # retrieve part meshes (call python script)
        # print(f"Retrieving part meshes for the object {i}...")
        # os.system(f"python scripts/mesh_retrieval/retrieve.py --src_dir {save_dir} --json_name object.json --gt_data_root {gt_data_root}")


    

def load_model(ckpt_path, config):
    print('Loading model from checkpoint...')
    model = models.make(config.name, config)
    state_dict = torch.load(ckpt_path)
    state_dict = {k.replace("model.", ""): v for k, v in state_dict.items()}
    model.load_state_dict(state_dict)
    model.eval()
    return model.cuda()

def convert_pred_graph(pred_graph):
    cond = {}
    B, K = pred_graph.shape[:2]
    adj = np.zeros((B, K, K), dtype=np.float32)
    padding = np.zeros((B, 5 * K, 5* K), dtype=bool)
    parents = np.zeros((B, K), dtype=np.int32)
    n_nodes = np.zeros((B,), dtype=np.int32)
    for b in range(B):
        node_len = 0
        for k in range(K):
            if pred_graph[b, k] == k and k > 0:
                node_len = k
                break
            node = pred_graph[b, k]
            adj[b, k, node] = 1
            adj[b, node, k] = 1
            parents[b, k] = node
        adj[b, node_len:] = 1
        padding[b, :, :5 * node_len] = 1
        parents[b, 0] = -1
        n_nodes[b] = node_len
    adj_mask = adj.astype(bool).repeat(5, axis=1).repeat(5, axis=2)
    attr_mask = np.eye(32, 32, dtype=bool)
    attr_mask = attr_mask.repeat(5, axis=0).repeat(5, axis=1)

    cond['adj_mask'] = adj_mask
    cond['attr_mask'] = attr_mask
    cond['key_pad_mask'] = padding

    cond['adj'] = adj
    cond['parents'] = parents
    cond['n_nodes'] = n_nodes
    cond['cat'] = 'StorageFurniture'

    data = np.zeros((32*5, 6), dtype=bool)

    return data, cond

def bfs_tree_simple(tree_list):
    order = [0] * len(tree_list)
    queue = []
    current_node_idx = 0
    for node_idx, node in enumerate(tree_list):
        if node['parent'] == -1:
            queue.append(node['id'])
            order[current_node_idx] = node_idx
            current_node_idx += 1
            break
    while len(queue) > 0:
        current_node = queue.pop(0)
        for node_idx, node in enumerate(tree_list):
            if node['parent'] == current_node:
                queue.append(node['id'])
                order[current_node_idx] = node_idx
                current_node_idx += 1

    return order

def get_graph_from_gpt(img_path_1, img_path_2):
    first_img = Image.open(img_path_1)
    first_img_data = first_img.resize((1024, 1024))
    buffer = BytesIO()
    first_img_data.save(buffer, format="PNG")
    buffer.seek(0)
    # encode the image as base64
    first_encoded_image = base64.b64encode(buffer.read()).decode("utf-8")


    second_img = Image.open(img_path_2)
    second_img_data = second_img.resize((1024, 1024))
    buffer = BytesIO()
    second_img_data.save(buffer, format="PNG")
    buffer.seek(0)
    # encode the image as base64
    second_encoded_image = base64.b64encode(buffer.read()).decode("utf-8")

    pred_gpt = predict_graph_twomode('', first_img_data=first_encoded_image, second_img_data=second_encoded_image)
    print(pred_gpt)
    pred_graph = pred_gpt['diffuse_tree']
    # order = bfs_tree_simple(pred_graph)
    # pred_graph = [pred_graph[i] for i in order]
    
    
    # generate array [0, 1, 2, ..., 31] for init
    graph_array = np.array([i for i in range(32)])
    for node_idx, node in enumerate(pred_graph):
        if node['parent'] == -1:
            graph_array[node_idx] = node_idx
        else:
            graph_array[node_idx] = node['parent']

    # new axis for batch
    graph_array = np.expand_dims(graph_array, axis=0)

    cat_str = gpt_infer_image_category(first_encoded_image, second_encoded_image)
    
    return torch.from_numpy(graph_array).cuda().repeat(3, 1), cat_str
        
@spaces.GPU
def run_demo(args):
    # extract DINOV2 feature from the input image
    t1 = time.time()
    feat = extract_dino_feature(args.img_path_1, args.img_path_2)
    t2 = time.time()
    print(f'Extracted DINO feature in {t2 - t1:.2f} seconds')
    scheduler = set_scheduler(args.n_denoise_steps)
    # load the checkpoint of the model
    model = load_model(args.ckpt_path, args.config.system.model)

    # inference
    with torch.no_grad():
        t3 = time.time()
        pred_graph, cat_str = get_graph_from_gpt(args.img_path_1, args.img_path_2)
        t4 = time.time()
        print(f'Got the predicted graph in {t4 - t3:.2f} seconds')
        print(pred_graph)
        data, cond = convert_pred_graph(pred_graph)
        inputs = prepare_model_input(data, cond, feat, n_samples=args.n_samples)

        # Update the object category
        cond['cat'] = cat_str
        inputs['cat'][:] = cat_ref[cat_str]
        print(f'Object category predicted by GPT: {cat_str}, {cat_ref[cat_str]}')

        output = forward(model, scheduler, inputs, omega=args.omega).cpu().numpy()
        t5 = time.time()
        print(f'Forwarded the model in {t5 - t4:.2f} seconds')

        # post-process
        post_process(output, cond, args.save_dir, args.gt_data_root, visualize=True)

    # retrieve
    for sample in os.listdir(args.save_dir):
        sample_dir = os.path.join(args.save_dir, sample)
        t6 = time.time()
        run_retrieve(sample_dir, 'object.json', args.gt_data_root)
        t7 = time.time()
        print(f'Retrieved part meshes for in {t7 - t6:.2f} seconds')
    
    save_json_path = os.path.join(args.save_dir, "0", "object.json")
    with open(save_json_path, 'r') as file:
        json_data = json.load(file)
    create_urdf_from_json(json_data, save_json_path.replace('.json', '.urdf'))
    pybullet_render(save_json_path.replace('.json', '.urdf'), os.path.join(args.save_dir, "0"), 8)
    

if __name__ == '__main__':
    '''
    Script for running the inference on an example image input.
    '''
    parser = argparse.ArgumentParser()
    parser.add_argument("--img_path_1", type=str, default='examples/1.png', help="path to the input image")
    parser.add_argument("--img_path_2", type=str, default='examples/1_open_1.png', help="path to the input image")
    parser.add_argument("--ckpt_path", type=str, default='exps/singapo/final/ckpts/last.ckpt', help="path to the checkpoint of the model")
    parser.add_argument("--config_path", type=str, default='exps/singapo/final/config/parsed.yaml', help="path to the config file")
    parser.add_argument("--use_example_graph", action="store_true", default=False, help="if you don't have the openai key yet, turn on to use the example graph for inference")
    parser.add_argument("--save_dir", type=str, default='results', help="path to save the output")
    parser.add_argument("--gt_data_root", type=str, default='./', help="the root directory of the original data, used for part mesh retrieval")
    parser.add_argument("--n_samples", type=int, default=3, help="number of samples to generate given the input")
    parser.add_argument("--omega", type=float, default=0.5, help="the weight of the condition-free mode in the inference")
    parser.add_argument("--n_denoise_steps", type=int, default=100, help="number of denoising steps")
    args = parser.parse_args()

    assert os.path.exists(args.img_path_1), "The input image does not exist"
    # assert os.path.exists(args.ckpt_path), "The checkpoint does not exist"
    assert os.path.exists(args.config_path), "The config file does not exist"
    os.makedirs(args.save_dir, exist_ok=True)

    config = load_config(args.config_path)
    args.config = config

    run_demo(args)