Add model.

main.py

@@ -0,0 +1,49 @@
+import gradio as gr
+
+from utils import load_example, run_nn, load_model_, next_, prev_
+
+demo = gr.Blocks()
+import models
+
+with demo:
+    headline = gr.Markdown("## Raven resolver ")
+    markdown = gr.Markdown("Below we show all 9 images from raven matrix. "
+                           "Model gets 8 images and predicts the properties of last one. "
+                           "Based on this properties the answer image is render in the  right panel. <br />"
+                           "Note that angle rotation is only used as a noise. "
+                           "There are not rules applied to angle property, so angle rotation of final output do not need to be the same as in example. "
+                           "Additionally there are cases that other properties could be used as noise.")
+    with gr.Row():
+        with gr.Column():
+            with gr.Row():
+                text = gr.Textbox(models.START_IMAGE,
+                                  label="Write the example number from validation dataset (0, 14,000). You can also paste here matrix representation from generator.")
+            with gr.Row():
+                prev = gr.Button("Prev")
+                show = gr.Button("Show")
+                next = gr.Button("Next")
+                # button = gr.Button("Run")
+            with gr.Row():
+                image = gr.Image(value=load_example(models.START_IMAGE)[0], label="Raven matrix")
+                desc = gr.Markdown(value=load_example(models.START_IMAGE)[1])
+
+        with gr.Column():
+            with gr.Row():
+                output = gr.Image(label="Generated image", shape=(200, 200))
+    with gr.Row():
+        button = gr.Button("Run")
+
+    # text.change(load_example, inputs=text, outputs=[image, desc])
+    show.click(load_example, inputs=text, outputs=[image, desc])
+    # button.click(run_nn, inputs=image, outputs=output)
+    button.click(run_nn, inputs=text, outputs=output)
+
+    # next.click(next_, inputs=text, outputs=text)
+    # next.click(load_example, inputs=text, outputs=[image, desc])
+    next.click(next_, inputs=text, outputs=[text, image, desc])
+
+    # prev.click(prev_, inputs=text, outputs=text)
+    # prev.click(load_example, inputs=text, outputs=[image, desc])
+    prev.click(prev_, inputs=text, outputs=[text, image, desc])
+
+demo.launch(debug=True)

models.py

@@ -0,0 +1,13 @@
+import os
+
+START_IMAGE = 12000
+
+from tensorflow.keras.models import load_model
+model = load_model("saved_model/1")
+
+from data_utils import nload, ims, DataSetFromFolder
+data = nload("/home/jkwiatkowski/all/dataset/arr/val.npy")
+indexes = nload("/home/jkwiatkowski/all/dataset/arr/val_target.npy")
+
+folders = DataSetFromFolder("/home/jkwiatkowski/all/dataset/arr/RAVEN-10000-release/RAVEN-10000", file_type="dir")
+properties = DataSetFromFolder(folders[:], file_type="xml", extension="val")

raven_utils/__init__.py

@@ -0,0 +1,10 @@
+import raven_utils.group as group
+import raven_utils.entity as entity
+import raven_utils.properties as properties
+import raven_utils.target as target
+import raven_utils.rules as rules
+import raven_utils.output as output
+import raven_utils.inference as inference
+import raven_utils.decode as decode
+import raven_utils.render_ as render_
+import raven_utils.draw as draw

raven_utils/config/constant.py

@@ -0,0 +1,54 @@
+RAVEN = "arr"
+RAVEN_BIG = "arrb"
+INDEX = "index"
+LABELS = "labels"
+TARGET_LABELS = "target_labels"
+FEATURES = "features"
+ACC_SAME = "acc_same"
+ACC_CHOOSE_LOWER = "acc_choose_lower"
+ACC_CHOOSE_UPPER = "acc_choose_upper"
+ACC_NO_GROUP = "acc_NO_group"
+CLASSIFICATION = "classification"
+INFERENCE = "inference"
+# PROPERTIES = "properties"
+PROPERTY = "property"
+MEMORY = "memory"
+CONTROL = "control"
+LATENT = "latent"
+TARGET = "target"
+INPUTS = "inputs"
+RES = "res"
+RESULT = "result"
+MERGE = "merge"
+MEMORY_STATE = "memory_state"
+CONTROL_STATE = "control_state"
+CONCAT = "concat"
+FLATTEN = "flatten"
+CROSS_ENTROPY = "cross_entropy"
+SLOT = "slot"
+PROPERTIES = "properties"
+ACC = "acc"
+GROUP = 'group'
+NUMBER = 'number'
+TRANS = 'trans'
+TAIL = "tail"
+MASK = "mask"
+
+RAV_METRICS = [
+    ACC_NO_GROUP,
+    ACC_SAME,
+    ACC_CHOOSE_UPPER,
+    ACC_CHOOSE_LOWER,
+    "acc",
+    "c_acc_NO_group",
+    "c_acc",
+    "loss",
+]
+
+IMP_RAV_METRICS = [
+    ACC_NO_GROUP,
+    ACC_SAME,
+    ACC_CHOOSE_UPPER,
+    ACC_CHOOSE_LOWER,
+    ACC,
+]

raven_utils/config/models.py

@@ -0,0 +1,9 @@
+
+AVAILABLE_MODELS = [
+    "197-0.31",
+    "53-0.48",
+    "74-0.50",
+    "21-0.48",
+    "10-0.52",
+    "179-0.50"
+]

raven_utils/const.py

@@ -0,0 +1,2 @@
+HORIZONTAL = "horizontal"
+VERTICAL = "vertical"

raven_utils/constant.py

@@ -0,0 +1,53 @@
+RAVEN = "arr"
+RAVEN_BIG = "arrb"
+INDEX = "index"
+LABELS = "labels"
+TARGET_LABELS = "target_labels"
+FEATURES = "features"
+ACC_SAME = "acc_same"
+ACC_CHOOSE_LOWER = "acc_choose_lower"
+ACC_CHOOSE_UPPER = "acc_choose_upper"
+ACC_NO_GROUP = "acc_NO_group"
+CLASSIFICATION = "classification"
+INFERENCE = "inference"
+# PROPERTIES = "properties"
+PROPERTY = "property"
+MEMORY = "memory"
+CONTROL = "control"
+LATENT = "latent"
+TARGET = "target"
+INPUTS = "inputs"
+RES = "res"
+RESULT = "result"
+MERGE = "merge"
+MEMORY_STATE = "memory_state"
+CONTROL_STATE = "control_state"
+CONCAT = "concat"
+FLATTEN = "flatten"
+CROSS_ENTROPY = "cross_entropy"
+SLOT = "slot"
+PROPERTIES = "properties"
+ACC = "acc"
+GROUP = 'group'
+NUMBER = 'number'
+TRANS = 'trans'
+TAIL = "tail"
+MASK = "mask"
+
+RAV_METRICS = [
+    ACC_NO_GROUP,
+    ACC_SAME,
+    ACC_CHOOSE_UPPER,
+    ACC_CHOOSE_LOWER,
+    "acc",
+    "c_acc_NO_group",
+    "c_acc",
+    "loss",
+]
+
+IMP_RAV_METRICS = [
+    ACC_NO_GROUP,
+    ACC_SAME,
+    ACC_CHOOSE_UPPER,
+    ACC_CHOOSE_LOWER,
+]

raven_utils/data.py

@@ -0,0 +1,46 @@
+import os
+
+import tensorflow as tf
+
+from models_utils import INPUTS, TARGET
+
+from raven_utils.config.constant import RAVEN, LABELS, INDEX, FEATURES, RAV_METRICS, IMP_RAV_METRICS, ACC_NO_GROUP
+
+
+from typing import Any
+
+from data_utils import pre, Data, gather, vec, resize
+from data_utils.data_generator import DataGenerator
+from funcy import identity
+
+
+def get_data(data, batch_size, steps=None, val_steps=None):
+    if val_steps is None:
+        val_steps = steps
+    fn = identity
+    train_target_index = data[4] + 8
+    train_generator = DataGenerator({
+        INPUTS: Data(data[0], fn),
+        TARGET: Data(data[2], identity),
+        LABELS: Data(data[2], identity),
+        INDEX: train_target_index[:, None],
+        # FEATURES: data[6]
+    },
+        batch=batch_size,
+        steps=steps
+    )
+    val_target_index = data[5] + 8
+    val_data = {
+        INPUTS: Data(data[1], fn),
+        TARGET: Data(data[3], identity),
+        LABELS: Data(data[3], identity),
+        INDEX: val_target_index[:, None],
+        # FEATURES: data[7]
+    }
+    val_generator = DataGenerator(
+        val_data,
+        batch=batch_size,
+        sampler="val",
+        steps=val_steps
+    )
+    return train_generator, val_generator

raven_utils/decode.py

@@ -0,0 +1,100 @@
+import numpy as np
+from data_utils import np_split
+from ml_utils import lw
+from models_utils.ops import ibin
+
+import raven_utils as rv
+
+
+def output(x, split_fn=np_split, predict_fn_1=np.argmax, predict_fn_2=ibin):
+    res = output_divide(x, split_fn=split_fn)
+    res = output_predict(res, predict_fn_1=predict_fn_1, predict_fn_2=predict_fn_2)
+    return (res[0], res[1]) + tuple(output_properties(res[2], predict_fn=predict_fn_1))
+
+
+def output_divide(output, split_fn=np_split):
+    group_output = output[..., rv.output.GROUP_SLICE_END]
+    slot_output = output[..., rv.output.SLOT_SLICE_END]
+    properties_output = output[..., rv.output.PROPERTIES_SLICE_END]
+    properties_output_splited = split_fn(properties_output, list(rv.properties.INDEX.values()), axis=-1)
+    return group_output, slot_output, properties_output_splited
+
+
+def output_predict(output, predict_fn_1=np.argmax, predict_fn_2=ibin):
+    return predict_fn_1(output[0]), predict_fn_2(output[1]), output[2]
+
+
+def output_properties(x, predict_fn=np.argmax):
+    out_reshaped = []
+    for i, out in enumerate(x):
+        shape = (-1, rv.entity.SUM, rv.properties.RAW_SIZE[i])
+        out_reshaped.append(predict_fn(out.reshape(shape)))
+    return out_reshaped
+
+
+def output_result(output, split_fn=np_split, arg_max=np.argmax):
+    result = output_properties(output, predict_fn=split_fn)
+    res = []
+    for i, r in enumerate(result):
+        if i == 1:
+            res.append(r)
+        else:
+            res.append(arg_max(r, axis=-1))
+    return tuple(res)
+
+
+def decode_inference(inference, reshape=np.reshape):
+    return reshape(inference[rv.inference.SLOT_SLICE],
+                   [-1, rv.group.NO, rv.inference.PROPERTY_TRANSFORMATION_NO]), reshape(
+        inference[rv.inference.PROPERTIES_SLICE],
+        [-1, rv.properties.NO, rv.entity.SUM, rv.inference.PROPERTY_TRANSFORMATION_NO])
+
+
+def decode_target(target):
+    target_group = target[..., 0]
+    target_slot = target[..., 1:rv.target.INDEX[0]]
+    target_properties = target[..., rv.target.INDEX[0]:rv.target.END_INDEX]
+    target_properties_splited = [
+        target_properties[..., ::rv.properties.NO],
+        target_properties[..., 1::rv.properties.NO],
+        target_properties[..., 2::rv.properties.NO]
+    ]
+    return target_group, target_slot, target_properties_splited
+
+
+
+
+def decode_target_flat(target):
+    t = decode_target(target)
+    return t[0], t[1], t[2][0], t[2][1], t[2][2]
+
+
+def demask(target, mask=None, group=None, zeroes=None):
+    if mask is None:
+        if group is None:
+            group = target[0]
+        # todo Use numpy range Mask
+        from models.uitls_ import RangeMask
+        mask = RangeMask()(group).numpy()
+    if zeroes is None:
+        return np.concatenate([t[mask] for t in lw(target[1:])])
+    return np.concatenate([target[0][None]] + [t * mask for t in lw(target[1:])],axis=-1)
+
+
+def target_mask(mask,right=1):
+    shape = mask.shape
+    return np.concatenate([np.ones([shape[0], 1]) ,mask, np.repeat(mask,3,axis=1), np.ones([shape[0], right])],axis=1)
+
+
+def get_full_range_mask(mask):
+    return np.concatenate([mask, np.repeat(mask, 3, axis=-1)], axis=-1)
+
+def compare(target, predict, mask):
+    target_comp = target[:, 1:rv.target.END_INDEX]
+    predict_comp = predict[:, 1:rv.target.END_INDEX]
+
+    mask = get_full_range_mask(mask)
+
+    target_masked = target_comp * mask
+    predict_masked = predict_comp * mask
+    return target_masked == predict_masked

raven_utils/depricated/old_raven.py

@@ -0,0 +1,490 @@
+from functools import partial
+
+import numpy as np
+from data_utils import take, EXIST, COR
+from data_utils.image import draw_images, add_text
+from data_utils.op import np_split
+from ml_utils import lu, dict_from_list2, filter_keys, none
+from data_utils import ops as K
+
+from config.constant import PROPERTY, TARGET, INPUTS
+# from raven_utils.render.rendering import render_panels
+
+RENDER_POSITIONS = [
+    [(0.5, 0.5, 1, 1)],
+    # ...
+    [(0.25, 0.25, 0.5, 0.5),
+     (0.25, 0.75, 0.5, 0.5),
+     (0.75, 0.25, 0.5, 0.5),
+     (0.75, 0.75, 0.5, 0.5)],
+    # ...
+    [(0.16, 0.16, 0.33, 0.33),
+     (0.16, 0.5, 0.33, 0.33),
+     (0.16, 0.83, 0.33, 0.33),
+     (0.5, 0.16, 0.33, 0.33),
+     (0.5, 0.5, 0.33, 0.33),
+     (0.5, 0.83, 0.33, 0.33),
+     (0.83, 0.16, 0.33, 0.33),
+     (0.83, 0.5, 0.33, 0.33),
+     (0.83, 0.83, 0.33, 0.33)],
+    # ...
+    [(0.5, 0.25, 0.5, 0.5)],
+    [(0.5, 0.75, 0.5, 0.5)],
+    # ...
+    [(0.25, 0.5, 0.5, 0.5)],
+    [(0.75, 0.5, 0.5, 0.5)],
+    # ...
+    [(0.5, 0.5, 1, 1)],
+    [(0.5, 0.5, 0.33, 0.33)],
+    # ...
+    [(0.5, 0.5, 1, 1)],
+    [(0.42, 0.42, 0.15, 0.15),
+     (0.42, 0.58, 0.15, 0.15),
+     (0.58, 0.42, 0.15, 0.15),
+     (0.58, 0.58, 0.15, 0.15)],
+    # ...
+
+]
+
+HORIZONTAL = "horizontal"
+VERTICAL = "vertical"
+
+NAMES = ['center_single',
+         'distribute_four',
+         'distribute_nine',
+         'in_center_single_out_center_single',
+         'in_distribute_four_out_center_single',
+         'left_center_single_right_center_single',
+         'up_center_single_down_center_single']
+
+PROPERTIES_NAMES = [
+    'Color',
+    'Size',
+    'Type',
+
+]
+PROPERTIES = dict_from_list2(PROPERTIES_NAMES, [10, 6, 5])
+ANGLE_MAX = 7
+
+PROPERTIES_NO = len(PROPERTIES)
+
+RULES_COMBINE = "Number/Position"
+
+RULES_ATTRIBUTES = [
+    "Number",
+    "Position",
+    "Color",
+    "Size",
+    "Type"
+]
+RULES_ATTRIBUTES_LEN = len(RULES_ATTRIBUTES)
+
+RULES_ATTRIBUTES_INDEX = dict_from_list2(RULES_ATTRIBUTES)
+
+RULES_TYPES = [
+    "Constant",
+    "Arithmetic",
+    "Progression",
+    "Distribute_Three"
+]
+RULES_TYPES_INDEX = dict_from_list2(RULES_TYPES)
+RULES_TYPES_LEN = len(RULES_ATTRIBUTES)
+
+GROUPS_NO = len(NAMES)
+ENTITY_NO = dict(zip(NAMES, [1, 4, 9, 2, 5, 2, 2]))
+ENTITY_SUM = sum(list(ENTITY_NO.values()))
+ENTITY_INDEX = np.concatenate([[0], np.cumsum(list(ENTITY_NO.values()))])
+ENTITY_INDEX_TARGET = ENTITY_INDEX + 1
+ENTITY_DICT = dict(zip(NAMES, ENTITY_INDEX_TARGET[:-1]))
+NAMES_ORDER = dict(zip(NAMES, np.arange(len(NAMES))))
+PROPERTIES_INDEXES = np.cumsum(np.array(list(ENTITY_NO.values())) * len(PROPERTIES))
+INDEX = np.concatenate([[0], PROPERTIES_INDEXES]) + ENTITY_SUM + 1  # +2 type and uniformity
+
+SECOND_LAYOUT = [i - 1 for i in [
+    ENTITY_DICT["in_center_single_out_center_single"] + 1,
+    ENTITY_DICT["in_distribute_four_out_center_single"] + 1,
+    ENTITY_DICT["in_distribute_four_out_center_single"] + 2,
+    ENTITY_DICT["in_distribute_four_out_center_single"] + 3,
+    ENTITY_DICT["left_center_single_right_center_single"] + 1,
+    ENTITY_DICT["up_center_single_down_center_single"] + 1
+]]
+
+FIRST_LAYOUT = list(set(range(ENTITY_SUM)) - set(SECOND_LAYOUT))
+LAYOUT_NO = 2
+
+START_INDEX = dict(zip(NAMES, INDEX[:-1]))
+END_INDEX = INDEX[-1]
+
+RULES_ATTRIBUTES_ALL_LEN = RULES_ATTRIBUTES_LEN * LAYOUT_NO
+UNIFORMITY_NO = 2
+UNIFORMITY_INDEX = END_INDEX + RULES_ATTRIBUTES_ALL_LEN
+
+FEATURE_NO = UNIFORMITY_INDEX + UNIFORMITY_NO
+MAPPING = {
+    "distribute_nine":
+        {0.16: 0,
+         0.5: 1,
+         0.83: 2},
+    "distribute_four":
+        {0.25: 0,
+         0.75: 1},
+    'in_distribute_four_out_center_single':
+        {0.42: 0,
+         0.58: 1}
+}
+MUL = {
+    "distribute_nine": 3,
+    "distribute_four": 2,
+    'in_distribute_four_out_center_single': 2
+}
+
+# SIZES = np.linspace(0.4, 0.9, 6)
+TYPES = ["triangle", "square", "pentagon", "hexagon", "circle"]
+# TYPES = ["triangle", "square", "pentagon", "circle", "circle"]
+SIZES = ["vs", "s", "m", "h", "vh", "e"]
+COLORS = ["vs", "s", "m", "h", "vh", "e"]
+# TYPES = ["", "", "circle", "hexagon", "square"]
+
+ENTITY_PROPERTIES_VALUES = list(PROPERTIES.values())
+ENTITY_PROPERTIES_KEYS = list(PROPERTIES.keys())
+ENTITY_PROPERTIES_NO = len(PROPERTIES)
+INDEX = dict(zip(PROPERTIES, np.array(ENTITY_PROPERTIES_VALUES) * ENTITY_SUM))
+ENTITY_PROPERTIES_SUM = sum(list(PROPERTIES.values()))
+
+OUTPUT_SIZE = ENTITY_SUM * ENTITY_PROPERTIES_SUM + GROUPS_NO + ENTITY_SUM
+
+SLOT_AND_GROUP = ENTITY_SUM + GROUPS_NO
+
+OUTPUT_GROUP_SLICE = np.s_[:, -GROUPS_NO:]
+OUTPUT_SLOT_SLICE = np.s_[:, -SLOT_AND_GROUP:-GROUPS_NO]
+OUTPUT_PROPERTIES_SLICE = np.s_[:, :-SLOT_AND_GROUP]
+
+OUTPUT_GROUP_SLICE_END = np.s_[-GROUPS_NO:]
+OUTPUT_SLOT_SLICE_END = np.s_[-SLOT_AND_GROUP:-GROUPS_NO]
+OUTPUT_PROPERTIES_SLICE_END = np.s_[:-SLOT_AND_GROUP]
+
+# Transformation
+# constant
+# progression -2, -1,1 ,2
+# arithmetic -/+ Position set arithmetic
+# distribute three
+
+# todo
+SLOTS_GROUPS = GROUPS_NO
+
+SLOT_TRANSFORMATION_NO = 4
+PROPERTY_TRANSFORMATION_NO = 8
+PROPERTIES_TRANSFORMATION_NO = PROPERTY_TRANSFORMATION_NO * PROPERTIES_NO
+PROPERTIES_TRANSFORMATION_SIZE = PROPERTIES_TRANSFORMATION_NO * ENTITY_SUM
+
+SLOT_TRANSFORMATION_SIZE = PROPERTY_TRANSFORMATION_NO * SLOTS_GROUPS
+INFERENCE_SIZE = SLOT_TRANSFORMATION_SIZE + PROPERTIES_TRANSFORMATION_SIZE
+
+INFERENCE_SLOT_SLICE = np.s_[:, :SLOT_TRANSFORMATION_SIZE]
+INFERENCE_PROPERTIES_SLICE = np.s_[:, -PROPERTIES_TRANSFORMATION_SIZE:]
+from operator import add
+
+
+# todo Refactor
+# Maybe properties should be on same level as rest.
+def decode_output(output, split_fn=np_split):
+    group_output = output[..., OUTPUT_GROUP_SLICE_END]
+    slot_output = output[..., OUTPUT_SLOT_SLICE_END]
+    properties_output = output[..., OUTPUT_PROPERTIES_SLICE_END]
+    properties_output_splited = split_fn(properties_output, list(rv.properties.INDEX.values()), axis=-1)
+    return group_output, slot_output, properties_output_splited
+
+
+def decode_inference(inference, reshape=np.reshape):
+    return reshape(inference[INFERENCE_SLOT_SLICE],
+                   [-1, SLOTS_GROUPS, PROPERTY_TRANSFORMATION_NO]), reshape(
+        inference[INFERENCE_PROPERTIES_SLICE],
+        [-1, PROPERTIES_NO, ENTITY_SUM, PROPERTY_TRANSFORMATION_NO])
+
+
+def decode_output_reshape(output, split_fn=np_split):
+    result = decode_output(output, split_fn=split_fn)
+    out_reshaped = []
+    for i, out in enumerate(result[2]):
+        shape = (-1, ENTITY_SUM, ENTITY_PROPERTIES_VALUES[i])
+        out_reshaped.append(out.reshape(shape))
+    return result[:2] + tuple(out_reshaped)
+
+
+def take_target(target):
+    return target[1], target[2]
+
+
+def create_target(images, index, pattern_index=(2, 5), full_index=False, arrange=np.arange, shape=lambda x: x.shape):
+    return [images[:, pattern_index[0]], images[:, pattern_index[1]],
+            images[arrange(shape(index)[0]), (0 if full_index else 8) + index[:, 0]]]
+
+
+def take_target_simple(target):
+    return target[1], target[0]
+
+
+def create_target_simple(images, target, index=slice(None), pattern_index=(2, 5)):
+    return [images[:, pattern_index[0]], images[:, pattern_index[1]], target][index]
+
+
+def decode_output_result(output, split_fn=np_split, arg_max=np.argmax):
+    result = decode_output_reshape(output, split_fn=split_fn)
+    res = []
+    for i, r in enumerate(result):
+        if i == 1:
+            res.append(r)
+        else:
+            res.append(arg_max(r, axis=-1))
+    return tuple(res)
+
+
+def decode_target(target):
+    target_group = target[..., 0]
+    target_slot = target[..., 1:INDEX[0]]
+    target_properties = target[..., INDEX[0]:END_INDEX]
+    target_properties_splited = [
+        target_properties[..., ::PROPERTIES_NO],
+        target_properties[..., 1::PROPERTIES_NO],
+        target_properties[..., 2::PROPERTIES_NO]
+    ]
+    return target_group, target_slot, target_properties_splited
+
+
+def decode_target_flat(target):
+    t = decode_target(target)
+    return t[0], t[1], t[2][0], t[2][1], t[2][2]
+
+
+def draw_board(images,  target=None, predict=None,image=None, desc=None, layout=None, break_=20):
+    if image != "target" and predict is not None:
+        image = images[predict:predict + 1]
+    elif images is None and target is not None:
+        image = images[target:target + 1]
+    # image = False to not draw anything
+    border = [{COR: target - 8, EXIST: (1, 3)}] + [{COR: p, EXIST: (0, 2)} for p in none(predict)]
+
+    boards = []
+    boards.append(draw_images(np.concatenate([images[:8], image[None] if len(image.shape)==3 else image]) if image is not None else images[:8]))
+    if layout == 1:
+        i = draw_images(images[8:], column=4, border=border)
+        if break_:
+            i = np.concatenate([np.zeros([ break_, i.shape[1],1]),i ],axis=0)
+        boards.append(i)
+
+    else:
+        boards.append(
+            draw_images(np.concatenate([images[8:], predict]) if predict is not None else images[8:], column=4,
+                        border=target - 8))
+    full_board = draw_images(boards, grid=False)
+    if desc:
+        full_board = add_text(full_board, desc)
+    return full_board
+
+
+def draw_boards(images, target=None, predict=None, image=None, desc=None, no=1, layout=None):
+    boards = []
+    for i, image in enumerate(images):
+        boards.append(draw_board(image, target[i][0] if target is not None else None,
+                                 predict[i] if predict is not None else None,
+                                 image[i] if image is not None else None,
+                                 desc[i] if desc is not None else None, layout=layout))
+    return boards
+
+
+def draw_raven(generator, predict=None, no=1, add_target_desc=True, indexes=None, types=TYPES,
+               layout=1):
+    if indexes is None:
+        indexes = val_sample(no)
+    data = generator.data[indexes]
+    if is_model(predict):
+        d = filter_keys(data, PROPERTY,reverse=True)
+        # tmp change
+        pro = predict(d)['predict']
+        print(pro)
+        predict = render_panels(pro, target=False)
+    # if target is not None:
+    target = data[TARGET]
+    target_index = data["index"]
+    images = data[INPUTS]
+
+    if hasattr(predict, "shape"):
+        if len(predict.shape) > 3:
+            # iamges
+            image = predict
+            # todo create index and output based on image
+            predict = None
+            predict_index = None
+        elif len(predict.shape) == 3:
+            image = render_panels(predict, target=False)
+            # Create index based on predict.
+            predict_index = None
+        else:
+            image = images[predict]
+            predict_index = predict
+            predict = target
+    else:
+        image = K.gather(images, target_index[:, 0])
+        predict_index = None
+        predict = None
+
+    # elif not(hasattr(target,"shape") and len(target.shape) > 3):
+    #     if hasattr(target,"shape") and target.shape[-1] == OUTPUT_SIZE:
+    #         pro = target
+    #         predict = render_panels(pro)
+    #     elif hasattr(target,"shape") and target.shape[-1] == FEATURE_NO:
+    #         # pro = target
+    #         pro = np.zeros([no, OUTPUT_SIZE], dtype="int")
+    #     else:
+    #         pro = np.zeros([no, OUTPUT_SIZE], dtype="int")
+    #         # predict = [None] * no
+    #         predict = render_panels(data[TARGET])
+
+    all_rules = []
+    for d in data[PROPERTY]:
+        rules = []
+        for j, rule_group in enumerate(d.findAll("Rule_Group")):
+            # rules_all.append(rule_group['id'])
+            for j, rule in enumerate(rule_group.findAll("Rule")):
+                rules.append(f"{rule['attr']} - {rule['name']}")
+            rules.append("")
+        all_rules.append(rules)
+    target_desc = get_desc(target)
+    if predict is not None:
+        predict_desc = decode_output_result(predict) if predict.shape[-1] == OUTPUT_SIZE else get_desc(predict)
+    else:
+        predict_desc = [None] * len(target_desc)
+    for a, po, to in zip(all_rules, predict_desc, target_desc):
+        # fl(predict_desc[-1])
+        if po is None:
+            po = [None] * len(to)
+        for p, t in zip(po, to):
+            a.extend(
+                [" ".join([str(i) for i in t])] + (
+                    [" ".join([str(i) for i in p]), ""] if p is not None else []
+                )
+            )
+        # a.extend([""] + [] + [""] + [" ".join(fl(p))])
+
+    # image = draw_boards(data[INPUTS],target=data["index"], predict=predict[:no], desc=all_rules, no=no,layer=layer)
+    image = draw_boards(images, target=target_index, predict=predict_index, image=image, desc=None, no=no,
+                        layout=layout)
+    return lu([(i, j) for i, j in zip(image, all_rules)])
+
+
+def val_sample(no=GROUPS_NO, base=3):
+    indexes = np.arange(no) * 2000 + base
+    return indexes
+
+
+def get_desc(target, exist=None, types=TYPES, sizes=SIZES):
+    decoded = decode_target(target)
+    exist = decoded[1] if exist is None else exist
+    taken = np.stack(take(decoded[2], np.array(exist, dtype=bool))).T
+
+    figures_no = np.sum(exist, axis=-1)
+    desc = np.split(taken, np.cumsum(figures_no))[:-1]
+    # figures_no = np.sum(exist, axis=-1)
+    # div = np.split(desc, np.cumsum(figures_no))[:-1]
+    result = []
+    for pd in desc:
+        r = []
+        for p in pd:
+            r.append([p[0], sizes[p[1]], types[p[2]]])
+        result.append(r)
+
+    return result
+
+
+# def get
+
+
+def get_description(inputs, predict, pro, no, types=TYPES, sizes=SIZES):
+    # target = inputs[1][2][:no]
+    target = inputs[TARGET]
+    target_group = target[:, 0]
+    target_exist = np.asarray(target[:, 1:ENTITY_SUM + 1], dtype="bool")
+    target_rest = target[:, ENTITY_SUM + 1:ENTITY_SUM + 1 + ENTITY_SUM * PROPERTIES_NO]
+    pro_reshaped = np.reshape(pro, (pro.shape[0], -1, PROPERTIES_NO))
+    target_reshaped = np.reshape(target_rest, (target_rest.shape[0], -1, PROPERTIES_NO))
+
+    # mask = np.repeat(target_exist, [4] * ENTITY_SUM, axis=-1)
+    # masked_result = np.repeat(target_exist, [4] * ENTITY_SUM, axis=-1)
+    pro_res = pro_reshaped[target_exist]
+    target_res = target_reshaped[target_exist]
+    figures_no = np.sum(target_exist, axis=-1)
+
+    pro_div = np.split(pro_res, np.cumsum(figures_no))[:-1]
+    target_div = np.split(target_res, np.cumsum(figures_no))[:-1]
+    pro_result_full = []
+    target_result_full = []
+    for pd, td in zip(pro_div, target_div):
+        pro_result = []
+        target_result = []
+        for p in pd:
+            pro_result.append([p[0], sizes[p[1]], types[p[2]]])
+        for t in td:
+            target_result.append([t[0], sizes[t[1]], types[t[2]]])
+        pro_result_full.append(pro_result)
+        target_result_full.append(target_result)
+
+    return pro_result_full, target_result_full
+
+
+def get_properties(target, types=TYPES, sizes=SIZES):
+    target_exist = np.asarray(target[:, 1:ENTITY_SUM + 1], dtype="bool")
+    target_rest = target[:, ENTITY_SUM + 1:ENTITY_SUM + 1 + ENTITY_SUM * PROPERTIES_NO]
+    target_reshaped = np.reshape(target_rest, (target_rest.shape[0], -1, PROPERTIES_NO))
+    target_res = target_reshaped[target_exist]
+    figures_no = np.sum(target_exist, axis=-1)
+    target_div = np.split(target_res, np.cumsum(figures_no))[:-1]
+    target_result_full = []
+    for td in target_div:
+        target_result = []
+        for t in td:
+            target_result.append([t[0], sizes[t[1]], types[t[2]]])
+        target_result_full.append(target_result)
+    return target_result_full
+
+
+def desc_properties(target, decode_fn=None, types=TYPES, sizes=SIZES):
+    if decode_fn is None:
+        if target.shape[1] == OUTPUT_SIZE:
+            decode_fn = decode_output_result
+        else:
+            decode_fn = decode_target
+
+    target_div = decode_fn(target)[2:]
+    target_result_full = []
+    for td in target_div:
+        target_result = []
+        for t in td:
+            target_result.append([t[0], sizes[t[1]], types[t[2]]])
+        target_result_full.append(target_result)
+    return target_result_full
+
+
+def get_pro(t, types=TYPES, sizes=SIZES):
+    return [int(t[0]), sizes[t[1]], types[t[2]]]
+
+
+def get_pro2(td, types=TYPES, sizes=SIZES):
+    target_result = []
+    for t in td:
+        target_result.append([int(t[0]), sizes[t[1]], types[t[2]]])
+    return target_result
+
+
+def get_pro3(target_div, types=TYPES, sizes=SIZES):
+    target_result_full = []
+    for td in target_div.to_list():
+        target_result = []
+        for t in td:
+            target_result.append([int(t[0]), sizes[t[1]], types[t[2]]])
+        target_result_full.append(target_result)
+    return target_result_full
+
+
+from models_utils import init_image as def_init_image, is_model
+
+init_image = partial(def_init_image, shape=(16, 8, 80, 80, 1))

