class Perception(object):
def __init__(self):
# input rgbd image in numpy array format [w h c]
self.sdmrcnn_model = get_model_instance_segmentation(2).to(
device, dtype=torch.float)
self.sdmrcnn_model.load_state_dict(torch.load(os.path.join('19.pth')))
self.sdmrcnn_model.eval()
self.siamese_model = SiameseNetwork().cuda()
self.siamese_model.load_state_dict(torch.load('siamese.pt'))
self.siamese_model.eval()
def segmentation(self, raw_rgb, raw_depth):
rgb_raw_img = np.zeros_like(raw_rgb)
for i in range(raw_rgb.shape[2]):
rgb_raw_img[:, :, i] = raw_rgb[:, :, 2 - i]
color_img = rgb_raw_img.astype(np.float) / 255.
img_mean = [0.485, 0.456, 0.406]
img_std = [0.229, 0.224, 0.225]
for c in range(color_img.shape[2]):
color_img[:, :,
c] = (color_img[:, :, c] - img_mean[c]) / img_std[c]
depth_img = raw_depth.astype(np.float)
x, y = np.shape(depth_img)
depth_img.shape = (x, y, 1)
# depth_img = depth_img / np.amax(depth_img)
img = np.concatenate((color_img[:, :, 0:2], depth_img), axis=2)
test_input = [
torch.from_numpy(np.transpose(img,
[2, 0, 1])).to(device,
dtype=torch.float)
]
output = self.sdmrcnn_model(test_input)
mask_list = output[0]['masks'].cpu().detach().numpy()
masks = np.reshape(mask_list, (len(mask_list), 480, 640))
# masks = []
# for mask in mask_list:
# mask = mask.reshape(960, 1280)
# mask = mask[240:720, 320:960]
# masks.append(mask)
return masks
def classification(self, masks, raw_rgb, anchor_img):
scores = []
img0 = anchor_img
img0 = torch.from_numpy(np.reshape(img0, (1, 3, 120, 120)))
masks = np.reshape(masks, (len(masks), 480, 640))
print('Number of objects detected: %d' % len(masks))
for mask in masks:
color_img = np.copy(raw_rgb)
color_img[np.where(mask < 0.5)] = 0
color_img = pad(color_img)
color_img = center_crop_150(color_img)
img1 = np.reshape(color_img, (1, 3, 120, 120))
img1 = torch.from_numpy(img1)
img0, img1 = img0.type(torch.FloatTensor), img1.type(
torch.FloatTensor)
img0, img1, = img0.cuda(), img1.cuda()
output = self.siamese_model(img0, img1)
output = output.detach().cpu().numpy()
scores.append(output[0])
scores = np.asarray(scores)
res_mask = masks[np.argmax(scores)]
res_mask = np.reshape(res_mask, (480, 640))
return res_mask
parser.add_argument(
'-c',
'--checkpoint',
type=str,
help="Path of model checkpoint to be used for inference.",
required=True
)
parser.add_argument(
'-o',
'--out_path',
type=str,
help="Path for saving tensorrt model.",
required=True
)
args = parser.parse_args()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
checkpoint = torch.load(args.checkpoint)
model = SiameseNetwork(backbone=checkpoint['backbone'])
model.to(device)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
torch.onnx.export(model, (torch.rand(1, 3, 224, 224).to(device), torch.rand(1, 3, 224, 224).to(device)), args.out_path, input_names=['input'],
output_names=['output'], export_params=True)
onnx_model = onnx.load(args.out_path)
onnx.checker.check_model(onnx_model)
if __name__ == '__main__':
lr = 0.001
batch_size = 32
epochs = 100
start_epoch = 0
train_data = FacesDataset(train=True, validation=False, base='resnet101')
train_loader = DataLoader(train_data, batch_size, False)
test = FacesDataset(train=False, validation=False, base='resnet101')
test_loader = DataLoader(test, batch_size, False)
siamese_network = SiameseNetwork(base='resnet101').cuda()
siamese_network.load('./models/triplet/resnet101.pt')
siamese_network.eval()
predictor_ = LinearPredictor().cuda()
sys.stdout.write('Training Linear predictor:\n')
losses, test_losses, train_accuracies_, test_accuracies_ = train(
siamese_network,
predictor_,
train_loader,
test_loader,
epochs,
lr,
'predictor-linear',
start_epoch_=start_epoch,
adam=False,
patience=3)