jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def testSeparableConv2DSameWithInputOddSize(self):
    n, n2 = 5, 3

    # Input image.
    x = create_test_input(1, n, n, 1)

    # Convolution kernel.
    dw = create_test_input(1, 3, 3, 1)
    dw = tf.reshape(dw, [3, 3, 1, 1])

    tf.get_variable('Conv/depthwise_weights', initializer=dw)
    tf.get_variable('Conv/pointwise_weights',
                    initializer=tf.ones([1, 1, 1, 1]))
    tf.get_variable('Conv/biases', initializer=tf.zeros([1]))
    tf.get_variable_scope().reuse_variables()

    y1 = slim.separable_conv2d(x, 1, [3, 3], depth_multiplier=1,
                               stride=1, scope='Conv')
    y1_expected = tf.cast([[14, 28, 43, 58, 34],
                           [28, 48, 66, 84, 46],
                           [43, 66, 84, 102, 55],
                           [58, 84, 102, 120, 64],
                           [34, 46, 55, 64, 30]], tf.float32)
    y1_expected = tf.reshape(y1_expected, [1, n, n, 1])

    y2 = resnet_utils.subsample(y1, 2)
    y2_expected = tf.cast([[14, 43, 34],
                           [43, 84, 55],
                           [34, 55, 30]], tf.float32)
    y2_expected = tf.reshape(y2_expected, [1, n2, n2, 1])

    y3 = xception.separable_conv2d_same(x, 1, 3, depth_multiplier=1,
                                        regularize_depthwise=True,
                                        stride=2, scope='Conv')
    y3_expected = y2_expected

    y4 = slim.separable_conv2d(x, 1, [3, 3], depth_multiplier=1,
                               stride=2, scope='Conv')
    y4_expected = y2_expected

    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      self.assertAllClose(y1.eval(), y1_expected.eval())
      self.assertAllClose(y2.eval(), y2_expected.eval())
      self.assertAllClose(y3.eval(), y3_expected.eval())
      self.assertAllClose(y4.eval(), y4_expected.eval())
//  [python]
def get_model(data, name="UNET-BASE", **params):
    def _get_log_weight_dirs(name, dt):
        _dt = dt.strftime("%Y_%m_%d_T%H_%M_%S")
        relative_dt_dir = os.path.join(name, _dt)
        dt_dir = os.path.join(config.MODEL_OUT_DIR, relative_dt_dir)
        return os.path.join(dt_dir, 'logs'), os.path.join(dt_dir, 'weights')


    start_filter = params.get('START_FILTER')
    middle_filter = params.get('MIDDLE_FILTERS')
    drop_rate = params.get('DROP_RATE')
    epochs = params.get('EPOCHS')
    batch_size = params.get('BATCH_SIZE')
    num_thresholds = params.get('NUM_THRESHOLDS')

    MODEL_OUT = '/home/elmar/Cigarette_buts_detection/cigarette_buts/cigarettNet/model/out'
    assert np.mod(np.log2(start_filter), 1) == 0.0
    assert np.mod(np.log2(middle_filter), 1) == 0.0
    assert middle_filter >= start_filter

    depth = np.int(np.log2(middle_filter) - np.log2(start_filter))
    
    (tr_ds, mask_ds),(val_ds, val_mask_ds) = split_train_val(data)
    
    
    X_train = tf.cast(np.array(tr_ds), tf.float32)
    y_train = tf.cast(mask_ds, tf.float32)
    X_val = tf.cast(np.array(val_ds), tf.float32)
    y_val = tf.cast(val_mask_ds, tf.float32)


    input_channels = X_train.shape[-1]
    output_channels = y_train.shape[-1]

    weight_dir = MODEL_OUT
    log_dir = MODEL_OUT
    weight_path = str(os.path.join(weight_dir, 'weight.h5'))

    inputs = Input(shape=(128, 128, 3), name="inputs")
    temp = Lambda(lambda x: x / config.IMG_MAX_VAL, name="normalize")(inputs)

    k_init = tf.initializers.glorot_normal(seed=313)
    b_init = tf.initializers.glorot_normal(seed=313)
    dconvs = list()

    x, dconvs = downsample(temp, dconvs, depth)
    x = bottleneck(x, depth)
    x = Upsample(x, dconvs, depth)


    outputs = Conv2D(output_channels,(1, 1), padding="same", activation="sigmoid", kernel_initializer=k_init,
                     bias_initializer=b_init, name="outputs")(x)


    model = Model(inputs=inputs, outputs=outputs)


    auc = AUC(num_thresholds=num_thresholds, curve='PR', name='auc')
    model.compile(loss="binary_crossentropy", optimizer="adam", metrics=[auc])

    tensorboard_callback = TensorBoard(log_dir, histogram_freq=1)
    early_stopping = EarlyStopping(patience=10, verbose=1)
    reduce_lr = ReduceLROnPlateau(factor=0.1, patience=5, min_lr=0.00001, verbose=1)
    model_checkpoint = ModelCheckpoint(weight_path, monitor='val_loss', mode='min', save_best_only=True, verbose=1,
                                       save_weights_only=True)
    custom_metrics = CustomMetrics(model, X_val, y_val, num_thresholds)

    history = model.fit(X_train, y_train,
                        validation_data=[X_val, y_val],
                        epochs=epochs,
                        batch_size=batch_size,
                        callbacks=[model_checkpoint, reduce_lr, custom_metrics, tensorboard_callback],
                        shuffle=False,
                        verbose=1)

    model.load_weights(weight_path)


    return history, model
//  [python]
def _preprocessing_fn_with_control_dependency(inputs):
  with tf.init_scope():
    initializer = tf.lookup.KeyValueTensorInitializer(['foo', 'bar'], [0, 1])
    table = tf.lookup.StaticHashTable(initializer, default_value=-1)
  # The table created here will add an automatic control dependency.
  s_int = table.lookup(inputs['s']) + 1

  # Perform some TF Ops to ensure x is part of the graph of dependencies for the
  # outputs.
  x_abs = tf.math.abs(inputs['x'])
  y_centered = (
      tf.sparse.add(
          tf.to_float(inputs['y']), -analyzers.mean(inputs['y'])))
  return {'s_int': s_int, 'x_abs': x_abs, 'y_centered': y_centered}