raven_utils/draw.py

@@ -0,0 +1,174 @@
+import numpy as np
+from data_utils import take, EXIST, COR
+from data_utils.image import draw_images, add_text
+from funcy import identity
+from ml_utils import none, filter_keys, lu
+from models_utils import is_model
+from models_utils import ops as K
+
+from raven_utils.constant import PROPERTY, TARGET, INPUTS
+from raven_utils.decode import decode_target, target_mask
+from raven_utils.render.rendering import render_panels
+from raven_utils.render_ import TYPES, SIZES
+from raven_utils.uitls import get_val_index
+
+
+def draw_board(images, target=None, predict=None, image=None, desc=None, layout=None, break_=20):
+    if image != "target" and predict is not None:
+        image = images[predict:predict + 1]
+    elif images is None and target is not None:
+        image = images[target:target + 1]
+    # image = False to not draw anything
+    border = [{COR: target - 8, EXIST: list(range(4)) if predict is None else (1, 3)}] + [{COR: p, EXIST: (0, 2)} for p
+                                                                                          in none(predict)]
+
+    boards = []
+    boards.append(draw_images(
+        np.concatenate([images[:8], image[None] if len(image.shape) == 3 else image]) if image is not None else images[
+                                                                                                                :8]))
+    if layout == 1:
+        i = draw_images(images[8:], column=4, border=border)
+        if break_:
+            i = np.concatenate([np.zeros([break_, i.shape[1], 1]), i], axis=0)
+        boards.append(i)
+
+    else:
+        boards.append(
+            draw_images(np.concatenate([images[8:], predict]) if predict is not None else images[8:], column=4,
+                        border=target - 8))
+    full_board = draw_images(boards, grid=False)
+    if desc:
+        full_board = add_text(full_board, desc)
+    return full_board
+
+
+def draw_boards(images, target=None, predict=None,image=None, desc=None,  layout=None):
+    boards = []
+    for i, im in enumerate(images):
+        boards.append(draw_board(im, target[i][0] if target is not None else None,
+                                 predict[i] if predict is not None else None,
+                                 image[i] if image is not None else None,
+                                 desc[i] if desc is not None else None, layout=layout))
+    return boards
+
+
+def draw_from_generator(generator, predict=None, no=1, indexes=None, layout=1):
+    data,_ = val_sample(generator, no, indexes)
+    return draw_raven(data, predict=predict, pre_fn=generator.data.data["inputs"].fn, layout=layout)
+
+
+def val_sample(generator, no=1, indexes=None):
+    if indexes is None:
+        indexes = get_val_index(base=no)
+    data = generator.data[indexes]
+    return data, indexes
+
+def render_from_model(data,predict,pre_fn=identity):
+    data = filter_keys(data, PROPERTY, reverse=True)
+    if is_model(predict):
+        predict = predict(data)
+    pro = np.array(target_mask(predict['predict_mask'].numpy()) * predict["predict"].numpy(), dtype=np.int8)
+    return pre_fn(render_panels(pro, target=False)[None])[0]
+
+def draw_raven(data, predict=None, pre_fn=identity, layout=1):
+    if is_model(predict):
+        d = filter_keys(data, PROPERTY, reverse=True)
+        # tmp change
+        res = predict(d)
+        pro = np.array(target_mask(res['mask'].numpy()) * res["predict"].numpy(),dtype=np.int8)
+        predict = pre_fn(render_panels(pro, target=False)[None])[0]
+        # from data_utils import ims
+        # ims(1 - predict[0])
+    # if target is not None:
+    target = data[TARGET]
+    target_index = data["index"]
+    images = data[INPUTS]
+    # np.equal(res['predict'], pro[:,:102]).sum()
+
+    if hasattr(predict, "shape"):
+        if len(predict.shape) > 3:
+            # iamges
+            image = predict
+            # todo create index and output based on image
+            predict = None
+            predict_index = None
+        elif len(predict.shape) == 3:
+            image = render_panels(predict, target=False)
+            # Create index based on predict.
+            predict_index = None
+        else:
+            image = images[predict]
+            predict_index = predict
+            predict = target
+    else:
+        image = K.gather(images, target_index[:, 0])
+        predict_index = None
+        predict = None
+
+    # elif not(hasattr(target,"shape") and len(target.shape) > 3):
+    #     if hasattr(target,"shape") and target.shape[-1] == OUTPUT_SIZE:
+    #         pro = target
+    #         predict = render_panels(pro)
+    #     elif hasattr(target,"shape") and target.shape[-1] == FEATURE_NO:
+    #         # pro = target
+    #         pro = np.zeros([no, OUTPUT_SIZE], dtype="int")
+    #     else:
+    #         pro = np.zeros([no, OUTPUT_SIZE], dtype="int")
+    #         # predict = [None] * no
+    #         predict = render_panels(data[TARGET])
+
+    image = draw_boards(images, target=target_index, predict=predict_index,image=image,  desc=None,
+                        layout=layout)
+
+    all_rules = extract_rules(data[PROPERTY])
+    target_desc = get_desc(target)
+    if predict is not None:
+        predict_desc = get_desc(predict)
+    else:
+        predict_desc = [None] * len(target_desc)
+    for a, po, to in zip(all_rules, predict_desc, target_desc):
+        # fl(predict_desc[-1])
+        if po is None:
+            po = [None] * len(to)
+        for p, t in zip(po, to):
+            a.extend(
+                [" ".join([str(i) for i in t])] + (
+                    [" ".join([str(i) for i in p]), ""] if p is not None else []
+                )
+            )
+        # a.extend([""] + [] + [""] + [" ".join(fl(p))])
+
+    # image = draw_boards(data[INPUTS],target=data["index"], predict=predict[:no], desc=all_rules, no=no,layer=layer)
+    return lu([(i, j) for i, j in zip(image, all_rules)])
+
+
+def extract_rules(data):
+    all_rules = []
+    for d in data:
+        rules = []
+        for j, rule_group in enumerate(d.findAll("Rule_Group")):
+            # rules_all.append(rule_group['id'])
+            for j, rule in enumerate(rule_group.findAll("Rule")):
+                rules.append(f"{rule['attr']} - {rule['name']}")
+            rules.append("")
+        all_rules.append(rules)
+    return all_rules
+
+
+def get_desc(target, exist=None, types=TYPES, sizes=SIZES):
+    decoded = decode_target(target)
+    exist = decoded[1] if exist is None else exist
+    taken = np.stack(take(decoded[2], np.array(exist, dtype=bool))).T
+
+    figures_no = np.sum(exist, axis=-1)
+    desc = np.split(taken, np.cumsum(figures_no))[:-1]
+    # figures_no = np.sum(exist, axis=-1)
+    # div = np.split(desc, np.cumsum(figures_no))[:-1]
+    result = []
+    for pd in desc:
+        r = []
+        for p in pd:
+            r.append([p[0], sizes[p[1]], types[p[2]]])
+        result.append(r)
+
+    return result

raven_utils/entity.py

@@ -0,0 +1,6 @@
+import raven_utils.group as group
+import numpy as np
+NO = dict(zip(group.NAMES, [1, 4, 9, 2, 5, 2, 2]))
+SUM = sum(list(NO.values()))
+
+INDEX = np.concatenate([[0], np.cumsum(list(NO.values()))])

raven_utils/group.py

@@ -0,0 +1,11 @@
+import numpy as np
+
+NAMES = ['center_single',
+         'distribute_four',
+         'distribute_nine',
+         'in_center_single_out_center_single',
+         'in_distribute_four_out_center_single',
+         'left_center_single_right_center_single',
+         'up_center_single_down_center_single']
+
+NO = len(NAMES)

raven_utils/inference.py

@@ -0,0 +1,15 @@
+import numpy as np
+import raven_utils.properties as properties
+import raven_utils.group as group
+
+
+SLOT_TRANSFORMATION_NO = 4
+PROPERTY_TRANSFORMATION_NO = 8
+PROPERTIES_TRANSFORMATION_NO = PROPERTY_TRANSFORMATION_NO * properties.NO
+PROPERTIES_TRANSFORMATION_SIZE = PROPERTIES_TRANSFORMATION_NO * group.NO
+
+SLOT_TRANSFORMATION_SIZE = PROPERTY_TRANSFORMATION_NO * group.NO
+SIZE = SLOT_TRANSFORMATION_SIZE + PROPERTIES_TRANSFORMATION_SIZE
+
+SLOT_SLICE = np.s_[:, :SLOT_TRANSFORMATION_SIZE]
+PROPERTIES_SLICE = np.s_[:, -PROPERTIES_TRANSFORMATION_SIZE:]

raven_utils/models/attn.py

@@ -0,0 +1,187 @@
+from __future__ import print_function
+
+import tensorflow as tf
+from tensorflow.keras import backend as K
+from tensorflow.keras.layers import LSTMCell
+from tensorflow.keras.models import Model
+from tensorflow.keras.layers import Conv2D, Dense
+from tensorflow.keras.losses import mse
+from tensorflow.keras.models import clone_model
+from tensorflow.layers.base import InputSpec, Layer
+
+from models.dense import create_conv_model
+from models.utils import broadcast
+
+
+class ReflectionPadding2D(Layer):
+    def __init__(self, padding=(1, 1), **kwargs):
+        self.padding = tuple(padding)
+        self.input_spec = [InputSpec(ndim=4)]
+        super(ReflectionPadding2D, self).__init__(**kwargs)
+
+    def compute_output_shape(self, s):
+        """ If you are using "channels_last" configuration"""
+        return (s[0], s[1] + 2 * self.padding[0], s[2] + 2 * self.padding[1], s[3])
+
+    def call(self, x, mask=None):
+        w_pad, h_pad = self.padding
+        return tf.pad(x, [[0, 0], [h_pad, h_pad], [w_pad, w_pad], [0, 0]], 'REFLECT')
+
+
+class Conv2Ref(Layer):
+    def __init__(self, padding=(1, 1), **kwargs):
+        self.padding = tuple(padding)
+        self.input_spec = [InputSpec(ndim=4)]
+        super(ReflectionPadding2D, self).__init__(**kwargs)
+
+    def compute_output_shape(self, s):
+        """ If you are using "channels_last" configuration"""
+        return (s[0], s[1] + 2 * self.padding[0], s[2] + 2 * self.padding[1], s[3])
+
+    def call(self, x, mask=None):
+        w_pad, h_pad = self.padding
+        return tf.pad(x, [[0, 0], [h_pad, h_pad], [w_pad, w_pad], [0, 0]], 'REFLECT')
+
+
+class SegmentationNetwork(Model):
+
+    def __init__(self, filters=64, kernels=(3, 3)):
+        super(RecAE, self).__init__()
+        self.conv_1 = Conv2D(filters, kernels, padding=SAME)
+        self.conv_2 = Conv2D(filters, kernels, padding=SAME)
+
+    def call(self, inputs):
+        x = K.relu(inputs)
+        x = self.conv_1(x)
+        x = K.relu(x)
+        x = self.conv_2(x)
+        return x + inputs
+
+
+class QueryNetwork(Model):
+
+    def __init__(self, units=64):
+        super(RecAE, self).__init__()
+        self.conv_1 = Dense(units)
+        self.conv_2 = Dense(units)
+
+    def call(self, inputs):
+        x = K.relu(inputs)
+        x = self.conv_1(x)
+        x = K.relu(x)
+        x = self.conv_2(x)
+        return x + inputs
+
+
+class RecAE(Model):
+
+    def __init__(self, head, bottle, decoder):
+        super(RecAE, self).__init__()
+        self.head = head
+        self.bottle = bottle
+        self.base = clone_model(bottle)
+        self.decoder = decoder
+        self.segmentation_network = SegmentationNetwork()
+        self.query_network = QueryNetwork()
+        self.control = LSTMCell(64)
+        self.memory = LSTMCell(64)
+
+    def call(self, inputs):
+        feature = self.head(inputs)
+        segmentation = self.segmentation_network(feature)
+        control_base = self.base(feature)
+        h_c = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
+        h_m = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
+        shape = K.shape(feature)[:-1]
+        full_attention = tf.zeros(shape)[..., tf.newaxis]
+        full_image = tf.zeros(K.shape(inputs))
+        masks = []
+        ff = tf.zeros(K.shape(inputs))
+        scope = tf.ones(shape)[..., tf.newaxis]
+        for i in range(4):
+            r_c, h_c = self.control(tf.concat([control_base, h_m[0]], 1), h_c)
+            query = self.query_network(h_c[0])
+            log_attention = image_attention(segmentation, query)
+            attention = K.sigmoid(log_attention)
+            mask = attention * scope
+            scope = scope - mask
+            im = feature * mask
+            # im = feature
+            latent = self.bottle(im)
+            decoded = self.decoder(latent)
+            # self.add_loss(K.mean(-mse(full_attention, attention)))
+            # self.add_loss(K.mean(-mse(tf.ones(attention.shape), attention)))
+            full_attention += attention
+            big_mask = tf.image.resize(mask, K.shape(inputs)[1:-1])
+            ff += K.sigmoid(decoded)
+            full_image += K.sigmoid(decoded) * big_mask
+            r_m, h_m = self.memory(latent, h_m)
+            masks.append(big_mask)
+        self.add_loss(K.mean(mse(inputs, full_image)))
+        return full_image, masks
+
+
+# def image_attention(image, query, scale=True):
+@tf.function
+def image_attention(image, query):
+    log_attention = K.sum(query[:, tf.newaxis, tf.newaxis, :] * image, axis=-1, keepdims=True)
+    # if scale is not None:
+    log_attention /= tf.sqrt(tf.cast(K.shape(image)[-1], dtype=float))
+    return log_attention
+
+
+class RecAE_2(Model):
+
+    def __init__(self, head, bottle, decoder):
+        super(RecAE_2, self).__init__()
+        self.head = head
+        self.bottle = bottle
+        # self.base = clone_model(bottle)
+        self.base = self.bottle
+        self.decoder = decoder
+        self.segmentation_network = create_conv_model((64, 64, 1))
+        self.control = LSTMCell(64)
+        self.memory = LSTMCell(64)
+
+    def call(self, inputs):
+        feature = self.head(inputs)
+        control_base = self.base(feature)
+        h_c = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
+        h_m = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
+        shape = K.shape(feature)[:-1]
+        full_attention = tf.zeros(shape)[..., tf.newaxis]
+        full_image = tf.zeros(K.shape(inputs))
+        big_masks = []
+        masks = []
+        ff = tf.zeros(K.shape(inputs))
+        scope = tf.ones(shape)[..., tf.newaxis]
+        for i in range(4):
+            if i ==3:
+                mask = scope
+            else:
+                r_c, h_c = self.control(tf.concat([control_base, h_m[0]], 1), h_c)
+                query = broadcast(h_c[0], feature.shape[1:])
+                log_attention = self.segmentation_network(tf.concat([feature, query], axis=-1))
+                attention = K.sigmoid(log_attention)
+                mask = attention * scope
+                scope = scope - mask
+            masks.append(mask)
+            im = feature * mask
+            # im = feature
+            latent = self.bottle(im)
+            decoded = self.decoder(latent)
+            # self.add_loss(K.mean(-mse(scope, mask)))
+            sum = K.sum(tf.ones(K.shape(mask)))
+            self.add_loss(K.abs((sum/4)-K.sum(mask))/sum)
+            # self.add_loss(K.mean(-mse(tf.zeros(K.shape(mask)), mask)))
+            for m in masks:
+                self.add_loss(K.mean(-mse(m,mask)))
+
+            full_attention += mask
+            big_mask = tf.image.resize(mask, K.shape(inputs)[1:-1])
+            ff += K.sigmoid(decoded)
+            full_image += K.sigmoid(decoded) * big_mask
+            r_m, h_m = self.memory(latent, h_m)
+            big_masks.append(big_mask)
+        self.add_loss(K.mean(mse(inputs, full_image)))
+        return full_image, big_masks

raven_utils/models/attn2.py

