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)