//  [python]
def _preprocessing_fn_with_control_dependency(inputs):
  with tf.init_scope():
    initializer = tf.lookup.KeyValueTensorInitializer(['foo', 'bar'], [0, 1])
    table = tf.lookup.StaticHashTable(initializer, default_value=-1)
  # The table created here will add an automatic control dependency.
  s_int = table.lookup(inputs['s']) + 1

  # Perform some TF Ops to ensure x is part of the graph of dependencies for the
  # outputs.
  x_abs = tf.math.abs(inputs['x'])
  y_centered = (
      tf.sparse.add(
          tf.cast(inputs['y'], tf.float32), -analyzers.mean(inputs['y'])))
  return {'s_int': s_int, 'x_abs': x_abs, 'y_centered': y_centered}
//  [python]
def call(
      self,
      inputs: Tuple[tfgnn.GraphTensor, Mapping[str, tf.RaggedTensor]],
  ) -> tfgnn.GraphTensor:
    graph_tensor, readout_features = inputs

    seed_node_set = graph_tensor.node_sets[self._seed_node_set]

    ids = seed_node_set[core.NODE_ID_NAME]
    link_source_idx = t.cast(
        ext_ops.ragged_lookup(
            readout_features[tfgnn.SOURCE_NAME], ids, global_indices=False
        ),
        graph_tensor.indices_dtype,
    )
    link_target_idx = t.cast(
        ext_ops.ragged_lookup(
            readout_features[tfgnn.TARGET_NAME], ids, global_indices=False
        ),
        graph_tensor.indices_dtype,
    )

    with tf.control_dependencies(
        [
            tf.debugging.assert_equal(
                link_source_idx.row_lengths(),
                link_target_idx.row_lengths(),
                message=(
                    'The number of link source and target nodes must be the'
                    ' same'
                ),
            )
        ]
    ):
      num_seeds = t.cast(
          link_source_idx.row_lengths(), graph_tensor.indices_dtype
      )
      readout_index = tf.ragged.range(num_seeds, dtype=num_seeds.dtype)
      sizes = tf.expand_dims(num_seeds, axis=-1)

    readout_node_sets = {
        self._readout_node_set: tfgnn.NodeSet.from_fields(
            features=readout_features, sizes=sizes
        )
    }
    readout_edge_sets = {
        f'{self._readout_node_set}/source': tfgnn.EdgeSet.from_fields(
            sizes=sizes,
            adjacency=tfgnn.Adjacency.from_indices(
                (self._seed_node_set, link_source_idx),
                (self._readout_node_set, readout_index),
            ),
        ),
        f'{self._readout_node_set}/target': tfgnn.EdgeSet.from_fields(
            sizes=sizes,
            adjacency=tfgnn.Adjacency.from_indices(
                (self._seed_node_set, link_target_idx),
                (self._readout_node_set, readout_index),
            ),
        ),
    }
    return tfgnn.GraphTensor.from_pieces(
        node_sets={**graph_tensor.node_sets, **readout_node_sets},
        edge_sets={**graph_tensor.edge_sets, **readout_edge_sets},
        context=graph_tensor.context,
    )
//  [python]
def call(
      self,
      inputs: Tuple[tfgnn.GraphTensor, Mapping[str, tf.RaggedTensor]],
  ) -> tfgnn.GraphTensor:
    graph_tensor, readout_features = inputs

    seed_node_set = graph_tensor.node_sets[self._seed_node_set]

    ids = seed_node_set[core.NODE_ID_NAME]
    link_source_idx = tf.cast(
        ext_ops.ragged_lookup(
            readout_features[tfgnn.SOURCE_NAME], ids, global_indices=False
        ),
        graph_tensor.indices_dtype,
    )
    link_target_idx = tf.cast(
        ext_ops.ragged_lookup(
            readout_features[tfgnn.TARGET_NAME], ids, global_indices=False
        ),
        graph_tensor.indices_dtype,
    )

    with tf.control_dependencies(
        [
            tf.debugging.assert_equal(
                link_source_idx.row_lengths(),
                link_target_idx.row_lengths(),
                message=(
                    'The number of link source and target nodes must be the'
                    ' same'
                ),
            )
        ]
    ):
      num_seeds = tf.cast(
          link_source_idx.row_lengths(), graph_tensor.indices_dtype
      )
      readout_index = tf.ragged.range(num_seeds, dtype=num_seeds.dtype)
      sizes = tf.expand_dims(num_seeds, axis=-1)

    readout_node_sets = {
        self._readout_node_set: tfgnn.NodeSet.from_fields(
            features=readout_features, sizes=sizes
        )
    }
    readout_edge_sets = {
        f'{self._readout_node_set}/source': tfgnn.EdgeSet.from_fields(
            sizes=sizes,
            adjacency=tfgnn.Adjacency.from_indices(
                (self._seed_node_set, link_source_idx),
                (self._readout_node_set, readout_index),
            ),
        ),
        f'{self._readout_node_set}/target': tfgnn.EdgeSet.from_fields(
            sizes=sizes,
            adjacency=tfgnn.Adjacency.from_indices(
                (self._seed_node_set, link_target_idx),
                (self._readout_node_set, readout_index),
            ),
        ),
    }
    return tfgnn.GraphTensor.from_pieces(
        node_sets={**graph_tensor.node_sets, **readout_node_sets},
        edge_sets={**graph_tensor.edge_sets, **readout_edge_sets},
        context=graph_tensor.context,
    )
//  [python]
def _predict_op(self, logits, **kwargs):
        group_idxs, group_probas = GroupRelationLabeler._predict_op(
            self, logits["group_logits"]
        )
        ner_idxs, ner_probas = SequenceLabeler._predict_op(
            self, logits["ner_logits"], **kwargs
        )

        # Stack to feed predictions as a single tensor
        # Will be unpacked in the target encoder
        # [batch_size, 2, seq_len]
        group_idxs = t.cast(group_idxs, t.int32)
        idxs = tf.stack([group_idxs, ner_idxs], axis=1)
        probas = ner_probas

        return idxs, probas