@@ -0,0 +1,187 @@
+from __future__ import print_function
+
+import tensorflow as tf
+from tensorflow.keras import backend as K
+from tensorflow.keras.layers import LSTMCell
+from tensorflow.keras.models import Model
+from tensorflow.keras.layers import Conv2D, Dense
+from tensorflow.keras.losses import mse
+from tensorflow.keras.models import clone_model
+from tensorflow.layers.base import InputSpec, Layer
+
+from models.dense import create_conv_model
+from models.utils import broadcast
+
+
+class ReflectionPadding2D(Layer):
+    def __init__(self, padding=(1, 1), **kwargs):
+        self.padding = tuple(padding)
+        self.input_spec = [InputSpec(ndim=4)]
+        super(ReflectionPadding2D, self).__init__(**kwargs)
+
+    def compute_output_shape(self, s):
+        """ If you are using "channels_last" configuration"""
+        return (s[0], s[1] + 2 * self.padding[0], s[2] + 2 * self.padding[1], s[3])
+
+    def call(self, x, mask=None):
+        w_pad, h_pad = self.padding
+        return tf.pad(x, [[0, 0], [h_pad, h_pad], [w_pad, w_pad], [0, 0]], 'REFLECT')
+
+
+class Conv2Ref(Layer):
+    def __init__(self, padding=(1, 1), **kwargs):
+        self.padding = tuple(padding)
+        self.input_spec = [InputSpec(ndim=4)]
+        super(ReflectionPadding2D, self).__init__(**kwargs)
+
+    def compute_output_shape(self, s):
+        """ If you are using "channels_last" configuration"""
+        return (s[0], s[1] + 2 * self.padding[0], s[2] + 2 * self.padding[1], s[3])
+
+    def call(self, x, mask=None):
+        w_pad, h_pad = self.padding
+        return tf.pad(x, [[0, 0], [h_pad, h_pad], [w_pad, w_pad], [0, 0]], 'REFLECT')
+
+
+class SegmentationNetwork(Model):
+
+    def __init__(self, filters=64, kernels=(3, 3)):
+        super(RecAE, self).__init__()
+        self.conv_1 = Conv2D(filters, kernels)
+        self.conv_2 = Conv2D(filters, kernels)
+
+    def call(self, inputs):
+        x = K.relu(inputs)
+        x = self.conv_1(x)
+        x = K.relu(x)
+        x = self.conv_2(x)
+        return x + inputs
+
+
+class QueryNetwork(Model):
+
+    def __init__(self, units=64):
+        super(RecAE, self).__init__()
+        self.conv_1 = Dense(units)
+        self.conv_2 = Dense(units)
+
+    def call(self, inputs):
+        x = K.relu(inputs)
+        x = self.conv_1(x)
+        x = K.relu(x)
+        x = self.conv_2(x)
+        return x + inputs
+
+
+class RecAE(Model):
+
+    def __init__(self, head, bottle, decoder):
+        super(RecAE, self).__init__()
+        self.head = head
+        self.bottle = bottle
+        self.base = clone_model(bottle)
+        self.decoder = decoder
+        self.segmentation_network = SegmentationNetwork()
+        self.query_network = QueryNetwork()
+        self.control = LSTMCell(64)
+        self.memory = LSTMCell(64)
+
+    def call(self, inputs):
+        feature = self.head(inputs)
+        segmentation = self.segmentation_network(feature)
+        control_base = self.base(feature)
+        h_c = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
+        h_m = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
+        shape = K.shape(feature)[:-1]
+        full_attention = tf.zeros(shape)[..., tf.newaxis]
+        full_image = tf.zeros(K.shape(inputs))
+        masks = []
+        ff = tf.zeros(K.shape(inputs))
+        scope = tf.ones(shape)[..., tf.newaxis]
+        for i in range(10):
+            r_c, h_c = self.control(tf.concat([control_base, h_m[0]], 1), h_c)
+            query = self.query_network(h_c[0])
+            log_attention = image_attention(segmentation, query)
+            attention = K.softmax(log_attention)
+            mask = attention * scope
+            scope = scope - mask
+            im = feature * mask
+            # im = feature
+            latent = self.bottle(im)
+            decoded = self.decoder(latent)
+            # self.add_loss(K.mean(-mse(full_attention, attention)))
+            # self.add_loss(K.mean(-mse(tf.ones(attention.shape), attention)))
+            full_attention += attention
+            big_mask = tf.image.resize(mask, K.shape(inputs)[1:-1])
+            ff += K.sigmoid(decoded)
+            full_image += K.sigmoid(decoded) * big_mask
+            r_m, h_m = self.memory(latent, h_m)
+            masks.append(big_mask)
+        self.add_loss(K.mean(mse(inputs, full_image)))
+        return full_image, masks
+
+
+# def image_attention(image, query, scale=True):
+@tf.function
+def image_attention(image, query):
+    log_attention = K.sum(query[:, tf.newaxis, tf.newaxis, :] * image, axis=-1, keepdims=True)
+    # if scale is not None:
+    log_attention /= tf.sqrt(tf.cast(K.shape(image)[-1], dtype=float))
+    return log_attention
+
+
+class RecAE_2(Model):
+
+    def __init__(self, head, bottle, decoder):
+        super(RecAE_2, self).__init__()
+        self.head = head
+        self.bottle = bottle
+        # self.base = clone_model(bottle)
+        self.base = self.bottle
+        self.decoder = decoder
+        self.segmentation_network = create_conv_model((64, 64, 1))
+        self.control = LSTMCell(64)
+        self.memory = LSTMCell(64)
+
+    def call(self, inputs):
+        feature = self.head(inputs)
+        control_base = self.base(feature)
+        h_c = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
+        h_m = [tf.random.normal([K.shape(inputs)[0], self.control.units])] * 2
+        shape = K.shape(feature)[:-1]
+        full_attention = tf.zeros(shape)[..., tf.newaxis]
+        full_image = tf.zeros(K.shape(inputs))
+        big_masks = []
+        masks = []
+        ff = tf.zeros(K.shape(inputs))
+        scope = tf.ones(shape)[..., tf.newaxis]
+        for i in range(4):
+            if i ==3:
+                mask = scope
+            else:
+                r_c, h_c = self.control(tf.concat([control_base, h_m[0]], 1), h_c)
+                query = broadcast(h_c[0], feature.shape[1:])
+                log_attention = self.segmentation_network(tf.concat([feature, query], axis=-1))
+                attention = K.sigmoid(log_attention)
+                mask = attention * scope
+                scope = scope - mask
+            masks.append(mask)
+            im = feature * mask
+            # im = feature
+            latent = self.bottle(im)
+            decoded = self.decoder(latent)
+            # self.add_loss(K.mean(-mse(scope, mask)))
+            sum = K.sum(tf.ones(K.shape(mask)))
+            self.add_loss(K.abs((sum/4)-K.sum(mask))/sum)
+            # self.add_loss(K.mean(-mse(tf.zeros(K.shape(mask)), mask)))
+            for m in masks:
+                self.add_loss(K.mean(-mse(m,mask)))
+
+            full_attention += mask
+            big_mask = tf.image.resize(mask, K.shape(inputs)[1:-1])
+            ff += K.sigmoid(decoded)
+            full_image += K.sigmoid(decoded) * big_mask
+            r_m, h_m = self.memory(latent, h_m)
+            big_masks.append(big_mask)
+        self.add_loss(K.mean(mse(inputs, full_image)))
+        return full_image, big_masks

raven_utils/models/body.py

@@ -0,0 +1,276 @@
+import itertools
+
+import tensorflow as tf
+from ml_utils import self_product, lw
+
+from models_utils import DictModel, ListModel, Flat, bm, Base, Cat, Res, Flat2, conv, KERNEL_SIZE, FILTERS, SAME, \
+    Get, SM, bs, RELU, ACTIVATION, dense, bd, HardBlock, MaxBlock
+import models_utils.ops as K
+from models_utils import Merge, SoftBlock
+from models_utils.build import build_multi_dense, build_multi_conv, build_conv_model, build_encoder
+from tensorflow.keras.layers import Lambda, Dense
+from tensorflow.keras.layers import Conv2D
+
+from config.constant import MEMORY, CONTROL, LATENT, MERGE, CONCAT, INFERENCE, FLATTEN
+from models_utils.config import config
+
+
+class RavRes(Res):
+    def __init__(self, model="v2", latent=256, act=RELU):
+        super().__init__(model=model)
+        self.latent = latent
+
+    def call(self, inputs):
+        return self.model(inputs) + inputs[0][:, ..., self.latent:]
+
+
+# not working
+class RavResConv(Res):
+    def __init__(self, model="v2", latent=256, act=RELU):
+        super().__init__(model=model)
+        self.latent = latent
+        self.conv = conv(latent, (1, 1), activation=act)
+
+    def call(self, inputs):
+        return self.model(inputs) + self.conv(inputs[0])
+
+
+class RavResDense(Res):
+    def __init__(self, model="v2", latent=256, act=config.DEF_DENSE.activation):
+        super().__init__(model=model)
+        self.latent = latent
+        self.conv = dense(latent, activation=act)
+
+    def call(self, inputs):
+        return self.model(inputs) + self.conv(inputs[0])
+
+
+def create_dense_block(latent=256, loop=1):
+    soft_block = Res(SoftBlock(build_multi_dense(latent), add_identity=None,
+                               score_activation=tf.sigmoid), latent=latent)
+    cells = [
+        (lambda x: K.cat([x[:, 0], x[:, 1]]), LATENT, CONCAT),
+        (None, CONCAT, MEMORY),
+        (Dense(latent), CONCAT, MERGE),
+        (Merge(latent), [INFERENCE, MERGE], CONTROL),
+        (soft_block, [MEMORY, CONTROL], MEMORY)
+    ]
+
+    return ListModel([DictModel(*cell) for cell in cells] * loop, [LATENT, INFERENCE], MEMORY)
+
+
+def build_multi_conv(filters=32, end_filters=64, padding="same",mul=1, norm=None, **kwargs):
+    base = [(1, 3), (3, 1), (3, 3)]
+    block = list(self_product(base))
+    block2 = [b + b[0:1] for b in block]
+    block3 = [b + b for b in block]
+    block4 = ([[(3, 3)]] + [[(3, 3), (3, 3)]] + [[(3, 3), (3, 3), (3, 3)]]) * 2
+    block5 = [[], []]
+    all_blocks = [s for b in [block, block2, block3, block4, block5] for s in b]
+    start = {
+        FILTERS: filters,
+        KERNEL_SIZE: (1, 1)
+    }
+
+    end = {
+        FILTERS: end_filters,
+        KERNEL_SIZE: (1, 1),
+        ACTIVATION: None
+    }
+
+    all_arch = []
+    for ab in all_blocks:
+        arch = [{
+            FILTERS: filters,
+            KERNEL_SIZE: a,
+            **kwargs
+        } for a in ab]
+        all_arch.append([start] + arch + [end])
+
+    all_arch = all_arch * mul
+
+    return [
+        build_encoder(a, add_norm=norm if norm else None, padding=padding, name=f"b{i}", order=(1, 0) if norm else None)
+        for i, a in enumerate(all_arch)]
+
+
+def create_block(latent=256, simpler=0, loop=1, padding=SAME, norm=None, trans_div=2, act="pass", type_="conv",
+                 block_=SoftBlock,max_k=16,
+                 **kwargs):
+    trans_size = int(latent / trans_div)
+    # if block_ == HardBlock:
+    #     mul = 2
+    # elif block_ == MaxBlock:
+    #     mul = int(38/max_k)
+    # else:
+    #     mul = 1
+
+    if act == "pass":
+        res_class = RavRes
+    else:
+        if type_ == "dense":
+            res_class = RavResDense
+        else:
+            res_class = RavResConv
+
+    if type_ == "dense":
+        build_res = lambda: Res(model="dv2")
+        # build_reduction = lambda: bm([dense(latent), "IN"])
+        build_reduction = lambda: dense(latent)
+        build_flatten = lambda: bd([latent] * 2)
+    else:
+        build_res = lambda: Res(padding=padding)
+        build_reduction = lambda: bm([conv(trans_size if simpler else latent, 1, padding=padding), "BN"])
+        # build_reduction = lambda: bm([conv(latent, 1, padding=padding), "BN"])
+        # build_reduction = lambda: bm([conv(trans_size, 1, padding=padding), "BN"])
+        # build_reduction = lambda: conv(trans_size, 1, padding=padding)
+        # build_flatten = lambda: Flat2(filters=trans_size,res_no=2, padding=padding, units=64)
+        build_flatten = lambda: Flat2(filters=trans_size,padding=padding, units=64)
+
+    if simpler == 1:
+        cells = [
+            (lambda x: K.cat([x[:, 0], x[:, 1]]), LATENT, CONCAT,"concatenation"),
+            # (None, CONCAT, MEMORY),
+            (build_reduction(), CONCAT, MERGE,"Start_resnet_block"),
+            # (Get(), INFERENCE, INFERENCE),
+            (K.cat, [INFERENCE, MERGE], CONTROL,"concatenation"),
+        ]
+    else:
+        cells = [
+            (lambda x: K.cat([x[:, 0], x[:, 1]]), LATENT, CONCAT),
+            (build_reduction(), CONCAT, MEMORY),
+            (build_reduction(), INFERENCE, CONTROL),
+        ]
+    for i, l in enumerate(lw(loop)):
+        if l:
+            concat = K.cat
+            control_reduction = build_reduction()
+            control_res = build_res()
+            control_flatten = build_flatten()
+            if i == 0 and simpler == 1:
+                rest_params = {
+                    "latent": latent,
+                    "act": act
+                }
+            else:
+                rest_params = {
+                    "latent": 0
+                }
+
+
+            if block_ == SoftBlock:
+                block_params = {
+                }
+            else:
+                block_params = {
+                    "trans_output_shape": latent
+                }
+                if block_ == MaxBlock:
+                    block_params['max_k'] = max_k
+
+
+            # todo change name
+            soft_block = res_class(
+                block_(
+                    build_multi_dense(latent) if type_ == "dense" else build_multi_conv(trans_size, end_filters=latent,
+                                                                                        norm=norm, padding=padding,
+                                                                                        **kwargs),
+                    add_identity=None,
+                    score_activation=tf.sigmoid,
+                    **block_params
+
+                ),
+                **rest_params)
+
+        if i == 0 and simpler == 1:
+            cells.extend([
+                (control_reduction, CONTROL, CONTROL,"Reduction"),
+                (control_res, CONTROL, CONTROL,"Control_resnet_block"),
+                (control_flatten, CONTROL, FLATTEN,"Weights"),
+                (soft_block, [CONCAT, FLATTEN], MEMORY,"Transformation"),
+                # (soft_block, [MEMORY, FLATTEN], MEMORY,"Transformation"),
+            ])
+        else:
+            if l:
+                memory_res = build_res()
+
+            cells.extend([
+                (memory_res, MEMORY, MEMORY,"Memory_resnet_block"),
+                (concat, [CONTROL, MEMORY], CONTROL,"concatenation"),
+                (control_reduction, CONTROL, CONTROL,"Reduction"),
+                (control_res, CONTROL, CONTROL,"Control_resnet_block"),
+                (control_flatten, CONTROL, FLATTEN,"Weights"),
+                (soft_block, [MEMORY, FLATTEN], MEMORY, "Transformation"),
+            ])
+    return ListModel([DictModel(*cell) for cell in cells], [LATENT, INFERENCE], MEMORY, debug_=False)
+
+#
+#
+# def test(x):
+#     np.zeros(4)
+#     self_product((1, 3))
+#
+#
+# list(itertools.product())
+# u.layers[0].layers[-1].model.layers[1]
+
+# class RecurrentBodyDict(Model):
+#     # def __init__(self, start=None, cell=None, output_network=None, output_activation="tanh", latent=64, loop_no=5):
+#     def __init__(self, start=None, cell=None, output_network=None, output_activation=None, latent=64, loop_no=5):
+#         super().__init__()
+#         self.start = sm(start, lambda: SubClassingModel([StartLSTMControl(latent), StartLSTMMemory(latent)]),
+#                         latent=latent)
+#         self.cell = sm(cell, lambda: SubClassingModel([LSTMControl(latent), LSTMMemory(latent)]), latent=latent)
+#         self.output_network = sm(output_network, lf(take_memory_states))
+#         self.loop_no = loop_no
+#         # tmp
+#         self.activation = Activation(output_activation)
+#
+#     def call(self, inputs):
+#         outputs = []
+#         for j in range(3):
+#             outputs.append(self.start({"latent": inputs[0][j], "inference": inputs[1]}))
+#         for i in range(self.loop_no):
+#             for j in range(3):
+#                 outputs[j] = self.cell(outputs[j])
+#
+#         return self.activation(self.output_network(outputs))
+#
+#
+# class RecurrentBodySimpleMix4Dict(RecurrentBodyDict):
+#     def __init__(self, latent=64, output_network=None, loop_no=5):
+#         super().__init__(
+#             start=SubClassingModel(
+#                 [ConcatCell(), DenseCell(latent), InfMergeCell(latent),
+#                  WeigthCell(latent, layer_no=np.repeat([1, 2, 3, 4, 5, 6, 7, 8], 4),
+#                             add_identity=Lambda(lambda x: x[:, latent:]))]),
+#             cell=False,
+#             output_network=output_network, loop_no=0)
+# class RecurrentBodySimpleMix4Conv(RecurrentBodyDict):
+#     def __init__(self, latent=64, output_network=None, loop_no=5):
+#         super().__init__(
+#             start=SubClassingModel(
+#                 [ConcatCell(), ConvCell(latent), ReduceCell(latent), InfMergeCell(latent),
+#                  ModelCell(latent=latent, layers_no=2, input_name=CONTROL, result_name=CONTROL),
+#                  WeigthCell(latent,
+#                             transformation_network=[build_conv_model2([latent] * i, kernels=(j, j)) for i in range(1, 7) for j in
+#                                                     range(1, 5) for _ in range(1)],
+#                             add_identity=Lambda(lambda x: x[:, ..., latent:]))
+#                  ]),
+#             cell=False,
+#             output_network=output_network, loop_no=0)
+#
+#
+# class RecurrentBodySimpleMix4Conv2(RecurrentBodyDict):
+#     def __init__(self, latent=64, output_network=None, loop_no=5):
+#         super().__init__(
+#             start=SubClassingModel(
+#                 [ConcatCell(), ConvCell(latent), ReduceCell2(latent), InfMergeCell(latent),
+#                  ModelCell(latent=latent, layers_no=2, input_name=CONTROL, result_name=CONTROL),
+#                  WeigthCell(latent,
+#                             transformation_network=[bc([latent] * i, kernels=(j, j)) for i in range(1, 7) for j in
+#                                                     range(1, 5) for _ in range(1)],
+#                             add_identity=Lambda(lambda x: x[:, ..., latent:]))
+#                  ]),
+#             cell=False,
+#             output_network=output_network, loop_no=0)

raven_utils/models/class_.py

@@ -0,0 +1,31 @@
+from ml_utils import lw
+from models_utils import SubClassingModel, ops as K, Base
+import tensorflow as tf
+
+
+class Merge(SubClassingModel):
+    def call(self, inputs):
+        results = []
+        for i, model in enumerate(self.model[:-1]):
+            results.append(model(inputs[i]))
+        # todo why K.cat not working
+        results = self.model[-1](tf.concat(results, axis=-1))
+        return results
+
+
+class RavenClass(Base):
+    def __init__(self, model, scales=None, no=3, name=None):
+        super().__init__(model=model, name=name)
+        self.scales = scales
+        self.no = no
+
+    def call(self, inputs):
+        inputs = lw(inputs)
+        class_res = []
+        # for i in range(inputs[0].shape[1]):
+        for i in range(self.no):
+            # d = [r[:, i] if r.ndim == 5 else r for r in inputs]
+            d = [inputs[s][:, i] if inputs[s].ndim > 2 else inputs for s in self.scales]
+            class_res.append(self.model(d))
+        # return tf.stack(class_res,axis=1)
+        return [class_res]

raven_utils/models/head.py

@@ -0,0 +1,159 @@
+import tensorflow as tf
+from ml_utils import set_default
+from models_utils import build_dense_model, bm, ActivationModel, sm, large_conv_dense_encoder, Pass
+from models_utils import res
+from tensorflow.keras import Model
+from models_utils import ops as K
+from tensorflow.keras.layers import Dense, Conv2D, Flatten
+from keras.backend import batch_flatten
+
+
+# todo Refactoring
+class HeadModel(Model):
+    def __init__(self, encoder=None, inference_network=None, output_size=64, inference_output_size=None,
+                 inference_activation="relu", stem=None, images_no=8, inference_image_no=None):
+        super().__init__()
+        # self.encoder = sm(encoder, bm([en.large_conv_dense_encoder(), Dense(output_size)], False))
+        self.encoder = encoder or bm([large_conv_dense_encoder(), Dense(output_size)])
+        # self.head = head or HeadBatch(encoder=encoder, output_size=output_size)
+        inference_output_size = inference_output_size or output_size
+        self.inference_network = inference_network or bm([
+            K.flat,
+            build_dense_model([1028, 512, 512, inference_output_size],
+                              last_activation=inference_activation)]
+        )
+        self.stem = stem or Pass()
+        self.images_no = images_no
+        self.inference_image_no = self.images_no if inference_image_no is None else inference_image_no
+
+
+class LatentHeadModel(HeadModel):
+    def call(self, inputs):
+        result = K.map_batch(inputs[:, :self.images_no], self.encoder)
+        inference = self.inference_network(result[:, :self.inference_image_no])
+        latents = self.stem(result)
+        return [latents, inference,result]
+
+
+# # todo use map_batch
+# class HeadBatch(Model):
+#     def __init__(self, encoder=None, output_size=64):
+#         super().__init__()
+#         self.encoder = sm(encoder, bm([large_conv_dense_encoder(), Dense(output_size)], False))
+#
+#     def call(self, inputs):
+#         shape = tf.shape(inputs)
+#         latents = self.encoder(tf.reshape(inputs, shape=tf.concat([[-1], shape[2:]], axis=-1)))
+#         latents = K.reshape(latents, tf.concat([[-1, shape[1]], latents.shape[1:]], axis=-1))
+#         return latents
+
+
+# Not working
+class DuoHeadModel(HeadModel):
+    def __init__(self, encoder=None, inference_network=None, images_no=8, filters=-4):
+        super().__init__(encoder=encoder, inference_network=inference_network, images_no=images_no)
+        self.encoder = ActivationModel(self.encoder, filters=filters, include_input=False)
+
+    def call(self, inputs):
+        shape = inputs.shape
+        result = reversed(self.encoder(K.reshape(inputs, shape=[-1] + list(shape[2:]))))
+        latents = K.reshape(result[0], [-1, self.images_no] + [result[0].shape[-1]])
+        inference = self.inference_network(K.flat(result[1]))
+        return [latents, inference]
+
+
+class MultiHeadModel(Model):
+    def __init__(self, encoder=None, images_no=8, filters=(1, 3, 6)):
+        super().__init__()
+        self.encoder = ActivationModel(encoder, filters=filters, include_input=False)
+        self.merge = MergeSacles()
+        self.images_no = images_no
+
+    def call(self, inputs):
+        shape = tf.shape(inputs)
+        results = self.encoder(tf.reshape(inputs, shape=tf.concat([[-1], shape[2:]], axis=-1)))
+        latents = [tf.reshape(result, shape=tf.concat([[-1, self.images_no], tf.shape(result)[1:]], axis=-1)) for result
+                   in results]
+
+        l1 = tf.transpose(latents[0], (0, 2, 3, 1, 4))
+        # l1 = tf.reshape(l1, tuple(list(l1.shape[:3]) + [l1.shape[-2] * l1.shape[-1]]))
+        shape = tf.shape(l1)
+        l1 = tf.reshape(l1, tf.concat([[-1], shape[1:3], [shape[-2] * shape[-1]]], axis=-1))
+
+        l2 = tf.transpose(latents[1], (0, 2, 3, 1, 4))
+        # l2 = tf.reshape(l2, [-1] + list(l2.shape[1:3]) + [l2.shape[-2] * l2.shape[-1]])
+        shape = tf.shape(l2)
+        l2 = tf.reshape(l2, tf.concat([[-1], shape[1:3], [shape[-2] * shape[-1]]], axis=-1))
+
+        l3 = latents[2]
+        shape = tf.shape(l3)
+        # l3 = tf.reshape(l3, [-1] + [l3.shape[-2] * l3.shape[-1]])
+        l3 = tf.reshape(l3, tf.concat([[-1], [shape[-2] * shape[-1]]], axis=-1))
+
+        inference = self.merge([l1, l2, l3])
+        return [latents, inference]
+
+
+class MergeSacles(Model):
+    def __init__(self):
+        super().__init__()
+        self.inf_1 = bm([Conv2D(64, 1, activation="relu"), res(64),
+                         Conv2D(64, 3, strides=2, padding=SAME, activation="relu"),
+                         res(64),
+                         Flatten(),
+                         Dense(256, "relu")])
+        self.inf_2 = bm([Conv2D(128, 1, activation="relu"),
+                         res(128),
+                         Flatten(),
+                         Dense(256, "relu")])
+        self.inf_3 = Dense(256, "relu")
+
+    def call(self, inputs):
+        il1 = self.inf_1(inputs[0])
+        il2 = self.inf_2(inputs[1])
+        il3 = self.inf_3(inputs[2])
+        inference = tf.concat([il1, il2, il3], axis=1)
+        return inference
+
+
+class MultiHeadModel2(Model):
+    def __init__(self, encoder=None, images_no=8, filters=(3, 6)):
+        super().__init__()
+        self.encoder = ActivationModel(encoder, filters=filters, include_input=False)
+        self.merge = MergeSacles2()
+        self.images_no = images_no
+
+    def call(self, inputs):
+        shape = tf.shape(inputs)
+        results = self.encoder(tf.reshape(inputs, shape=tf.concat([[-1], shape[2:]], axis=-1)))
+        latents = [tf.reshape(result, shape=tf.concat([[-1, self.images_no], tf.shape(result)[1:]], axis=-1)) for result
+                   in results]
+
+        l2 = tf.transpose(latents[0], (0, 2, 3, 1, 4))
+        # l2 = tf.reshape(l2, [-1] + list(l2.shape[1:3]) + [l2.shape[-2] * l2.shape[-1]])
+        shape = tf.shape(l2)
+        l2 = tf.reshape(l2, tf.concat([[-1], shape[1:3], [shape[-2] * shape[-1]]], axis=-1))
+
+        l3 = latents[1]
+        shape = tf.shape(l3)
+        # l3 = tf.reshape(l3, [-1] + [l3.shape[-2] * l3.shape[-1]])
+        l3 = tf.reshape(l3, tf.concat([[-1], [shape[-2] * shape[-1]]], axis=-1))
+
+        inference = self.merge([l2, l3])
+        return [latents, inference]
+
+
+class MergeSacles2(Model):
+    def __init__(self):
+        super().__init__()
+        self.inf_1 = bm([Conv2D(128, 1, activation="relu"),
+                         res(128),
+                         Flatten(),
+                         Dense(256, "relu")])
+        self.inf_2 = Dense(256, "relu")
+
+    def call(self, inputs):
+        il1 = self.inf_1(inputs[0])
+        il2 = self.inf_2(inputs[1])
+        inference = tf.concat([il1, il2], axis=1)
+        return inference

raven_utils/models/loss.py

