Esempio n. 1
0
    def __init__(self,
                 patch_size=512,
                 saved_model_path='./model_logs/release_places2_256'):
        '''
        Saved model weights url:
            https://drive.google.com/drive/folders/1y7Irxm3HSHGvp546hZdAZwuNmhLUVcjO
        '''
        self.FLAGS = ng.Config('inpaint.yml')

        self.model = InpaintCAModel()
        self.checkpoint_dir = saved_model_path

        self.patch_size = patch_size

        self.sess_config = tf.ConfigProto()
        self.sess_config.gpu_options.allow_growth = True
        self._inpaint_input_placeholder = tf.placeholder(
            shape=(1, self.patch_size, self.patch_size * 2, 3),
            dtype=tf.float32)
        self.sess = tf.Session(config=self.sess_config)

        output = self.model.build_server_graph(self.FLAGS,
                                               self._inpaint_input_placeholder,
                                               reuse=tf.AUTO_REUSE)
        output = (output + 1.) * 127.5
        output = tf.reverse(output, [-1])
        output = tf.saturate_cast(output, tf.uint8)
        self._cached_inpaint_output = output

        self.load_model()
Esempio n. 2
0
def test_save_and_run_model(tmpdir):
    checkpoint_dir = tmpdir.join('ckpt')  # a 'local' object.
    with tf.Graph().as_default():
        # Setup.
        # Create model.
        model = deep_fill_model.InpaintCAModel()
        dummy_img = tf.zeros([batch_size, 256, 256, 3], tf.float32)
        dummy_mask = tf.ones([batch_size, 256, 256, 1], tf.float32)
        stage_1, stage_2, offset_flow = model.build_inpaint_net(
            x=dummy_img, mask=dummy_mask, training=False, name=net_name)
        with tf.Session() as sess:
            # Initialize variables.
            sess.run(tf.global_variables_initializer())

            # Test.
            # 1. Save random model. There should be no errors.
            global_vars = tf.global_variables()
            saver = tf.train.Saver(global_vars)
            saver.save(sess, str(checkpoint_dir))

    # 2. Use random model model to fill.
    #input_image_path = './test/resources/case1_input.png'
    #mask_path = './test/resources/case1_mask.png'
    input_image_path = './waseda_fill/statue1.jpg'
    mask_path = './waseda_fill/statue1_mask2.png'
    out = deep_fill_app.fill(
        image_path=input_image_path,
        mask_path=mask_path,
        checkpoint_dir=str(checkpoint_dir),
        FLAGS=ng.Config('./test/resources/inpaint_test.yml'))
    reference_out = str(tmpdir.join('case1_filled_using_untrained_model.png'))
    cv2.imwrite('waseda_rand.png', out[0][:, :, ::-1])
    (score, diff) = cv2.compare_ssim(out, reference_out, full=True)
    assert score > 0.9