//  [python]
def _predict_op(self, logits, **kwargs):
        group_idxs, group_probas = GroupRelationLabeler._predict_op(
            self, logits["group_logits"]
        )
        ner_idxs, ner_probas = SequenceLabeler._predict_op(
            self, logits["ner_logits"], **kwargs
        )

        # Stack to feed predictions as a single tensor
        # Will be unpacked in the target encoder
        # [batch_size, 2, seq_len]
        group_idxs = tf.cast(group_idxs, tf.int32)
        idxs = tf.stack([group_idxs, ner_idxs], axis=1)
        probas = ner_probas

        return idxs, probas
//  [python]
def _predict_op(self, logits, **kwargs):
        group_idxs, group_probas = BROSLabeler._predict_op(
            self, logits, **kwargs
        )
        ner_idxs, ner_probas = SequenceLabeler._predict_op(
            self, logits["ner_logits"], **kwargs
        )

        # Produces [batch_size, 3, seq_len]
        # Required to pass predictions through Finetune
        ner_idxs = tf.cast(ner_idxs, tf.int32)
        group_idxs = tf.cast(group_idxs, tf.int32)
        start_token_idxs, next_token_idxs = group_idxs[:, 0, :], group_idxs[:, 1, :]
        idxs = tf.stack([ner_idxs, start_token_idxs, next_token_idxs], axis=1)
        probas = ner_probas

        return idxs, probas
//  [python]
def disentangled_att_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):
        if relative_pos is None:
            q = shape_list(query_layer)[-2]
            relative_pos = build_relative_position(q, shape_list(key_layer)[-2])
        shape_list_pos = shape_list(relative_pos)
        if len(shape_list_pos) == 2:
            relative_pos = tf.expand_dims(tf.expand_dims(relative_pos, 0), 0)
        elif len(shape_list_pos) == 3:
            relative_pos = tf.expand_dims(relative_pos, 1)
        # bxhxqxk
        elif len(shape_list_pos) != 4:
            raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {len(shape_list_pos)}")

        att_span = tf.minimum(
                tf.maximum(shape_list(query_layer)[-2], shape_list(key_layer)[-2]), self.max_relative_positions
            )
        att_span = tf.cast(att_span, tf.int64)

        rel_embeddings = tf.expand_dims(
            rel_embeddings[self.max_relative_positions - att_span : self.max_relative_positions + att_span, :], 0
        )

        score = 0

        # content->position
        if "c2p" in self.pos_att_type:
            pos_key_layer = self.pos_proj(rel_embeddings)
            pos_key_layer = self.transpose_for_scores(pos_key_layer)
            c2p_att = tf.matmul(query_layer, tf.transpose(pos_key_layer, [0, 1, 3, 2]))
            c2p_pos = tf.clip_by_value(relative_pos + att_span, 0, att_span * 2 - 1)
            c2p_att = torch_gather(c2p_att, c2p_dynamic_expand(c2p_pos, query_layer, relative_pos), -1)
            score += c2p_att

        # position->content
        if "p2c" in self.pos_att_type:
            pos_query_layer = self.pos_q_proj(rel_embeddings)
            pos_query_layer = self.transpose_for_scores(pos_query_layer)
            pos_query_layer /= tf.math.sqrt(tf.cast(shape_list(pos_query_layer)[-1] * scale_factor, dtype=tf.float32))
            if shape_list(query_layer)[-2] != shape_list(key_layer)[-2]:
                r_pos = build_relative_position(shape_list(key_layer)[-2], shape_list(key_layer)[-2])
            else:
                r_pos = relative_pos
            p2c_pos = tf.clip_by_value(-r_pos + att_span, 0, att_span * 2 - 1)
            p2c_att = tf.matmul(key_layer, tf.transpose(pos_query_layer, [0, 1, 3, 2]))
            p2c_att = tf.transpose(
                torch_gather(p2c_att, p2c_dynamic_expand(p2c_pos, query_layer, key_layer), -1), [0, 1, 3, 2]
            )
            if shape_list(query_layer)[-2] != shape_list(key_layer)[-2]:
                pos_index = tf.expand_dims(relative_pos[:, :, :, 0], -1)
                p2c_att = torch_gather(p2c_att, pos_dynamic_expand(pos_index, p2c_att, key_layer), -2)
            score += p2c_att

        return score
//  [python]
def disentangled_att_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):
        if relative_pos is None:
            q = shape_list(query_layer)[-2]
            relative_pos = build_relative_position(q, shape_list(key_layer)[-2])
        shape_list_pos = shape_list(relative_pos)
        if len(shape_list_pos) == 2:
            relative_pos = tf.expand_dims(tf.expand_dims(relative_pos, 0), 0)
        elif len(shape_list_pos) == 3:
            relative_pos = tf.expand_dims(relative_pos, 1)
        # bxhxqxk
        elif len(shape_list_pos) != 4:
            raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {len(shape_list_pos)}")

        att_span = tf.cast(
            tf.minimum(
                tf.maximum(shape_list(query_layer)[-2], shape_list(key_layer)[-2]), self.max_relative_positions
            ),
            tf.int64,
        )
        rel_embeddings = tf.expand_dims(
            rel_embeddings[self.max_relative_positions - att_span : self.max_relative_positions + att_span, :], 0
        )

        score = 0

        # content->position
        if "c2p" in self.pos_att_type:
            pos_key_layer = self.pos_proj(rel_embeddings)
            pos_key_layer = self.transpose_for_scores(pos_key_layer)
            c2p_att = tf.matmul(query_layer, tf.transpose(pos_key_layer, [0, 1, 3, 2]))
            c2p_pos = tf.clip_by_value(relative_pos + att_span, 0, att_span * 2 - 1)
            c2p_att = torch_gather(c2p_att, c2p_dynamic_expand(c2p_pos, query_layer, relative_pos), -1)
            score += c2p_att

        # position->content
        if "p2c" in self.pos_att_type:
            pos_query_layer = self.pos_q_proj(rel_embeddings)
            pos_query_layer = self.transpose_for_scores(pos_query_layer)
            pos_query_layer /= tf.math.sqrt(tf.cast(shape_list(pos_query_layer)[-1] * scale_factor, dtype=tf.float32))
            if shape_list(query_layer)[-2] != shape_list(key_layer)[-2]:
                r_pos = build_relative_position(shape_list(key_layer)[-2], shape_list(key_layer)[-2])
            else:
                r_pos = relative_pos
            p2c_pos = tf.clip_by_value(-r_pos + att_span, 0, att_span * 2 - 1)
            p2c_att = tf.matmul(key_layer, tf.transpose(pos_query_layer, [0, 1, 3, 2]))
            p2c_att = tf.transpose(
                torch_gather(p2c_att, p2c_dynamic_expand(p2c_pos, query_layer, key_layer), -1), [0, 1, 3, 2]
            )
            if shape_list(query_layer)[-2] != shape_list(key_layer)[-2]:
                pos_index = tf.expand_dims(relative_pos[:, :, :, 0], -1)
                p2c_att = torch_gather(p2c_att, pos_dynamic_expand(pos_index, p2c_att, key_layer), -2)
            score += p2c_att

        return score