@@ -0,0 +1,630 @@
+from functools import partial
+
+import tensorflow as tf
+import tensorflow.experimental.numpy as tnp
+from models_utils import OUTPUT, TARGET, PREDICT, DictModel, add_loss, LOSS, Predict
+from models_utils import SubClassingModel
+from models_utils.models.utils import interleave
+from models_utils.op import reshape
+from tensorflow.keras import Model
+# from tensorflow.keras import backend as K
+from tensorflow.keras.layers import Lambda
+from tensorflow.keras.losses import SparseCategoricalCrossentropy, mse
+from tensorflow.keras.metrics import SparseCategoricalAccuracy, Accuracy, BinaryAccuracy
+import models_utils.ops as K
+
+import raven_utils.decode
+import raven_utils as rv
+from raven_utils.config.constant import LABELS, INDEX, ACC_SAME, ACC_CHOOSE_LOWER, ACC_CHOOSE_UPPER, CLASSIFICATION, \
+    SLOT, \
+    PROPERTIES, ACC, GROUP, NUMBER, MASK
+from raven_utils.models.uitls_ import RangeMask
+from raven_utils.const import VERTICAL, HORIZONTAL
+
+
+def get_properties_mask(target):
+    return target[:, rv.target.END_INDEX:rv.target.UNIFORMITY_INDEX] > 0
+
+
+def create_change_mask(target):
+    properties_mask = get_properties_mask(target)
+    return [create_mask(properties_mask, i) for i, _ in enumerate(rv.rules.ATTRIBUTES)]
+
+
+def create_uniform_mask(target):
+    u_mask = lambda i: tf.tile(target[:, rv.target.UNIFORMITY_INDEX + i, None] == 3, [1, rv.rules.ATTRIBUTES_LEN])
+    properties_mask = tf.concat([u_mask(0), u_mask(1)], axis=-1) | get_properties_mask(target)
+    return [create_mask(properties_mask, i) for i, _ in enumerate(rv.rules.ATTRIBUTES)]
+
+
+def create_all_mask(target):
+    return [
+        tf.cast(tf.ones(tf.stack([tf.shape(target)[0], rv.entity.SUM])), dtype=tf.bool) for i, _ in
+        enumerate(rv.rules.ATTRIBUTES)]
+
+
+class BaselineClassificationLossModel(Model):
+    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True):
+        super().__init__()
+        self.predict_fn = SubClassingModel([lambda x: x[0], PredictModel()])
+        self.loss_fn = ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
+                                       group_loss=group_loss)
+        self.metric_fn = SimilarityRaven(mode=mode)
+
+    def call(self, inputs):
+        losses = []
+        output = inputs[1]
+        losses.append(self.loss_fn([inputs[0][0], output]))
+        losses.append(self.metric_fn([inputs[0][2], inputs[3][0], inputs[0][1][:, 8:]]))
+        return losses
+
+
+class RavenLoss(Model):
+    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True, lw=(1.0, 0.3),
+                 classification=False, trans=True, anneal=False):
+        super().__init__()
+        if anneal:
+            self.weight_scheduler
+        self.classification = classification
+        self.trans = trans
+        self.predict_fn = DictModel(SubClassingModel([lambda x: x[-1], PredictModel()]), in_=OUTPUT,
+                                    out=[PREDICT, MASK], name="pred")
+        if self.trans:
+            self.loss_fn = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
+                                                    group_loss=group_loss, enable_metrics=False, lw=lw[0]),
+                                    name="main_loss")
+            self.loss_fn_2 = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
+                                                      group_loss=group_loss), name="add_loss")
+            self.metric_fn = SimilarityRaven(mode=mode)
+        if self.classification:
+            self.loss_fn_3 = add_loss(
+                ClassRavenModel(mode=create_all_mask, number_loss=number_loss, slot_loss=slot_loss,
+                                group_loss=group_loss, enable_metrics="c" if self.trans else True), lw=lw[1],
+                name="class_loss")
+
+    def call(self, inputs):
+        losses = []
+        output = inputs[OUTPUT]
+        target = inputs[TARGET]
+        labels = inputs[LABELS]
+
+        if self.trans:
+            losses.append(self.loss_fn([labels[:, 2], output[0]]))
+            losses.append(self.loss_fn([labels[:, 5], output[1]]))
+            losses.append(self.loss_fn_2([target, output[2]]))
+            losses.append(self.metric_fn([inputs[INDEX], inputs[PREDICT], labels]))
+        if self.classification:
+            for i in range(8):
+                losses.append(self.loss_fn_3([labels[:, i], inputs[CLASSIFICATION][i]]))
+        return {**inputs, LOSS: losses}
+
+
+class VTRavenLoss(Model):
+    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True, lw=(1.0, 0.1),
+                 classification=False, trans=True, anneal=False, plw=None):
+        super().__init__()
+        if anneal:
+            self.weight_scheduler
+        self.classification = classification
+        self.trans = trans
+        self.predict_fn = DictModel(SubClassingModel([lambda x: x[:, -1], PredictModel()]), in_=OUTPUT,
+                                    out=[PREDICT, MASK], name="pred")
+        self.loss_fn = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
+                                                group_loss=group_loss, plw=plw), lw=lw[0] , name="add_loss")
+        self.metric_fn = SimilarityRaven(mode=mode)
+        if self.classification:
+            self.loss_fn_2 = add_loss(
+                ClassRavenModel(mode=create_all_mask, number_loss=number_loss, slot_loss=slot_loss,
+                                group_loss=group_loss, enable_metrics="c", plw=plw), lw=lw[1],  name="class_loss")
+
+    def call(self, inputs):
+        losses = []
+        output = inputs[OUTPUT]
+        target = inputs[TARGET]
+        labels = inputs[LABELS]
+
+        for i in range(9):
+            losses.append(self.loss_fn_2([labels[:, i], output[:, i]]))
+        losses.append(self.loss_fn([target, output[:, 8]]))
+        losses.append(self.metric_fn([inputs[INDEX], inputs[PREDICT], labels]))
+        return {**inputs, LOSS: losses}
+
+
+class SingleVTRavenLoss(Model):
+    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True, lw=(1.0, 0.1),
+                 classification=False, trans=True, anneal=False):
+        super().__init__()
+        if anneal:
+            self.weight_scheduler
+        self.classification = classification
+        self.trans = trans
+        self.predict_fn = DictModel(PredictModel(), in_=OUTPUT, out=[PREDICT, MASK], name="pred")
+        self.loss_fn = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
+                                                group_loss=group_loss), lw=lw[0], name="add_loss")
+        self.metric_fn = SimilarityRaven(mode=mode)
+
+    def call(self, inputs):
+        losses = []
+        output = inputs[OUTPUT]
+        target = inputs[TARGET]
+        labels = inputs[LABELS]
+
+        losses.append(self.loss_fn([target, output]))
+        losses.append(self.metric_fn([inputs[INDEX], inputs[PREDICT], labels]))
+        return {**inputs, LOSS: losses}
+
+
+class ClassRavenModel(Model):
+    def __init__(self, mode=create_all_mask,plw=None, number_loss=False, slot_loss=True, group_loss=True, enable_metrics=True,
+                 lw=1.0):
+        super().__init__()
+        self.number_loss = number_loss
+        self.group_loss = group_loss
+        self.enable_metrics = enable_metrics
+        self.slot_loss = slot_loss
+        self.predict_fn = PredictModel()
+        self.loss_fn = SparseCategoricalCrossentropy(from_logits=True)
+        if self.slot_loss:
+            self.loss_fn_2 = tf.nn.sigmoid_cross_entropy_with_logits
+        if self.enable_metrics:
+            self.enable_metrics = f"{self.enable_metrics}_" if isinstance(self.enable_metrics, str) else ""
+            self.metric_fn = [
+                SparseCategoricalAccuracy(name=f"{self.enable_metrics}{ACC}_{property_}") for property_ in
+                rv.properties.NAMES]
+            if self.group_loss:
+                self.metric_fn_group = SparseCategoricalAccuracy(name=f"{self.enable_metrics}{ACC}_{GROUP}")
+            if self.slot_loss:
+                self.metric_fn_2 = BinaryAccuracy(name=f"{self.enable_metrics}{ACC}_{SLOT}")
+        self.range_mask = RangeMask()
+        self.mode = mode
+        self.lw = lw
+        if not plw:
+            plw = [1., 95.37352927, 2.83426987, 0.85212836, 1.096005, 1.21943385]
+        elif isinstance(plw, int) or isinstance(plw, float):
+            plw = [1., plw, 2.83426987, 0.85212836, 1.096005, 1.21943385]
+            # plw = [plw] * 6
+        self.plw = plw
+
+    # self.predict_fn = partial(tf.argmax, axis=-1)
+
+    def call(self, inputs):
+        losses = []
+        metrics = {}
+        target = inputs[0]
+        output = inputs[1]
+
+        target_group, target_slot, target_all = raven_utils.decode.decode_target(target)
+
+        group_output, output_slot, outputs = raven_utils.decode.output_divide(output, split_fn=tf.split)
+
+        # group
+        if self.group_loss:
+            group_loss = self.lw * self.plw[0] *  self.loss_fn(target_group, group_output)
+            losses.append(group_loss)
+
+            if isinstance(self.enable_metrics, str):
+                group_metric = self.metric_fn_group(target_group, group_output)
+                # metrics[GROUP] = group_metric
+                self.add_metric(group_metric)
+                self.add_metric(tf.reduce_sum(group_metric), f"{self.enable_metrics}{ACC}")
+
+        # setting uniformity mask
+        full_properties_musks = self.mode(target)
+
+        range_mask = self.range_mask(target_group)
+
+        if self.slot_loss:
+            # number
+            number_mask = range_mask & full_properties_musks[0]
+            number_mask = tf.cast(number_mask, tf.float32)
+            target_number = tf.reduce_sum(
+                tf.cast(target_slot, "float32") * number_mask, axis=-1)
+            output_number = tf.reduce_sum(
+                tf.cast(tf.sigmoid(output_slot) >= 0.5, "float32") * number_mask, axis=-1)
+
+            # output_number = tf.reduce_sum(tf.sigmoid(output_slot) * number_mask, axis=-1)
+            if self.number_loss:
+                scale = 1 / 9
+                if self.number_loss == 2:
+                    output_number_2 = tf.reduce_sum(tf.sigmoid(output_slot) * number_mask, axis=-1)
+                else:
+                    output_number_2 = output_number
+                number_loss = self.lw * self.plw[1] * mse(tf.stop_gradient(target_number) * scale, output_number_2 * scale)
+                losses.append(number_loss)
+
+            # metrics[NUMBER] = number_acc
+
+            if isinstance(self.enable_metrics, str):
+                number_acc = tf.reduce_mean(
+                    tf.cast(tf.cast(target_number, "int8") == tf.cast(output_number, "int8"), "float32"))
+                self.add_metric(tf.reduce_sum(number_acc), f"{self.enable_metrics}{ACC}_{NUMBER}")
+                self.add_metric(tf.reduce_sum(number_acc), f"{self.enable_metrics}{ACC}")
+                self.add_metric(tf.reduce_sum(number_acc), f"{self.enable_metrics}{ACC}_NO_{GROUP}")
+
+            # position/slot
+            slot_mask = range_mask & full_properties_musks[1]
+            # tf.boolean_mask(target_slot,slot_mask)
+
+            if tf.reduce_any(slot_mask):
+                # if tf.reduce_mean(tf.cast(slot_mask, dtype=tf.int32)) > 0:
+                target_slot_masked = tf.boolean_mask(target_slot, slot_mask)[:, None]
+                output_slot_masked = tf.boolean_mask(output_slot, slot_mask)[:, None]
+                loss_slot = self.lw * self.plw[2] * tf.reduce_mean(
+                    self.loss_fn_2(tf.cast(target_slot_masked, "float32"), output_slot_masked))
+                if isinstance(self.enable_metrics, str):
+                    acc_slot = self.metric_fn_2(target_slot_masked, output_slot_masked)
+                    self.add_metric(acc_slot)
+                    self.add_metric(tf.reduce_sum(acc_slot), f"{self.enable_metrics}{ACC}")
+                    self.add_metric(tf.reduce_sum(acc_slot), f"{self.enable_metrics}{ACC}_NO_{GROUP}")
+            else:
+                loss_slot = 0.0
+                acc_slot = -1.0
+
+            losses.append(loss_slot)
+            # metrics[SLOT] = acc_slot
+        # if loss_slot != 0:
+
+        # if tf.reduce_any(slot_mask):
+
+        # self.add_metric(acc_slot, f"{self.enable_metrics}{ACC}_{NUMBER}")
+        # self.add_metric(acc_slot, f"{self.enable_metrics}{ACC}")
+        # self.add_metric(acc_slot, f"{self.enable_metrics}{ACC}_NO_{GROUP}")
+
+        # properties
+        for i, out in enumerate(outputs):
+            shape = (-1, rv.entity.SUM, rv.properties.RAW_SIZE[i])
+            out_reshaped = tf.reshape(out, shape)
+            properties_mask = tf.cast(target_slot, "bool") & full_properties_musks[i + 2]
+
+            if tf.reduce_any(properties_mask):
+                out_masked = tf.boolean_mask(out_reshaped, properties_mask)
+                out_target = tf.boolean_mask(target_all[i], properties_mask)
+                loss = self.lw * self.plw[3+i] * self.loss_fn(out_target, out_masked)
+                if isinstance(self.enable_metrics, str):
+                    metric = self.metric_fn[i](out_target, out_masked)
+                    self.add_metric(metric)
+                    # self.add_metric(metric, f"{self.enable_metrics}{ACC}")
+                    self.add_metric(tf.reduce_sum(metric), f"{self.enable_metrics}{ACC}")
+                    self.add_metric(tf.reduce_sum(metric), f"{self.enable_metrics}{ACC}_{PROPERTIES}")
+                    self.add_metric(tf.reduce_sum(metric), f"{self.enable_metrics}{ACC}_NO_{GROUP}")
+            else:
+                loss = 0.0
+                metric = -1.0
+
+            losses.append(loss)
+        return losses
+
+
+class FullMask(Model):
+    def __init__(self, mode=create_uniform_mask):
+        super().__init__()
+        self.range_mask = RangeMask()
+        self.mode = mode
+
+    def call(self, inputs):
+        target_group, target_slot, _ = raven_utils.decode.decode_target(inputs)
+        full_properties_musks = self.mode(inputs)
+        range_mask = self.range_mask(target_group)
+
+        number_mask = range_mask & full_properties_musks[0]
+
+        slot_mask = range_mask & full_properties_musks[1]
+        properties_mask = []
+        for property_mask in full_properties_musks[2:]:
+            properties_mask.append(tf.cast(target_slot, "bool") & property_mask)
+        return [slot_mask, properties_mask, number_mask]
+
+
+def create_mask(rules, i):
+    mask_1 = tf.tile(rules[:, i][None], [len(rv.target.FIRST_LAYOUT), 1])
+    mask_2 = tf.tile(rules[:, i + 5][None], [len(rv.target.SECOND_LAYOUT), 1])
+    shape = tf.shape(rules)
+    full_mask_1 = tf.scatter_nd(tnp.array(rv.target.FIRST_LAYOUT)[:, None], mask_1, shape=(rv.entity.SUM, shape[0]))
+    full_mask_2 = tf.tensor_scatter_nd_update(full_mask_1, tnp.array(rv.target.SECOND_LAYOUT)[:, None], mask_2)
+    return tf.transpose(full_mask_2)
+
+
+# class PredictModel(Model):
+#     def __init__(self):
+#         super().__init__()
+#         self.predict_fn = Lambda(partial(tf.argmax, axis=-1))
+#         self.predict_fn_2 = Lambda(lambda x: tf.sigmoid(x) > 0.5)
+#         self.range_mask = RangeMask()
+#
+#     # self.predict_fn = partial(tf.argmax, axis=-1)
+#
+#     def call(self, inputs):
+#         group_output = inputs[rv.OUTPUT_GROUP_SLICE]
+#         group_loss = self.predict_fn(group_output)[:, None]
+#
+#         output_slot = inputs[rv.OUTPUT_SLOT_SLICE]
+#         range_mask = self.range_mask(group_loss[:, 0])
+#         loss_slot = tf.cast(self.predict_fn_2(output_slot), dtype=tf.int64)
+#
+#         properties_output = inputs[rv.OUTPUT_PROPERTIES_SLICE]
+#         properties = []
+#         outputs = tf.split(properties_output, list(rv.ENTITY_PROPERTIES_INDEX.values()), axis=-1)
+#         for i, out in enumerate(outputs):
+#             shape = (-1, rv.ENTITY_SUM, rv.ENTITY_PROPERTIES_VALUES[i])
+#             out_reshaped = tf.reshape(out, shape)
+#             properties.append(self.predict_fn(out_reshaped))
+#         number_loss = tf.reduce_sum(loss_slot, axis=-1, keepdims=True)
+#
+#         result = tf.concat([group_loss, loss_slot, interleave(properties), number_loss], axis=-1)
+#
+#         return [result, range_mask, range_mask, range_mask, range_mask]
+
+class PredictModel(Model):
+    def __init__(self):
+        super().__init__()
+        self.predict_fn = Predict()
+        self.predict_fn_2 = Lambda(lambda x: tf.sigmoid(x) > 0.5)
+        self.range_mask = RangeMask()
+
+    # self.predict_fn = partial(tf.argmax, axis=-1)
+
+    def call(self, inputs):
+        group_output, output_slot, *properties = rv.decode.output(inputs, tf.split, self.predict_fn, self.predict_fn_2)
+        number_loss = K.int64(K.sum(output_slot))
+        result = tf.concat(
+            [group_output[:, None], tf.cast(output_slot, dtype=tf.int64), interleave(properties), number_loss[:, None]],
+            axis=-1)
+
+        range_mask = self.range_mask(group_output)
+        return [result, range_mask]
+        # return [result, range_mask, range_mask, range_mask, range_mask]
+
+
+# todo change slices
+class PredictModelMasked(Model):
+    def __init__(self):
+        super().__init__()
+        self.predict_fn = Lambda(partial(tf.argmax, axis=-1))
+        self.loss_fn_2 = Lambda(lambda x: tf.sigmoid(x) > 0.5)
+        self.range_mask = RangeMask()
+
+    # self.predict_fn = partial(tf.argmax, axis=-1)
+
+    def call(self, inputs):
+        group_output = inputs[:, -rv.GROUPS_NO:]
+        group_loss = self.predict_fn(group_output)[:, None]
+
+        output_slot = inputs[:, :rv.ENTITY_SUM]
+        range_mask = self.range_mask(group_loss[:, 0])
+        loss_slot = tf.cast(self.predict_fn_2(output_slot * range_mask), dtype=tf.int64)
+
+        properties_output = inputs[:, rv.ENTITY_SUM:-rv.GROUPS_NO]
+
+        properties = []
+        outputs = tf.split(properties_output, list(rv.ENTITY_PROPERTIES_INDEX.values()), axis=-1)
+        for i, out in enumerate(outputs):
+            shape = (-1, rv.ENTITY_SUM, rv.ENTITY_PROPERTIES_VALUES[i])
+            out_reshaped = tf.reshape(out, shape)
+            out_masked = out_reshaped * loss_slot[..., None]
+            properties.append(self.predict_fn(out_masked))
+            # out_masked[0].numpy()
+        number_loss = tf.reduce_sum(loss_slot, axis=-1, keepdims=True)
+
+        result = tf.concat([group_loss, loss_slot, interleave(properties), number_loss], axis=-1)
+
+        return result
+
+
+def final_predict_mask(x, mask):
+    r = reshape(x[0][:, rv.INDEX[0]:-1], [-1, 3])
+    return tf.ragged.boolean_mask(r, mask)
+
+
+def final_predict(x, mode=False):
+    m = x[1] if mode else tf.cast(x[0][:, 1:rv.INDEX[0]], tf.bool)
+    return final_predict_mask(x[0], m)
+
+
+def final_predict_2(x):
+    ones = tf.cast(tf.ones(tf.shape(x[0])[0]), tf.bool)[:, None]
+    mask = tf.concat([ones, tf.tile(x[1], [1, 4]), ones], axis=-1)
+    return tf.ragged.boolean_mask(x[0], mask)
+
+
+class PredictModelOld(Model):
+
+    def call(self, inputs):
+        output = inputs[-2]
+
+        rest_output = output[:, :-rv.GROUPS_NO]
+
+        result_all = []
+        outputs = tf.split(rest_output, list(rv.ENTITY_PROPERTIES_INDEX.values()), axis=-3)
+        for i, out in enumerate(outputs):
+            shape = (-3, rv.ENTITY_SUM, rv.ENTITY_PROPERTIES_VALUES[i])
+            out_reshaped = tf.reshape(out, shape)
+
+            result = tf.cast(tf.argmax(out_reshaped, axis=-3), dtype="int8")
+            result_all.append(result)
+
+        result_all = interleave(result_all)
+        return result_all
+
+
+def get_matches(diff, target_index):
+    diff_sum = K.sum(diff)
+    db_argsort = tf.argsort(diff_sum, axis=-1)
+    db_sorted = tf.sort(diff_sum)
+    db_mask = db_sorted[:, 0, None] == db_sorted
+    db_same = tf.where(db_mask, db_argsort, -1 * tf.ones_like(db_argsort))
+    matched_index = db_same == target_index
+    # setting shape needed for TensorFlow graph
+    matched_index.set_shape(db_same.shape)
+    matches = K.any(matched_index)
+    more_matches = K.sum(db_mask) > 1
+    once_matches = K.sum(matches & tf.math.logical_not(more_matches))
+    return matches, more_matches, once_matches
+
+
+class SimilarityRaven(Model):
+    def __init__(self, mode=create_all_mask, number_loss=False):
+        super().__init__()
+        self.range_mask = RangeMask()
+        self.mode = mode
+
+    # self.predict_fn = partial(tf.argmax, axis=-1)
+
+    # INDEX, PREDICT, LABELS
+    def call(self, inputs):
+        metrics = []
+        target_index = inputs[0] - 8
+        predict = inputs[1]
+        answers = inputs[2][:, 8:]
+        shape = tf.shape(predict)
+
+        target = K.gather(answers, target_index[:, 0])
+
+        target_group = target[:, 0]
+
+        # comp_slice = np.
+        target_comp = target[:, 1:rv.target.END_INDEX]
+        predict_comp = predict[:, 1:rv.target.END_INDEX]
+        answers_comp = answers[:, :, 1:rv.target.END_INDEX]
+
+        full_properties_musks = self.mode(target)
+        fpm = K.cat([full_properties_musks[0], interleave(full_properties_musks[2:])])
+
+        range_mask = self.range_mask(target_group)
+        full_range_mask = K.cat([range_mask, tf.repeat(range_mask, 3, axis=-1)], axis=-1)
+
+        final_mask = fpm & full_range_mask
+
+        target_masked = target_comp * final_mask
+        predict_masked = predict_comp * final_mask
+        answers_masked = answers_comp * tf.tile(final_mask[:, None], [1, 8, 1])
+
+        acc_same = K.mean(K.all(target_masked == predict_masked))
+        self.add_metric(acc_same, ACC_SAME)
+        metrics.append(acc_same)
+
+        diff = tf.abs(predict_masked[:, None] - answers_masked)
+        diff_bool = diff != 0
+
+        matches, more_matches, once_matches = get_matches(tf.cast(diff_bool, dtype=tf.int32), target_index)
+
+        second_phase_mask = (more_matches & matches)
+        diff_second_phase = tf.boolean_mask(diff, second_phase_mask)
+        target_index_2 = tf.boolean_mask(target_index, second_phase_mask, axis=0)
+
+        matches_2, more_matches_2, once_matches_2 = get_matches(diff_second_phase, target_index_2)
+        matches_2_no = K.sum(matches_2)
+
+        acc_choose_upper = (once_matches + matches_2_no) / shape[0]
+        self.add_metric(acc_choose_upper, ACC_CHOOSE_UPPER)
+        metrics.append(acc_choose_upper)
+
+        acc_choose_lower = (once_matches + once_matches_2) / shape[0]
+        self.add_metric(acc_choose_lower, ACC_CHOOSE_LOWER)
+        metrics.append(acc_choose_lower)
+
+        return metrics
+
+
+class SimilarityRaven2(Model):
+    def __init__(self, mode=create_all_mask, number_loss=False):
+        super().__init__()
+        self.range_mask = RangeMask()
+        self.mode = mode
+
+    # self.predict_fn = partial(tf.argmax, axis=-1)
+
+    # INDEX, PREDICT, LABELS
+    def call(self, inputs):
+        metrics = []
+        target_index = inputs[0] - 8
+        predict = inputs[1]
+        answers = inputs[2][:, 8:]
+        shape = tf.shape(predict)
+
+        target = K.gather(answers, target_index[:, 0])
+
+        target_group = target[:, 0]
+
+        # comp_slice = np.
+        target_comp = target[:, 1:rv.target.END_INDEX]
+        predict_comp = predict[:, 1:rv.target.END_INDEX]
+        answers_comp = answers[:, :, 1:rv.target.END_INDEX]
+
+        full_properties_musks = self.mode(target)
+        fpm = K.cat([full_properties_musks[0], interleave(full_properties_musks[2:])])
+
+        range_mask = self.range_mask(target_group)
+        full_range_mask = K.cat([range_mask, tf.repeat(range_mask, 3, axis=-1)], axis=-1)
+
+        final_mask = fpm & full_range_mask
+
+        target_masked = target_comp * final_mask
+        predict_masked = predict_comp * final_mask
+        answers_masked = answers_comp * tf.tile(final_mask[:, None], [1, 8, 1])
+
+        acc_same = K.mean(K.all(target_masked == predict_masked))
+        self.add_metric(acc_same, ACC_SAME)
+        metrics.append(acc_same)
+
+        diff = tf.abs(predict_masked[:, None] - answers_masked)
+        diff_bool = diff != 0
+
+        matches, more_matches, once_matches = get_matches(tf.cast(diff_bool, dtype=tf.int32), target_index)
+
+        second_phase_mask = (more_matches & matches)
+        diff_second_phase = tf.boolean_mask(diff, second_phase_mask)
+        target_index_2 = tf.boolean_mask(target_index, second_phase_mask, axis=0)
+
+        matches_2, more_matches_2, once_matches_2 = get_matches(diff_second_phase, target_index_2)
+        matches_2_no = K.sum(matches_2)
+
+        acc_choose_upper = (once_matches + matches_2_no) / shape[0]
+        self.add_metric(acc_choose_upper, ACC_CHOOSE_UPPER)
+        metrics.append(acc_choose_upper)
+
+        acc_choose_lower = (once_matches + once_matches_2) / shape[0]
+        self.add_metric(acc_choose_lower, ACC_CHOOSE_LOWER)
+        metrics.append(acc_choose_lower)
+
+        metrics.append(K.sum(target_masked != predict_masked))
+
+        return metrics
+
+
+class LatentLossModel(Model):
+    def __init__(self, dir_=HORIZONTAL):
+        super().__init__()
+        # self.sum_metrics = []
+        # for i in range(8):
+        #     self.sum_metrics.append(Sum(name=f"no_{i}"))
+        self.metric_fn = Accuracy(name="acc_latent")
+        if dir_ == VERTICAL:
+            self.dir = (6, 7)
+        else:
+            self.dir = (2, 5)
+
+    def call(self, inputs):
+        target_image = tf.reshape(inputs[0][2], [-1])
+        output = inputs[1]
+        latents = tnp.asarray(inputs[2])
+
+        target_hor = tf.concat([
+            latents[:, self.dir],
+            latents[tf.range(latents.shape[0]), target_image + 8][:, None]
+        ],
+            axis=1)
+
+        loss_hor = mse(K.stop_gradient(target_hor), output)
+        self.add_loss(loss_hor)
+
+        self.add_metric(self.metric_fn(inputs[3], target_image))
+
+        return loss_hor
+
+
+class PredRav(Model):
+
+    def call(self, inputs):
+        output = inputs[0][:, -1]
+        answers = inputs[1][:, 8:]
+        return tf.argmin(tf.reduce_sum(tf.abs(output[:, None] - answers), axis=-1), axis=-1)