Esempio n. 3
0
def deepfill_inpaint(basedata, imagepath, maskinfo, status, checkpointdir,
                     inputimgpath, outputpath):
    # basedata = BaseData()
    FLAGS = ng.Config(join(basedata.DEEPFILL_BASE_DIR, 'inpaint.yml'))
    # FLAGS = ng.Config('./inpaint.yml')
    # FLAGS = ng.Config('/home/zzy/work/dnnii_web/dnnii_web/App/deepfill/inpaint.yml')
    # ng.get_gpus(1)
    # args, unknown = parser.parse_known_args()

    model = InpaintCAModel()
    image = cv2.imread(imagepath)
    h, w, _ = image.shape
    if status == 0:
        mask = np.zeros((h, w, 3)).astype(np.uint8)
        for rect in maskinfo:
            mask[rect[1]:rect[1] + rect[3], rect[0]:rect[0] + rect[2], :] = 255
    else:
        mask = cv2.imread(maskinfo)

    mask = cv2.resize(mask, (w, h), fx=0.5, fy=0.5)
    assert image.shape == mask.shape

    #把原始图片划分成grid*grid个格子区域,'//'表示向下取整的除法
    grid = 8
    image = image[:h // grid * grid, :w // grid * grid, :]
    mask = mask[:h // grid * grid, :w // grid * grid, :]
    print('Shape of image: {}'.format(image.shape))

    inputimage = image * ((255 - mask) // 255) + mask
    cv2.imwrite(inputimgpath, inputimage.astype(np.uint8))

    image = np.expand_dims(image, 0)
    mask = np.expand_dims(mask, 0)
    input_image = np.concatenate([image, mask], axis=2)

    sess_config = tf.ConfigProto()
    sess_config.gpu_options.allow_growth = True
    deepfill_graph = tf.Graph()
    with tf.Session(config=sess_config, graph=deepfill_graph) as deepfill_sess:
        input_image = tf.constant(input_image, dtype=tf.float32)
        output = model.build_server_graph(FLAGS, input_image)
        output = (output + 1.) * 127.5
        output = tf.reverse(output, [-1])
        output = tf.saturate_cast(output, tf.uint8)
        # load pretrained model
        vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
        assign_ops = []
        for var in vars_list:
            vname = var.name
            from_name = vname
            var_value = tf.contrib.framework.load_variable(
                checkpointdir, from_name)
            assign_ops.append(tf.assign(var, var_value))
        deepfill_sess.run(assign_ops)
        print('deepfill Model loaded.')
        result = deepfill_sess.run(output)
        cv2.imwrite(outputpath, result[0][:, :, ::-1])
        deepfill_sess.close()
def Output(img_in, mask_in, img_out):
    print("import from test:", img_in, mask_in, img_out)

    #if __name__ == "__main__":
    FLAGS = ng.Config('inpaint.yml')
    # ng.get_gpus(1)

    #args, unknown = parser.parse_known_args()

    model = InpaintCAModel()
    #image = cv2.imread(args.image)
    image = cv2.imread("examples/places356/" + img_in)
    #mask = cv2.imread(args.mask)
    mask = cv2.imread("places356_mask/" + mask_in)
    # mask = cv2.resize(mask, (0,0), fx=0.5, fy=0.5)

    assert image.shape == mask.shape

    h, w, _ = image.shape
    grid = 8
    image = image[:h // grid * grid, :w // grid * grid, :]
    mask = mask[:h // grid * grid, :w // grid * grid, :]
    print('Shape of image: {}'.format(image.shape))

    image = np.expand_dims(image, 0)
    mask = np.expand_dims(mask, 0)
    input_image = np.concatenate([image, mask], axis=2)

    sess_config = tf.ConfigProto()
    sess_config.gpu_options.allow_growth = True
    with tf.Session(config=sess_config) as sess:
        input_image = tf.constant(input_image, dtype=tf.float32)
        output = model.build_server_graph(FLAGS, input_image)
        output = (output + 1.) * 127.5
        output = tf.reverse(output, [-1])
        output = tf.saturate_cast(output, tf.uint8)
        # load pretrained model
        vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
        assign_ops = []
        for var in vars_list:
            vname = var.name
            from_name = vname

            #var_value = tf.contrib.framework.load_variable(args.checkpoint_dir, from_name)
            var_value = tf.contrib.framework.load_variable(
                "model_logs/release_places2_256", from_name)
            assign_ops.append(tf.assign(var, var_value))
        sess.run(assign_ops)
        print('Model loaded.')
        result = sess.run(output)
        #cv2.imwrite(args.output, result[0][:, :, ::-1])
        #cv2.imshow("result", result[0][:, :, ::-1])
        cv2.imwrite("examples/places356/" + img_out, result[0][:, :, ::-1])
        show1 = cv2.imread("examples/places356/" + img_in)
        show2 = cv2.imread("examples/places356/" + img_out)
        show = np.hstack([show1, show2])
        cv2.imshow("result", show)
def inpaint(arg_image_dir, arg_mask_dir, arg_checkpoint_dir, arg_output_dir):
    tf.reset_default_graph()
    FLAGS = ng.Config('inpaint.yml')
    # ng.get_gpus(1)

    model = InpaintCAModel()
    for arg_image in os.listdir(arg_image_dir):
        arg_mask = arg_image  # assume the mask has the same name as the image
        if os.path.exists(arg_output_dir + arg_image):
            print("note |", arg_image, "already inpainted.")
            continue
        if os.path.exists(arg_image_dir +
                          arg_image) and os.path.exists(arg_mask_dir +
                                                        arg_mask):
            pass
        else:
            continue

        image = cv2.imread(os.path.join(arg_image_dir, arg_image))
        mask = cv2.imread(os.path.join(arg_mask_dir, arg_mask))
        name = arg_image

        assert image.shape == mask.shape

        h, w, _ = image.shape
        grid = 8
        image = image[:h // grid * grid, :w // grid * grid, :]
        mask = mask[:h // grid * grid, :w // grid * grid, :]
        print('Shape of image: {}'.format(image.shape))

        image = np.expand_dims(image, 0)
        mask = np.expand_dims(mask, 0)
        input_image = np.concatenate([image, mask], axis=2)

        tf.reset_default_graph()
        sess_config = tf.ConfigProto()
        sess_config.gpu_options.allow_growth = True
        with tf.Session(config=sess_config) as sess:
            input_image = tf.constant(input_image, dtype=tf.float32)
            output = model.build_server_graph(FLAGS, input_image)
            output = (output + 1.) * 127.5
            output = tf.reverse(output, [-1])
            output = tf.saturate_cast(output, tf.uint8)
            # load pretrained model
            vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
            assign_ops = []
            for var in vars_list:
                vname = var.name
                from_name = vname
                var_value = tf.contrib.framework.load_variable(
                    arg_checkpoint_dir, from_name)
                assign_ops.append(tf.assign(var, var_value))
            sess.run(assign_ops)
            print('Model loaded.')
            result = sess.run(output)
            cv2.imwrite(arg_output_dir + name, result[0][:, :, ::-1])
Esempio n. 6
0
 def __init__(self,
              checkpoint_dir=None,
              get_checkpoint_dir=download_radish,
              config=abspath('inpaint.yml'),
              max_size=(1024, 1024)):
     self.FLAGS = ng.Config(config)
     self.max_size = max_size
     if checkpoint_dir is None:
         checkpoint_dir = get_checkpoint_dir()
     self.checkpoint_dir = checkpoint_dir
     self.model_loaded = False
Esempio n. 7
0
def inpaintP(name):
    CUR_DIR = os.path.join(WORK_DIR, name)
    TMASK_DIR = os.path.join(WORK_DIR, name + "//tmask")
    INPAINT_MODEL_PATH = os.path.join(ROOT_DIR,
                                      "model_logs/release_places2_256")

    FLAGS = ng.Config('inpaint.yml')
    model = InpaintCAModel()

    image = cv.imread(os.path.join(CUR_DIR, f"{name}.png"))
    mask = cv.imread(os.path.join(TMASK_DIR, "mask.png"))
    filename = f'4#_{name}.png'

    assert image.shape == mask.shape

    h, w, _ = image.shape
    grid = 8
    image = image[:h // grid * grid, :w // grid * grid, :]
    mask = mask[:h // grid * grid, :w // grid * grid, :]
    print('Shape of image: {}'.format(image.shape))

    image = np.expand_dims(image, 0)
    mask = np.expand_dims(mask, 0)
    input_image = np.concatenate([image, mask], axis=2)

    sess_config = tf.ConfigProto()
    sess_config.gpu_options.per_process_gpu_memory_fraction = 0.5
    tf2 = tf.Graph()
    with tf2.as_default():
        with tf.Session(config=sess_config) as sess:
            input_image = tf.constant(input_image, dtype=tf.float32)

            output = model.build_server_graph(FLAGS, input_image)
            output = (output + 1.) * 127.5
            output = tf.reverse(output, [-1])
            output = tf.saturate_cast(output, tf.uint8)

            # load pretrained model
            vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
            assign_ops = []
            for var in vars_list:
                vname = var.name
                from_name = vname
                var_value = tf.contrib.framework.load_variable(
                    INPAINT_MODEL_PATH, from_name)
                assign_ops.append(tf.assign(var, var_value))

            sess.run(assign_ops)
            print('Model loaded.')
            result = sess.run(output)
            cv.imwrite(os.path.join(CUR_DIR, filename), result[0][:, :, ::-1])
            print('Image has been made')

    return 0
Esempio n. 8
0
def generate_counterfactual(image_fp,
                            mask_fp,
                            output_fp,
                            checkpoint_dir,
                            model_id=None):
    try:
        FLAGS = ng.Config('config/inpaint.yml')
    except AssertionError:
        raise ValueError('check directory above')
    # ng.get_gpus(1)
    # args, unknown = parser.parse_known_args()
    model = InpaintCAModel()
    image = cv2.imread(image_fp)
    mask = cv2.imread(mask_fp)
    # mask = cv2.resize(mask, (0,0), fx=0.5, fy=0.5)

    assert image.shape == mask.shape

    h, w, _ = image.shape
    grid = 8
    image = image[:h // grid * grid, :w // grid * grid, :]
    mask = mask[:h // grid * grid, :w // grid * grid, :]

    image = np.expand_dims(image, 0)
    mask = np.expand_dims(mask, 0)
    input_image = np.concatenate([image, mask], axis=2)

    sess_config = tf.ConfigProto()
    sess_config.gpu_options.allow_growth = True
    with tf.Session(config=sess_config) as sess:
        input_image = tf.constant(input_image, dtype=tf.float32)
        output = model.build_server_graph(FLAGS, input_image)
        output = (output + 1.) * 127.5
        output = tf.reverse(output, [-1])
        output = tf.saturate_cast(output, tf.uint8)
        # load pretrained model
        vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
        assign_ops = []
        for var in vars_list:
            vname = var.name
            from_name = vname
            var_value = tf.contrib.framework.load_variable(
                checkpoint_dir, from_name)
            assign_ops.append(tf.assign(var, var_value))
        sess.run(assign_ops)
        print('Model loaded.')
        result = sess.run(output)
        cv2.imwrite(output_fp, result[0][:, :, ::-1])
        print(f'IMAGE WROTE TO {output_fp}\n\n\n')
    tf.reset_default_graph()
Esempio n. 9
0
def pixel_fill(image, mask):
    import tensorflow as tf2
    import neuralgym as ng
    from inpaint_model import InpaintCAModel

    if image.ndim > mask.ndim:
        mask = np.dstack([mask] * image.shape[2])
    assert image.shape == mask.shape
    model = InpaintCAModel()
    FLAGS = ng.Config('inpaint.yml')

    h, w, _ = image.shape
    grid = 8
    image = image[:h // grid * grid, :w // grid * grid, :]
    mask = mask[:h // grid * grid, :w // grid * grid, :]
    print('Shape of image: {}'.format(image.shape))

    image = np.expand_dims(image, 0)
    mask = np.expand_dims(mask, 0)
    input_image = np.concatenate([image, mask], axis=2)

    sess_config = tf2.ConfigProto()
    sess_config.gpu_options.allow_growth = True
    with tf2.Session(config=sess_config) as sess:
        input_image = tf2.constant(input_image, dtype=tf2.float32)
        output = model.build_server_graph(FLAGS, input_image)
        output = (output + 1.) * 127.5
        output = tf2.reverse(output, [-1])
        output = tf2.saturate_cast(output, tf2.uint8)
        # load pretrained model
        vars_list = tf2.get_collection(tf2.GraphKeys.GLOBAL_VARIABLES)
        assign_ops = []
        print("checkpoint_dir = ", checkpoint_dir)
        for var in vars_list:
            vname = var.name
            from_name = vname
            if "inpaint_net" in var.name:  # or else is going to mix with mrcnn
                var_value = tf2.contrib.framework.load_variable(
                    checkpoint_dir, from_name)
                assign_ops.append(tf2.assign(var, var_value))
        sess.run(assign_ops)
        print('Model loaded.' * 10)
        result = sess.run(output)
    sess.close()
    tf2.reset_default_graph()
    return result[0][:, :, ::-1]
Esempio n. 10
0
    def build_server_graph(self, batch_data, reuse=False, is_training=False):
        """
        """
        # generate mask, 1 represents masked point
        batch_raw, masks_raw = tf.split(batch_data, 2, axis=2)
        masks = tf.cast(masks_raw[0:1, :, :, 0:1] > 127.5, tf.float32)

        batch_pos = batch_raw / 127.5 - 1.
        batch_incomplete = batch_pos * (1. - masks)
        # inpaint
        config = ng.Config('inpaint.yml')
        x1, x2, flow = self.build_inpaint_net(
            batch_incomplete, masks, reuse=reuse, training=is_training,
            config=config)
        batch_predict = x2
        # apply mask and reconstruct
        batch_complete = batch_predict*masks + batch_incomplete*(1-masks)
        return batch_complete
Esempio n. 11
0
def inpainting_api(image, mask):
    FLAGS = ng.Config('inpaint.yml')
    tf.reset_default_graph()

    model = InpaintCAModel()
    # image = cv2.imread(img_path)
    # mask = cv2.imread(mask_path)
    # cv2.imwrite('new.png', image - mask)
    # mask = cv2.resize(mask, (0,0), fx=0.5, fy=0.5)

    assert image.shape == mask.shape

    image = crop(image)
    mask = crop(mask)
    print('Shape of image: {}'.format(image.shape))

    image = np.expand_dims(image, 0)
    mask = np.expand_dims(mask, 0)
    input_image = np.concatenate([image, mask], axis=2)

    sess_config = tf.ConfigProto()
    sess_config.gpu_options.allow_growth = True
    with tf.Session(config=sess_config) as sess:
        input_image = tf.constant(input_image, dtype=tf.float32)
        output = model.build_server_graph(FLAGS, input_image)
        output = (output + 1.) * 127.5
        output = tf.reverse(output, [-1])
        output = tf.saturate_cast(output, tf.uint8)
        # load pretrained model
        vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
        assign_ops = []
        for var in vars_list:
            vname = var.name
            from_name = vname
            var_value = tf.contrib.framework.load_variable(
                './model_logs/inpaint', from_name)
            assign_ops.append(tf.assign(var, var_value))
        sess.run(assign_ops)
        print('Model loaded.')
        result = sess.run(output)
        return result[0][:, :, ::-1]
Esempio n. 12
0
    def __init__(self,
                 checkpoint_dir,
                 use_gpu=False,
                 preview_socket_name=None):
        self.video_creator = VideoCreator(fps=5)
        if preview_socket_name is not None:
            self.preview_socket = socket.socket(socket.AF_UNIX,
                                                socket.SOCK_STREAM)
            self.preview_socket.connect(preview_socket_name)
        self.recording = False

        FLAGS = ng.Config('generative_inpainting/inpaint.yml')
        if use_gpu:
            ng.get_gpus(1)

        sess_config = tf.compat.v1.ConfigProto()
        sess_config.gpu_options.allow_growth = True
        self.sess = tf.compat.v1.Session(config=sess_config)

        model = InpaintCAModel()
        self.input_image_ph = tf.compat.v1.placeholder(tf.float32,
                                                       shape=(1, 640, 480 * 2,
                                                              3))
        output = model.build_server_graph(FLAGS, self.input_image_ph)
        output = (output + 1.) * 127.5
        output = tf.reverse(output, [-1])
        self.output = tf.saturate_cast(output, tf.uint8)
        vars_list = tf.compat.v1.get_collection(
            tf.compat.v1.GraphKeys.GLOBAL_VARIABLES)
        assign_ops = []
        for var in vars_list:
            vname = var.name
            from_name = vname
            var_value = tf.contrib.framework.load_variable(
                checkpoint_dir, from_name)
            assign_ops.append(tf.compat.v1.assign(var, var_value))
        self.sess.run(assign_ops)

        self.on_image_received(
            ThermalImage(0, Image.new('RGB', (480, 640), color='black'),
                         np.full((480 * 640), False)))
Esempio n. 13
0
def inpaintP(name):
    import os
    import time
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    import cv2 as cv
    import numpy as np
    import tensorflow as tf
    tf.compat.v1.logging.set_verbosity("ERROR")
    import neuralgym as ng
    from inpaint_model import InpaintCAModel

    ROOT_DIR = os.path.abspath("./")
    WAIT_DIR = os.path.abspath("./waiting")
    WORK_DIR = os.path.abspath("./workspace")

    CUR_DIR = os.path.join(WORK_DIR, name)
    TMASK_DIR = os.path.join(WORK_DIR, name+"//tmask")
    INPAINT_MODEL_PATH = os.path.join(ROOT_DIR, "model_logs/release_places2_256")

    FLAGS = ng.Config('inpaint.yml')
    model = InpaintCAModel()

    image = cv.imread(os.path.join(CUR_DIR, f"{name}.png"))
    mask = cv.imread(os.path.join(TMASK_DIR, "mask.png"))
    filename = f'4#_{name}.png'

    assert image.shape == mask.shape

    h, w, _ = image.shape
    grid = 8
    image = image[:h//grid*grid, :w//grid*grid, :]
    mask = mask[:h//grid*grid, :w//grid*grid, :]
    print('Shape of image: {}'.format(image.shape))

    image = np.expand_dims(image, 0)
    mask = np.expand_dims(mask, 0)
    input_image = np.concatenate([image, mask], axis=2)

    sess_config = tf.ConfigProto()
    sess_config.gpu_options.per_process_gpu_memory_fraction = 0.5

    with tf.Session(config=sess_config) as sess:
        input_image = tf.constant(input_image, dtype=tf.float32)

        output = model.build_server_graph(FLAGS, input_image)
        output = (output + 1.) * 127.5
        output = tf.reverse(output, [-1])
        output = tf.saturate_cast(output, tf.uint8)

        # load pretrained model
        vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
        assign_ops = []
        for var in vars_list:
            vname = var.name
            from_name = vname
            var_value = tf.contrib.framework.load_variable(INPAINT_MODEL_PATH, from_name)
            assign_ops.append(tf.assign(var, var_value))

        sess.run(assign_ops)
        print('Model loaded.')
        result = sess.run(output)
        cv.imwrite(os.path.join(CUR_DIR, filename), result[0][:, :, ::-1])
Esempio n. 14
0
                                                     summary=True,
                                                     reuse=True)
    else:
        _, _, losses = model.build_graph_with_losses(images,
                                                     config,
                                                     reuse=True)
    if loss_type == 'g':
        return losses['g_loss']
    elif loss_type == 'd':
        return losses['d_loss']
    else:
        raise ValueError('loss type is not supported.')


if __name__ == "__main__":
    config = ng.Config('inpaint.yml')
    if config.GPU_ID != -1:
        ng.set_gpus(config.GPU_ID)
    else:
        ng.get_gpus(config.NUM_GPUS)

    # training data
    with open(config.DATA_FLIST[config.DATASET][0]) as f:
        fnames = f.read().splitlines()
    data = ng.data.DataFromFNames(fnames,
                                  config.IMG_SHAPES,
                                  random_crop=config.RANDOM_CROP)
    images = data.data_pipeline(config.BATCH_SIZE)

    # main model
    model = InpaintCAModel()
Esempio n. 15
0
 def __init__(self, checkpoint_dir: os.PathLike,
              ng_config_path: os.PathLike) -> None:
     self.checkpoint_dir = checkpoint_dir
     self.FLAGS = ng.Config(ng_config_path)
Esempio n. 16
0
    '--image_height', default=256, type=int,
    help='The height of images should be defined, otherwise batch mode is not'
    ' supported.')
parser.add_argument(
    '--image_width', default=256, type=int,
    help='The width of images should be defined, otherwise batch mode is not'
    ' supported.')
parser.add_argument(
    '--checkpoint_dir', default='', type=str,
    help='The directory of tensorflow checkpoint.')
parser.add_argument(
    '--outlist', default='', type=str,
    help='The directory of putting out image.')

if __name__ == "__main__":
    FLAGS = ng.Config('./inpaint_dem.yml')
    ng.get_gpus(1)
    # os.environ['CUDA_VISIBLE_DEVICES'] =''
    args = parser.parse_args()

    sess_config = tf.ConfigProto()
    sess_config.gpu_options.allow_growth = True
    sess = tf.Session(config=sess_config)

    model = InpaintCAModel()
    input_image_ph = tf.placeholder(
        tf.float32, shape=(1, args.image_height, args.image_width*2, 1))
    output = model.build_server_graph(FLAGS, input_image_ph)
    output = tf.reverse(output, [-1])
    vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
    assign_ops = []
        _, _, losses = model.build_graph_with_losses(images,
                                                     config,
                                                     reuse=True)
    if loss_type == 'g':
        return losses['g_loss']
    elif loss_type == 'd':
        return losses['d_loss']
    else:
        raise ValueError('loss type is not supported.')


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument('config', default='inpaint.yml')
    args = parser.parse_args()
    config = ng.Config(args.config)
    if config.GPU_ID != -1:
        ng.set_gpus(config.GPU_ID)
    else:
        ng.get_gpus(config.NUM_GPUS, dedicated=False)
    # training data
    with open(config.DATA_FLIST[config.DATASET][0]) as f:
        fnames = f.read().splitlines()
    data = ng.data.DataFromFNames(fnames,
                                  config.IMG_SHAPES,
                                  random_crop=config.RANDOM_CROP,
                                  random_flip=config.RANDOM_FLIP,
                                  gamma=config.GAMMA,
                                  exposure=config.EXPOSURE)
    images = data.data_pipeline(config.BATCH_SIZE)
    # main model
Esempio n. 18
0
    else:
        _, _, losses = model.build_graph_with_losses(images,
                                                     masks,
                                                     guides,
                                                     config,
                                                     reuse=True)
    if loss_type == 'g':
        return losses['g_loss']
    elif loss_type == 'd':
        return losses['d_loss']
    else:
        raise ValueError('loss type is not supported.')


if __name__ == "__main__":
    config = ng.Config(sys.argv[1])
    if config.GPU_ID != -1:
        ng.set_gpus(config.GPU_ID)
    else:
        ng.get_gpus(config.NUM_GPUS)
    # training data
    # Image Data
    with open(config.DATA_FLIST[config.DATASET][0]) as f:
        fnames = f.read().splitlines()
    # # Mask Data

    if config.MASKFROMFILE:
        with open(config.DATA_FLIST[config.MASKDATASET][0]) as f:
            mask_fnames = f.read().splitlines()
        data_mask_data = DataMaskFromFNames(list(zip(
            fnames, mask_fnames)), [config.IMG_SHAPES, config.MASK_SHAPES],
            'loss_type': 'g'
        },
        spe=config.TRAIN_SPE,
        max_iters=config.MAX_ITERS,
        log_dir=log_prefix,
    )

    trainer.add_callbacks([
        ng.callbacks.WeightsViewer(),
        # ng.callbacks.ModelRestorer(trainer.context['saver'], dump_prefix=config.RESTORE_PREFIX, optimistic=True),
        discriminator_training_callback,
        ng.callbacks.ModelSaver(config.TRAIN_SPE, trainer.context['saver'],
                                log_prefix + '/snap'),
        ng.callbacks.SummaryWriter((config.VAL_PSTEPS // 1),
                                   trainer.context['summary_writer'],
                                   tf.summary.merge_all()),
    ])

    trainer.train()


if __name__ == "__main__":
    config = ng.Config('progressive_gan.yml')
    if config.GPU_ID != -1:
        ng.set_gpus(config.GPU_ID)
    else:
        ng.get_gpus(config.NUM_GPUS)
    np.random.seed(config.RANDOM_SEED)

    eval(config.TRAIN.func + '(config)')
Esempio n. 20
0
import cv2
import numpy as np
import tensorflow as tf
import neuralgym as ng

from .inpaint_model import InpaintCAModel

checkpoint_dir = 'generative_inpainting/models'
FLAGS = ng.Config('generative_inpainting/inpaint.yml')


def run_fill(file_test, file_mask):
    model = InpaintCAModel()
    image = cv2.imread(file_test)
    mask = cv2.imread(file_mask)

    h, w, _ = image.shape
    grid = 8
    image = image[:h // grid * grid, :w // grid * grid, :]
    mask = mask[:h // grid * grid, :w // grid * grid, :]

    image = np.expand_dims(image, 0)
    mask = np.expand_dims(mask, 0)
    input_image = np.concatenate([image, mask], axis=2)

    sess_config = tf.ConfigProto()
    sess_config.gpu_options.allow_growth = True
    with tf.Session(config=sess_config) as sess:
        input_image = tf.constant(input_image, dtype=tf.float32)
        output = model.build_server_graph(FLAGS, input_image)
        output = (output + 1.) * 127.5
Esempio n. 21
0
        vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
        assign_ops = []
        idx_list = [54, 56, 58, 60, 62] # random
        for idx, var in enumerate(vars_list):
            vname = var.name
            from_name = vname
            var_value = tf.contrib.framework.load_variable(checkpoint_dir,
                                                           from_name)
            assign_ops.append(tf.assign(var, var_value))
        sess.run(assign_ops)
        print('Model loaded.')
        result = sess.run(output)
        return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--image', default='', type=str,
                        help='The filename of image to be completed.')
    parser.add_argument('--mask', default='', type=str,
                        help='The filename of mask, value 255 indicates mask.')
    parser.add_argument('--output', default='output.png', type=str,
                        help='Where to write output.')
    parser.add_argument('--checkpoint_dir', default='', type=str,
                        help='The directory of tensorflow checkpoint.')

    FLAGS = ng.Config('./resources/inpaint.yml')
    args, _ = parser.parse_known_args()
    result = fill(args.image, args.mask, args.checkpoint_dir, FLAGS)
    cv2.imwrite(args.output, result[0][:, :, ::-1])
Esempio n. 22
0
                                                     reuse=True)
    if loss_type == 'g':
        return losses['g_loss']
    elif loss_type == 'd':
        return losses['d_loss']
    else:
        raise ValueError('loss type is not supported.')


if __name__ == "__main__":
    if len(sys.argv) > 1:
        yml_path = sys.argv[1]
    else:
        yml_path = 'inpaint.yml'

    config = ng.Config(yml_path)
    if config.GPU_ID != -1:
        ng.set_gpus(config.GPU_ID)
    else:
        ng.get_gpus(config.NUM_GPUS)
    # training data
    with open(config.DATA_FLIST[config.DATASET][0]) as f:
        fnames = f.read().splitlines()
    data = ng.data.DataFromFNames(fnames,
                                  config.IMG_SHAPES,
                                  random_crop=config.RANDOM_CROP)
    images = data.data_pipeline(config.BATCH_SIZE)
    # main model
    model = InpaintCAModel()
    g_vars, d_vars, losses = model.build_graph_with_losses(images,
                                                           config=config)
Esempio n. 23
0
def cagc_inp_batch():
    baserecutpath = '/home/zzy/TrainData/MITPlace2Dataset/val_recut_512x680'
    imagesnums = 132
    checkpointdir = "./checkpoints/places2_512x680"
    maskpath = "./examples/places2_680x512/wooden_mask.png"
    recutoutput = './examples/recutoutput/'
    recutmaskedoutput = './examples/recutmasked/'

    FLAGS = ng.Config('inpaint.yml')
    # ng.get_gpus(1)
    # args, unknown = parser.parse_known_args()
    model = InpaintCAModel()

    sess_config = tf.ConfigProto()
    sess_config.gpu_options.allow_growth = True
    with tf.Session(config=sess_config) as sess:
        for imageindex in range(1, imagesnums):
            imagename = 'Places365_' + str(imageindex).zfill(8) + '.png'
            maskedimagename = 'Places365_' + str(imageindex).zfill(
                8) + '_masked.png'
            outputimagename = 'Places365_' + str(imageindex).zfill(
                8) + '_output.png'
            imagepath = os.path.join(baserecutpath, imagename)
            outputpath = os.path.join(recutoutput, outputimagename)
            maskedimagepath = os.path.join(recutmaskedoutput, maskedimagename)

            image = cv2.imread(imagepath)
            mask = cv2.imread(maskpath)
            h, w, _ = image.shape
            mask = cv2.resize(mask, (w, h), fx=0.5, fy=0.5)

            inputimage = image * ((255 - mask) // 255) + mask
            cv2.imwrite(maskedimagepath, inputimage.astype(np.uint8))

            assert image.shape == mask.shape
            # 把原始图片划分成grid*grid个格子区域,'//'表示向下取整的除法
            grid = 8
            image = image[:h // grid * grid, :w // grid * grid, :]
            mask = mask[:h // grid * grid, :w // grid * grid, :]
            print('Shape of image: {}'.format(image.shape))

            image = np.expand_dims(image, 0)
            mask = np.expand_dims(mask, 0)
            input_image = np.concatenate([image, mask], axis=2)
            input_image = tf.constant(input_image, dtype=tf.float32)
            output = model.build_server_graph(FLAGS, input_image)
            output = (output + 1.) * 127.5
            output = tf.reverse(output, [-1])
            output = tf.saturate_cast(output, tf.uint8)
            # load pretrained model
            vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
            assign_ops = []
            for var in vars_list:
                vname = var.name
                from_name = vname
                var_value = tf.contrib.framework.load_variable(
                    checkpointdir, from_name)
                assign_ops.append(tf.assign(var, var_value))
            sess.run(assign_ops)
            print('Model loaded.')
            result = sess.run(output)
            cv2.imwrite(outputpath, result[0][:, :, ::-1])
    else:
        _, _, losses = model.build_graph_with_losses(FLAGS,
                                                     images,
                                                     FLAGS,
                                                     reuse=True)
    if loss_type == 'g':
        return losses['g_loss']
    elif loss_type == 'd':
        return losses['d_loss']
    else:
        raise ValueError('loss type is not supported.')


if __name__ == "__main__":
    # training data
    FLAGS = ng.Config('inpaint.yml')
    img_shapes = FLAGS.img_shapes
    with open(FLAGS.data_flist[FLAGS.dataset][0]) as f:
        fnames = f.read().splitlines()
    if FLAGS.guided:
        fnames = [(fname, fname[:-4] + '_edge.jpg') for fname in fnames]
        img_shapes = [img_shapes, img_shapes]
    data = ng.data.DataFromFNames(fnames,
                                  img_shapes,
                                  random_crop=FLAGS.random_crop,
                                  nthreads=FLAGS.num_cpus_per_job)
    images = data.data_pipeline(FLAGS.batch_size)
    # main model
    model = InpaintCAModel()
    g_vars, d_vars, losses = model.build_graph_with_losses(FLAGS, images)
    # validation images
Esempio n. 25
0
def deepfillbatch(image_height, image_width, checkpoint_dir, img_mask_txt,
                  outputdir):
    FLAGS = ng.Config('./deepfill/inpaint.yml')
    ng.get_gpus(1)
    # os.environ['CUDA_VISIBLE_DEVICES'] =''

    sess_config = tf.ConfigProto()
    sess_config.gpu_options.allow_growth = True
    sess = tf.Session(config=sess_config)

    model = InpaintCAModel()
    input_image_ph = tf.placeholder(tf.float32,
                                    shape=(1, image_height, image_width * 2,
                                           3))
    output = model.build_server_graph(FLAGS, input_image_ph)
    output = (output + 1.) * 127.5
    output = tf.reverse(output, [-1])
    output = tf.saturate_cast(output, tf.uint8)
    vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
    assign_ops = []
    for var in vars_list:
        vname = var.name
        from_name = vname
        var_value = tf.contrib.framework.load_variable(checkpoint_dir,
                                                       from_name)
        assign_ops.append(tf.assign(var, var_value))
    sess.run(assign_ops)
    print('Model loaded.')
    t = time.time()
    with open(img_mask_txt, 'r') as f:
        while True:
            line = f.readline()  # 整行读取数据
            if not line:
                break

            imagepath, maskpath = line.replace('\n', '').replace('\r',
                                                                 '').split(',')
            print(imagepath)
            print(maskpath)
            if not os.path.exists(outputdir):
                os.mkdir(outputdir)
            outputpath = os.path.join(outputdir, os.path.basename(imagepath))

            image = cv2.imread(imagepath)
            mask = cv2.imread(maskpath)
            print(image.shape)
            print(mask.shape)
            image = cv2.resize(image, (image_width, image_height))
            mask = cv2.resize(mask, (image_width, image_height))
            # cv2.imwrite(maskedimg, image*(1-mask/255.) + mask)
            # # continue
            # image = np.zeros((128, 256, 3))
            # mask = np.zeros((128, 256, 3))

            assert image.shape == mask.shape

            h, w, _ = image.shape
            grid = 4
            image = image[:h // grid * grid, :w // grid * grid, :]
            mask = mask[:h // grid * grid, :w // grid * grid, :]
            print('Shape of image: {}'.format(image.shape))

            image = np.expand_dims(image, 0)
            mask = np.expand_dims(mask, 0)
            input_image = np.concatenate([image, mask], axis=2)

            # load pretrained model
            result = sess.run(output, feed_dict={input_image_ph: input_image})
            print('Processed: {}'.format(outputpath))
            cv2.imwrite(outputpath, result[0][:, :, ::-1])

    print('Time total: {}'.format(time.time() - t))