//  [python]
def process_data():   
    current_dir = os.getcwd()
    # parent = os.path.dirname(current_dir)
    pokemon_dir = os.path.join(current_dir, 'data')
    images = []
    for each in os.listdir(pokemon_dir):
        images.append(os.path.join(pokemon_dir,each))
    # print images    
    all_images = tf.convert_to_tensor(images, dtype = tf.string)
    
    images_queue = tf.train.slice_input_producer(
                                        [all_images])
                                        
    content = tf.read_file(images_queue[0])
    image = tf.image.decode_jpeg(content, channels = CHANNEL)
    # sess1 = tf.Session()
    # print sess1.run(image)
    image = tf.image.random_flip_left_right(image)
    image = tf.image.random_brightness(image, max_delta = 0.1)
    image = tf.image.random_contrast(image, lower = 0.9, upper = 1.1)
    # noise = tf.Variable(tf.truncated_normal(shape = [HEIGHT,WIDTH,CHANNEL], dtype = tf.float32, stddev = 1e-3, name = 'noise')) 
    # print image.get_shape()
    size = [HEIGHT, WIDTH]
    image = tf.image.resize_images(image, size)
    image.set_shape([HEIGHT,WIDTH,CHANNEL])
    # image = image + noise
    # image = tf.transpose(image, perm=[2, 0, 1])
    # print image.get_shape()
    
    image = tf.cast(image, tf.float32)
    image = image / 255.0
    
    iamges_batch = tf.train.shuffle_batch(
                                    [image], batch_size = BATCH_SIZE,
                                    num_threads = 4, capacity = 200 + 3* BATCH_SIZE,
                                    min_after_dequeue = 200)
    num_images = len(images)

    return iamges_batch, num_images
//  [python]
def _top_k_tactics(tactic_logits: tf.Tensor,  # [batch, tactics]
                       tactic_mask: tf.Tensor,  # [batch, tactics]
                       tactic_expand_bound: int
                       ) -> Tuple[tf.Tensor, tf.Tensor]:  #([batch, top_k_tactics], [batch, top_k_tactics])
        tactic_logits = tf.math.log_softmax(tactic_logits + tf.math.log(tf.to_float(tactic_mask)), axis=-1)

        # Sometimes the number of tactics is less than the tactic_expand_bound
        # In that case return only the number of tactics.  It is ok if some have probability 0.
        top_k = tf.math.top_k(tactic_logits, k=tf.minimum(tactic_expand_bound, tf.shape(tactic_logits)[1]))
        # ([batch, top_k_tactics], [batch, top_k_tactics])
        return top_k.indices, top_k.values
//  [python]
def _top_k_tactics(tactic_logits: tf.Tensor,  # [batch, tactics]
                       tactic_mask: tf.Tensor,  # [batch, tactics]
                       tactic_expand_bound: int
                       ) -> Tuple[tf.Tensor, tf.Tensor]:  #([batch, top_k_tactics], [batch, top_k_tactics])
        tactic_logits = tf.math.log_softmax(tactic_logits + tf.math.log(tf.cast(tactic_mask, tf.float32)), axis=-1)

        # Sometimes the number of tactics is less than the tactic_expand_bound
        # In that case return only the number of tactics.  It is ok if some have probability 0.
        top_k = tf.math.top_k(tactic_logits, k=tf.minimum(tactic_expand_bound, tf.shape(tactic_logits)[1]))
        # ([batch, top_k_tactics], [batch, top_k_tactics])
        return top_k.indices, top_k.values
//  [python]
def _global_arguments_logits_mask(scalar_proofstate_graph: tfgnn.GraphTensor,
                                      global_context_size: int) -> tf.Tensor:
        """
        @param scalar_proofstate_graph: the proofstate graph containing ids for the available global context definitions
        @param global_context_size: the size of the full global context
        @return: a mask for logits of the global context, taking into account the definitions that are actually available
        """
        global_context_ids = scalar_proofstate_graph.context['global_context_ids']

        indices = tf.stack([
            tf.to_int64(global_context_ids.value_rowids()),
            global_context_ids.values
        ], axis = -1)
        updates = tf.ones_like(global_context_ids.values, dtype=tf.float32)
        shape = [global_context_ids.nrows(), global_context_size]

        return tf.math.log(tf.scatter_nd(indices, updates, shape))
//  [python]
def _global_arguments_logits_mask(scalar_proofstate_graph: tfgnn.GraphTensor,
                                      global_context_size: int) -> tf.Tensor:
        """
        @param scalar_proofstate_graph: the proofstate graph containing ids for the available global context definitions
        @param global_context_size: the size of the full global context
        @return: a mask for logits of the global context, taking into account the definitions that are actually available
        """
        global_context_ids = scalar_proofstate_graph.context['global_context_ids']

        indices = tf.stack([
            tf.cast(global_context_ids.value_rowids(), tf.int64),
            global_context_ids.values
        ], axis = -1)
        updates = tf.ones_like(global_context_ids.values, dtype=tf.float32)
        shape = [global_context_ids.nrows(), global_context_size]

        return tf.math.log(tf.scatter_nd(indices, updates, shape))