raven_utils/models/loss_3.py

@@ -0,0 +1,638 @@
+from functools import partial
+
+import tensorflow as tf
+import tensorflow.experimental.numpy as tnp
+from models_utils import OUTPUT, TARGET, PREDICT, DictModel, add_loss, LOSS, Predict
+from models_utils import SubClassingModel
+from models_utils.models.utils import interleave
+from models_utils.op import reshape
+from tensorflow.keras import Model
+# from tensorflow.keras import backend as K
+from tensorflow.keras.layers import Lambda
+from tensorflow.keras.losses import SparseCategoricalCrossentropy, mse
+from tensorflow.keras.metrics import SparseCategoricalAccuracy, Accuracy, BinaryAccuracy
+import models_utils.ops as K
+
+import raven_utils.decode
+import raven_utils as rv
+from raven_utils.config.constant import LABELS, INDEX, ACC_SAME, ACC_CHOOSE_LOWER, ACC_CHOOSE_UPPER, CLASSIFICATION, \
+    SLOT, \
+    PROPERTIES, ACC, GROUP, NUMBER, MASK
+from raven_utils.models.uitls_ import RangeMask
+from raven_utils.const import VERTICAL, HORIZONTAL
+
+
+def get_properties_mask(target):
+    return target[:, rv.target.END_INDEX:rv.target.UNIFORMITY_INDEX] > 0
+
+
+def create_change_mask(target):
+    properties_mask = get_properties_mask(target)
+    return [create_mask(properties_mask, i) for i, _ in enumerate(rv.rules.ATTRIBUTES)]
+
+
+def create_uniform_mask(target):
+    u_mask = lambda i: tf.tile(target[:, rv.target.UNIFORMITY_INDEX + i, None] == 3, [1, rv.rules.ATTRIBUTES_LEN])
+    properties_mask = tf.concat([u_mask(0), u_mask(1)], axis=-1) | get_properties_mask(target)
+    return [create_mask(properties_mask, i) for i, _ in enumerate(rv.rules.ATTRIBUTES)]
+
+
+def create_all_mask(target):
+    return [
+        tf.cast(tf.ones(tf.stack([tf.shape(target)[0], rv.entity.SUM])), dtype=tf.bool) for i, _ in
+        enumerate(rv.rules.ATTRIBUTES)]
+
+
+class BaselineClassificationLossModel(Model):
+    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True):
+        super().__init__()
+        self.predict_fn = SubClassingModel([lambda x: x[0], PredictModel()])
+        self.loss_fn = ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
+                                       group_loss=group_loss)
+        self.metric_fn = SimilarityRaven(mode=mode)
+
+    def call(self, inputs):
+        losses = []
+        output = inputs[1]
+        losses.append(self.loss_fn([inputs[0][0], output]))
+        losses.append(self.metric_fn([inputs[0][2], inputs[3][0], inputs[0][1][:, 8:]]))
+        return losses
+
+
+class RavenLoss(Model):
+    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True, lw=(1.0, 0.3),
+                 classification=False, trans=True, anneal=False):
+        super().__init__()
+        if anneal:
+            self.weight_scheduler
+        self.classification = classification
+        self.trans = trans
+        self.predict_fn = DictModel(SubClassingModel([lambda x: x[-1], PredictModel()]), in_=OUTPUT,
+                                    out=[PREDICT, MASK], name="pred")
+        if self.trans:
+            self.loss_fn = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
+                                                    group_loss=group_loss, enable_metrics=False, lw=lw[0]),
+                                    name="main_loss")
+            self.loss_fn_2 = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
+                                                      group_loss=group_loss), name="add_loss")
+            self.metric_fn = SimilarityRaven(mode=mode)
+        if self.classification:
+            self.loss_fn_3 = add_loss(
+                ClassRavenModel(mode=create_all_mask, number_loss=number_loss, slot_loss=slot_loss,
+                                group_loss=group_loss, enable_metrics="c" if self.trans else True), lw=lw[1],
+                name="class_loss")
+
+    def call(self, inputs):
+        losses = []
+        output = inputs[OUTPUT]
+        target = inputs[TARGET]
+        labels = inputs[LABELS]
+
+        if self.trans:
+            losses.append(self.loss_fn([labels[:, 2], output[0]]))
+            losses.append(self.loss_fn([labels[:, 5], output[1]]))
+            losses.append(self.loss_fn_2([target, output[2]]))
+            losses.append(self.metric_fn([inputs[INDEX], inputs[PREDICT], labels]))
+        if self.classification:
+            for i in range(8):
+                losses.append(self.loss_fn_3([labels[:, i], inputs[CLASSIFICATION][i]]))
+        return {**inputs, LOSS: losses}
+
+
+class VTRavenLoss(Model):
+    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True, lw=(2.0, 1.0),
+                 classification=False, trans=True, anneal=False, plw=None):
+        super().__init__()
+        if anneal:
+            self.weight_scheduler
+        self.classification = classification
+        self.trans = trans
+        self.predict_fn = DictModel(SubClassingModel([lambda x: x[:, -1], PredictModel()]), in_=OUTPUT,
+                                    out=[PREDICT, "predict_mask"], name="pred")
+        self.loss_fn = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
+                                                group_loss=group_loss, plw=plw), lw=lw[0], name="add_loss")
+        self.metric_fn = SimilarityRaven(mode=mode)
+        if self.classification:
+            self.loss_fn_2 = add_loss(
+                ClassRavenModel(mode=create_all_mask, number_loss=number_loss, slot_loss=slot_loss,
+                                group_loss=group_loss, enable_metrics="c", plw=plw), lw=lw[1], name="class_loss")
+
+    def call(self, inputs):
+        losses = []
+        output = inputs[OUTPUT]
+        target = inputs[TARGET]
+        labels = inputs[LABELS]
+        mask = inputs[MASK]
+
+        target_masked = target[mask]
+        output_masked = output[mask]
+        losses.append(self.loss_fn([target_masked, output_masked]))
+
+        target_unmasked = target[~mask]
+        output_unmasked = output[~mask]
+        losses.append(self.loss_fn_2([target_unmasked, output_unmasked]))
+
+        losses.append(self.metric_fn([inputs[INDEX], inputs[PREDICT], labels]))
+        return {**inputs, LOSS: losses}
+
+
+class SingleVTRavenLoss(Model):
+    def __init__(self, mode=create_all_mask, number_loss=False, slot_loss=True, group_loss=True, lw=(1.0, 0.1),
+                 classification=False, trans=True, anneal=False):
+        super().__init__()
+        if anneal:
+            self.weight_scheduler
+        self.classification = classification
+        self.trans = trans
+        self.predict_fn = DictModel(PredictModel(), in_=OUTPUT, out=[PREDICT, MASK], name="pred")
+        self.loss_fn = add_loss(ClassRavenModel(mode=mode, number_loss=number_loss, slot_loss=slot_loss,
+                                                group_loss=group_loss), lw=lw[0], name="add_loss")
+        self.metric_fn = SimilarityRaven(mode=mode)
+
+    def call(self, inputs):
+        losses = []
+        output = inputs[OUTPUT]
+        target = inputs[TARGET]
+        labels = inputs[LABELS]
+
+        losses.append(self.loss_fn([target, output]))
+        losses.append(self.metric_fn([inputs[INDEX], inputs[PREDICT], labels]))
+        return {**inputs, LOSS: losses}
+
+
+class ClassRavenModel(Model):
+    def __init__(self, mode=create_all_mask, plw=None, number_loss=False, slot_loss=True, group_loss=True,
+                 enable_metrics=True,
+                 lw=1.0):
+        super().__init__()
+        self.number_loss = number_loss
+        self.group_loss = group_loss
+        self.enable_metrics = enable_metrics
+        self.slot_loss = slot_loss
+        self.predict_fn = PredictModel()
+        self.loss_fn = SparseCategoricalCrossentropy(from_logits=True)
+        if self.slot_loss:
+            self.loss_fn_2 = tf.nn.sigmoid_cross_entropy_with_logits
+        if self.enable_metrics:
+            self.enable_metrics = f"{self.enable_metrics}_" if isinstance(self.enable_metrics, str) else ""
+            self.metric_fn = [
+                SparseCategoricalAccuracy(name=f"{self.enable_metrics}{ACC}_{property_}") for property_ in
+                rv.properties.NAMES]
+            if self.group_loss:
+                self.metric_fn_group = SparseCategoricalAccuracy(name=f"{self.enable_metrics}{ACC}_{GROUP}")
+            if self.slot_loss:
+                self.metric_fn_2 = BinaryAccuracy(name=f"{self.enable_metrics}{ACC}_{SLOT}")
+        self.range_mask = RangeMask()
+        self.mode = mode
+        self.lw = lw
+        if not plw:
+            plw = [1., 95.37352927, 2.83426987, 0.85212836, 1.096005, 1.21943385]
+        elif isinstance(plw, int) or isinstance(plw, float):
+            plw = [1., plw, 2.83426987, 0.85212836, 1.096005, 1.21943385]
+            # plw = [plw] * 6
+        self.plw = plw
+
+    # self.predict_fn = partial(tf.argmax, axis=-1)
+
+    def call(self, inputs):
+        losses = []
+        metrics = {}
+        target = inputs[0]
+        output = inputs[1]
+
+        target_group, target_slot, target_all = raven_utils.decode.decode_target(target)
+
+        group_output, output_slot, outputs = raven_utils.decode.output_divide(output, split_fn=tf.split)
+
+        # group
+        if self.group_loss:
+            group_loss = self.lw * self.plw[0] * self.loss_fn(target_group, group_output)
+            losses.append(group_loss)
+
+            if isinstance(self.enable_metrics, str):
+                group_metric = self.metric_fn_group(target_group, group_output)
+                # metrics[GROUP] = group_metric
+                self.add_metric(group_metric)
+                self.add_metric(tf.reduce_sum(group_metric), f"{self.enable_metrics}{ACC}")
+
+        # setting uniformity mask
+        full_properties_musks = self.mode(target)
+
+        range_mask = self.range_mask(target_group)
+
+        if self.slot_loss:
+            # number
+            number_mask = range_mask & full_properties_musks[0]
+            number_mask = tf.cast(number_mask, tf.float32)
+            target_number = tf.reduce_sum(
+                tf.cast(target_slot, "float32") * number_mask, axis=-1)
+            output_number = tf.reduce_sum(
+                tf.cast(tf.sigmoid(output_slot) >= 0.5, "float32") * number_mask, axis=-1)
+
+            # output_number = tf.reduce_sum(tf.sigmoid(output_slot) * number_mask, axis=-1)
+            if self.number_loss:
+                scale = 1 / 9
+                if self.number_loss == 2:
+                    output_number_2 = tf.reduce_sum(tf.sigmoid(output_slot) * number_mask, axis=-1)
+                else:
+                    output_number_2 = output_number
+                number_loss = self.lw * self.plw[1] * mse(tf.stop_gradient(target_number) * scale,
+                                                          output_number_2 * scale)
+                losses.append(number_loss)
+
+            # metrics[NUMBER] = number_acc
+
+            if isinstance(self.enable_metrics, str):
+                number_acc = tf.reduce_mean(
+                    tf.cast(tf.cast(target_number, "int8") == tf.cast(output_number, "int8"), "float32"))
+                self.add_metric(tf.reduce_sum(number_acc), f"{self.enable_metrics}{ACC}_{NUMBER}")
+                self.add_metric(tf.reduce_sum(number_acc), f"{self.enable_metrics}{ACC}")
+                self.add_metric(tf.reduce_sum(number_acc), f"{self.enable_metrics}{ACC}_NO_{GROUP}")
+
+            # position/slot
+            slot_mask = range_mask & full_properties_musks[1]
+            # tf.boolean_mask(target_slot,slot_mask)
+
+            if tf.reduce_any(slot_mask):
+                # if tf.reduce_mean(tf.cast(slot_mask, dtype=tf.int32)) > 0:
+                target_slot_masked = tf.boolean_mask(target_slot, slot_mask)[:, None]
+                output_slot_masked = tf.boolean_mask(output_slot, slot_mask)[:, None]
+                loss_slot = self.lw * self.plw[2] * tf.reduce_mean(
+                    self.loss_fn_2(tf.cast(target_slot_masked, "float32"), output_slot_masked))
+                if isinstance(self.enable_metrics, str):
+                    acc_slot = self.metric_fn_2(target_slot_masked, output_slot_masked)
+                    self.add_metric(acc_slot)
+                    self.add_metric(tf.reduce_sum(acc_slot), f"{self.enable_metrics}{ACC}")
+                    self.add_metric(tf.reduce_sum(acc_slot), f"{self.enable_metrics}{ACC}_NO_{GROUP}")
+            else:
+                loss_slot = 0.0
+                acc_slot = -1.0
+
+            losses.append(loss_slot)
+            # metrics[SLOT] = acc_slot
+        # if loss_slot != 0:
+
+        # if tf.reduce_any(slot_mask):
+
+        # self.add_metric(acc_slot, f"{self.enable_metrics}{ACC}_{NUMBER}")
+        # self.add_metric(acc_slot, f"{self.enable_metrics}{ACC}")
+        # self.add_metric(acc_slot, f"{self.enable_metrics}{ACC}_NO_{GROUP}")
+
+        # properties
+        for i, out in enumerate(outputs):
+            shape = (-1, rv.entity.SUM, rv.properties.RAW_SIZE[i])
+            out_reshaped = tf.reshape(out, shape)
+            properties_mask = tf.cast(target_slot, "bool") & full_properties_musks[i + 2]
+
+            if tf.reduce_any(properties_mask):
+                out_masked = tf.boolean_mask(out_reshaped, properties_mask)
+                out_target = tf.boolean_mask(target_all[i], properties_mask)
+                loss = self.lw * self.plw[3 + i] * self.loss_fn(out_target, out_masked)
+                if isinstance(self.enable_metrics, str):
+                    metric = self.metric_fn[i](out_target, out_masked)
+                    self.add_metric(metric)
+                    # self.add_metric(metric, f"{self.enable_metrics}{ACC}")
+                    self.add_metric(tf.reduce_sum(metric), f"{self.enable_metrics}{ACC}")
+                    self.add_metric(tf.reduce_sum(metric), f"{self.enable_metrics}{ACC}_{PROPERTIES}")
+                    self.add_metric(tf.reduce_sum(metric), f"{self.enable_metrics}{ACC}_NO_{GROUP}")
+            else:
+                loss = 0.0
+                metric = -1.0
+
+            losses.append(loss)
+        return losses
+
+
+class FullMask(Model):
+    def __init__(self, mode=create_uniform_mask):
+        super().__init__()
+        self.range_mask = RangeMask()
+        self.mode = mode
+
+    def call(self, inputs):
+        target_group, target_slot, _ = raven_utils.decode.decode_target(inputs)
+        full_properties_musks = self.mode(inputs)
+        range_mask = self.range_mask(target_group)
+
+        number_mask = range_mask & full_properties_musks[0]
+
+        slot_mask = range_mask & full_properties_musks[1]
+        properties_mask = []
+        for property_mask in full_properties_musks[2:]:
+            properties_mask.append(tf.cast(target_slot, "bool") & property_mask)
+        return [slot_mask, properties_mask, number_mask]
+
+
+def create_mask(rules, i):
+    mask_1 = tf.tile(rules[:, i][None], [len(rv.target.FIRST_LAYOUT), 1])
+    mask_2 = tf.tile(rules[:, i + 5][None], [len(rv.target.SECOND_LAYOUT), 1])
+    shape = tf.shape(rules)
+    full_mask_1 = tf.scatter_nd(tnp.array(rv.target.FIRST_LAYOUT)[:, None], mask_1, shape=(rv.entity.SUM, shape[0]))
+    full_mask_2 = tf.tensor_scatter_nd_update(full_mask_1, tnp.array(rv.target.SECOND_LAYOUT)[:, None], mask_2)
+    return tf.transpose(full_mask_2)
+
+
+# class PredictModel(Model):
+#     def __init__(self):
+#         super().__init__()
+#         self.predict_fn = Lambda(partial(tf.argmax, axis=-1))
+#         self.predict_fn_2 = Lambda(lambda x: tf.sigmoid(x) > 0.5)
+#         self.range_mask = RangeMask()
+#
+#     # self.predict_fn = partial(tf.argmax, axis=-1)
+#
+#     def call(self, inputs):
+#         group_output = inputs[rv.OUTPUT_GROUP_SLICE]
+#         group_loss = self.predict_fn(group_output)[:, None]
+#
+#         output_slot = inputs[rv.OUTPUT_SLOT_SLICE]
+#         range_mask = self.range_mask(group_loss[:, 0])
+#         loss_slot = tf.cast(self.predict_fn_2(output_slot), dtype=tf.int64)
+#
+#         properties_output = inputs[rv.OUTPUT_PROPERTIES_SLICE]
+#         properties = []
+#         outputs = tf.split(properties_output, list(rv.ENTITY_PROPERTIES_INDEX.values()), axis=-1)
+#         for i, out in enumerate(outputs):
+#             shape = (-1, rv.ENTITY_SUM, rv.ENTITY_PROPERTIES_VALUES[i])
+#             out_reshaped = tf.reshape(out, shape)
+#             properties.append(self.predict_fn(out_reshaped))
+#         number_loss = tf.reduce_sum(loss_slot, axis=-1, keepdims=True)
+#
+#         result = tf.concat([group_loss, loss_slot, interleave(properties), number_loss], axis=-1)
+#
+#         return [result, range_mask, range_mask, range_mask, range_mask]
+
+class PredictModel(Model):
+    def __init__(self):
+        super().__init__()
+        self.predict_fn = Predict()
+        self.predict_fn_2 = Lambda(lambda x: tf.sigmoid(x) > 0.5)
+        self.range_mask = RangeMask()
+
+    # self.predict_fn = partial(tf.argmax, axis=-1)
+
+    def call(self, inputs):
+        group_output, output_slot, *properties = rv.decode.output(inputs, tf.split, self.predict_fn, self.predict_fn_2)
+        number_loss = K.int64(K.sum(output_slot))
+        result = tf.concat(
+            [group_output[:, None], tf.cast(output_slot, dtype=tf.int64), interleave(properties), number_loss[:, None]],
+            axis=-1)
+
+        range_mask = self.range_mask(group_output)
+        return [result, range_mask]
+        # return [result, range_mask, range_mask, range_mask, range_mask]
+
+
+# todo change slices
+class PredictModelMasked(Model):
+    def __init__(self):
+        super().__init__()
+        self.predict_fn = Lambda(partial(tf.argmax, axis=-1))
+        self.loss_fn_2 = Lambda(lambda x: tf.sigmoid(x) > 0.5)
+        self.range_mask = RangeMask()
+
+    # self.predict_fn = partial(tf.argmax, axis=-1)
+
+    def call(self, inputs):
+        group_output = inputs[:, -rv.GROUPS_NO:]
+        group_loss = self.predict_fn(group_output)[:, None]
+
+        output_slot = inputs[:, :rv.ENTITY_SUM]
+        range_mask = self.range_mask(group_loss[:, 0])
+        loss_slot = tf.cast(self.predict_fn_2(output_slot * range_mask), dtype=tf.int64)
+
+        properties_output = inputs[:, rv.ENTITY_SUM:-rv.GROUPS_NO]
+
+        properties = []
+        outputs = tf.split(properties_output, list(rv.ENTITY_PROPERTIES_INDEX.values()), axis=-1)
+        for i, out in enumerate(outputs):
+            shape = (-1, rv.ENTITY_SUM, rv.ENTITY_PROPERTIES_VALUES[i])
+            out_reshaped = tf.reshape(out, shape)
+            out_masked = out_reshaped * loss_slot[..., None]
+            properties.append(self.predict_fn(out_masked))
+            # out_masked[0].numpy()
+        number_loss = tf.reduce_sum(loss_slot, axis=-1, keepdims=True)
+
+        result = tf.concat([group_loss, loss_slot, interleave(properties), number_loss], axis=-1)
+
+        return result
+
+
+def final_predict_mask(x, mask):
+    r = reshape(x[0][:, rv.INDEX[0]:-1], [-1, 3])
+    return tf.ragged.boolean_mask(r, mask)
+
+
+def final_predict(x, mode=False):
+    m = x[1] if mode else tf.cast(x[0][:, 1:rv.INDEX[0]], tf.bool)
+    return final_predict_mask(x[0], m)
+
+
+def final_predict_2(x):
+    ones = tf.cast(tf.ones(tf.shape(x[0])[0]), tf.bool)[:, None]
+    mask = tf.concat([ones, tf.tile(x[1], [1, 4]), ones], axis=-1)
+    return tf.ragged.boolean_mask(x[0], mask)
+
+
+class PredictModelOld(Model):
+
+    def call(self, inputs):
+        output = inputs[-2]
+
+        rest_output = output[:, :-rv.GROUPS_NO]
+
+        result_all = []
+        outputs = tf.split(rest_output, list(rv.ENTITY_PROPERTIES_INDEX.values()), axis=-3)
+        for i, out in enumerate(outputs):
+            shape = (-3, rv.ENTITY_SUM, rv.ENTITY_PROPERTIES_VALUES[i])
+            out_reshaped = tf.reshape(out, shape)
+
+            result = tf.cast(tf.argmax(out_reshaped, axis=-3), dtype="int8")
+            result_all.append(result)
+
+        result_all = interleave(result_all)
+        return result_all
+
+
+def get_matches(diff, target_index):
+    diff_sum = K.sum(diff)
+    db_argsort = tf.argsort(diff_sum, axis=-1)
+    db_sorted = tf.sort(diff_sum)
+    db_mask = db_sorted[:, 0, None] == db_sorted
+    db_same = tf.where(db_mask, db_argsort, -1 * tf.ones_like(db_argsort))
+    matched_index = db_same == target_index
+    # setting shape needed for TensorFlow graph
+    matched_index.set_shape(db_same.shape)
+    matches = K.any(matched_index)
+    more_matches = K.sum(db_mask) > 1
+    once_matches = K.sum(matches & tf.math.logical_not(more_matches))
+    return matches, more_matches, once_matches
+
+
+class SimilarityRaven(Model):
+    def __init__(self, mode=create_all_mask, number_loss=False):
+        super().__init__()
+        self.range_mask = RangeMask()
+        self.mode = mode
+
+    # self.predict_fn = partial(tf.argmax, axis=-1)
+
+    # INDEX, PREDICT, LABELS
+    def call(self, inputs):
+        metrics = []
+        target_index = inputs[0] - 8
+        predict = inputs[1]
+        answers = inputs[2][:, 8:]
+        shape = tf.shape(predict)
+
+        target = K.gather(answers, target_index[:, 0])
+
+        target_group = target[:, 0]
+
+        # comp_slice = np.
+        target_comp = target[:, 1:rv.target.END_INDEX]
+        predict_comp = predict[:, 1:rv.target.END_INDEX]
+        answers_comp = answers[:, :, 1:rv.target.END_INDEX]
+
+        full_properties_musks = self.mode(target)
+        fpm = K.cat([full_properties_musks[0], interleave(full_properties_musks[2:])])
+
+        range_mask = self.range_mask(target_group)
+        full_range_mask = K.cat([range_mask, tf.repeat(range_mask, 3, axis=-1)], axis=-1)
+
+        final_mask = fpm & full_range_mask
+
+        target_masked = target_comp * final_mask
+        predict_masked = predict_comp * final_mask
+        answers_masked = answers_comp * tf.tile(final_mask[:, None], [1, 8, 1])
+
+        acc_same = K.mean(K.all(target_masked == predict_masked))
+        self.add_metric(acc_same, ACC_SAME)
+        metrics.append(acc_same)
+
+        diff = tf.abs(predict_masked[:, None] - answers_masked)
+        diff_bool = diff != 0
+
+        matches, more_matches, once_matches = get_matches(tf.cast(diff_bool, dtype=tf.int32), target_index)
+
+        second_phase_mask = (more_matches & matches)
+        diff_second_phase = tf.boolean_mask(diff, second_phase_mask)
+        target_index_2 = tf.boolean_mask(target_index, second_phase_mask, axis=0)
+
+        matches_2, more_matches_2, once_matches_2 = get_matches(diff_second_phase, target_index_2)
+        matches_2_no = K.sum(matches_2)
+
+        acc_choose_upper = (once_matches + matches_2_no) / shape[0]
+        self.add_metric(acc_choose_upper, ACC_CHOOSE_UPPER)
+        metrics.append(acc_choose_upper)
+
+        acc_choose_lower = (once_matches + once_matches_2) / shape[0]
+        self.add_metric(acc_choose_lower, ACC_CHOOSE_LOWER)
+        metrics.append(acc_choose_lower)
+
+        return metrics
+
+
+class SimilarityRaven2(Model):
+    def __init__(self, mode=create_all_mask, number_loss=False):
+        super().__init__()
+        self.range_mask = RangeMask()
+        self.mode = mode
+
+    # self.predict_fn = partial(tf.argmax, axis=-1)
+
+    # INDEX, PREDICT, LABELS
+    def call(self, inputs):
+        metrics = []
+        target_index = inputs[0] - 8
+        predict = inputs[1]
+        answers = inputs[2][:, 8:]
+        shape = tf.shape(predict)
+
+        target = K.gather(answers, target_index[:, 0])
+
+        target_group = target[:, 0]
+
+        # comp_slice = np.
+        target_comp = target[:, 1:rv.target.END_INDEX]
+        predict_comp = predict[:, 1:rv.target.END_INDEX]
+        answers_comp = answers[:, :, 1:rv.target.END_INDEX]
+
+        full_properties_musks = self.mode(target)
+        fpm = K.cat([full_properties_musks[0], interleave(full_properties_musks[2:])])
+
+        range_mask = self.range_mask(target_group)
+        full_range_mask = K.cat([range_mask, tf.repeat(range_mask, 3, axis=-1)], axis=-1)
+
+        final_mask = fpm & full_range_mask
+
+        target_masked = target_comp * final_mask
+        predict_masked = predict_comp * final_mask
+        answers_masked = answers_comp * tf.tile(final_mask[:, None], [1, 8, 1])
+
+        acc_same = K.mean(K.all(target_masked == predict_masked))
+        self.add_metric(acc_same, ACC_SAME)
+        metrics.append(acc_same)
+
+        diff = tf.abs(predict_masked[:, None] - answers_masked)
+        diff_bool = diff != 0
+
+        matches, more_matches, once_matches = get_matches(tf.cast(diff_bool, dtype=tf.int32), target_index)
+
+        second_phase_mask = (more_matches & matches)
+        diff_second_phase = tf.boolean_mask(diff, second_phase_mask)
+        target_index_2 = tf.boolean_mask(target_index, second_phase_mask, axis=0)
+
+        matches_2, more_matches_2, once_matches_2 = get_matches(diff_second_phase, target_index_2)
+        matches_2_no = K.sum(matches_2)
+
+        acc_choose_upper = (once_matches + matches_2_no) / shape[0]
+        self.add_metric(acc_choose_upper, ACC_CHOOSE_UPPER)
+        metrics.append(acc_choose_upper)
+
+        acc_choose_lower = (once_matches + once_matches_2) / shape[0]
+        self.add_metric(acc_choose_lower, ACC_CHOOSE_LOWER)
+        metrics.append(acc_choose_lower)
+
+        metrics.append(K.sum(target_masked != predict_masked))
+
+        return metrics
+
+
+class LatentLossModel(Model):
+    def __init__(self, dir_=HORIZONTAL):
+        super().__init__()
+        # self.sum_metrics = []
+        # for i in range(8):
+        #     self.sum_metrics.append(Sum(name=f"no_{i}"))
+        self.metric_fn = Accuracy(name="acc_latent")
+        if dir_ == VERTICAL:
+            self.dir = (6, 7)
+        else:
+            self.dir = (2, 5)
+
+    def call(self, inputs):
+        target_image = tf.reshape(inputs[0][2], [-1])
+        output = inputs[1]
+        latents = tnp.asarray(inputs[2])
+
+        target_hor = tf.concat([
+            latents[:, self.dir],
+            latents[tf.range(latents.shape[0]), target_image + 8][:, None]
+        ],
+            axis=1)
+
+        loss_hor = mse(K.stop_gradient(target_hor), output)
+        self.add_loss(loss_hor)
+
+        self.add_metric(self.metric_fn(inputs[3], target_image))
+
+        return loss_hor
+
+
+class PredRav(Model):
+
+    def call(self, inputs):
+        output = inputs[0][:, -1]
+        answers = inputs[1][:, 8:]
+        return tf.argmin(tf.reduce_sum(tf.abs(output[:, None] - answers), axis=-1), axis=-1)

