Spaces:
Running
Running
File size: 1,509 Bytes
bcec73a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.
from __future__ import absolute_import, division, print_function
import torch
import torch.nn as nn
class PoseCNN(nn.Module):
def __init__(self, num_input_frames):
super(PoseCNN, self).__init__()
self.num_input_frames = num_input_frames
self.convs = {}
self.convs[0] = nn.Conv2d(3 * num_input_frames, 16, 7, 2, 3)
self.convs[1] = nn.Conv2d(16, 32, 5, 2, 2)
self.convs[2] = nn.Conv2d(32, 64, 3, 2, 1)
self.convs[3] = nn.Conv2d(64, 128, 3, 2, 1)
self.convs[4] = nn.Conv2d(128, 256, 3, 2, 1)
self.convs[5] = nn.Conv2d(256, 256, 3, 2, 1)
self.convs[6] = nn.Conv2d(256, 256, 3, 2, 1)
self.pose_conv = nn.Conv2d(256, 6 * (num_input_frames - 1), 1)
self.num_convs = len(self.convs)
self.relu = nn.ReLU(True)
self.net = nn.ModuleList(list(self.convs.values()))
def forward(self, out):
for i in range(self.num_convs):
out = self.convs[i](out)
out = self.relu(out)
out = self.pose_conv(out)
out = out.mean(3).mean(2)
out = 0.01 * out.view(-1, self.num_input_frames - 1, 1, 6)
axisangle = out[..., :3]
translation = out[..., 3:]
return axisangle, translation
|