//  [python]
def compute_metrics(self, x, y, y_pred, sample_weight):
        metric_results = super().compute_metrics(x, y, y_pred, sample_weight)

        tactic = y[LocalArgumentPrediction.TACTIC_LOGITS]
        tactic_logits = y_pred[LocalArgumentPrediction.TACTIC_LOGITS]
        tactic_accuracy = tf.keras.metrics.sparse_categorical_accuracy(tactic, tactic_logits)

        local_arguments = y[LocalArgumentPrediction.LOCAL_ARGUMENTS_LOGITS]
        local_arguments_logits = y_pred[LocalArgumentPrediction.LOCAL_ARGUMENTS_LOGITS]
        local_arguments_true = tf.squeeze(local_arguments.to_tensor(default_value=0), axis=1)
        local_arguments_pred = _local_arguments_pred(local_arguments_true, local_arguments_logits)
        arguments_seq_accuracy = t.reduce_all(local_arguments_true == local_arguments_pred, axis=-1)
        arguments_seq_mask = t.reduce_min(local_arguments_true, axis=-1) > -1

        self.arguments_seq_accuracy.update_state(t.cast(arguments_seq_mask, dtype=tf.float32) * t.cast(arguments_seq_accuracy, dtype=tf.float32),
                                                 sample_weight=sample_weight)
        self.strict_accuracy.update_state(t.cast(arguments_seq_mask, dtype=tf.float32) * t.cast(arguments_seq_accuracy, dtype=tf.float32) * t.cast(tactic_accuracy, dtype=tf.float32),
                                          sample_weight=sample_weight)

        metric_results[LocalArgumentPrediction.ARGUMENTS_SEQ_ACCURACY] = self.arguments_seq_accuracy.result()
        metric_results[LocalArgumentPrediction.STRICT_ACCURACY] = self.strict_accuracy.result()
        return metric_results
//  [python]
def compute_metrics(self, x, y, y_pred, sample_weight):
        metric_results = super().compute_metrics(x, y, y_pred, sample_weight)

        tactic = y[LocalArgumentPrediction.TACTIC_LOGITS]
        tactic_logits = y_pred[LocalArgumentPrediction.TACTIC_LOGITS]
        tactic_accuracy = tf.keras.metrics.sparse_categorical_accuracy(tactic, tactic_logits)

        local_arguments = y[LocalArgumentPrediction.LOCAL_ARGUMENTS_LOGITS]
        local_arguments_logits = y_pred[LocalArgumentPrediction.LOCAL_ARGUMENTS_LOGITS]
        local_arguments_true = tf.squeeze(local_arguments.to_tensor(default_value=0), axis=1)
        local_arguments_pred = _local_arguments_pred(local_arguments_true, local_arguments_logits)
        arguments_seq_accuracy = tf.cast(tf.reduce_all(local_arguments_true == local_arguments_pred, axis=-1), dtype=tf.float32)
        arguments_seq_mask = tf.cast(tf.reduce_min(local_arguments_true, axis=-1) > -1, dtype=tf.float32)

        self.arguments_seq_accuracy.update_state(arguments_seq_mask * arguments_seq_accuracy,
                                                 sample_weight=sample_weight)
        self.strict_accuracy.update_state(arguments_seq_mask * arguments_seq_accuracy * tactic_accuracy,
                                          sample_weight=sample_weight)

        metric_results[LocalArgumentPrediction.ARGUMENTS_SEQ_ACCURACY] = self.arguments_seq_accuracy.result()
        metric_results[LocalArgumentPrediction.STRICT_ACCURACY] = self.strict_accuracy.result()
        return metric_results
//  [python]
def _local_context_hidden(
        scalar_proofstate_graph: tfgnn.GraphTensor,
        hidden_graph: tfgnn.GraphTensor,
    ) -> tf.RaggedTensor:  # [batch_size, None(local_context), hidden_size]
        """
        Computes logits for local arguments from the hidden states and the local context node ids.
        """
        # the sizes of the components of this graph
        # [ batch_size, ]
        sizes = t.cast(scalar_proofstate_graph.node_sets['node'].sizes, dtype=tf.int64)

        # the offsets of the node ids for each graph component
        # [ batch_size, 1 ]
        cumulative_sizes = tf.expand_dims(tf.cumsum(sizes, exclusive=True), axis=-1)

        # the node ids for the local context nodes, shifted per components
        # [ batch_size, None(num_context_nodes) ]
        local_context_ids = cumulative_sizes + scalar_proofstate_graph.context['local_context_ids']

        # the hidden states for the nodes in the local context
        # [ batch_size, None(num_context_nodes), hidden_size ]
        return tf.gather(hidden_graph.node_sets['node']['hidden_state'],
                                               local_context_ids).with_row_splits_dtype(tf.int64)
//  [python]
def _local_context_hidden(
        scalar_proofstate_graph: tfgnn.GraphTensor,
        hidden_graph: tfgnn.GraphTensor,
    ) -> tf.RaggedTensor:  # [batch_size, None(local_context), hidden_size]
        """
        Computes logits for local arguments from the hidden states and the local context node ids.
        """
        # the sizes of the components of this graph
        # [ batch_size, ]
        sizes = tf.cast(scalar_proofstate_graph.node_sets['node'].sizes, dtype=tf.int64)

        # the offsets of the node ids for each graph component
        # [ batch_size, 1 ]
        cumulative_sizes = tf.expand_dims(tf.cumsum(sizes, exclusive=True), axis=-1)

        # the node ids for the local context nodes, shifted per components
        # [ batch_size, None(num_context_nodes) ]
        local_context_ids = cumulative_sizes + scalar_proofstate_graph.context['local_context_ids']

        # the hidden states for the nodes in the local context
        # [ batch_size, None(num_context_nodes), hidden_size ]
        return tf.gather(hidden_graph.node_sets['node']['hidden_state'],
                                               local_context_ids).with_row_splits_dtype(tf.int64)
