Beispiel #1
0
    def train(self, training_data, target_station, verbose=1):
        print(self.session_name)
        # print(training_data, target_station, self.vol_size, self.strides)

        dataset = Dataset(self._builder_class, training_data, target_station, self.vol_size, self.strides)
        (X_train, Y_train) = dataset.construct()

        model = self.model(summary=False)
        model.fit(X_train, Y_train, epochs=self.epochs, batch_size=self.batch_size, callbacks=self.callbacks(),
                  validation_split=0.05, verbose=verbose, shuffle=True)
Beispiel #2
0
    def test(self, test_data, target_station):
        """
        :param target_station:
        :param test_data:
        """

        model_path = '/home/jamesneill/vessel_seg/session/{}/models/{}.h5'.format(self.__class__.__name__,
                                                                                  self.session_name)

        save_path = '/home/jamesneill/vessel_seg/session/{}/results/{}.pkl'.format(self.__class__.__name__,
                                                                                   self.session_name)

        custom_objects = {
            'dsc': metrics.dsc,
            'iou': metrics.iou,
            'dsc_loss': losses.dsc,
            'prelu': activations.prelu
        }

        print('Testing: {}'.format(model_path))

        # Load model with custom objects suchas dice loss, mIoU etc.
        model = load_model(model_path, custom_objects=custom_objects)

        self.strides = (None, None, 1)
        dataset = Dataset(self._builder_class, test_data, target_station, self.vol_size, self.strides, testing=True)
        (X_test, _) = dataset.construct()
        output_shape = X_test.shape

        padding = []
        padding.append((0, 0))
        for axis in X_test.shape[1:-1]:
            padding.append((0, 16 - (axis % 16)))
        padding.append((0, 0))
        padding = tuple(padding)
        print(padding)

        X_test = np.pad(X_test, padding, 'constant', constant_values=(0, 0))

        rebuilt = model.predict(X_test)
        rebuilt = rebuilt[0:output_shape[0], 0:output_shape[1], 0:output_shape[2]]
        rebuilt = rebuilt[0:output_shape[0], 0:output_shape[1], 0:output_shape[2]]
        rebuilt = rebuilt[0:output_shape[0], 0:output_shape[1], 0:output_shape[2]]
        rebuilt = np.swapaxes(rebuilt, 2, 0)
        print('rebuilt.shape', rebuilt.shape)
        # index = 0
        # for x in range(0, (output_shape[0] / strides[0]) / 2):
        #     for y in range(0, (output_shape[1] / strides[1]) / 2):
        #         for z in range(0, (output_shape[2] / strides[2]) / 2):
        #             rebuilt[index] = predict_test[index]
        #
        #             # rebuilt[strides[0] * x:vol_size[0] + (strides[0] * x),
        #             # strides[1] * y:vol_size[1] + (strides[1] * y),
        #             # strides[2] * z:vol_size[2] + (strides[2] * z)] = predict_test[index]
        #
        #
        #             # strides[2] * z:vol_size[2] + (strides[2] * z)] += predict_test[index]
        #             index += 1

        # rebuilt -= np.min(rebuilt)
        # rebuilt /= np.max(rebuilt)
        # rebuilt -= np.mean(rebuilt) - 0.5

        # print('np.max(rebuilt)', np.max(rebuilt))
        # rebuilt /= np.max(rebuilt)

        classification_thresh = 0.5
        rebuilt[rebuilt < classification_thresh] = 0
        rebuilt[rebuilt >= classification_thresh] = 1

        print('Saving output..')
        pickle.dump(rebuilt, open(save_path, 'wb'))
        imsave('{}.tiff'.format(save_path), rebuilt.astype(np.int16))
        print('Done!')