raven_utils/models/multi_transformer.py

@@ -0,0 +1,274 @@
+import tensorflow as tf
+from functools import partial
+from tensorflow.keras.layers import Lambda
+from tensorflow.keras.layers import Dense
+from tensorflow.keras import Input, Model
+from tensorflow.python.keras import Sequential
+
+from config.constant import TRANS
+from ml_utils import filter_init
+from models.loss import VTRavenLoss, create_uniform_mask, SingleVTRavenLoss
+from models_utils import pmodel, DictModel, bt, INPUTS, bm, OUTPUT, LATENTS, transformer, BatchModel, get_extractor, \
+    build_seq_model, BUILD, build_train_list, InitialWeight
+from models_utils import SumPositionEmbedding, TransformerBlock, CatPositionEmbedding, transformer, BatchInitialWeight
+import models_utils.ops as K
+from models_utils.image import inverse_fn
+from models_utils.ops_core import IndexReshape
+from models_utils.random_ import EpsilonGreedy, EpsilonSoft
+from models_utils.step import StepDict
+
+
+def init_weights(shape, dtype=None):
+    return tf.cast(K.var.image(shape=shape, pre=True), dtype=tf.float32)
+
+
+def conversion(x, max_=45):
+    shape = tf.shape(x)
+    return tf.reshape(x[:, :max_], tf.stack([shape[0], 9, -1]))
+
+
+def take_left(x):
+    return x[..., 7:8]
+
+
+def take_by_index(x, i=8):
+    return x[..., i:i + 1]
+
+
+def mix(x):
+    return (x[..., 7:8] + x[..., 5:6]) / 2
+
+
+def empty_last(x):
+    return tf.zeros_like(x[..., 7:8])
+
+
+class Conversion(Model):
+    def __init__(self):
+        super().__init__()
+        self.model = IndexReshape((0, "9", None))
+
+    def call(self, inputs):
+        return self.model(inputs[:, :45])
+
+
+class RandomImageMask(Model):
+    def __init__(self, last, last_index=9):
+        super().__init__()
+        self.get_last = last
+        self.last_index = last_index
+
+    def call(self, inputs):
+        shape = tf.shape(inputs)
+        indexes = tf.random.uniform(shape=shape[0:1], maxval=self.last_index, dtype=tf.int32)
+        mask = tf.one_hot(indexes, self.last_index)[:, None, None]
+
+        return (1 - mask) * inputs[..., :self.last_index] + mask * tf.tile(self.get_last(inputs),
+                                                                           (1, 1, 1, self.last_index))
+
+
+# res = (1 - mask) * inputs[..., :self.last_index] + mask * tf.tile(self.get_last(inputs),
+#                                                                            (1, 1, 1, self.last_index))
+
+
+# from data_utils import ims
+# for i in range(50):
+#     ims(res[i].numpy().swapaxes(0, 2))
+# res[12].numpy()
+# self.get_last(inputs).numpy()
+# import tensorflow as tf
+# tf.random.uniform(shape=shape[0:1], maxval=255, dtype=tf.int32)
+# from ml_utils import print_error
+# ims(mask[0].numpy())
+# print_error(lambda :ims(mask[0]))
+# from models_utils import ops as K
+
+
+class ImageMask(Model):
+    def __init__(self, last, index=8, last_index=9):
+        super().__init__()
+        self.get_last = last
+        self.index = index
+        self.last_index = last_index
+
+    def call(self, inputs):
+        return tf.concat([inputs[..., :8], self.get_last(inputs)], axis=-1)
+
+
+class CreateGrid(Model):
+    def __init__(self,
+                 no=4,
+                 extractor="ef",
+                 type_=3,
+                 base="seq",
+                 last=take_left,
+                 epsilon=None,
+                 pooling=None,
+                 mask_fn=None,
+                 model=None,
+                 **kwargs
+                 ):
+        super().__init__()
+        self.type_ = type_
+        if type_ == 9:
+            self.start_shape = 75
+            data = (224, 224, 3)
+            conv = lambda: Conversion()
+        else:
+            self.start_shape = 84
+            data = (84, 84, 3)
+            extractor = BUILD[base]([
+                BatchModel(get_extractor(data=data, model=extractor)),
+                lambda x: tf.transpose(x, (1, 0, 2, 3, 4))
+                # lambda x: tf.tile(x[:, :224, :224], (1, 1, 1, 3))
+            ])
+            conv = lambda: conversion
+
+        self.epsilon = epsilon
+        if mask_fn == "random":
+            mask_fn = RandomImageMask(last=last)
+        elif mask_fn is None:
+            mask_fn = ImageMask(last=last)
+
+        self.mask_fn = mask_fn
+
+
+    def call(self, inputs):
+        transposed = tf.image.resize(tf.transpose(inputs, (0, 2, 3, 1)), (self.start_shape, self.start_shape))
+        re = self.mask_fn(transposed)
+
+        # re = tf.concat([transposed[..., :8], self.get_last(transposed)], axis=-1)
+        if self.type_ == 9:
+            x = tf.transpose(re, [0, 3, 1, 2])[..., None]
+            x = K.create_image_grid(x, 3, 3)
+            x = x[:, :224, :224]
+            x = tf.tile(x, [1, 1, 1, 3])
+        else:
+
+            x = tf.stack([
+                re[..., :3],
+                re[..., 3:6],
+                re[..., 6:9],
+            ])
+        return self.model(x)
+
+
+# self.model.layers[0](x)
+
+
+def grid_transformer(
+        *args,
+        type_=9,
+        no=4,
+        extractor="ef",
+        loss_mode=create_uniform_mask,
+        output_size=10,
+        loss_weight=1.0,
+        out_layers=(1000, 1000, 1000),
+        pos_emd="cat",
+        base="seq",
+        inverse_image=True,
+        last="left",
+        mask_fn=None,
+        model=None,
+        trans=None,
+        **kwargs):
+
+    if last == "left":
+        last = take_left
+    elif last == "mix":
+        last = mix
+    elif last == "empty":
+        last = empty_last
+    elif last == "start":
+        last = Sequential([Lambda(empty_last), BatchInitialWeight(initializer=init_weights)])
+
+    create_grid = CreateGrid(
+        type_=type_,
+        no=no,
+        extractor=extractor,
+        model=model,
+        output_size=output_size,
+        out_layer=out_layers,
+        pos_emd=pos_emd,
+        base=base,
+        last=last,
+        mask_fn=mask_fn,
+        **kwargs
+    )
+
+    if model is None:
+        trans = transformer(
+            extractor=extractor,
+            pos_emd=pos_emd,
+            data=data,
+            output_size=output_size,
+            out_layers=out_layer,
+            pooling=conv,
+            no=no,
+            base=base,
+            **kwargs
+            # **as_dict(p.trans)
+        )
+    else:
+        trans = trans
+
+
+
+def get_rav_trans(
+        *args,
+        type_=9,
+        no=4,
+        extractor="ef",
+        loss_mode=create_uniform_mask,
+        output_size=10,
+        loss_weight=1.0,
+        out_layers=(1000, 1000, 1000),
+        pos_emd="cat",
+        base="seq",
+        inverse_image=True,
+        last="left",
+        epsilon="greedy",
+        epsilon_step=500,
+        mask_fn=None,
+        model=None,
+        loss="multi",
+        **kwargs):
+    if last == "left":
+        last = take_left
+    elif last == "mix":
+        last = mix
+    elif last == "empty":
+        last = empty_last
+    elif last == "start":
+        last = Sequential([Lambda(empty_last), BatchInitialWeight(initializer=init_weights)])
+
+    trans_raven = CreateGrid(
+        type_=type_,
+        no=no,
+        extractor=extractor,
+        model=model,
+        output_size=output_size,
+        out_layer=out_layers,
+        pos_emd=pos_emd,
+        base=base,
+        last=last,
+        epsilon=epsilon,
+        mask_fn=mask_fn,
+        **kwargs
+    )
+
+    if loss == "single":
+        loss = SingleVTRavenLoss
+    else:
+        loss = VTRavenLoss
+
+    return bt(
+        DictModel(
+            Sequential([Lambda(lambda x: 255 - x), trans_raven]) if inverse_image else trans_raven,
+            in_=INPUTS,
+            name="Body"
+        ),
+        loss=loss(mode=loss_mode, classification=True, lw=(loss_weight, 1.0)),
+        loss_wrap=False
+    )

raven_utils/models/raven.py

@@ -0,0 +1,239 @@
+from ml_utils import lw, lu
+from models_utils import bm, Base, res, bt, DictModel, dense_drop, drop, build_encoder, MODEL_ARCH, ListModel, short, \
+    dense, Flatten, Cat, CatDenseBefore, \
+    CatDense, CatBefore, Drop, Flat2, down, Pass, conv, Flat, Get, bs, Res, SoftBlock
+from models_utils import SubClassingModel
+from models_utils.config.constants import *
+import numpy as np
+
+from config.constant import *
+from tensorflow.keras.layers import Dense, Activation, BatchNormalization
+import tensorflow as tf
+
+import raven_utils as rv
+
+from models.body import create_block
+from models.class_ import Merge, RavenClass
+from models.head import LatentHeadModel
+
+from models.loss import RavenLoss
+from models.trans import TransModel, FullTrans
+from raven_utils.const import HORIZONTAL
+
+
+def raven_model(scales,
+                out_layers,
+                latent=(64, 128, 256),
+                output_size=None,
+                padding=SAME,
+                body_layers=1,
+                encoder=None,
+                loop=1,
+                model=None,
+                act=None,
+                simpler=0,
+                loss_mode=None,
+                loss_weight=0.3,
+                dir_=HORIZONTAL,
+                global_context=False,
+                images_no=8,
+                context_mul=2,
+                res_act="pass",
+                drop_latent=0,
+                drop_inference=0,
+                drop_end=0,
+                ga=False,
+                trans_norm=None,
+                trans_act="relu",
+                arch=HEAD3,
+                encoder_norm=False,
+                encoder_pool=False,
+                encoder_global="GM",
+                encoder_before=False,
+                tail_units=256,
+                tail_flatten=None,
+                # for now by default
+                tail_down="MP",
+                trans_no=1,
+                trans_score_activation=tf.nn.softmax,
+                block_=SoftBlock,
+                **kwargs):
+    if isinstance(latent, int):
+        latent = (latent, 128, 256)
+    scales = lw(scales)
+
+    context_size = np.array(latent) * context_mul
+    # context_size = latent[scales] * context_mul
+
+    # if scales == 2:
+    #     arch = HEAD
+    # elif scales == 1:
+    #     arch = HEAD2
+    # else:
+    #     arch = VERY2
+
+    if encoder_pool:
+        strides = (1, 1)
+    else:
+        strides = (2, 2)
+    if not isinstance(encoder_before, tuple):
+        encoder_before = [encoder_before] * 3
+
+    # if trans == 1:
+    #     trans_model = TransModel2
+    # else:
+    #     trans_model = TransModel
+
+    # if scales == 3:
+    #     head = MultiHeadModel(encoder=encoder)
+    arch = MODEL_ARCH[arch]
+    heads = []
+    for s in list(range(0, max(scales) + 1)):
+        if s in (0, 1):
+            if s == 0:
+                encoder = build_encoder(arch[:3], add_norm=encoder_norm, add_pool=encoder_pool, kerner_size=(4, 4),
+                                        strides=strides)
+            else:
+                encoder = build_encoder(arch[3:4], add_norm=encoder_norm, add_pool=encoder_pool, kerner_size=(4, 4),
+                                        strides=strides)
+            head = LatentHeadModel(
+                encoder=encoder,
+                inference_network=(
+                    bm([
+                           CatBefore(filters=int(context_size[s] / 8)) if encoder_before[s] else Cat(
+                               filters=context_size[s]),
+                           # todo activation?
+                           Res(filters=context_size[s], padding=padding)
+                       ] + ([drop(drop_inference)] if drop_inference else []),
+                       name="inference")
+                ) if s in scales else Pass(),
+                stem=Base(
+                    bm(
+                        # ok we choose by parameters anyway
+                        [res(filters=latent[s], padding=padding, act=act)] + (
+                            [drop(drop_latent)] if drop_latent else [])
+                    ),
+                    name="stem")
+            )
+        else:
+            encoder = bm([
+                Res(),
+                build_encoder(arch[4:], add_norm=encoder_norm, add_pool=encoder_pool, kerner_size=(4, 4),
+                              strides=strides),
+                short(encoder_global) if encoder_global else Flatten(),
+                dense(latent[s])
+            ])
+            head = LatentHeadModel(
+                encoder=encoder,
+                inference_network=bm([
+                                         # todo Echeck Cat
+                                         CatDenseBefore(filters=int(context_size[s] / 8)) if encoder_before[
+                                             s] else CatDense(filters=context_size[s]),
+                                         # todo activation?
+                                         Res(model="dv2", filters=context_size[s], padding=padding)
+                                     ] + ([dense_drop(drop_inference)] if drop_inference else []),
+                                     name="inference"),
+                stem=Base(
+                    bm(
+                        # ok we choose by parameters anyway
+                        [res(model="dv2", units=latent[s], padding=padding, act=act)] + (
+                            [dense_drop(drop_latent)] if drop_latent else [])
+                    ),
+                    name="stem")
+            )
+        heads.append(head)
+
+    concat_input = [f"{LATENT}_{i}" for i, _ in enumerate(heads)] + [f"{INFERENCE}_{i}" for i, _ in enumerate(heads)]
+    concat_output = ["LATENTS", "INFERENCES"]
+
+    def head_concat(inputs):
+        latents = inputs[:len(heads)]
+        inferences = inputs[len(heads):]
+        return latents, inferences
+
+    head = ListModel([(h, (INPUTS if i == 0 else OUTPUT), [f"{LATENT}_{i}", f"{INFERENCE}_{i}", OUTPUT]) for i, h in
+                      enumerate(heads)] + [
+                         (head_concat, concat_input, concat_output)], out=concat_output)
+    # from rav_utils.raven import init_image
+    # a = init_image()
+    # head(a)
+
+    if model is None:
+        model = []
+        for i in scales:
+            trans_models = []
+            for t in range(trans_no):
+                trans_models.append(
+                    bm(
+                        [create_block(latent=latent[i], simpler=simpler, padding=padding, norm=trans_norm, act=res_act,
+                                      loop=loop, type_="dense" if i == 2 else "conv", block_=block_)] +
+                        [Activation(trans_act)] + [
+                            res(filters=latent[i],
+                                padding=padding,
+                                act=act,
+                                name="body_out",
+                                model="dv2" if i == 2 else "v2") for _ in
+                            range(body_layers)] + ([Drop(drop_latent)] if drop_latent else []),
+                        base_class=SubClassingModel)
+                )
+            trans_models = lu(trans_models)
+            if trans_no > 1:
+                trans_models = bm([
+                    lambda x: [[x[0], x[1]], x[1]],
+                    SoftBlock(
+                        model=trans_models,
+                        score_model=bm([
+                            Flat2(filters=latent[i], units=256, res_no=2),
+                            Dense(trans_no, trans_score_activation)
+                        ])
+                    )
+                ],
+                    base_class=SubClassingModel
+                )
+
+            model.append(
+                TransModel(
+                    body=trans_models,
+                    dir_=dir_,
+                    images_no=images_no
+                )
+            )
+
+    tail = []
+    for i, s in enumerate(scales):
+        flatting = lambda: Flat2(filters=latent[s + 1], base_class=tail_flatten, units=tail_units)
+        if s == 0:
+            if tail_flatten is None:
+                branch = bm([res(filters=latent[s], padding=padding),
+                             conv(filters=latent[s], padding=padding),
+                             BatchNormalization(),
+                             conv(filters=latent[s], padding=padding),
+                             Flatten()])
+            else:
+                branch = bm([down(base_class=tail_down), flatting()])
+        elif s == 1:
+            if tail_flatten is None:
+                branch = bm([res(filters=latent[s], padding=padding),
+                             Flatten()])
+            else:
+                branch = flatting()
+        else:
+            branch = bm([tail_units] * 2, add_flatten=False)
+        tail.append(branch)
+
+    tail.append(
+        bm([dense(tail_units)] + ([dense_drop(drop_end)] if drop_end else []) + [Dense(output_size)], add_flatten=False,
+           name=TAIL))
+    class_input = []
+
+    return bt([
+        DictModel(head, in_=INPUTS, out=[LATENT, INFERENCE], name="Head"),
+        DictModel(FullTrans(model, scales=scales), in_=[LATENT, INFERENCE], out=TRANS, name="Body"),
+        DictModel(RavenClass(Merge(tail), scales=scales, no=8), in_=[LATENT] + class_input, out=CLASSIFICATION,
+                  name="Classificator"),
+        DictModel(RavenClass(Merge(tail), scales=list(range(len(scales))), no=3), in_=[TRANS] + class_input,
+                  out=OUTPUT, name="Classificator_trans"),
+    ],
+        loss=RavenLoss(mode=loss_mode, classification=True, trans=True, lw=(1.0, loss_weight)),
+        loss_wrap=False
+    )

raven_utils/models/trans.py

@@ -0,0 +1,74 @@
+import tensorflow as tf
+from ml_utils import lw
+from models_utils import ops as K, SubClassingModel
+from tensorflow.keras import Model
+
+from models.body import create_dense_block
+import raven_utils as rv
+from raven_utils.const import HORIZONTAL, VERTICAL
+
+
+class TransModel(Model):
+    def __init__(self, body=None, dir_=HORIZONTAL, images_no=8, latent=64):
+        super().__init__()
+        self.model = body or create_dense_block(latent=latent)
+        if dir_ == VERTICAL:
+            self.dir = (0, 3, 1, 4, 3, 5)
+        else:
+            self.dir = (0, 1, 3, 4, 6, 7)
+        self.images_no = images_no
+        self.latent = latent
+
+    def call(self, inputs):
+        # latents = tnp.asarray(inputs[0])
+        latents = inputs[0]
+        inference = inputs[1]
+        shape = tf.shape(latents)
+        new_shape = K.cat([[-1, 3, 2], shape[2:]])
+        horizontal = latents[:, self.dir].reshape(new_shape)
+        res = tf.TensorArray(tf.float32, size=3)
+        for i in range(3):
+            res = res.write(i, self.model([horizontal[:, i], inference]))
+        result = K.tran(res.stack())
+        return result
+
+
+class TransModel2(Model):
+    def __init__(self, body=None, dir_=HORIZONTAL, images_no=8, latent=64):
+        super().__init__()
+        self.body = body or create_dense_block(latent=latent)
+        if dir_ == VERTICAL:
+            self.dir = (0, 3, 1, 4, 3, 5)
+        else:
+            self.dir = (0, 1, 3, 4, 6, 7)
+        self.images_no = images_no
+        self.latent = latent
+
+    def call(self, inputs):
+        # latents = tnp.asarray(inputs[0])
+        latents = inputs[0]
+        inference = inputs[1]
+        shape = tf.shape(latents)
+        new_shape = K.cat([[-1, 3, 2], shape[2:]])
+        horizontal = latents[:, self.dir].reshape(new_shape)
+        res = tf.TensorArray(tf.float32, size=3)
+        for i in tf.range(3):
+            res = res.write(i, self.body([horizontal[:, i], inference[:,i]]))
+        result = K.tran(res.stack())
+        return result
+
+
+class FullTrans(SubClassingModel):
+    def __init__(self, model,scales,name=None):
+        super().__init__(model=model,name=name)
+        self.scales = scales
+
+    def call(self, inputs):
+        latent = lw(inputs[0])
+        inference = lw(inputs[1])
+        results = []
+        # todo merging inference?
+        for i,s in enumerate(self.scales):
+            # results.append(model([latent[::-1][i], inference]))
+            results.append(self.model[i]([latent[s], inference[s]]))
+        return results,

raven_utils/models/transformer.py

@@ -0,0 +1,133 @@
+import tensorflow as tf
+from tensorflow.keras.layers import Lambda
+from tensorflow.python.keras import Sequential
+
+# from models_utils.models.loss import VTRavenLoss, create_uniform_mask, SingleVTRavenLoss
+from models_utils import DictModel, bt, INPUTS, BatchInitialWeight
+import models_utils.ops as K
+from models_utils.models.transformer.img_seq import init_weights, take_left, mix, empty_last
+from models_utils.models.transformer.img_seq2 import init_weights, take_left, mix, empty_last, img_sec_trans
+from models_utils.ops_core import IndexReshape
+from models_utils.random_ import EpsilonGreedy, EpsilonSoft
+from models_utils.step import StepDict
+
+# res = (1 - mask) * inputs[..., :self.last_index] + mask * tf.tile(self.get_last(inputs),
+#                                                                            (1, 1, 1, self.last_index))
+
+
+# from data_utils import ims
+# for i in range(50):
+#     ims(res[i].numpy().swapaxes(0, 2))
+# res[12].numpy()
+# self.get_last(inputs).numpy()
+# import tensorflow as tf
+# tf.random.uniform(shape=shape[0:1], maxval=255, dtype=tf.int32)
+# from ml_utils import print_error
+# ims(mask[0].numpy())
+# print_error(lambda :ims(mask[0]))
+# from models_utils import ops as K
+
+
+# self.model.layers[0](x)
+from raven_utils.models.loss import VTRavenLoss, SingleVTRavenLoss, create_uniform_mask
+
+
+def get_rav_trans(
+        data,
+        type_=9,
+        no=4,
+        extractor="ef",
+        loss_mode=create_uniform_mask,
+        output_size=10,
+        loss_weight=1.0,
+        out_layers=(1000, 1000, 1000),
+        pos_emd="cat",
+        base="seq",
+        inverse_image=True,
+        last="left",
+        epsilon="greedy",
+        epsilon_step=500,
+        mask_fn=None,
+        model=None,
+        loss="multi",
+        **kwargs):
+    if last == "left":
+        last = take_left
+    elif last == "mix":
+        last = mix
+    elif last == "empty":
+        last = empty_last
+    elif last == "start":
+        last = Sequential([Lambda(empty_last), BatchInitialWeight(initializer=init_weights)])
+
+    if epsilon == "greedy":
+        epsilon = EpsilonGreedy(step=epsilon_step)
+    elif epsilon == "soft":
+        epsilon = EpsilonSoft(step=epsilon_step)
+    elif epsilon is False:
+        epsilon = None
+
+    if epsilon:
+        trans_raven = TransRavenwithStep(
+            type_=type_,
+            no=no,
+            extractor=extractor,
+            output_size=output_size,
+            out_layer=out_layers,
+            pos_emd=pos_emd,
+            base=base,
+            last=last,
+            epsilon=epsilon,
+            **kwargs
+        )
+        return StepDict(bt(
+            DictModel(
+                Sequential([Lambda(lambda x: (255 - x[0], x[1])), trans_raven]) if inverse_image else trans_raven,
+                in_=[INPUTS, "step"],
+                name="Body"
+            ),
+            loss=VTRavenLoss(mode=loss_mode, classification=True, lw=(loss_weight, 1.0)),
+            loss_wrap=False),
+            add_step=epsilon_step,
+        )
+
+    trans_raven = img_sec_trans(
+        type_=type_,
+        no=no,
+        extractor=extractor,
+        model=model,
+        output_size=output_size,
+        out_layer=out_layers,
+        pos_emd=pos_emd,
+        base=base,
+        last=last,
+        epsilon=epsilon,
+        mask_fn=mask_fn,
+        **kwargs
+    )
+    if loss == "single":
+        loss = SingleVTRavenLoss
+    else:
+        loss = VTRavenLoss
+
+    # return bt(
+    #     DictModel(
+    #         Sequential([Lambda(lambda x: 255 - x), trans_raven]) if inverse_image else trans_raven,
+    #         inputs=INPUTS,
+    #         name="Body"
+    #     ),
+    #     loss=loss(mode=loss_mode, classification=True, lw=(loss_weight, 1.0)),
+    #     loss_wrap=False
+    # )
+
+    return bt([
+        DictModel(
+            Sequential([Lambda(lambda x: 255 - x), trans_raven]) if inverse_image else trans_raven,
+            in_=INPUTS,
+            name="Body"
+        ),
+
+    ],
+        loss=loss(mode=loss_mode, classification=True, lw=(loss_weight, 1.0)),
+        loss_wrap=False
+    )