//  [python]
def _mask_defined_embeddings(scalar_definition_graph: tfgnn.GraphTensor, embedded_graph: tfgnn.GraphTensor):
        num_definitions = t.cast(scalar_definition_graph.context['num_definitions'], dtype=t.int32)
        is_defined = tf.ragged.range(scalar_definition_graph.node_sets['node'].sizes) < tf.expand_dims(num_definitions, axis=-1)
        mask = tf.expand_dims(1 - t.cast(is_defined.flat_values, dtype=t.float32), axis=-1)
        masked_hidden_state = embedded_graph.node_sets['node']['hidden_state'] * mask
        return embedded_graph.replace_features(node_sets={'node': {'hidden_state': masked_hidden_state}})
//  [python]
def _mask_defined_embeddings(scalar_definition_graph: tfgnn.GraphTensor, embedded_graph: tfgnn.GraphTensor):
        num_definitions = tf.cast(scalar_definition_graph.context['num_definitions'], dtype=tf.int32)
        is_defined = tf.ragged.range(scalar_definition_graph.node_sets['node'].sizes) < tf.expand_dims(num_definitions, axis=-1)
        mask = tf.expand_dims(1 - tf.cast(is_defined.flat_values, dtype=tf.float32), axis=-1)
        masked_hidden_state = embedded_graph.node_sets['node']['hidden_state'] * mask
        return embedded_graph.replace_features(node_sets={'node': {'hidden_state': masked_hidden_state}})
//  [python]
def main(args):
  setproctitle.setproctitle('DeepUPE')

  inputs = get_input_list(args.input)

  # -------- Load params ----------------------------------------------------
  config = tf.ConfigProto()
  config.gpu_options.allow_growth = True
  with tf.Session(config=config) as sess:
    checkpoint_path = tf.train.latest_checkpoint(args.checkpoint_dir)
    if checkpoint_path is None:
      log.error('Could not find a checkpoint in {}'.format(args.checkpoint_dir))
      return

    metapath = ".".join([checkpoint_path, "meta"])
    log.info('Loading graph from {}'.format(metapath))
    tf.train.import_meta_graph(metapath)

    model_params = utils.get_model_params(sess)

  # -------- Setup graph ----------------------------------------------------
  if not hasattr(models, model_params['model_name']):
    log.error("Model {} does not exist".format(params.model_name))
    return
  mdl = getattr(models, model_params['model_name'])

  tf.reset_default_graph()
  net_shape = model_params['net_input_size']
  t_fullres_input = tf.placeholder(tf.float32, (1, None, None, 3))
  t_lowres_input = tf.placeholder(tf.float32, (1, net_shape, net_shape, 3))

  with tf.variable_scope('inference'):
    prediction = mdl.inference(
        t_lowres_input, t_fullres_input, model_params, is_training=False)
  output = tf.to_int32(255.0*tf.squeeze(tf.clip_by_value(prediction, 0, 1)))
  saver = tf.train.Saver()


  with tf.Session(config=config) as sess:
    log.info('Restoring weights from {}'.format(checkpoint_path))
    saver.restore(sess, checkpoint_path)

    for idx, input_path in enumerate(inputs):

      log.info("Processing {}".format(input_path))
      im_input = cv2.imread(input_path, -1)  # -1 means read as is, no conversions.
      if im_input.shape[2] == 4:
        log.info("Input {} has 4 channels, dropping alpha".format(input_path))
        im_input = im_input[:, :, :3]

      im_input = np.flip(im_input, 2)  # OpenCV reads BGR, convert back to RGB.



      im_input = skimage.img_as_float(im_input)

      lowres_input = skimage.transform.resize(
          im_input, [net_shape, net_shape], order = 0)


      fname = os.path.splitext(os.path.basename(input_path))[0]
      output_path = os.path.join(args.output, fname+".png")
      basedir = os.path.dirname(output_path)

      im_input = im_input[np.newaxis, :, :, :]
      lowres_input = lowres_input[np.newaxis, :, :, :]

      feed_dict = {
          t_fullres_input: im_input,
          t_lowres_input: lowres_input
      }

      out_ =  sess.run(output, feed_dict=feed_dict)

      if not os.path.exists(basedir):
        os.makedirs(basedir)
      

      out_ =  (out_/255.0)**1.0
      out_ = np.clip(out_, im_input, 1)
      
      final_result = np.divide(im_input.astype(np.float), out_.astype(np.float))
     
      skimage.io.imsave(output_path, final_result)
//  [python]
def main(args):
  setproctitle.setproctitle('DeepUPE')

  inputs = get_input_list(args.input)

  # -------- Load params ----------------------------------------------------
  config = tf.ConfigProto()
  config.gpu_options.allow_growth = True
  with tf.Session(config=config) as sess:
    checkpoint_path = tf.train.latest_checkpoint(args.checkpoint_dir)
    if checkpoint_path is None:
      log.error('Could not find a checkpoint in {}'.format(args.checkpoint_dir))
      return

    metapath = ".".join([checkpoint_path, "meta"])
    log.info('Loading graph from {}'.format(metapath))
    tf.train.import_meta_graph(metapath)

    model_params = utils.get_model_params(sess)

  # -------- Setup graph ----------------------------------------------------
  if not hasattr(models, model_params['model_name']):
    log.error("Model {} does not exist".format(params.model_name))
    return
  mdl = getattr(models, model_params['model_name'])

  tf.reset_default_graph()
  net_shape = model_params['net_input_size']
  t_fullres_input = tf.placeholder(tf.float32, (1, None, None, 3))
  t_lowres_input = tf.placeholder(tf.float32, (1, net_shape, net_shape, 3))

  with tf.variable_scope('inference'):
    prediction = mdl.inference(
        t_lowres_input, t_fullres_input, model_params, is_training=False)
  output = tf.cast(255.0*tf.squeeze(tf.clip_by_value(prediction, 0, 1)), tf.uint8)
  saver = tf.train.Saver()


  with tf.Session(config=config) as sess:
    log.info('Restoring weights from {}'.format(checkpoint_path))
    saver.restore(sess, checkpoint_path)

    for idx, input_path in enumerate(inputs):

      log.info("Processing {}".format(input_path))
      im_input = cv2.imread(input_path, -1)  # -1 means read as is, no conversions.
      if im_input.shape[2] == 4:
        log.info("Input {} has 4 channels, dropping alpha".format(input_path))
        im_input = im_input[:, :, :3]

      im_input = np.flip(im_input, 2)  # OpenCV reads BGR, convert back to RGB.



      im_input = skimage.img_as_float(im_input)

      lowres_input = skimage.transform.resize(
          im_input, [net_shape, net_shape], order = 0)


      fname = os.path.splitext(os.path.basename(input_path))[0]
      output_path = os.path.join(args.output, fname+".png")
      basedir = os.path.dirname(output_path)

      im_input = im_input[np.newaxis, :, :, :]
      lowres_input = lowres_input[np.newaxis, :, :, :]

      feed_dict = {
          t_fullres_input: im_input,
          t_lowres_input: lowres_input
      }

      out_ =  sess.run(output, feed_dict=feed_dict)

      if not os.path.exists(basedir):
        os.makedirs(basedir)
      

      out_ =  (out_/255.0)**1.0
      out_ = np.clip(out_, im_input, 1)
      
      final_result = np.divide(im_input.astype(np.float), out_.astype(np.float))
     
      skimage.io.imsave(output_path, final_result)
