main_test.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function
import argparse
import os
import time
import torch
import torch.utils.data
import torch.nn as nn
import numpy as np
from scipy.stats import mode
import sys
from os import path

sys.path.append(path.dirname(path.dirname(path.abspath("__file__"))))
parser = argparse.ArgumentParser(description="pointnetGPD")
parser.add_argument("--cuda", action="store_true", default=False)
parser.add_argument("--gpu", type=int, default=0)
parser.add_argument("--load-model", type=str,
                    default="../data/pointgpd_3class.model")
parser.add_argument("--show_final_grasp", action="store_true", default=False)
parser.add_argument("--tray_grasp", action="store_true", default=False)
parser.add_argument("--using_mp", action="store_true", default=False)
parser.add_argument("--model_type", type=str)

args = parser.parse_args()

args.cuda = args.cuda if torch.cuda.is_available else False

if args.cuda:
    torch.cuda.manual_seed(1)

np.random.seed(int(time.time()))

if args.model_type == "100":
    args.load_model = "../data/pointgpd_chann3_local.model"
elif args.model_type == "50":
    args.load_model = "../data/pointgpd_50_points.model"
elif args.model_type == "3class":  # input points number is 500
    args.load_model = "../data/pointgpd_3class.model"
else:
    print("Using default model file")
model = torch.load(args.load_model, map_location="cpu")
model.device_ids = [args.gpu]
print("load model {}".format(args.load_model))

if args.cuda:
    model = torch.load(args.load_model, map_location="cuda:{}".format(args.gpu))
    if args.gpu != -1:
        torch.cuda.set_device(args.gpu)
        model = model.cuda()
    else:
        device_id = [0, 1]
        torch.cuda.set_device(device_id[0])
        model = nn.DataParallel(model, device_ids=device_id).cuda()
if isinstance(model, torch.nn.DataParallel):
    model = model.module


def test_network(model_, local_pc):
    local_pc = local_pc.T
    local_pc = local_pc[np.newaxis, ...]
    local_pc = torch.FloatTensor(local_pc)
    if args.cuda:
        local_pc = local_pc.cuda()
    output, _ = model_(local_pc)  # N*C
    output = output.softmax(1)
    pred = output.data.max(1, keepdim=True)[1]
    output = output.cpu()
    return pred[0], output.data.numpy()


def main():
    repeat = 10
    num_point = 500
    model.eval()
    torch.set_grad_enabled(False)

    # load pc(should be in local gripper coordinate)
    # local_pc: (N, 3)
    # local_pc = np.load("test.npy")
    local_pc = np.random.random([500, 3])  # test only
    predict = []
    for _ in range(repeat):
        if len(local_pc) >= num_point:
            local_pc = local_pc[np.random.choice(len(local_pc), num_point, replace=False)]
        else:
            local_pc = local_pc[np.random.choice(len(local_pc), num_point, replace=True)]

        # run model
        predict.append(test_network(model, local_pc)[0])
    print("voting: ", predict)
    predict = mode(predict).mode[0]

    # output
    print("Test result:", predict)


if __name__ == "__main__":
    main()