raven_utils/models/transformer_2.py

@@ -0,0 +1,146 @@
+from functools import partial
+
+import tensorflow as tf
+from tensorflow.keras.layers import Lambda
+from tensorflow.python.keras import Sequential
+from models_utils import ops as K, SubClassing
+from models_utils.models.transformer import aug
+
+# from models_utils.models.loss import VTRavenLoss, create_uniform_mask, SingleVTRavenLoss
+from data_utils import DataGenerator, LOSS, TARGET, IMAGES
+from models_utils import DictModel, bt, INPUTS, BatchInitialWeight, build_functional_model, get_input_layer
+import models_utils.ops as K
+from models_utils.models.transformer.img_seq import init_weights, take_left, mix, empty_last
+from models_utils.models.transformer.img_seq2 import init_weights, take_left, mix, empty_last, img_sec_trans
+from models_utils.ops_core import IndexReshape
+from models_utils.random_ import EpsilonGreedy, EpsilonSoft
+from models_utils.step import StepDict
+
+from models_utils.models.transformer import aug
+
+# res = (1 - mask) * inputs[..., :self.last_index] + mask * tf.tile(self.get_last(inputs),
+#                                                                            (1, 1, 1, self.last_index))
+
+
+# from data_utils import ims
+# for i in range(50):
+#     ims(res[i].numpy().swapaxes(0, 2))
+# res[12].numpy()
+# self.get_last(inputs).numpy()
+# import tensorflow as tf
+# tf.random.uniform(shape=shape[0:1], maxval=255, dtype=tf.int32)
+# from ml_utils import print_error
+# ims(mask[0].numpy())
+# print_error(lambda :ims(mask[0]))
+# from models_utils import ops as K
+
+
+# self.model.layers[0](x)
+from raven_utils.constant import INDEX, LABELS
+from raven_utils.models.loss import VTRavenLoss, SingleVTRavenLoss, create_uniform_mask
+
+
+def get_matrix(inputs, index):
+    return tf.concat([inputs[:, :8], K.gather(inputs, index[:, 0])[:, None]], axis=1)
+
+
+def get_images(inputs):
+    return get_matrix(inputs[0], inputs[1])
+
+
+def random_last(inputs, max_=8):
+    index = K.init.label(max=max_, shape=[tf.shape(inputs[0])[0]])[..., None]
+    return get_matrix(inputs[0], index)
+
+
+def get_images_no_answer(inputs):
+    return inputs[0][:, :9]
+
+
+def repeat_last(inputs):
+    return inputs[0][:, list(range(8)) + [7]]
+
+
+def get_rav_trans(
+        data,
+        inverse_image=True,
+        loss_mode=create_uniform_mask,
+        loss_weight=1.0,
+        loss="multi",
+        number_loss=False,
+        plw=None,
+        pre="auto",
+        augmentation=None,
+        **kwargs):
+    if isinstance(data, DataGenerator):
+        data = data[0]['inputs'], data[0]['index']
+    # u = img_sec_trans(**kwargs)(get_images(data) if kwargs['mask'] == "random" else get_images_no_answer(data))
+    # u.shape
+    from keras import Model
+    if pre == "auto":
+        pre = get_images if kwargs['mask'] == "random" else get_images_no_answer
+    elif pre == "no_answer":
+        pre = get_images_no_answer
+    elif pre == "last":
+        pre = repeat_last
+    elif pre == "images":
+        pre = get_images
+    elif pre == "random_last":
+        pre = random_last
+    elif pre == "noise":
+        pre = SubClassing([get_matrix, partial(aug.noise, max_=8)])
+    elif pre == "batch_noise":
+        pre = SubClassing([get_matrix, partial(aug.batch_noise, max_=8)])
+
+    if augmentation == "transpose":
+        augmentation = aug.Transpose(axis=(0, 2, 1))
+        augmentation_label = aug.Transpose(axis=(0, 2, 1))
+    elif augmentation == "shuffle_col":
+        augmentation = aug.shuffle_col
+        augmentation_label = aug.shuffle_col
+    elif augmentation == "shuffle":
+        augmentation = aug.shuffle
+        augmentation_label = aug.shuffle
+    if augmentation:
+        augmentation = [
+            # DictModel(augmentation, IMAGES, IMAGES),
+            # DictModel(aug.reshape_static(pre(data),augmentation), IMAGES, IMAGES),
+            DictModel(aug.ReshapeStatic(augmentation), IMAGES, IMAGES),
+            DictModel(
+                aug.PartialModel(
+                    aug.ReshapeStatic(augmentation_label),
+                    last_axis=9)
+                , LABELS, LABELS)
+        ]
+    else:
+        augmentation = []
+
+    trans_raven = build_functional_model(
+        img_sec_trans(**kwargs),
+        # get_images(data) if kwargs['mask'] == "random" else get_images_no_answer(data)
+        pre(data)
+        # data[0]
+    )
+    if loss == "single":
+        loss = SingleVTRavenLoss
+    else:
+        loss = VTRavenLoss
+    if isinstance(loss_weight, float):
+        loss_weight = (loss_weight, 1.0)
+
+    return bt([
+        # DictModel(get_images if kwargs['mask'] == "random" else get_images_no_answer, [INPUTS, INDEX], IMAGES),
+        DictModel(pre, [INPUTS, INDEX], IMAGES),
+        *augmentation,
+        DictModel(
+            Sequential([Lambda(lambda x: 255 - x), trans_raven]) if inverse_image else trans_raven,
+            in_=IMAGES,
+            # inputs=INPUTS,
+            name="Body"
+        ),
+
+    ],
+        loss=loss(mode=loss_mode, classification=True, number_loss=number_loss, lw=loss_weight, plw=plw),
+        predict=LOSS,
+        loss_wrap=False
+    )

raven_utils/models/transformer_3.py

@@ -0,0 +1,206 @@
+import logging
+
+from loguru import logger
+from tensorflow.keras.layers import Lambda
+from tensorflow.keras.layers import Activation
+
+from grid_transformer import aug_trans
+from raven_utils.models.loss_3 import VTRavenLoss, SingleVTRavenLoss, create_uniform_mask
+from data_utils import get_shape, TakeDict
+
+from data_utils import DataGenerator, LOSS, TARGET, IMAGES
+from models_utils import DictModel, bt, INPUTS, BatchInitialWeight, build_functional, get_input_layer, Last, bm, \
+    add_end, AUGMENTATION
+# from report.select_ import SelectModel2, SelectModel, SelectModel9
+from experiment_utils.keras_model import load_weights as model_load_weights
+
+
+def get_rav_trans(
+        data,
+        loss_mode=create_uniform_mask,
+        loss_weight=2.0,
+        number_loss=False,
+        dry_run="auto",
+        plw=None,
+        **kwargs):
+    if isinstance(loss_weight, float):
+        loss_weight = (loss_weight, 1.0)
+
+    # seq_trans(**kwargs)(data[0])
+    # trans_raven = build_functional_model2(
+    #     seq_trans(**kwargs),
+    #     data[0],
+    #     batch=None
+    # )
+    trans_raven = build_functional(
+        model=aug_trans,
+        inputs_=data[0] if isinstance(data, DataGenerator) else data,
+        batch_=None,
+        dry_run=dry_run,
+        **kwargs
+    )
+
+    return bt(
+        model=trans_raven,
+        loss=VTRavenLoss(mode=loss_mode, classification=True, number_loss=number_loss, lw=loss_weight, plw=plw),
+        model_wrap=False,
+        predict=LOSS,
+        loss_wrap=False
+    )
+
+
+def rav_select_model(
+        data,
+        load_weights=None,
+        loss_weight=(0.01, 0.0),
+        plw=5.0,
+        result_metric="sparse_categorical_accuracy",
+        select_type=2,
+        select_out=0,
+        additional_out=0,
+        additional_copy=True,
+        tail_out=(1000, 1000),
+        **kwargs
+):
+    out_layers = Last()
+    if additional_out > 0:
+        model3 = get_rav_trans(
+            data,
+            plw=plw,
+            loss_weight=loss_weight,
+            **kwargs
+        )
+
+        model_load_weights(
+            model3,
+            load_weights,
+            # sample_data,
+            None,
+            template="weights_{epoch:02d}-{val_loss:.2f}",
+            key=result_metric,
+        )
+
+        if AUGMENTATION in kwargs and kwargs[AUGMENTATION] is not None:
+            index = -1
+        else:
+            index = -2
+
+        out = model3[0, index, :additional_out]
+        logger.info(f"Additional out from: {model3[0, index]}.")
+
+        if additional_out > 2:
+            out += [Activation("gelu")]
+        out_layers = bm([out_layers] + out, add_flatten=False)
+    model = get_rav_trans(
+        TakeDict(data[0])[:, 8:],
+        plw=plw,
+        loss_weight=loss_weight,
+        **{
+            **kwargs,
+            "out_layers": out_layers,
+        }
+        # **{**as_dict(p.mp), "show_shape": True, "save_shape": f"output/shapes/type_{p.mp.type_}.json"},
+    )
+    # from data_utils.ops import Equal
+    # o = []
+    # for i in range(1, 3):
+    #     for j in range(2):
+    #         o.append(
+    #             Equal(
+    #                 # model[0,:,-2, i].variables[j],
+    #                 model2[0, :, -2, i].variables[j],
+    #                 # out_layers[i].variables[j]
+    #                 second_pooling[i].variables[j]
+    #             ).equal
+    #         )
+    # assert all(o)
+    # model = get_rav_trans(
+    #     # TakeDict(val_generator[0])[:, 8:],
+    #     # TakeDict(val_generator[0])[:, 8:],
+    #     val_generator[0],
+    #     plw=p.plw,
+    #     loss_weight=p.loss_weight,
+    #     **{**as_dict(p.mp),
+    #        # "out_layers": out_layers,
+    #        }
+    #     # **{**as_dict(p.mp), "show_shape": True, "save_shape": f"output/shapes/type_{p.mp.type_}.json"},
+    # )
+    model_load_weights(model,
+                       load_weights,
+                       # sample_data,
+                       None,
+                       template="weights_{epoch:02d}-{val_loss:.2f}",
+                       key=result_metric,
+                       )
+    # model.compile()
+    # model.evaluate(val_generator.data[:1000])
+    # model(TakeDict(val_generator[0])[:, 8:])
+    trans_raven = model[0]
+    # s = trans_raven(TakeDict(val_generator[0])[:, 8:])
+    if select_type == 2:
+        second_pooling = Lambda(lambda x: x[:, :-1])
+    else:
+        second_pooling = Last()
+    if additional_out > 0:
+        if additional_copy:
+            model4 = get_rav_trans(
+                data,
+                plw=plw,
+                loss_weight=loss_weight,
+                **kwargs
+            )
+            model_load_weights(model4,
+                               load_weights,
+                               # sample_data,
+                               None,
+                               template="weights_{epoch:02d}-{val_loss:.2f}",
+                               key=result_metric,
+                               )
+
+            if AUGMENTATION in kwargs and kwargs[AUGMENTATION] is not None:
+                index = -1
+            else:
+                index = -2
+            out2 = model4[0, index, :additional_out]
+            logger.info(f"Additional out from: {model4[0, index]}.")
+
+            if additional_out > 2:
+                out2 += [Activation("gelu")]
+        else:
+            out2 = out
+
+        second_pooling = bm([second_pooling] + out2, add_flatten=False)
+
+    model2 = get_rav_trans(
+        TakeDict(data[0])[:, 8:],
+        plw=plw,
+        loss_weight=loss_weight,
+        **{
+            **kwargs,
+            "out_layers": second_pooling,
+        }
+        # **{**as_dict(p.mp), "show_shape": True, "save_shape": f"output/shapes/type_{p.mp.type_}.json"},
+    )
+    model_load_weights(
+        model2,
+        load_weights,
+        # sample_data,
+        None,
+        template="weights_{epoch:02d}-{val_loss:.2f}",
+        key=result_metric,
+    )
+    if select_type == 0:
+        # not working
+        trans_raven2 = model2[0]
+    else:
+        trans_raven2 = model2[0]
+    tail = add_end(out_layers=tail_out, output_size=8 if select_out else 1)
+    # trans_raven2.mask_fn = ImageMask(last=take_by_index)
+    if select_type == 2:
+        select_model_class = SelectModel2
+    elif select_type == 1:
+        select_model_class = SelectModel
+    else:
+        select_model_class = SelectModel9
+    select_model = select_model_class(trans_raven, model2=trans_raven2, tail=tail, select_out=select_out)
+    return select_model

raven_utils/models/uitls_.py

@@ -0,0 +1,16 @@
+import tensorflow as tf
+import tensorflow.experimental.numpy as tnp
+from tensorflow.keras import Model
+import raven_utils as rv
+
+
+class RangeMask(Model):
+    def __init__(self):
+        super().__init__()
+        ranges = tf.tile(tf.range(rv.entity.INDEX[-1])[None], [rv.group.NO, 1])
+        start_index = rv.entity.INDEX[:-1][:, None]
+        end_index = rv.entity.INDEX[1:][:, None]
+        self.mask = tnp.array((start_index <= ranges) & (ranges < end_index))
+
+    def call(self, inputs):
+        return self.mask[inputs]

raven_utils/output.py

@@ -0,0 +1,16 @@
+import numpy as np
+import raven_utils.entity as entity
+import raven_utils.properties as properties
+import raven_utils.group as group
+
+SIZE = entity.SUM * properties.SUM + group.NO + entity.SUM
+
+SLOT_AND_GROUP = group.NO + entity.SUM
+
+PROPERTIES_SLICE = np.s_[:, :-SLOT_AND_GROUP]
+SLOT_SLICE = np.s_[:, -SLOT_AND_GROUP:-group.NO]
+GROUP_SLICE = np.s_[:, -group.NO:]
+
+GROUP_SLICE_END = np.s_[-group.NO:]
+SLOT_SLICE_END = np.s_[-SLOT_AND_GROUP:-group.NO]
+PROPERTIES_SLICE_END = np.s_[:-SLOT_AND_GROUP]

raven_utils/params.py

@@ -0,0 +1,110 @@
+from dataclasses import dataclass
+from typing import Any, Tuple
+
+from ml_utils import get_str_name
+from grid_transformer.params import ImgSeqTransformerParameters
+from raven_utils import output
+
+from experiment_utils.parameters.nn_default import TP, EP
+
+
+@dataclass
+class SudokuParameters(ImgSeqTransformerParameters):
+    mask: str = "input"
+    col: int = 3
+    row: int = 3
+    pooling: int = 81
+    output_size: int = 9
+    size: int = 384
+
+
+@dataclass
+class RavenTransParameters(ImgSeqTransformerParameters):
+    mask: str = "last"
+    last_index: int = 8
+    col: int = 1
+    row: int = 1
+    output_size: int = output.SIZE
+    number_loss: bool = 0
+    pre: str = "images"
+    num_heads: int = 8
+
+
+MP = RavenTransParameters
+
+
+@dataclass
+class RavenSelectTransParameters(RavenTransParameters):
+    select_type: int = 2
+    select_out: int = 0
+    additional_out: int = 0
+    additional_copy: bool = True
+    tail_out: Tuple = (1000, 1000)
+    pre: str = "index"
+
+
+SMP = RavenSelectTransParameters
+
+from raven_utils.config.models import AVAILABLE_MODELS
+from experiment_utils.parameters.nn_clean import Parameters as BaseParameters
+from raven_utils.config.constant import RAVEN, LABELS, INDEX, FEATURES, RAV_METRICS, IMP_RAV_METRICS, ACC_NO_GROUP, \
+    ACC_SAME
+
+MODEL_NO = -1
+
+
+@dataclass
+class RavenParameters(BaseParameters):
+    dataset_name: str = RAVEN
+    data: Any = (
+        f"{dataset_name}/train.npy",
+        f"{dataset_name}/val.npy",
+        f"{dataset_name}/train_labels.npy",
+        f"{dataset_name}/val_labels.npy",
+        f"{dataset_name}/train_target.npy",
+        f"{dataset_name}/val_target.npy",
+        f"arr/train_features_{AVAILABLE_MODELS[MODEL_NO]}.npy",
+        f"arr/val_features_{AVAILABLE_MODELS[MODEL_NO]}.npy",
+        f"{dataset_name}/val_index.npy"
+        # DataParameters2()
+    )
+
+    # core_metrics: tuple = tuple(RAV_METRICS)
+    filter_metrics: tuple = tuple(IMP_RAV_METRICS)
+    # result_metric: str = ACC_NO_GROUP
+    result_metric: str = ACC_SAME
+
+    lw: float = 0.0001  # Autoencoder
+    loss_weight: float = 2.0
+    plw: int = 5.0
+    mp: RavenTransParameters = RavenTransParameters()
+
+    @property
+    def experiment(self):
+        # return "rav/trans"
+        return "rav/best_test3"
+        # return "rav/trans_weight"
+
+    # @property
+    # def name(self):
+    #     # return f"i{self.extractor}_{len(self.tail)}{self.tail[0]}_{self.type_}_{self.epsilon}_{self.last}_{self.epsilon_step}"
+    #     return f"{get_str_name(self.mp.pre)[0]}_{str(self.plw)[0]}_{str(self.mp.number_loss)[0]}_{self.mp.extractor}_{self.mp.noise if self.mp.noise else ''}_{self.mp.augmentation if self.mp.augmentation else ''}_{self.mp.extractor_shape}_{self.mp.no}_{self.mp.num_heads}_{self.mp.size}_{self.mp.pos_emd}_{self.mp.ff_mul}_{self.tp.batch}"
+
+
+@dataclass
+class BaselineRavenParameters(RavenParameters):
+
+    @property
+    def experiment(self):
+        # return "rav/best_test3"
+        return "rav/baseline"
+        # return "rav/trans_weight"
+
+    @property
+    def name(self):
+        # return f"i{self.extractor}_{len(self.tail)}{self.tail[0]}_{self.type_}_{self.epsilon}_{self.last}_{self.epsilon_step}"
+        return f"{get_str_name(self.mp.pre)[0]}_{str(self.plw)[0]}_{str(self.mp.number_loss)[0]}_{self.mp.extractor}_{self.mp.noise if self.mp.noise else ''}_{self.mp.augmentation if self.mp.augmentation else ''}_{self.mp.extractor_shape}_{self.mp.no}_{self.mp.num_heads}_{self.mp.size}_{self.mp.pos_emd}_{self.mp.ff_mul}_{self.tp.batch}"
+
+
+if __name__ == '__main__':
+    params = PreRavenTransParameters()

raven_utils/properties.py

@@ -0,0 +1,16 @@
+import raven_utils as rv
+from ml_utils import dict_from_list2, CalcDict
+import raven_utils.entity as entity
+
+NAMES = [
+    'Color',
+    'Size',
+    'Type',
+]
+RAW_SIZE = [10, 6, 5]
+SIZE = dict_from_list2(NAMES, RAW_SIZE)
+ANGLE_SIZE = 7
+NO = len(NAMES)
+
+INDEX = (CalcDict(SIZE) * entity.SUM).to_dict()
+SUM = sum(list(SIZE.values()))

raven_utils/range_mask.py

@@ -0,0 +1,16 @@
+import tensorflow as tf
+import tensorflow.experimental.numpy as tnp
+from tensorflow.keras import Model
+import raven_utils as rv
+
+
+class RangeMask(Model):
+    def __init__(self):
+        super().__init__()
+        ranges = tf.tile(tf.range(rv.entity.INDEX[-1])[None], [rv.group.NO, 1])
+        start_index = rv.entity.INDEX[:-1][:, None]
+        end_index = rv.entity.INDEX[1:][:, None]
+        self.mask = tnp.array((start_index <= ranges) & (ranges < end_index))
+
+    def call(self, inputs):
+        return self.mask[inputs]

raven_utils/render/const.py