//  [python]
def call(self, inputs, y_true, y_pred, training=None):
        """
        Get the softmax cross entropy loss

        Parameters
        ----------
        inputs: dict of dict of tensors
            Dictionary of input feature tensors
        y_true: tensor
            True labels
        y_pred: tensor
            Predicted scores
        training: boolean
            Boolean indicating whether the layer is being used in training mode

        Returns
        -------
        tensor
            Scalar softmax cross entropy loss tensor

        Notes
        -----
            Uses `mask` field to exclude padded records from contributing
            to the loss
        """
        mask = tf.cast(inputs[FeatureTypeKey.MASK], y_pred.dtype)
        y_true = tf.cast(y_true, y_pred.dtype)

        return self.loss_fn(y_true=tf.math.multiply(y_true, mask),
                            y_pred=tf.math.multiply(y_pred, mask))
//  [python]
def call(self, inputs, y_true, y_pred, training=None):
        """
        Get the softmax cross entropy loss

        Parameters
        ----------
        inputs: dict of dict of tensors
            Dictionary of input feature tensors
        y_true: tensor
            True labels
        y_pred: tensor
            Predicted scores
        training: boolean
            Boolean indicating whether the layer is being used in training mode

        Returns
        -------
        tensor
            Scalar basic cross entropy loss tensor

        Notes
        -----
        - Uses `mask` field to exclude padded records from contributing to the loss
        """
        mask = tf.cast(inputs[FeatureTypeKey.MASK], y_pred.dtype)
        y_true = tf.cast(y_true, y_pred.dtype)

        # Convert y_true to a probability distribution
        y_true_softmax = self.final_activation_op({
            FeatureTypeKey.METADATA: {
                FeatureTypeKey.MASK: mask
            },
            FeatureTypeKey.LOGITS: y_true
        }, training=training)

        return self.loss_fn(y_true=tf.math.multiply(y_true_softmax, mask),
                            y_pred=tf.math.multiply(y_pred, mask))
//  [python]
def iou_coe(output, target, threshold=0.5, epsilon=1e-10):
    """Non-differentiable Intersection over Union, usually be used for evaluating binary image segmentation.
    The coefficient = [0, 1], 1 means totally match.

    Parameters
    -----------
    output : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    target : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    threshold : float
        The threshold value to be true.
    epsilon : float
        A small value to avoid zero denominator when both output and target output nothing.

    Examples
    ---------
    >>> outputs = tl.act.pixel_wise_softmax(network.outputs)
    >>> iou = tl.cost.iou_coe(outputs[:,:,:,0], y_[:,:,:,0])

    Notes
    ------
    - IOU cannot be used as training loss, people usually use dice coefficient for training, and IOU for evaluating.
    """
    pre = t.cast(output > threshold, dtype=t.float32)
    truth = t.cast(target > threshold, dtype=t.float32)
    intersection = tf.reduce_sum(pre * truth)
    union = tf.reduce_sum(t.cast((pre + truth) > threshold, dtype=t.float32))
    return tf.reduce_sum(intersection) / (tf.reduce_sum(union) + epsilon)
//  [python]
def iou_coe(output, target, threshold=0.5, epsilon=1e-10):
    """Non-differentiable Intersection over Union, usually be used for evaluating binary image segmentation.
    The coefficient = [0, 1], 1 means totally match.

    Parameters
    -----------
    output : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    target : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    threshold : float
        The threshold value to be true.
    epsilon : float
        A small value to avoid zero denominator when both output and target output nothing.

    Examples
    ---------
    >>> outputs = tl.act.pixel_wise_softmax(network.outputs)
    >>> iou = tl.cost.iou_coe(outputs[:,:,:,0], y_[:,:,:,0])

    Notes
    ------
    - IOU cannot be used as training loss, people usually use dice coefficient for training, and IOU for evaluating.
    """
    pre = tf.cast(output > threshold, dtype=tf.float32)
    truth = tf.cast(target > threshold, dtype=tf.float32)
    intersection = tf.reduce_sum(pre * truth)
    union = tf.reduce_sum(tf.cast((pre + truth) > threshold, dtype=tf.float32))
    return tf.reduce_sum(intersection) / (tf.reduce_sum(union) + epsilon)
//  [python]
def call(self, inputs, y_true, y_pred, training=None):
        """
        Get the softmax cross entropy loss on the auxiliary label

        Parameters
        ----------
        inputs: dict of dict of tensors
            Dictionary of input feature tensors
        y_true: tensor
            True labels
        y_pred: tensor
            Predicted scores
        training: boolean
            Boolean indicating whether the layer is being used in training mode

        Returns
        -------
        tensor
            Scalar softmax cross entropy loss tensor

        Notes
        -----
        - Uses `mask` field to exclude padded records from contributing
        to the loss
        - Queries with ties in the highest scores would have multiple one's in the 1-hot vector.
        - Queries with all zeros for y_true would have all ones as their 1-hot vector.
        - A simple remedy is to scale down the loss by the number of ties per query.
        """
        mask = t.cast(inputs[FeatureTypeKey.MASK], y_pred.dtype)
        y_true = t.cast(y_true, y_pred.dtype)

        # Convert y_true to 1-hot labels
        y_true_one_hot = tf.equal(y_true, tf.expand_dims(tf.math.reduce_max(y_true, axis=1), axis=1))
        y_true_one_hot = t.cast(y_true_one_hot, dtype=y_pred.dtype)

        # Scale down the loss of a query by 1 / (number of ties)
        sample_weight = tf.math.divide(tf.constant(1, dtype=tf.float32), tf.reduce_sum(y_true_one_hot, axis=1))

        return self.loss_fn(y_true=tf.math.multiply(y_true_one_hot, mask),
                            y_pred=tf.math.multiply(y_pred, mask),
                            sample_weight=sample_weight)
