def __init__(self, name, task):
self.name = name
self.task = task
self.total_iter = 0
self.num_rotations = 24
self.descriptor_dim = 16
self.pixel_size = 0.003125
self.input_shape = (320, 160, 6)
self.camera_config = cameras.RealSenseD415.CONFIG
self.models_dir = os.path.join('checkpoints', self.name)
self.bounds = np.array([[0.25, 0.75], [-0.5, 0.5], [0, 0.28]])
self.pick_model = Attention(input_shape=self.input_shape,
num_rotations=1,
preprocess=self.preprocess,
lite=True)
self.place_model = Attention(input_shape=self.input_shape,
num_rotations=1,
preprocess=self.preprocess,
lite=True)
self.match_model = Matching(input_shape=self.input_shape,
descriptor_dim=self.descriptor_dim,
num_rotations=self.num_rotations,
preprocess=self.preprocess,
lite=True)
def __init__(self, name, task, n_rotations=36):
super().__init__(name, task, n_rotations)
self.attention = Attention(in_shape=self.in_shape,
n_rotations=1,
preprocess=utils.preprocess)
self.transport = Attention(in_shape=self.in_shape,
n_rotations=self.n_rotations,
preprocess=utils.preprocess)
def __init__(self, name, task, n_rotations=36):
super().__init__(name, task, n_rotations)
self.attention = Attention(in_shape=self.in_shape,
n_rotations=1,
preprocess=utils.preprocess)
self.transport = TransportPerPixelLoss(in_shape=self.in_shape,
n_rotations=self.n_rotations,
crop_size=self.crop_size,
preprocess=utils.preprocess)
def __init__(self, name, task, n_rotations=12):
super().__init__(name, task, n_rotations)
self.attention = Attention(in_shape=self.in_shape,
n_rotations=1,
preprocess=utils.preprocess)
self.transport = TransportGoal(input_shape=self.in_shape,
num_rotations=self.n_rotations,
crop_size=self.crop_size,
preprocess=utils.preprocess)
def __init__(self, name, task, n_rotations=36):
super().__init__(name, task, n_rotations)
# Stack the goal image for the vanilla Transport module.
t_shape = (self.in_shape[0], self.in_shape[1],
int(self.in_shape[2] * 2))
self.attention = Attention(in_shape=self.in_shape,
n_rotations=1,
preprocess=utils.preprocess)
self.transport = Transport(in_shape=t_shape,
n_rotations=self.n_rotations,
crop_size=self.crop_size,
preprocess=utils.preprocess)