@@ -0,0 +1,86 @@
+# -*- coding: utf-8 -*-
+
+
+# Maximum number of components in a RPM
+MAX_COMPONENTS = 2
+
+# Canvas parameters
+IMAGE_SIZE = 160
+CENTER = (IMAGE_SIZE / 2, IMAGE_SIZE / 2)
+DEFAULT_RADIUS = IMAGE_SIZE / 4
+DEFAULT_WIDTH = 2
+
+# Attribute parameters
+# Number
+NUM_VALUES = [1, 2, 3, 4, 5, 6, 7, 8, 9]
+NUM_MIN = 0
+NUM_MAX = len(NUM_VALUES) - 1
+
+# Uniformity
+UNI_VALUES = [False, False, False, True]
+UNI_MIN = 0
+UNI_MAX = len(UNI_VALUES) - 1
+
+# Type
+TYPE_VALUES = ["none", "triangle", "square", "pentagon", "hexagon", "circle"]
+TYPE_MIN = 0
+TYPE_MAX = len(TYPE_VALUES) - 1
+
+# Size
+SIZE_VALUES = [0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+SIZE_MIN = 0
+SIZE_MAX = len(SIZE_VALUES) - 1
+
+# Color
+COLOR_VALUES = [255, 224, 196, 168, 140, 112, 84, 56, 28, 0]
+COLOR_MIN = 0
+COLOR_MAX = len(COLOR_VALUES) - 1
+
+# Angle: self-rotation
+ANGLE_VALUES = [-135, -90, -45, 0, 45, 90, 135, 180]
+ANGLE_MIN = 0
+ANGLE_MAX = len(ANGLE_VALUES) - 1
+
+META_TARGET_FORMAT = ["Constant", "Progression", "Arithmetic", "Distribute_Three", "Number", "Position", "Type", "Size", "Color"]
+META_STRUCTURE_FORMAT = ["Singleton", "Left_Right", "Up_Down", "Out_In", "Left", "Right", "Up", "Down", "Out", "In", "Grid", "Center_Single", "Distribute_Four", "Distribute_Nine", "Left_Center_Single", "Right_Center_Single", "Up_Center_Single", "Down_Center_Single", "Out_Center_Single", "In_Center_Single", "In_Distribute_Four"]
+
+# Rule, Attr, Param
+# The design encodes rule priority order: Number/Position always comes first
+# Number and Position could not both be sampled
+# Progression on Number: Number on each Panel +1/2 or -1/2
+# Progression on Position: Entities on each Panel roll over the layout
+# Arithmetic on Number: Numeber on the third Panel = Number on first +/- Number on second (1 for + and -1 for -)
+# Arithmetic on Position: 1 for SET_UNION and -1 for SET_DIFF
+# Distribute_Three on Number: Three numbers through each row
+# Distribute_Three on Position: Three positions (same number) through each row
+# Constant on Number/Position: Nothing changes
+# Progression on Type: Type progression defined as the number of edges on each entity (Triangle, Square, Pentagon, Hexagon, Circle)
+# Distribute_Three on Type: Three types through each row
+# Constant on Type: Nothing changes
+# Progression on Size: Size on each entity +1/2 or -1/2
+# Arithmetic on Size: Size on the third Panel = Size on the first +/- Size on the second (1 for + and -1 for -)
+# Distribute_Three on Size: Three sizes through each row
+# Constant on Size: Nothing changes
+# Progression on Color: Color +1/2 or -1/2
+# Arithmetic on Color: Color on the third Panel = Color on the first +/- Color on the second (1 for + and -1 for -)
+# Distribute_Three on Color: Three colors through each row
+# Constant on Color: Nothing changes
+# Note that all rules on Type, Size and Color enforce value consistency in a panel
+RULE_ATTR = [[["Progression", "Number", [-2, -1, 1, 2]], 
+              ["Progression", "Position", [-2, -1, 1, 2]], 
+              ["Arithmetic", "Number", [1, -1]],
+              ["Arithmetic", "Position", [1, -1]],
+              ["Distribute_Three", "Number", None],
+              ["Distribute_Three", "Position", None],
+              ["Constant", "Number/Position", None]],
+             [["Progression", "Type", [-2, -1, 1, 2]],
+              ["Distribute_Three", "Type", None], 
+              ["Constant", "Type", None]],
+             [["Progression", "Size", [-2, -1, 1, 2]],
+              ["Arithmetic", "Size", [1, -1]],
+              ["Distribute_Three", "Size", None],
+              ["Constant", "Size", None]],
+             [["Progression", "Color", [-2, -1, 1, 2]],
+              ["Arithmetic", "Color", [1, -1]],
+              ["Distribute_Three", "Color", None],
+              ["Constant", "Color", None]]]

raven_utils/render/rendering.py

@@ -0,0 +1,304 @@
+# -*- coding: utf-8 -*-
+
+
+import cv2
+import numpy as np
+from PIL import Image
+#
+# from AoT import Root
+import raven_utils.decode
+from raven_utils.render.const import CENTER, DEFAULT_WIDTH, IMAGE_SIZE
+
+from data_utils import Bag
+
+from raven_utils.render_ import COLOR_VALUES, SIZE_VALUES, TYPE_VALUES, ANGLE_VALUES, RENDER_POSITIONS
+
+
+def imshow(array):
+    image = Image.fromarray(array)
+    image.show()
+
+
+def imsave(array, filepath):
+    image = Image.fromarray(array)
+    image.save(filepath)
+
+
+def generate_matrix(array_list):
+    # row-major array_list
+    assert len(array_list) <= 9
+    img_grid = np.zeros((IMAGE_SIZE * 3, IMAGE_SIZE * 3), np.uint8)
+    for idx in range(len(array_list)):
+        i, j = divmod(idx, 3)
+        img_grid[i * IMAGE_SIZE:(i + 1) * IMAGE_SIZE, j * IMAGE_SIZE:(j + 1) * IMAGE_SIZE] = array_list[idx]
+    # draw grid
+    for x in [0.33, 0.67]:
+        img_grid[int(x * IMAGE_SIZE * 3) - 1:int(x * IMAGE_SIZE * 3) + 1, :] = 0
+    for y in [0.33, 0.67]:
+        img_grid[:, int(y * IMAGE_SIZE * 3) - 1:int(y * IMAGE_SIZE * 3) + 1] = 0
+    return img_grid
+
+
+def generate_answers(array_list):
+    assert len(array_list) <= 8
+    img_grid = np.zeros((IMAGE_SIZE * 2, IMAGE_SIZE * 4), np.uint8)
+    for idx in range(len(array_list)):
+        i, j = divmod(idx, 4)
+        img_grid[i * IMAGE_SIZE:(i + 1) * IMAGE_SIZE, j * IMAGE_SIZE:(j + 1) * IMAGE_SIZE] = array_list[idx]
+    # draw grid
+    for x in [0.5]:
+        img_grid[int(x * IMAGE_SIZE * 2) - 1:int(x * IMAGE_SIZE * 2) + 1, :] = 0
+    for y in [0.25, 0.5, 0.75]:
+        img_grid[:, int(y * IMAGE_SIZE * 4) - 1:int(y * IMAGE_SIZE * 4) + 1] = 0
+    return img_grid
+
+
+def generate_matrix_answer(array_list):
+    # row-major array_list
+    assert len(array_list) <= 18
+    img_grid = np.zeros((IMAGE_SIZE * 6, IMAGE_SIZE * 3), np.uint8)
+    for idx in range(len(array_list)):
+        i, j = divmod(idx, 3)
+        img_grid[i * IMAGE_SIZE:(i + 1) * IMAGE_SIZE, j * IMAGE_SIZE:(j + 1) * IMAGE_SIZE] = array_list[idx]
+    # draw grid
+    for x in [0.33, 0.67, 1.00, 1.33, 1.67]:
+        img_grid[int(x * IMAGE_SIZE * 3), :] = 0
+    for y in [0.33, 0.67]:
+        img_grid[:, int(y * IMAGE_SIZE * 3)] = 0
+    return img_grid
+
+
+def merge_matrix_answer(matrix, answer):
+    matrix_image = generate_matrix(matrix)
+    answer_image = generate_answers(answer)
+    img_grid = np.ones((IMAGE_SIZE * 5 + 20, IMAGE_SIZE * 4), np.uint8) * 255
+    img_grid[:IMAGE_SIZE * 3, int(0.5 * IMAGE_SIZE):int(3.5 * IMAGE_SIZE)] = matrix_image
+    img_grid[-(IMAGE_SIZE * 2):, :] = answer_image
+    return img_grid
+
+
+def render_panels(feature, target=True,angle=None):
+    # Decompose the panel into a structure and its entities
+    # root
+    # rv.decode_output(root)
+    # rv.decode_output_reshape(root)
+    # decoded =
+    # panel = decoded[0]
+    panels = []
+    for group, exist, color, size, type_ in Bag(raven_utils.decode.decode_target_flat(feature)):
+        canvas = np.ones((IMAGE_SIZE, IMAGE_SIZE), np.uint8) * 255
+        structure_img = render_structure(group)
+        background = np.zeros((IMAGE_SIZE, IMAGE_SIZE), np.uint8)
+        # note left components entities are in the lower layer
+        for i, entity in enumerate(exist):
+            if entity:
+                entity_img = render_entity(RENDER_POSITIONS[i], color[i], size[i], type_[i] + 1, angle=angle)
+                background = layer_add(background, entity_img)
+        background = layer_add(background, structure_img)
+        panels.append(canvas - background)
+    return np.stack(panels)
+
+
+def render_structure(structure):
+    if structure == 5:
+        ret = np.zeros((IMAGE_SIZE, IMAGE_SIZE), np.uint8)
+        ret[:, int(0.5 * IMAGE_SIZE)] = 255.0
+    elif structure == 6:
+        ret = np.zeros((IMAGE_SIZE, IMAGE_SIZE), np.uint8)
+        ret[int(0.5 * IMAGE_SIZE), :] = 255.0
+    else:
+        ret = np.zeros((IMAGE_SIZE, IMAGE_SIZE), np.uint8)
+    return ret
+
+
+def render_entity(bbox, color, size, type_, angle=None):
+    color = COLOR_VALUES[color]
+    size = SIZE_VALUES[size]
+    type_ = TYPE_VALUES[type_]
+    if angle is None:
+        angle = np.random.randint(0, 7, 1)[0]
+    angle = ANGLE_VALUES[angle]
+    img = np.zeros((IMAGE_SIZE, IMAGE_SIZE), np.uint8)
+
+    # planar position: [x, y, w, h]
+    # angular position: [x, y, w, h, x_c, y_c, omega]
+    # center: (columns, rows)
+    center = (int(bbox[1] * IMAGE_SIZE), int(bbox[0] * IMAGE_SIZE))
+    if type_ == "triangle":
+        unit = min(bbox[2], bbox[3]) * IMAGE_SIZE / 2
+        dl = int(unit * size)
+        pts = np.array([[center[0], center[1] - dl],
+                        [center[0] + int(dl / 2.0 * np.sqrt(3)), center[1] + int(dl / 2.0)],
+                        [center[0] - int(dl / 2.0 * np.sqrt(3)), center[1] + int(dl / 2.0)]],
+                       np.int32)
+        pts = pts.reshape((-1, 1, 2))
+        color = 255 - color
+        width = DEFAULT_WIDTH
+        draw_triangle(img, pts, color, width)
+    elif type_ == "square":
+        unit = min(bbox[2], bbox[3]) * IMAGE_SIZE / 2
+        dl = int(unit / 2 * np.sqrt(2) * size)
+        pt1 = (center[0] - dl, center[1] - dl)
+        pt2 = (center[0] + dl, center[1] + dl)
+        color = 255 - color
+        width = DEFAULT_WIDTH
+        draw_square(img, pt1, pt2, color, width)
+    elif type_ == "pentagon":
+        unit = min(bbox[2], bbox[3]) * IMAGE_SIZE / 2
+        dl = int(unit * size)
+        pts = np.array([[center[0], center[1] - dl],
+                        [center[0] - int(dl * np.cos(np.pi / 10)), center[1] - int(dl * np.sin(np.pi / 10))],
+                        [center[0] - int(dl * np.sin(np.pi / 5)), center[1] + int(dl * np.cos(np.pi / 5))],
+                        [center[0] + int(dl * np.sin(np.pi / 5)), center[1] + int(dl * np.cos(np.pi / 5))],
+                        [center[0] + int(dl * np.cos(np.pi / 10)), center[1] - int(dl * np.sin(np.pi / 10))]],
+                       np.int32)
+        pts = pts.reshape((-1, 1, 2))
+        color = 255 - color
+        width = DEFAULT_WIDTH
+        draw_pentagon(img, pts, color, width)
+    elif type_ == "hexagon":
+        unit = min(bbox[2], bbox[3]) * IMAGE_SIZE / 2
+        dl = int(unit * size)
+        pts = np.array([[center[0], center[1] - dl],
+                        [center[0] - int(dl / 2.0 * np.sqrt(3)), center[1] - int(dl / 2.0)],
+                        [center[0] - int(dl / 2.0 * np.sqrt(3)), center[1] + int(dl / 2.0)],
+                        [center[0], center[1] + dl],
+                        [center[0] + int(dl / 2.0 * np.sqrt(3)), center[1] + int(dl / 2.0)],
+                        [center[0] + int(dl / 2.0 * np.sqrt(3)), center[1] - int(dl / 2.0)]],
+                       np.int32)
+        pts = pts.reshape((-1, 1, 2))
+        color = 255 - color
+        width = DEFAULT_WIDTH
+        draw_hexagon(img, pts, color, width)
+    elif type_ == "circle":
+        # Minus because of the way we show the image. See: render_panel's return
+        color = 255 - color
+        unit = min(bbox[2], bbox[3]) * IMAGE_SIZE / 2
+        radius = int(unit * size)
+        width = DEFAULT_WIDTH
+        draw_circle(img, center, radius, color, width)
+    elif type_ == "none":
+        pass
+    # angular
+    if len(bbox) > 4:
+        # [x, y, w, h, x_c, y_c, omega]
+        angle = bbox[6]
+        center = (int(bbox[5] * IMAGE_SIZE), int(bbox[4] * IMAGE_SIZE))
+        img = rotate(img, angle, center=center)
+    # planar 
+    else:
+        img = rotate(img, angle, center=center)
+    # img = shift(img, *entity_position)
+
+    return img
+
+
+def shift(img, dx, dy):
+    M = np.array([[1, 0, dx], [0, 1, dy]], np.float32)
+    img = cv2.warpAffine(img, M, (IMAGE_SIZE, IMAGE_SIZE), flags=cv2.INTER_LINEAR)
+    return img
+
+
+def rotate(img, angle, center=CENTER):
+    M = cv2.getRotationMatrix2D(center, angle, 1)
+    img = cv2.warpAffine(img, M, (IMAGE_SIZE, IMAGE_SIZE), flags=cv2.INTER_LINEAR)
+    return img
+
+
+def scale(img, tx, ty, center=CENTER):
+    M = np.array([[tx, 0, center[0] * (1 - tx)], [0, ty, center[1] * (1 - ty)]], np.float32)
+    img = cv2.warpAffine(img, M, (IMAGE_SIZE, IMAGE_SIZE), flags=cv2.INTER_LINEAR)
+    return img
+
+
+def layer_add(lower_layer_np, higher_layer_np):
+    # higher_layer_np is superimposed on lower_layer_np
+    # new_np = lower_layer_np.copy()
+    # lower_layer_np is modified
+    lower_layer_np[higher_layer_np > 0] = 0
+    return lower_layer_np + higher_layer_np
+
+
+# Draw primitives
+def draw_triangle(img, pts, color, width):
+    # if filled
+    if color != 0:
+        # fill the interior
+        cv2.fillConvexPoly(img, pts, color)
+        # draw the edge
+        cv2.polylines(img, [pts], True, 255, width)
+    # if not filled
+    else:
+        cv2.polylines(img, [pts], True, 255, width)
+
+
+def draw_square(img, pt1, pt2, color, width):
+    # if filled
+    if color != 0:
+        # fill the interior
+        cv2.rectangle(img,
+                      pt1,
+                      pt2,
+                      color,
+                      -1)
+        # draw the edge
+        cv2.rectangle(img,
+                      pt1,
+                      pt2,
+                      255,
+                      width)
+    # if not filled
+    else:
+        cv2.rectangle(img,
+                      pt1,
+                      pt2,
+                      255,
+                      width)
+
+
+def draw_pentagon(img, pts, color, width):
+    # if filled
+    if color != 0:
+        # fill the interior
+        cv2.fillConvexPoly(img, pts, color)
+        # draw the edge
+        cv2.polylines(img, [pts], True, 255, width)
+    # if not filled
+    else:
+        cv2.polylines(img, [pts], True, 255, width)
+
+
+def draw_hexagon(img, pts, color, width):
+    # if filled
+    if color != 0:
+        # fill the interior
+        cv2.fillConvexPoly(img, pts, color)
+        # draw the edge
+        cv2.polylines(img, [pts], True, 255, width)
+    # if not filled
+    else:
+        cv2.polylines(img, [pts], True, 255, width)
+
+
+def draw_circle(img, center, radius, color, width):
+    # if filled
+    if color != 0:
+        # fill the interior
+        cv2.circle(img,
+                   center,
+                   radius,
+                   color,
+                   -1)
+        # draw the edge
+        cv2.circle(img,
+                   center,
+                   radius,
+                   255,
+                   width)
+    # if not filled
+    else:
+        cv2.circle(img,
+                   center,
+                   radius,
+                   255,
+                   width)

raven_utils/render_.py

@@ -0,0 +1,104 @@
+COLOR_VALUES = [255, 224, 196, 168, 140, 112, 84, 56, 28, 0]
+SIZE_VALUES = [0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+TYPE_VALUES = ["none", "triangle", "square", "pentagon", "hexagon", "circle"]
+ANGLE_VALUES = [-135, -90, -45, 0, 45, 90, 135, 180]
+RENDER_POSITIONS_GROUPED = [
+    [(0.5, 0.5, 1, 1)],
+    # ...
+    [(0.25, 0.25, 0.5, 0.5),
+     (0.25, 0.75, 0.5, 0.5),
+     (0.75, 0.25, 0.5, 0.5),
+     (0.75, 0.75, 0.5, 0.5)],
+    # ...
+    [(0.16, 0.16, 0.33, 0.33),
+     (0.16, 0.5, 0.33, 0.33),
+     (0.16, 0.83, 0.33, 0.33),
+     (0.5, 0.16, 0.33, 0.33),
+     (0.5, 0.5, 0.33, 0.33),
+     (0.5, 0.83, 0.33, 0.33),
+     (0.83, 0.16, 0.33, 0.33),
+     (0.83, 0.5, 0.33, 0.33),
+     (0.83, 0.83, 0.33, 0.33)],
+    # ...
+    [(0.5, 0.5, 1, 1)],
+    [(0.5, 0.5, 0.33, 0.33)],
+    # ...
+    [(0.5, 0.5, 1, 1)],
+    [(0.42, 0.42, 0.15, 0.15),
+     (0.42, 0.58, 0.15, 0.15),
+     (0.58, 0.42, 0.15, 0.15),
+     (0.58, 0.58, 0.15, 0.15)],
+    # ....
+    [(0.5, 0.25, 0.5, 0.5)],
+    [(0.5, 0.75, 0.5, 0.5)],
+    # ...
+    [(0.25, 0.5, 0.5, 0.5)],
+    [(0.75, 0.5, 0.5, 0.5)],
+    # ...
+]
+RENDER_POSITIONS = [pos_ for group_pos_ in RENDER_POSITIONS_GROUPED for pos_ in group_pos_]
+MAPPING = {
+    "distribute_nine":
+        {0.16: 0,
+         0.5: 1,
+         0.83: 2},
+    "distribute_four":
+        {0.25: 0,
+         0.75: 1},
+    'in_distribute_four_out_center_single':
+        {0.42: 0,
+         0.58: 1}
+}
+MUL = {
+    "distribute_nine": 3,
+    "distribute_four": 2,
+    'in_distribute_four_out_center_single': 2
+}
+TYPES = ["triangle", "square", "pentagon", "hexagon", "circle"]
+TYPES_NONE = ["none", "triangle", "square", "pentagon", "hexagon", "circle"]
+SIZES = ["vs", "s", "m", "h", "vh", "e"]
+SIZES_NAME = ["Very Small", "Small", "Medium", "High", "Very High", "Enormous"]
+COLORS = ["vs", "s", "m", "h", "vh", "e"]
+
+SAMPLE_TARGET = [[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 9, 5, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 2, 0, 0, 0, 0,
+        0, 1, 3],
+       [1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 2, 0, 3, 2, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 2, 3, 0, 0, 0, 0,
+        0, 3, 3],
+       [2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 3, 2, 0, 0, 0, 0,
+        0, 0, 3],
+       [3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 5, 2, 1, 2, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3,
+        3, 2, 1],
+       [4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 0, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 3,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 3, 2, 1,
+        3, 0, 1],
+       [5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 2, 0, 0, 7, 5, 4, 0, 0, 0, 0, 0, 0, 0, 0, 2, 1, 3, 0, 0, 3, 3,
+        3, 1, 0],
+       [6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+        0, 0, 0, 0, 0, 0, 0, 6, 5, 0, 8, 5, 1, 0, 0, 0, 3, 2, 0, 0, 1, 0,
+        3, 3, 3]]

raven_utils/rules.py

@@ -0,0 +1,21 @@
+from ml_utils import dict_from_list
+COMBINE = "Number/Position"
+
+ATTRIBUTES = [
+    "Number",
+    "Position",
+    "Color",
+    "Size",
+    "Type"
+]
+ATTRIBUTES_LEN = len(ATTRIBUTES)
+ATTRIBUTES_INDEX = dict_from_list(ATTRIBUTES)
+
+TYPES = [
+    "Constant",
+    "Arithmetic",
+    "Progression",
+    "Distribute_Three"
+]
+TYPES_INDEX = dict_from_list(TYPES)
+TYPES_LEN = len(TYPES)

raven_utils/target.py

@@ -0,0 +1,50 @@
+import numpy as np
+import raven_utils.group as group
+import raven_utils.entity as entity
+import raven_utils.rules as rules
+import raven_utils.properties as properties
+
+
+ENTITY_INDEX = entity.INDEX + 1
+ENTITY_DICT = dict(zip(group.NAMES, ENTITY_INDEX[:-1]))
+NAMES_ORDER = dict(zip(group.NAMES, np.arange(len(group.NAMES))))
+PROPERTIES_INDEXES = np.cumsum(np.array(list(entity.NO.values())) * properties.NO)
+INDEX = np.concatenate([[0], PROPERTIES_INDEXES]) + entity.SUM + 1  # +2 type and uniformity
+
+SECOND_LAYOUT = [i - 1 for i in [
+    ENTITY_DICT["in_center_single_out_center_single"] + 1,
+    ENTITY_DICT["in_distribute_four_out_center_single"] + 1,
+    ENTITY_DICT["in_distribute_four_out_center_single"] + 2,
+    ENTITY_DICT["in_distribute_four_out_center_single"] + 3,
+    ENTITY_DICT["left_center_single_right_center_single"] + 1,
+    ENTITY_DICT["up_center_single_down_center_single"] + 1
+]]
+
+FIRST_LAYOUT = list(set(range(entity.SUM)) - set(SECOND_LAYOUT))
+LAYOUT_NO = 2
+
+START_INDEX = dict(zip(group.NAMES, INDEX[:-1]))
+END_INDEX = INDEX[-1]
+
+RULES_ATTRIBUTES_ALL_LEN = rules.ATTRIBUTES_LEN * LAYOUT_NO
+UNIFORMITY_NO = 2
+UNIFORMITY_INDEX = END_INDEX + RULES_ATTRIBUTES_ALL_LEN
+
+SIZE = UNIFORMITY_INDEX + UNIFORMITY_NO
+
+def take(target):
+    return target[1], target[2]
+
+
+def create(images, index, pattern_index=(2, 5), full_index=False, arrange=np.arange, shape=lambda x: x.shape):
+    return [images[:, pattern_index[0]], images[:, pattern_index[1]],
+            images[arrange(shape(index)[0]), (0 if full_index else 8) + index[:, 0]]]
+
+
+
+def take_simple(target):
+    return target[1], target[0]
+
+
+def create_simple(images, target, index=slice(None), pattern_index=(2, 5)):
+    return [images[:, pattern_index[0]], images[:, pattern_index[1]], target][index]

raven_utils/uitls.py

@@ -0,0 +1,64 @@
+from functools import partial
+from itertools import product
+
+import numpy as np
+from funcy import identity
+
+from data_utils import gather, DataGenerator, Data
+from data_utils.sampling import DataSampler
+from models_utils import init_image as def_init_image, INPUTS, TARGET
+
+import raven_utils.group as group
+
+from data_utils import ops as D
+
+init_image = partial(def_init_image, shape=(16, 8, 80, 80, 1))
+
+
+def get_val_index(no=group.NO, base=3,add_end=False):
+    indexes = np.arange(no) * 2000 + base
+    if add_end:
+        indexes = np.concatenate([indexes, no*2000])
+    return indexes
+
+
+def get_matrix(inputs, index):
+    return np.concatenate([inputs[:, :8], gather(inputs, index[:, 0])[:, None]], axis=1)
+
+
+def get_matrix_from_data(x):
+    inputs = x["inputs"]
+    index = x["index"]
+    return get_matrix(inputs, index)
+
+
+def get_data_class(data, batch_size=128):
+    fn = identity
+    shape = data[0].shape
+    train_generator = DataGenerator(
+        {
+            INPUTS: Data(data[0], fn),
+            TARGET: Data(data[2], fn),
+        },
+        sampler=DataSampler(np.array(list(product(np.arange(shape[0]), np.arange(shape[1]))))),
+        batch=batch_size
+    )
+    shape = data[1].shape
+    val_generator = DataGenerator(
+        {
+            INPUTS: Data(data[1], fn),
+            TARGET: Data(data[3], fn),
+        },
+        sampler=DataSampler(np.array(list(product(np.arange(shape[0]), np.arange(shape[1])))), shuffle=False),
+        batch=batch_size
+    )
+    return train_generator, val_generator
+
+
+def compare_from_result(result, data):
+    data = data.data.data
+    answer = D.gather(data['target'].data, data['index'].data[:, 0])
+    import raven_utils as rv
+    predict = result['predict']
+    predict_mask = result['predict_mask']
+    return np.all(rv.decode.compare(answer[:len(predict)], predict, predict_mask), axis=-1)

saved_model/1/keras_metadata.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3065f1580247f096711cd61201a17a730a1e5a3d719f2c2778030dea78bb17b4
+size 730275

saved_model/1/saved_model.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:48d10d74324a5e993ceacc0a4bffc1fcb232d7e2f708a2ebbeabd864650baeeb
+size 12159312

saved_model/1/variables/variables.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e3c44b273228c834166b40b8a062a53dce76cc21d4cce42f65df2edc53533a7
+size 43002413

saved_model/1/variables/variables.index

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e69bbffa5d415625538b629762f1aaeeb355a83d676242110af3d633e31017dd
+size 24958

utils.py

@@ -0,0 +1,84 @@
+import numpy as np
+from data_utils.image import draw_images
+from ml_utils import il
+
+import raven_utils as rv
+from raven_utils.uitls import get_matrix
+from tensorflow.keras.models import load_model
+from raven_utils.draw import render_from_model
+import models
+import ast
+
+
+def load_example(index=0):
+    index = ast.literal_eval(str(index))
+    if il(index):
+        example = rv.draw.render_panels(np.array(index))
+        desc = "Custom matrix"
+    else:
+        if not index:
+            index = 0
+        index = int(index)
+
+        desc = rv.draw.extract_rules(models.properties[index])
+        desc = "<br /><br />".join(["<br />".join(d) for d in desc])
+
+        example = get_matrix(models.data[index:index + 1], models.indexes[index:index + 1, None] + 8)
+    result = np.tile(draw_images(example[:9], row=3), reps=(1, 1, 3))
+    return result, desc
+
+
+def load_model_(name):
+    if name == "Transformer":
+        path = "/home/jkwiatkowski/all/best/rav/full_trans/6e8e6bad403e4171ad10daa1a518ba09"
+    else:
+        path = name
+    models.model = load_model(path)
+    return f"Success loading: {name}"
+
+
+def run_nn(index=0):
+    index = ast.literal_eval(str(index))
+    if il(index):
+        data = rv.draw.render_panels(np.array(index))
+        data = np.concatenate([data, data[:7]])[None]
+    else:
+        if not index:
+            index = models.START_IMAGE
+        index = int(index)
+        data = models.data[index:index + 1]
+
+    # model = load_model("/home/jkwiatkowski/all/best/rav/full_trans/6e8e6bad403e4171ad10daa1a518ba09")
+    data = {
+        'inputs': data,
+        'index': np.zeros(shape=(1, 1), dtype="uint8"),
+        'labels': np.zeros(shape=(1, 16, 113), dtype="int8"),
+        'target': np.zeros(shape=(1, 16, 113), dtype="int8"),
+        # 'features': np.zeros(shape=(1, 16, 64), dtype="float32")
+    }
+    res = np.tile(render_from_model(data, models.model)[0, ..., None], reps=(1, 1, 3))
+
+    # res = model({'inputs': data[0:1]})
+
+    return res
+
+
+def next_(index=0):
+    index = ast.literal_eval(str(index))
+    if not isinstance(index, int):
+        index = models.START_IMAGE
+    index = int(index) + 1
+    return (index,) + load_example(index)
+
+
+def prev_(index=0):
+    index = ast.literal_eval(str(index))
+    if not isinstance(index, int):
+        index = models.START_IMAGE
+    index = int(index) - 1
+    return (index,) + load_example(index)
+
+
+if __name__ == '__main__':
+    image, _ = load_example(5)
+    run_nn(image)