//  [python]
def call(self, inputs, y_true, y_pred, training=None):
        """
        Get the softmax cross entropy loss on the auxiliary label

        Parameters
        ----------
        inputs: dict of dict of tensors
            Dictionary of input feature tensors
        y_true: tensor
            True labels
        y_pred: tensor
            Predicted scores
        training: boolean
            Boolean indicating whether the layer is being used in training mode

        Returns
        -------
        tensor
            Scalar softmax cross entropy loss tensor

        Notes
        -----
        - Uses `mask` field to exclude padded records from contributing
        to the loss
        - Queries with ties in the highest scores would have multiple one's in the 1-hot vector.
        - Queries with all zeros for y_true would have all ones as their 1-hot vector.
        - A simple remedy is to scale down the loss by the number of ties per query.
        """
        mask = tf.cast(inputs[FeatureTypeKey.MASK], y_pred.dtype)
        y_true = tf.cast(y_true, y_pred.dtype)

        # Convert y_true to 1-hot labels
        y_true_one_hot = tf.equal(y_true, tf.expand_dims(tf.math.reduce_max(y_true, axis=1), axis=1))
        y_true_one_hot = tf.cast(y_true_one_hot, dtype=y_pred.dtype)

        # Scale down the loss of a query by 1 / (number of ties)
        sample_weight = tf.math.divide(tf.constant(1, dtype=tf.float32), tf.reduce_sum(y_true_one_hot, axis=1))

        return self.loss_fn(y_true=tf.math.multiply(y_true_one_hot, mask),
                            y_pred=tf.math.multiply(y_pred, mask),
                            sample_weight=sample_weight)
//  [python]
def dice_hard_coe(output, target, epsilon=1e-10):
    """Non-differentiable Sørensen–Dice coefficient for comparing the similarity of two distributions,
    usually be used for binary image segmentation i.e. labels are binary.
    The coefficient = [0, 1], 1 if totally match.

    Parameters
    -----------
    output : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    target : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    epsilon : float
        An optional name to attach to this layer.

    Examples
    ---------
    >>> outputs = pixel_wise_softmax(network.outputs)
    >>> dice_loss = 1 - dice_coe(outputs, y_, epsilon=1e-5)

    References
    -----------
    - `wiki-dice <https://en.wikipedia.org/wiki/Sørensen–Dice_coefficient>`_
    """
    output = tf.to_float(tf.cast(output > 0.5, dtype=tf.bool))
    target = tf.to_float(tf.cast(target > 0.5, dtype=tf.bool))
    inse = tf.reduce_sum( output * target )
    l = tf.reduce_sum( output * output )
    r = tf.reduce_sum( target * target )
    dice = 2 * (inse) / (l + r)
    if epsilon == 0:
        return dice
    else:
        return tf.clip_by_value(dice, 0, 1.0-epsilon)
//  [python]
def dice_hard_coe(output, target, epsilon=1e-10):
    """Non-differentiable Sørensen–Dice coefficient for comparing the similarity of two distributions,
    usually be used for binary image segmentation i.e. labels are binary.
    The coefficient = [0, 1], 1 if totally match.

    Parameters
    -----------
    output : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    target : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    epsilon : float
        An optional name to attach to this layer.

    Examples
    ---------
    >>> outputs = pixel_wise_softmax(network.outputs)
    >>> dice_loss = 1 - dice_coe(outputs, y_, epsilon=1e-5)

    References
    -----------
    - `wiki-dice <https://en.wikipedia.org/wiki/Sørensen–Dice_coefficient>`_
    """
    output = tf.cast(output > 0.5, dtype=tf.float32)
    target = tf.cast(target > 0.5, dtype=tf.float32)
    inse = tf.reduce_sum( output * target )
    l = tf.reduce_sum( output * output )
    r = tf.reduce_sum( target * target )
    dice = 2 * (inse) / (l + r)
    if epsilon == 0:
        return dice
    else:
        return tf.clip_by_value(dice, 0, 1.0-epsilon)
//  [python]
def call(self, inputs, y_true, y_pred, training=None):
        """
        Define a masked rank 1 ListNet loss.
        This loss is useful for multi-label classification when we have multiple
        click labels per document. This is because the loss breaks down the comparison
        between y_pred and y_true into individual binary assessments.
        Ref -> https://www.microsoft.com/en-us/research/wp-content/uploads/2016/02/tr-2007-40.pdf

        Parameters
        ----------
        inputs: dict of dict of tensors
            Dictionary of input feature tensors
        y_true: tensor
            True labels
        y_pred: tensor
            Predicted scores
        training: boolean
            Boolean indicating whether the layer is being used in training mode

        Returns
        -------
        tensor
            Scalar sigmoid cross entropy loss tensor

        Notes
        -----
            Uses `mask` field to exclude padded records from contributing
            to the loss
        """
        mask = tf.to_float(inputs[FeatureTypeKey.MASK])
        y_true = tf.to_float(y_true)
        batch_size = tf.to_float(tf.shape(y_true)[0])

        # Mask the padded records
        y_true = tf.gather_nd(y_true, tf.where(tf.equal(mask, tf.constant(1.0))))
        y_pred = tf.gather_nd(y_pred, tf.where(tf.equal(mask, tf.constant(1.0))))

        # Reshape the tensors so that we sum the losses from each record
        y_true = tf.expand_dims(tf.squeeze(y_true), axis=-1)
        y_pred = tf.expand_dims(tf.squeeze(y_pred), axis=-1)

        # Scale the sum of losses down by number of queries in the batch
        return tf.math.divide(self.loss_fn(y_true, y_pred), batch_size)