Merge pull request #433 from cpaxton/feature/mhp

MHP predictions and GAN updates

costar_models/python/costar_models/callbacks.py

@@ -312,6 +312,85 @@ def on_epoch_end(self, epoch, logs={}):
             fig.savefig(name, bbox_inches="tight")
             plt.close(fig)
 
+class PredictorShowImageOnlyMultiStep(keras.callbacks.Callback):
+    '''
+    Save an image showing what some number of frames and associated predictions
+    will look like at the end of an epoch.
+    '''
+
+    def __init__(self, predictor, features, targets,
+            model_directory=DEFAULT_MODEL_DIRECTORY,
+            num_hypotheses=4,
+            verbose=False,
+            features_name=None,
+            noise_dim=64,
+            use_noise=False,
+            name="model",
+            use_prev_option=True,
+            min_idx=0, max_idx=66, step=11):
+        '''
+        Set up a data set we can use to output validation images.
+
+        Parameters:
+        -----------
+        predictor: model used to generate predictions
+        targets: training target info, in compressed form
+        num_hypotheses: how many outputs to expect
+        verbose: print out extra information
+        '''
+
+        if features_name is None:
+            self.features_name = "def"
+        else:
+            self.features_name = features_name
+        self.verbose = verbose
+        self.predictor = predictor
+        self.idxs = range(min_idx, max_idx, step)
+        self.num = len(self.idxs)
+        self.features = [f[self.idxs] for f in features]
+        self.targets = [np.squeeze(t[self.idxs]) for t in targets]
+        self.epoch = 0
+        self.num_hypotheses = num_hypotheses
+        self.directory = os.path.join(model_directory,'debug')
+        self.noise_dim = noise_dim
+        self.use_noise = use_noise
+        if not os.path.exists(self.directory):
+            os.makedirs(self.directory)
+
+    def on_epoch_end(self, epoch, logs={}):
+        # take the model and print it out
+        self.epoch += 1
+        data = self.predictor.predict(self.features)
+        plt.ioff()
+        if self.verbose:
+            print("============================")
+        for j in range(self.num):
+            name = os.path.join(self.directory,
+                    "%s_predictor_epoch%03d_result%d.png"%(self.features_name,
+                        self.epoch, j))
+            fig = plt.figure()#figsize=(3+int(1.5*self.num_hypotheses),2))
+
+            plt.subplot(2,2+self.num_hypotheses,1)
+            plt.title('Input Image')
+            plt.imshow(self.features[0][j])
+            for k in range(2):
+                # This counts off rows
+                rand_offset = (k*(2+self.num_hypotheses))
+                plt.subplot(2,2+self.num_hypotheses,2+self.num_hypotheses+rand_offset)
+                plt.title('Observed Goal')
+                plt.imshow(np.squeeze(self.targets[k][j]))
+                for i in range(self.num_hypotheses):
+                    plt.subplot(2,2+self.num_hypotheses,i+2+rand_offset)
+                    plt.imshow(np.squeeze(data[k][j][i]))
+                    plt.title('Hypothesis %d'%(i+1))
+
+            if self.verbose:
+                print(name)
+            fig.savefig(name, bbox_inches="tight")
+            plt.close(fig)
+
+
+
 class PredictorShowImageOnly(keras.callbacks.Callback):
     '''
     Save an image showing what some number of frames and associated predictions

costar_models/python/costar_models/conditional_image.py

@@ -104,12 +104,6 @@ def _makePredictor(self, features):
         #value_out = value_model([h,label_in])
         #next_option_out = next_model([h,label_in])
 
-        # create input for controlling noise output if that's what we decide
-        # that we want to do
-        if self.use_noise:
-            z = Input((self.num_hypotheses, self.noise_dim))
-            ins += [z]
-
         next_option_in = Input((1,), name="next_option_in")
         next_option_in2 = Input((1,), name="next_option_in2")
         ins += [next_option_in, next_option_in2]
@@ -159,7 +153,8 @@ def _makePredictor(self, features):
             train_predictor.compile(
                     loss=[lfn, lfn, "binary_crossentropy", val_loss,
                         lfn2, lfn2, "categorical_crossentropy"],
-                    loss_weights=[1., 1., 0.1, 0.1, 1., 0.2, 1e-4],
+                    #loss_weights=[1., 1., 0.1, 0.1, 1., 0.2, 1e-3],
+                    loss_weights=[1., 1., 0., 0., 0., 0., 1e-3],
                     optimizer=self.getOptimizer())
         else:
             train_predictor = Model(ins + [label_in],

costar_models/python/costar_models/conditional_image_gan.py

@@ -137,8 +137,6 @@ def _makePredictor(self, features):
                 loss=["mae"]*2 + ["binary_crossentropy"],
                 loss_weights=[100., 100., 1.],
                 optimizer=self.getOptimizer())
-        model.summary()
-        self.discriminator.summary()
         self.model = model
 
         return predictor, model, model, ins, h
@@ -200,10 +198,9 @@ def _makeImageDiscriminator(self, img_shape):
         #x = Concatenate()([x1, x2])
         x = x2
         x = AddConv2D(x, 128, [4,4], 2, dr, "same", lrelu=True)
-        #x = AddConv2D(x, 128, [4,4], 1, dr, "same", lrelu=True)
         x= AddConv2D(x, 256, [4,4], 2, dr, "same", lrelu=True)
-        #x = AddConv2D(x, 256, [4,4], 1, dr, "same", lrelu=True)
-        x = AddConv2D(x, 1, [4,4], 1, 0., "same", activation="sigmoid")
+        x = AddConv2D(x, 1, [1,1], 1, 0., "same", activation="sigmoid",
+                bn=False)
 
         #x = MaxPooling2D(pool_size=(8,8))(x)
         x = AveragePooling2D(pool_size=(8,8))(x)

costar_models/python/costar_models/conditional_image_gan_jigsaws.py

@@ -94,7 +94,6 @@ def _makeModel(self, image, *args, **kwargs):
                 loss_weights=[100., 100., 1.],
                 optimizer=self.getOptimizer())
         model.summary()
-        self.discriminator.summary()
         self.model = model
 
         self.predictor = generator
@@ -137,31 +136,39 @@ def _makeImageDiscriminator(self, img_shape):
         xg1 = AddConv2D(img_goal, 64, [4,4], 1, dr, "same", lrelu=True, bn=False)
         xg2 = AddConv2D(img_goal2, 64, [4,4], 1, dr, "same", lrelu=True, bn=False)
 
-        x1 = Add()([x0, xobs, xg1])
-        x2 = Add()([x0, xg1, xg2])
+        #x1 = Add()([x0, xobs, xg1])
+        #x2 = Add()([x0, xg1, xg2])
+        x1 = Add()([xobs, xg1])
+        x2 = Add()([xg1, xg2])
+        #x1 = Concatenate(axis=-1)([img, img_goal])
+        #x2 = Concatenate(axis=-1)([img_goal, img_goal2])
 
         # -------------------------------------------------------------
         y = OneHot(self.num_options)(option)
         y = AddDense(y, 64, "lrelu", dr)
-        x1 = TileOnto(x1, y, 64, img_size, add=True)
+        #x1 = TileOnto(x1, y, 64, img_size, add=True)
         x1 = AddConv2D(x1, 64, [4,4], 2, dr, "same", lrelu=True, bn=False)
-        x1 = AddConv2D(x1, 128, [4,4], 2, dr, "same", lrelu=True)
-        #x = AddConv2D(x, 256, [4,4], 2, dr, "same", lrelu=True)
+        x1 = AddConv2D(x1, 128, [4,4], 2, dr, "same", lrelu=True, bn=True)
+        #x1 = AddConv2D(x1, 256, [4,4], 2, dr, "same", lrelu=True, bn=True)
         #x1 = AddConv2D(x1, 1, [4,4], 1, 0., "same", activation="sigmoid")
 
         # -------------------------------------------------------------
         y = OneHot(self.num_options)(option2)
         y = AddDense(y, 64, "lrelu", dr)
         x2 = TileOnto(x2, y, 64, img_size, add=True)
         x2 = AddConv2D(x2, 64, [4,4], 2, dr, "same", lrelu=True, bn=False)
-        x2 = AddConv2D(x2, 128, [4,4], 2, dr, "same", lrelu=True)
-        #x = AddConv2D(x, 256, [4,4], 2, dr, "same", lrelu=True)
 
-        x = Concatenate(axis=-1)([x1, x2])
+        # Final block
+        x = x2
+        x2 = AddConv2D(x2, 128, [4,4], 2, dr, "same", lrelu=True, bn=True)
+        x2 = AddConv2D(x2, 256, [4,4], 2, dr, "same", lrelu=True, bn=True)
+        #x = Concatenate(axis=-1)([x1, x2])
         #x = Add()([x1, x2])
-        x = AddConv2D(x, 1, [4,4], 1, 0., "same", activation="sigmoid")
-        #x = AveragePooling2D(pool_size=(12,16))(x)
-        x = AveragePooling2D(pool_size=(24,32))(x)
+        x = AddConv2D(x2, 1, [1,1], 1, 0., "same", activation="sigmoid", bn=False)
+
+        # Combine
+        x = AveragePooling2D(pool_size=(12,16))(x)
+        #x = AveragePooling2D(pool_size=(24,32))(x)
         x = Flatten()(x)
         discrim = Model(ins, x, name="image_discriminator")
         self.lr *= 2.

costar_models/python/costar_models/conditional_image_jigsaws.py

@@ -15,10 +15,7 @@
 from keras.optimizers import Adam
 from matplotlib import pyplot as plt
 
-from .abstract import *
-from .callbacks import *
 from .robot_multi_models import *
-from .split import *
 from .mhp_loss import *
 from .loss import *
 from .sampler2 import *
@@ -33,10 +30,14 @@ def __init__(self, *args, **kwargs):
         super(ConditionalImageJigsaws, self).__init__(*args, **kwargs)
 
         self.num_options = SuturingNumOptions()
+        self.PredictorCb = PredictorShowImageOnlyMultiStep
 
     def _makeModel(self, image, *args, **kwargs):
 
         img_shape = image.shape[1:]
+        img_size = 1.
+        for dim in img_shape:
+            img_size *= dim
 
         img0_in = Input(img_shape, name="predictor_img0_in")
         img_in = Input(img_shape, name="predictor_img_in")
@@ -52,7 +53,7 @@ def _makeModel(self, image, *args, **kwargs):
 
         # =====================================================================
         # Load weights and stuff
-        LoadEncoderWeights(self, encoder, decoder)
+        LoadEncoderWeights(self, encoder, decoder, gan=True)
         image_discriminator = LoadGoalClassifierWeights(self,
                 make_classifier_fn=MakeJigsawsImageClassifier,
                 img_shape=img_shape)
@@ -82,36 +83,50 @@ def _makeModel(self, image, *args, **kwargs):
         option_in2 = Input((1,), name="option_in2")
         ins += [option_in, option_in2]
 
+        # --------------------------------------------------------------------
+        # Create multiple hypothesis loss
+        lfn = MhpLossWithShape(
+                num_hypotheses=self.num_hypotheses,
+                outputs=[img_size],
+                weights=[1.0],
+                loss=[self.loss],
+                avg_weight=0.05,
+                )
+
+        # --------------------------------------------------------------------
         # Image model
+        h_dim = (12, 16)
+        multi_decoder = MakeJigsawsMultiDecoder(self, decoder,
+                self.num_hypotheses, h_dim)
         y = Flatten()(OneHot(self.num_options)(option_in))
         y2 = Flatten()(OneHot(self.num_options)(option_in2))
-        x = h
-        tform = MakeJigsawsTransform(self, h_dim=(12,16))
-        x = tform([h0, h, y])
+        x = MakeJigsawsExpand(self, h, h_dim)
+        tform = MakeJigsawsTransform(self, h_dim)
+        x = tform([h0, x, y])
         x2 = tform([h0, x, y2])
-        image_out, image_out2 = decoder([x]), decoder([x2])
-        disc_out2 = image_discriminator(image_out2)
+        image_out, image_out2 = multi_decoder([x]), multi_decoder([x2])
+        #disc_out2 = image_discriminator(image_out2)
 
-        lfn = self.loss
         lfn2 = "logcosh"
 
         # =====================================================================
         # Create models to train
         predictor = Model(ins + [prev_option_in],
                 [image_out, image_out2, next_option_out])
         predictor.compile(
-                loss=[lfn, lfn, "binary_crossentropy"],
+                loss=[self.loss, self.loss, "binary_crossentropy"],
                 loss_weights=[1., 1., 0.1],
                 optimizer=self.getOptimizer())
         model = Model(ins + [prev_option_in],
-                [image_out, image_out2, next_option_out, disc_out2])
+                [image_out, image_out2, next_option_out])#, disc_out2])
         model.compile(
-                loss=[lfn, lfn, "binary_crossentropy", "categorical_crossentropy"],
-                loss_weights=[1., 1., 0.1, 1e-4],
+                loss=[lfn, lfn, "binary_crossentropy"],# "categorical_crossentropy"],
+                loss_weights=[1., 1., 0.1],#, 1e-3],
                 optimizer=self.getOptimizer())
 
         self.predictor = predictor
         self.model = model
+        self.model.summary()
 
     def _getData(self, image, label, goal_image, goal_label,
             prev_label, *args, **kwargs):
@@ -128,5 +143,7 @@ def _getData(self, image, label, goal_image, goal_label,
 
         label_1h = np.squeeze(ToOneHot2D(label, self.num_options))
         label2_1h = np.squeeze(ToOneHot2D(label2, self.num_options))
-        return [image0, image, label, goal_label, prev_label], [goal_image, goal_image2, label_1h, label2_1h]
+        return ([image0, image, label, goal_label, prev_label],
+                [np.expand_dims(goal_image, axis=1),
+                 np.expand_dims(goal_image2, axis=1), label_1h])#, label2_1h]
 

costar_models/python/costar_models/conditional_sampler2.py

@@ -67,35 +67,25 @@ def _makePredictor(self, features):
         label_in = Input((1,))
         ins = [img_in, arm_in, gripper_in, label_in]
 
-        encoder = self._makeImageEncoder(img_shape)
-        try:
-            encoder.load_weights(self._makeName(
-                "pretrain_image_encoder_model",
-                "image_encoder.h5f"))
-            encoder.trainable = self.retrain
-        except Exception as e:
-            raise e
-
         if self.skip_connections:
-            decoder = self._makeImageDecoder(self.hidden_shape,self.skip_shape)
+            encoder = self._makeImageEncoder2(img_shape)
+            decoder = self._makeImageDecoder2(self.hidden_shape)
         else:
+            encoder = self._makeImageEncoder(img_shape)
             decoder = self._makeImageDecoder(self.hidden_shape)
-        try:
-            decoder.load_weights(self._makeName(
-                "pretrain_image_encoder_model",
-                "image_decoder.h5f"))
-            decoder.trainable = self.retrain
-        except Exception as e:
-            raise e
 
+        LoadEncoderWeights(self, encoder, decoder)
+        image_discriminator = LoadGoalClassifierWeights(self,
+                make_classifier_fn=MakeImageClassifier,
+                img_shape=img_shape)
+
+        # =====================================================================
+        # Load the arm and gripper representation
         rep_channels = self.encoder_channels
         sencoder = self._makeStateEncoder(arm_size, gripper_size, False)
         sdecoder = self._makeStateDecoder(arm_size, gripper_size,
                 rep_channels)
 
-        # =====================================================================
-        # Load the arm and gripper representation
-
         # =====================================================================
         # combine these models together with state information and label
         # information
@@ -104,12 +94,12 @@ def _makePredictor(self, features):
         hidden_decoder = self._makeFromHidden(rep_channels)
 
         try:
-            hidden_encoder.load_weights(self._makeName(
-                "pretrain_sampler_model",
-                "hidden_encoder.h5f"))
-            hidden_decoder.load_weights(self._makeName(
-                "pretrain_sampler_model",
-                "hidden_decoder.h5f"))
+            hidden_encoder.load_weights(self.makeName(
+                "pretrain_sampler",
+                "hidden_encoder"))
+            hidden_decoder.load_weights(self.makeName(
+                "pretrain_sampler",
+                "hidden_decoder"))
             hidden_encoder.trainable = self.retrain
             hidden_decoder.trainable = self.retrain
         except Exception as e:
@@ -162,7 +152,7 @@ def _makePredictor(self, features):
         return predictor, predictor, actor, ins, h
 
     def _getData(self, *args, **kwargs):
-        features, targets = self._getAllData(*args, **kwargs)
+        features, targets = GetAllMultiData(self.num_options, *args, **kwargs)
         [I, q, g, oin, q_target, g_target,] = features
         tt, o1, v, qa, ga, I_target = targets
         if self.use_noise:

costar_models/python/costar_models/datasets/h5f_generator.py

@@ -15,15 +15,8 @@ class H5fGeneratorDataset(NpzGeneratorDataset):
     takes the load function so all we need to do is implement things so they'll
     load a particular class.
     '''
-    def __init__(self, name, split=0.1, ):
-        '''
-        Set name of directory to load files from
-
-        '''
-        self.name = name 
-        self.split = split
-        self.train = []
-        self.test = []
+    def __init__(self, *args, **kwargs):
+        super(H5fGeneratorDataset, self).__init__(*args, **kwargs)
 
     def _load(self, filename):
         '''