train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Script for training EBMs for discovering concepts, relations and operators.\n",
    "\"\"\"\n",
    "try:\n",
    "    %load_ext autoreload\n",
    "    %autoreload 2\n",
    "except:\n",
    "    pass\n",
    "\n",
    "import argparse\n",
    "from collections import OrderedDict, Iterable\n",
    "from copy import deepcopy\n",
    "from datetime import datetime\n",
    "import itertools\n",
    "import logging\n",
    "logging.getLogger('matplotlib.font_manager').disabled = True\n",
    "from numbers import Number\n",
    "import os\n",
    "import pdb\n",
    "import pickle\n",
    "import pprint as pp\n",
    "import random\n",
    "from typing import Union\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.distributions.normal import Normal\n",
    "from torch.distributions.one_hot_categorical import OneHotCategorical as cat\n",
    "from torchvision import datasets, transforms, utils\n",
    "from torchvision.transforms import functional as F_tr\n",
    "from torchvision.transforms.functional import InterpolationMode\n",
    "from einops import rearrange, repeat\n",
    "import torch.multiprocessing\n",
    "torch.multiprocessing.set_sharing_strategy('file_system')\n",
    "\n",
    "import sys, os\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "from zeroc.datasets.arc_image import ARCDataset\n",
    "from zeroc.datasets.BabyARC.code.dataset.dataset import *\n",
    "from zeroc.utils import ClevrImagePreprocessor\n",
    "# from zeroc.clevr_dataset_gen.dataset import ClevrRelationDataset, create_easy_dataset\n",
    "# from zeroc.clevr_dataset_gen.generate_concept_dataset import get_clevr_concept_data\n",
    "from zeroc.argparser import get_args_EBM\n",
    "# from zeroc.slot_attention.clevr import CLEVR\n",
    "# from zeroc.slot_attention.multi_dsprites import MultiDsprites\n",
    "# from zeroc.slot_attention.tetrominoes import Tetrominoes\n",
    "from zeroc.concept_library.models import get_model_energy, load_model_energy, neg_mask_sgd, neg_mask_sgd_with_kl, id_to_tensor, requires_grad\n",
    "from zeroc.concept_library.settings import REPR_DIM, DEFAULT_OBJ_TYPE\n",
    "from zeroc.utils import REA_PATH, EXP_PATH, get_root_dir\n",
    "from zeroc.concept_transfer import convert_babyarc\n",
    "from zeroc.concept_library.util import to_cpu_recur, try_call, Printer, transform_dict, MineDataset, is_diagnose, reduce_tensor, get_hashing, pdump, pload, remove_elements, loss_op_core, filter_kwargs, to_Variable, gather_broadcast, get_pdict, COLOR_LIST, set_seed, Zip, Early_Stopping, init_args, make_dir, str2bool, get_filename, get_filename_short, get_machine_name, get_device, record_data, plot_matrices, filter_filename, get_next_available_key, to_np_array, to_Variable, get_filename_short, write_to_config, Dictionary, Batch, to_cpu\n",
    "from zeroc.concept_library.util import model_parallel, color_dict, clip_grad, identity_fun, seperate_concept, to_one_hot, onehot_to_RGB, get_module_parameters, assign_embedding_value, get_hashing, to_device_recur, visualize_matrices, repeat_n, mask_iou_score, shrink, get_obj_from_mask\n",
    "p = Printer()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Dataset:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 ConceptDataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConceptDataset(Dataset):\n",
    "    \"\"\"Concept Dataset for learning basic concepts for ARC.\n",
    "\n",
    "    mode:\n",
    "        Concepts:  E(x; a; c)\n",
    "            \"Pixel\": one or many pixels\n",
    "            \"Line\": one or many lines\n",
    "            \"Rect\": hollow rectangles\n",
    "            \"{}+{}+...\": each \"{}\" can be a concept.\n",
    "\n",
    "        Symmetries: E(x; a; c)\n",
    "            \"hFlip\", \"vFlip\": one image where some object has property of symmetry w.r.t. hflip\n",
    "            \"Rotate\": one image where some object has property of symmetry w.r.t. rotation.\n",
    "\n",
    "        Relations: E(x; a1, a2; c)\n",
    "            \"Vertical\": lines where some of them are vertical\n",
    "            \"Parallel\": lines where some of them are parallel\n",
    "            \"Vertical+Parallel\": lines where some of them are vertical or parallel\n",
    "            \"IsInside\": obj_1 is inside obj_2\n",
    "            \"SameRow\": obj_1 and obj_2 are at the same row\n",
    "            \"SameCol\": obj_1 and obj_2 are at the same column\n",
    "\n",
    "        Operations: E(x1,x2; a1,a2; c1,c2)\n",
    "            \"RotateA+vFlip(Line+Rect)\": two images where some object1 in image1 is rotated or vertically-flipped w.r.t. some object2 in image2, and the objects are chosen from Line or Rect.\n",
    "            \"hFlip(Lshape)\", \"vFlip(Lshape+Line)\": two images where some object1 in image1 is flipped w.r.t. some object2 in image2.\n",
    "\n",
    "        ARC+:\n",
    "            \"arc^{}\": ARC images with property \"{}\" masked as above.\n",
    "        \"\"\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        mode,\n",
    "        canvas_size=8,\n",
    "        n_examples=10000,\n",
    "        rainbow_prob=0.,\n",
    "        data=None,\n",
    "        idx_list=None,\n",
    "        concept_collection=None,\n",
    "        allowed_shape_concept=None,\n",
    "        w_type=\"image+mask\",\n",
    "        color_avail=\"-1\",\n",
    "        min_n_distractors=0,\n",
    "        max_n_distractors=-1,\n",
    "        n_operators=1,\n",
    "        allow_connect=True,\n",
    "        parsing_check=False,\n",
    "        focus_type=None,\n",
    "        transform=None,\n",
    "        save_interval=-1,\n",
    "        save_filename=None,\n",
    "    ):\n",
    "        if allowed_shape_concept is None:\n",
    "            allowed_shape_concept=[\n",
    "                \"Line\", \"Rect\", \"RectSolid\", \"Lshape\", \"Randshape\", \"ARCshape\",\n",
    "                \"Tshape\", \"Eshape\",\n",
    "                \"Hshape\", \"Cshape\", \"Ashape\", \"Fshape\",\n",
    "                \"RectE1a\", \"RectE1b\", \"RectE1c\", \n",
    "                \"RectE2a\", \"RectE2b\", \"RectE2c\",\n",
    "                \"RectE3a\", \"RectE3b\", \"RectE3c\", \n",
    "                \"RectF1a\", \"RectF1b\", \"RectF1c\", \n",
    "                \"RectF2a\", \"RectF2b\", \"RectF2c\",\n",
    "                \"RectF3a\", \"RectF3b\", \"RectF3c\",\n",
    "            ]\n",
    "        self.mode = mode\n",
    "        self.canvas_size = canvas_size\n",
    "        self.rainbow_prob = rainbow_prob\n",
    "        self.n_examples = n_examples\n",
    "        self.allowed_shape_concept = allowed_shape_concept\n",
    "        self.min_n_distractors = min_n_distractors\n",
    "        self.max_n_distractors = max_n_distractors\n",
    "        self.n_operators = n_operators\n",
    "        self.w_type = w_type\n",
    "        self.allow_connect = allow_connect\n",
    "        self.parsing_check = parsing_check\n",
    "        self.focus_type = focus_type\n",
    "        if isinstance(color_avail, str):\n",
    "            if color_avail == \"-1\":\n",
    "                self.color_avail = None\n",
    "            else:\n",
    "                self.color_avail = [int(c) for c in color_avail.split(\",\")]\n",
    "                for c in self.color_avail:\n",
    "                    assert c >= 1 and c <= 9\n",
    "        else:\n",
    "            self.color_avail = color_avail\n",
    "        if idx_list is None:\n",
    "            assert data is None\n",
    "            if mode.startswith(\"arc^\"):\n",
    "                if \"(\" in mode:\n",
    "                    self.data = []\n",
    "                    # Operator:\n",
    "                    concept_raw = mode.split(\"(\")[0].split(\"+\")\n",
    "                    concept_collection = []\n",
    "                    for c in concept_raw:\n",
    "                        if \"^\" in c:\n",
    "                            concept_collection.append(c.split(\"^\")[1])\n",
    "                        else:\n",
    "                            concept_collection.append(c)\n",
    "                    self.concept_collection = concept_collection\n",
    "                    arcDataset = ARCDataset(\n",
    "                        n_examples=n_examples*2,\n",
    "                        canvas_size=canvas_size,\n",
    "                    )\n",
    "                    babyArcDataset = BabyARCDataset(\n",
    "                        pretrained_obj_cache=os.path.join(get_root_dir(), 'datasets/arc_objs.pt'),\n",
    "                        save_directory=get_root_dir() + \"/datasets/BabyARCDataset/\",\n",
    "                        object_limit=None,\n",
    "                        noise_level=0,\n",
    "                        canvas_size=canvas_size,\n",
    "                    )\n",
    "                    if set(self.concept_collection).issubset({\"RotateA\", \"RotateB\", \"RotateC\", \n",
    "                                                              \"hFlip\", \"vFlip\", \"DiagFlipA\", \"DiagFlipB\", \n",
    "                                                              \"Identity\", \n",
    "                                                              \"Move\"}):\n",
    "                        for arc_example_one_hot in arcDataset:\n",
    "                            arc_image = torch.zeros_like(arc_example_one_hot[0])\n",
    "                            for i in range(0, 10):\n",
    "                                arc_image += arc_example_one_hot[i]*i\n",
    "                            arc_image = arc_image.type(torch.int32)\n",
    "\n",
    "                            repre_dict = babyArcDataset.sample_task_canvas_from_arc(\n",
    "                                arc_image,\n",
    "                                color=np.random.choice([True, False], p=[0.6, 0.4]),\n",
    "                                is_plot=False,\n",
    "                            )\n",
    "                            if repre_dict == -1:\n",
    "                                continue\n",
    "                            in_canvas = Canvas(repre_dict=repre_dict)\n",
    "\n",
    "                            # Operate on the input:\n",
    "                            if len(list(repre_dict[\"node_id_map\"].keys())) == 0:\n",
    "                                continue # empty arc canvas\n",
    "                            chosen_obj_key = np.random.choice(list(repre_dict[\"node_id_map\"].keys()))\n",
    "                            chosen_obj_id = repre_dict[\"node_id_map\"][chosen_obj_key]\n",
    "                            chosen_op = np.random.choice(self.concept_collection)\n",
    "                            if chosen_op in [\"Identity\"]:\n",
    "                                inplace = True if random.random() < 0.5 else False\n",
    "                                out_canvas_list, concept = OperatorEngine().operator_identity(\n",
    "                                    [in_canvas],\n",
    "                                    [[chosen_obj_key]],\n",
    "                                    inplace=inplace,\n",
    "                                )\n",
    "                                if out_canvas_list == -1:\n",
    "                                    continue\n",
    "                            elif chosen_op in [\"Move\"]:\n",
    "                                # create operator spec as move is a complex operator\n",
    "                                move_spec = OperatorMoveSpec(\n",
    "                                                autonomous=False,\n",
    "                                                direction=random.randint(0,3), \n",
    "                                                distance=-1, \n",
    "                                                hit_type=None, # either wall, agent or None\n",
    "                                                linkage_move=False, \n",
    "                                                linkage_move_distance_ratio=None,\n",
    "                                            )\n",
    "                                out_canvas_list, concept = OperatorEngine().operator_move(\n",
    "                                    [in_canvas],\n",
    "                                    [[chosen_obj_key]],\n",
    "                                    [[move_spec]], \n",
    "                                    allow_overlap=False, \n",
    "                                    allow_shape_break=False,\n",
    "                                    allow_connect=self.allow_connect,\n",
    "                                    allow_stay=False,\n",
    "                                )\n",
    "                                if out_canvas_list == -1:\n",
    "                                    continue\n",
    "                            elif chosen_op in [\"RotateA\", \"RotateB\", \"RotateC\", \"hFlip\", \"vFlip\", \"DiagFlipA\", \"DiagFlipB\"]:\n",
    "                                out_canvas_list, concept = OperatorEngine().operate_rotate(\n",
    "                                    [in_canvas],\n",
    "                                    [[chosen_obj_key]],\n",
    "                                    operator_tag=f\"#{chosen_op}\",\n",
    "                                    allow_connect=self.allow_connect,\n",
    "                                    allow_shape_break=False,\n",
    "                                )\n",
    "                                if out_canvas_list == -1:\n",
    "                                    continue\n",
    "                            else:\n",
    "                                raise Exception(f\"operator={chosen_op} is not supported!\")\n",
    "\n",
    "                            # Add to self.data:\n",
    "                            in_canvas_dict = in_canvas.repr_as_dict()\n",
    "                            out_canvas_dict = out_canvas_list[0].repr_as_dict()\n",
    "                            in_mask = in_canvas_dict[\"id_object_mask\"][chosen_obj_id][None]\n",
    "                            out_mask = out_canvas_dict[\"id_object_mask\"][chosen_obj_id][None]\n",
    "                            self.data.append(\n",
    "                                ((to_one_hot(in_canvas_dict[\"image_t\"]), to_one_hot(out_canvas_dict[\"image_t\"])),\n",
    "                                 (in_mask, out_mask),\n",
    "                                 chosen_op,\n",
    "                                 Dictionary({}),\n",
    "                                )\n",
    "                            )\n",
    "                            if len(self.data) >= n_examples:\n",
    "                                break\n",
    "                            if i > n_examples * 2 and len(self.data) < n_examples * 0.05:\n",
    "                                raise Exception(\"Sampled {} times and only {} of them satisfies the specified condition. Try relaxing the condition!\".format(i, len(self.data)))\n",
    "                else:\n",
    "                    mode_core = mode.split(\"^\")[1]\n",
    "                    self.concept_collection = mode_core.split(\"+\")\n",
    "                    dataset = ARCDataset(\n",
    "                        n_examples=n_examples,\n",
    "                        canvas_size=canvas_size,\n",
    "                    )\n",
    "                    examples_all = []\n",
    "                    masks_all = []\n",
    "                    concepts_all = []\n",
    "                    examples = dataset.data\n",
    "                    examples_argmax = examples.argmax(1)\n",
    "                    self.data = []\n",
    "                    for i in range(len(examples)):\n",
    "                        concept_dict = seperate_concept(examples_argmax[i])\n",
    "                        masks, concepts = get_masks(concept_dict, allowed_concepts=self.concept_collection, canvas_size=canvas_size)\n",
    "                        if masks is not None:\n",
    "                            for mask, concept in zip(masks, concepts):\n",
    "                                self.data.append((\n",
    "                                        examples[i],\n",
    "                                        (mask,),\n",
    "                                        concept,\n",
    "                                        Dictionary({}),\n",
    "                                    )\n",
    "                                )\n",
    "            else:\n",
    "                if \"(\" in mode:\n",
    "                    # Operator:\n",
    "                    self.concept_collection = mode.split(\"(\")[0].split(\"+\")\n",
    "                    input_concepts = mode.split(\"(\")[1][:-1].split(\"+\")\n",
    "                else:\n",
    "                    self.concept_collection = mode.split(\"-\")[-1].split(\"+\")\n",
    "                    input_concepts = [\"\"]\n",
    "                dataset = BabyARCDataset(\n",
    "                    pretrained_obj_cache=os.path.join(get_root_dir(), 'datasets/arc_objs.pt'),\n",
    "                    save_directory=get_root_dir() + \"/datasets/BabyARCDataset/\",\n",
    "                    object_limit=None,\n",
    "                    noise_level=0,\n",
    "                    canvas_size=canvas_size,\n",
    "                )\n",
    "                concept_str_mapping = {\n",
    "                    \"line\": \"Line\", \n",
    "                    \"rectangle\": \"Rect\", \n",
    "                    \"rectangleSolid\": \"RectSolid\",\n",
    "                    \"Lshape\": \"Lshape\", \n",
    "                    \"Tshape\": \"Tshape\", \n",
    "                    \"Eshape\": \"Eshape\", \n",
    "                    \"Hshape\": \"Hshape\", \n",
    "                    \"Cshape\": \"Cshape\", \n",
    "                    \"Ashape\": \"Ashape\", \n",
    "                    \"Fshape\": \"Fshape\",\n",
    "                    \"randomShape\": \"Randshape\",\n",
    "                    \"arcShape\": \"ARCshape\"}  # Mapping between two conventions\n",
    "                concept_str_reverse_mapping = {\n",
    "                    \"Line\": \"line\", \n",
    "                    \"Rect\": \"rectangle\", \n",
    "                    \"RectSolid\": \"rectangleSolid\", \n",
    "                    \"Lshape\": \"Lshape\", \n",
    "                    \"Tshape\": \"Tshape\", \n",
    "                    \"Eshape\": \"Eshape\", \n",
    "                    \"Hshape\": \"Hshape\", \n",
    "                    \"Cshape\": \"Cshape\", \n",
    "                    \"Ashape\": \"Ashape\", \n",
    "                    \"Fshape\": \"Fshape\",\n",
    "                    \"Randshape\": \"randomShape\",\n",
    "                    \"ARCshape\": \"arcShape\"}  # Mapping between two conventions\n",
    "                composite_concepts = [                \n",
    "                    \"RectE1a\", \"RectE1b\", \"RectE1c\", \n",
    "                    \"RectE2a\", \"RectE2b\", \"RectE2c\",\n",
    "                    \"RectE3a\", \"RectE3b\", \"RectE3c\", \n",
    "                    \"RectF1a\", \"RectF1b\", \"RectF1c\", \n",
    "                    \"RectF2a\", \"RectF2b\", \"RectF2c\",\n",
    "                    \"RectF3a\", \"RectF3b\", \"RectF3c\",\n",
    "                ]\n",
    "                for c in composite_concepts:\n",
    "                    concept_str_mapping[c] = c\n",
    "                    concept_str_reverse_mapping[c] = c\n",
    "\n",
    "                if set(get_c_core(self.concept_collection)).issubset({\n",
    "                    \"Image\"\n",
    "                }):\n",
    "                    # Image is a collection of all shapes.\n",
    "                    if max_n_distractors == -1:\n",
    "                        max_n_objs = 3\n",
    "                    else:\n",
    "                        max_n_objs = 1 + max_n_distractors # 1 is for the core concept itself.\n",
    "                    self.data = generate_samples(\n",
    "                        dataset=dataset,\n",
    "                        obj_spec_fun=obj_spec_fun,\n",
    "                        n_examples=n_examples,\n",
    "                        mode=\"concept-image\",\n",
    "                        concept_collection=[\"Line\", \"Rect\", \"Lshape\", \n",
    "                                            \"RectSolid\", \"Randshape\", \"ARCshape\", \n",
    "                                            \"Tshape\", \"Eshape\", \n",
    "                                            \"Hshape\", \"Cshape\", \"Ashape\", \"Fshape\"],\n",
    "                        min_n_objs=1+self.min_n_distractors,\n",
    "                        max_n_objs=max_n_objs,\n",
    "                        canvas_size=canvas_size,\n",
    "                        rainbow_prob=rainbow_prob,\n",
    "                        concept_str_mapping=concept_str_mapping,\n",
    "                        concept_str_reverse_mapping=concept_str_reverse_mapping,\n",
    "                        allowed_shape_concept=[\"Line\", \"Rect\", \"Lshape\", \n",
    "                                               \"RectSolid\", \"Randshape\", \"ARCshape\", \n",
    "                                               \"Tshape\", \"Eshape\", \n",
    "                                               \"Hshape\", \"Cshape\", \"Ashape\", \"Fshape\"],\n",
    "                        color_avail=self.color_avail,\n",
    "                        allow_connect=self.allow_connect,\n",
    "                        parsing_check=self.parsing_check,\n",
    "                        save_interval=10,\n",
    "                        save_filename=save_filename,\n",
    "                    )\n",
    "                elif set(get_c_core(self.concept_collection)).issubset({\n",
    "                    \"Line\", \"Rect\", \"Lshape\",\n",
    "                    \"RectSolid\", \"Randshape\", \"ARCshape\",\n",
    "                    \"Tshape\", \"Eshape\",\n",
    "                    \"Hshape\", \"Cshape\", \"Ashape\", \"Fshape\",\n",
    "                    \"RectE1a\", \"RectE1b\", \"RectE1c\", \n",
    "                    \"RectE2a\", \"RectE2b\", \"RectE2c\",\n",
    "                    \"RectE3a\", \"RectE3b\", \"RectE3c\", \n",
    "                    \"RectF1a\", \"RectF1b\", \"RectF1c\", \n",
    "                    \"RectF2a\", \"RectF2b\", \"RectF2c\",\n",
    "                    \"RectF3a\", \"RectF3b\", \"RectF3c\",\n",
    "                }):\n",
    "                    if max_n_distractors == -1:\n",
    "                        max_n_objs = 3\n",
    "                    else:\n",
    "                        max_n_objs = 1 + max_n_distractors # 1 is for the core concept itself.\n",
    "                    self.data = generate_samples(\n",
    "                        dataset=dataset,\n",
    "                        obj_spec_fun=obj_spec_fun,\n",
    "                        n_examples=n_examples,\n",
    "                        mode=\"concept\",\n",
    "                        concept_collection=self.concept_collection,\n",
    "                        min_n_objs=1+self.min_n_distractors,\n",
    "                        max_n_objs=max_n_objs,\n",
    "                        canvas_size=canvas_size,\n",
    "                        rainbow_prob=rainbow_prob,\n",
    "                        concept_str_mapping=concept_str_mapping,\n",
    "                        concept_str_reverse_mapping=concept_str_reverse_mapping,\n",
    "                        allowed_shape_concept=self.allowed_shape_concept,\n",
    "                        color_avail=self.color_avail,\n",
    "                        allow_connect=self.allow_connect,\n",
    "                        parsing_check=self.parsing_check,\n",
    "                        focus_type=self.focus_type,\n",
    "                        save_interval=10,\n",
    "                        save_filename=save_filename,\n",
    "                    )\n",
    "                elif set(self.concept_collection).issubset({\n",
    "                    \"RectE1a\", \"RectE1b\", \"RectE1c\", \n",
    "                    \"RectE2a\", \"RectE2b\", \"RectE2c\",\n",
    "                    \"RectE3a\", \"RectE3b\", \"RectE3c\", \n",
    "                    \"RectF1a\", \"RectF1b\", \"RectF1c\", \n",
    "                    \"RectF2a\", \"RectF2b\", \"RectF2c\",\n",
    "                    \"RectF3a\", \"RectF3b\", \"RectF3c\",\n",
    "                }):\n",
    "                    max_n_objs = 1 # we currently don't allow distractors to be sampled.\n",
    "                    self.data = generate_samples(\n",
    "                        dataset=dataset,\n",
    "                        obj_spec_fun=obj_spec_fun,\n",
    "                        n_examples=n_examples,\n",
    "                        mode=\"compositional-concept\",\n",
    "                        concept_collection=self.concept_collection,\n",
    "                        min_n_objs=1+self.min_n_distractors,\n",
    "                        max_n_objs=max_n_objs,\n",
    "                        canvas_size=canvas_size,\n",
    "                        rainbow_prob=rainbow_prob,\n",
    "                        concept_str_mapping=concept_str_mapping,\n",
    "                        concept_str_reverse_mapping=concept_str_reverse_mapping,\n",
    "                        allowed_shape_concept=self.allowed_shape_concept,\n",
    "                        color_avail=self.color_avail,\n",
    "                        allow_connect=self.allow_connect,\n",
    "                        parsing_check=self.parsing_check,\n",
    "                        focus_type=self.focus_type,\n",
    "                        save_interval=10,\n",
    "                        save_filename=save_filename,\n",
    "                    )\n",
    "                elif set(self.concept_collection).issubset({\"Vertical\", \"Parallel\"}):\n",
    "                    if max_n_distractors == -1:\n",
    "                        max_n_objs = 3\n",
    "                    else:\n",
    "                        max_n_objs = 2 + max_n_distractors # 2 is for the core concept itself.\n",
    "                    def obj_spec_fun_re(concept_collection, min_n_objs, max_n_objs, canvas_size, allowed_shape_concept=None, color_avail=None, focus_type=None):\n",
    "                        n_objs = np.random.randint(min_n_objs, max_n_objs+1)\n",
    "                        obj_spec = [(('obj_{}'.format(k), 'line_[-1,1,-1]'), 'Attr') for k in range(n_objs)]\n",
    "                        return obj_spec\n",
    "                    self.data = generate_samples(\n",
    "                        dataset=dataset,\n",
    "                        obj_spec_fun=obj_spec_fun_re,\n",
    "                        n_examples=n_examples,\n",
    "                        mode=\"relation\",\n",
    "                        concept_collection=self.concept_collection,\n",
    "                        min_n_objs=2+self.min_n_distractors,\n",
    "                        max_n_objs=max_n_objs,\n",
    "                        canvas_size=canvas_size,\n",
    "                        rainbow_prob=rainbow_prob,\n",
    "                        concept_str_mapping=concept_str_mapping,\n",
    "                        concept_str_reverse_mapping=concept_str_reverse_mapping,\n",
    "                        allowed_shape_concept=self.allowed_shape_concept,\n",
    "                        color_avail=self.color_avail,\n",
    "                        allow_connect=self.allow_connect,\n",
    "                        parsing_check=self.parsing_check,\n",
    "                        save_interval=10,\n",
    "                        save_filename=save_filename,\n",
    "                    )\n",
    "                elif set(self.concept_collection).issubset({\"VerticalMid\", \"VerticalEdge\", \"VerticalSepa\", \"Parallel\"}):\n",
    "                    if max_n_distractors == -1:\n",
    "                        max_n_objs = 3\n",
    "                    else:\n",
    "                        max_n_objs = 2 + max_n_distractors # 2 is for the core concept itself.\n",
    "                    self.data = generate_lines_full_vertical_parallel(\n",
    "                        n_examples=n_examples,\n",
    "                        concept_collection=self.concept_collection,\n",
    "                        min_n_objs=2+self.min_n_distractors,\n",
    "                        max_n_objs=max_n_objs,\n",
    "                        canvas_size=canvas_size,\n",
    "                        min_size=3,\n",
    "                        max_size=canvas_size-2,\n",
    "                        color_avail=self.color_avail,\n",
    "                        isplot=False,\n",
    "                    )\n",
    "                elif set(self.concept_collection).issubset({\n",
    "                    \"SameAll\", \"SameShape\", \"SameColor\", \n",
    "                    \"SameRow\", \"SameCol\", \"IsInside\", \n",
    "                    \"IsTouch\", \"IsNonOverlapXY\",\n",
    "                    \"IsEnclosed\",\n",
    "                }):\n",
    "                    if max_n_distractors == -1:\n",
    "                        max_n_objs = 3\n",
    "                    else:\n",
    "                        max_n_objs = 2 + max_n_distractors # 2 is for the core relation itself.\n",
    "                    def obj_spec_fun_re(\n",
    "                        concept_collection, min_n_objs, max_n_objs, \n",
    "                        canvas_size, allowed_shape_concept=None, \n",
    "                        color_avail=None,\n",
    "                        focus_type=None,\n",
    "                    ):\n",
    "                        assert allowed_shape_concept != None\n",
    "                        n_objs = np.random.randint(min_n_objs, max_n_objs+1)\n",
    "                        # two slots are for the relation\n",
    "                        sampled_relation = np.random.choice(concept_collection)\n",
    "                        obj_spec = [(('obj_0', 'obj_1'), sampled_relation)]\n",
    "                        max_rect_size = canvas_size//2\n",
    "                        # choose distractors\n",
    "                        for k in range(2, n_objs):\n",
    "                            # choose a distractor shape\n",
    "                            distractor_shape = np.random.choice(allowed_shape_concept)\n",
    "                            if distractor_shape == \"Line\":\n",
    "                                obj_spec += [(('obj_{}'.format(k), 'line_[-1,-1,-1]'), 'Attr')]\n",
    "                            elif distractor_shape == \"Rect\":\n",
    "                                obj_spec += [(('obj_{}'.format(k), 'rectangle_[-1,-1]'), 'Attr')]\n",
    "                            elif distractor_shape == \"RectSolid\":\n",
    "                                obj_spec += [(('obj_{}'.format(k), 'rectangleSolid_[-1,-1]'), 'Attr')]\n",
    "                            elif distractor_shape == \"Lshape\":\n",
    "                                obj_spec += [(('obj_{}'.format(k), 'Lshape_[-1,-1,-1]'), 'Attr')]\n",
    "                            elif distractor_shape == \"Tshape\":\n",
    "                                w = np.random.randint(3, max_rect_size+2)\n",
    "                                h = np.random.randint(3, max_rect_size+2)\n",
    "                                obj_spec += [(('obj_{}'.format(k), f'Tshape_[{w},{h}]'), 'Attr')]\n",
    "                            elif distractor_shape == \"Eshape\":\n",
    "                                w = np.random.randint(3, max_rect_size+1)\n",
    "                                h = np.random.randint(5, max_rect_size+3)\n",
    "                                obj_spec += [(('obj_{}'.format(k), f'Eshape_[{w},{h}]'), 'Attr')]\n",
    "                            elif distractor_shape == \"Hshape\":\n",
    "                                w = np.random.randint(3, max_rect_size+2)\n",
    "                                h = np.random.randint(3, max_rect_size+2)\n",
    "                                obj_spec += [(('obj_{}'.format(k), f'Hshape_[{w},{h}]'), 'Attr')]\n",
    "                            elif distractor_shape == \"Cshape\":\n",
    "                                w = np.random.randint(3, max_rect_size+1)\n",
    "                                h = np.random.randint(3, max_rect_size+2)\n",
    "                                obj_spec += [(('obj_{}'.format(k), f'Cshape_[{w},{h}]'), 'Attr')]\n",
    "                            elif distractor_shape == \"Ashape\":\n",
    "                                w = np.random.randint(3, max_rect_size+2)\n",
    "                                h = np.random.randint(4, max_rect_size+3)\n",
    "                                obj_spec += [(('obj_{}'.format(k), f'Ashape_[{w},{h}]'), 'Attr')]\n",
    "                            elif distractor_shape == \"Fshape\":\n",
    "                                w = np.random.randint(3, max_rect_size+1)\n",
    "                                h = np.random.randint(4, max_rect_size+3)\n",
    "                                obj_spec += [(('obj_{}'.format(k), f'Fshape_[{w},{h}]'), 'Attr')]   \n",
    "                            elif distractor_shape == \"Randshape\":\n",
    "                                max_rect_size = canvas_size // 2\n",
    "                                w, h = np.random.randint(2, max_rect_size+1, size=2) # hard-code for the size\n",
    "                                obj_spec += [(('obj_{}'.format(k), f'randomShape_[{w},{h}]'), 'Attr')]\n",
    "                            elif distractor_shape == \"ARCshape\":\n",
    "                                max_rect_size = canvas_size // 2\n",
    "                                w, h = np.random.randint(2, max_rect_size+1, size=2) # hard-code for the size\n",
    "                                obj_spec += [(('obj_{}'.format(k), f'arcShape_[{w},{h}]'), 'Attr')]\n",
    "                        return obj_spec\n",
    "                    if len(input_concepts) == 1 and input_concepts[0] == \"\":\n",
    "                        _shape_concept=[c for c in self.allowed_shape_concept]\n",
    "                    else:\n",
    "                        _shape_concept=[c for c in input_concepts]\n",
    "\n",
    "                    self.data = generate_samples(\n",
    "                        dataset=dataset,\n",
    "                        obj_spec_fun=obj_spec_fun_re,\n",
    "                        n_examples=n_examples,\n",
    "                        mode=\"relation\",\n",
    "                        concept_collection=self.concept_collection,\n",
    "                        min_n_objs=2+self.min_n_distractors,\n",
    "                        max_n_objs=max_n_objs,\n",
    "                        canvas_size=canvas_size,\n",
    "                        rainbow_prob=rainbow_prob,\n",
    "                        concept_str_mapping=concept_str_mapping,\n",
    "                        concept_str_reverse_mapping=concept_str_reverse_mapping,\n",
    "                        allowed_shape_concept=_shape_concept,\n",
    "                        color_avail=self.color_avail,\n",
    "                        allow_connect=self.allow_connect,\n",
    "                        parsing_check=self.parsing_check,\n",
    "                    )\n",
    "                elif set(self.concept_collection).issubset({\n",
    "                    \"RotateA\", \"RotateB\", \"RotateC\", \n",
    "                    \"hFlip\", \"vFlip\", \"DiagFlipA\", \n",
    "                    \"DiagFlipB\", \"Identity\", \"Move\"\n",
    "                }):\n",
    "                    if max_n_distractors == -1:\n",
    "                        max_n_objs = 3\n",
    "                    else:\n",
    "                        max_n_objs = 1 + max_n_distractors # 1 is for the core operator itself.\n",
    "                    self.data = []\n",
    "                    for i in range(self.n_examples * 5):\n",
    "                        # Initialize input concept instance:\n",
    "                        obj_spec = obj_spec_fun(\n",
    "                            concept_collection=input_concepts,\n",
    "                            min_n_objs=1+self.min_n_distractors,\n",
    "                            max_n_objs=max_n_objs,\n",
    "                            canvas_size=canvas_size,\n",
    "                        )\n",
    "                        # get the number of the objects\n",
    "                        operatable_obj_set = set([])\n",
    "                        for spec in obj_spec:\n",
    "                            if spec[1] == \"Attr\":\n",
    "                                operatable_obj_set.add(spec[0][0])\n",
    "                            else:\n",
    "                                operatable_obj_set.add(spec[0][0])\n",
    "                                operatable_obj_set.add(spec[0][1])\n",
    "                        operatable_obj_set = list(operatable_obj_set)\n",
    "                        # let us enable distractors\n",
    "                        if set(input_concepts).issubset({\"SameColor\", \"IsTouch\"}):\n",
    "                            n_distractors = np.random.randint(0, max_n_distractors+1)\n",
    "                            max_rect_size = canvas_size//2\n",
    "                            for i in range(n_distractors):\n",
    "                                k = i+len(operatable_obj_set)\n",
    "                                distractor_shape = np.random.choice(self.allowed_shape_concept)\n",
    "                                if distractor_shape == \"Line\":\n",
    "                                    obj_spec += [(('obj_{}'.format(k), 'line_[-1,-1,-1]'), 'Attr')]\n",
    "                                elif distractor_shape == \"Rect\":\n",
    "                                    obj_spec += [(('obj_{}'.format(k), 'rectangle_[-1,-1]'), 'Attr')]\n",
    "                                elif distractor_shape == \"RectSolid\":\n",
    "                                    obj_spec += [(('obj_{}'.format(k), 'rectangleSolid_[-1,-1]'), 'Attr')]\n",
    "                                elif distractor_shape == \"Lshape\":\n",
    "                                    obj_spec += [(('obj_{}'.format(k), 'Lshape_[-1,-1,-1]'), 'Attr')]\n",
    "                                elif distractor_shape == \"Tshape\":\n",
    "                                    w = np.random.randint(3, max_rect_size+2)\n",
    "                                    h = np.random.randint(3, max_rect_size+2)\n",
    "                                    obj_spec += [(('obj_{}'.format(k), f'Tshape_[{w},{h}]'), 'Attr')]\n",
    "                                elif distractor_shape == \"Eshape\":\n",
    "                                    w = np.random.randint(3, max_rect_size+1)\n",
    "                                    h = np.random.randint(5, max_rect_size+3)\n",
    "                                    obj_spec += [(('obj_{}'.format(k), f'Eshape_[{w},{h}]'), 'Attr')]\n",
    "                                elif distractor_shape == \"Hshape\":\n",
    "                                    w = np.random.randint(3, max_rect_size+2)\n",
    "                                    h = np.random.randint(3, max_rect_size+2)\n",
    "                                    obj_spec += [(('obj_{}'.format(k), f'Hshape_[{w},{h}]'), 'Attr')]\n",
    "                                elif distractor_shape == \"Cshape\":\n",
    "                                    w = np.random.randint(3, max_rect_size+1)\n",
    "                                    h = np.random.randint(3, max_rect_size+2)\n",
    "                                    obj_spec += [(('obj_{}'.format(k), f'Cshape_[{w},{h}]'), 'Attr')]\n",
    "                                elif distractor_shape == \"Ashape\":\n",
    "                                    w = np.random.randint(3, max_rect_size+2)\n",
    "                                    h = np.random.randint(4, max_rect_size+3)\n",
    "                                    obj_spec += [(('obj_{}'.format(k), f'Ashape_[{w},{h}]'), 'Attr')]\n",
    "                                elif distractor_shape == \"Fshape\":\n",
    "                                    w = np.random.randint(3, max_rect_size+1)\n",
    "                                    h = np.random.randint(4, max_rect_size+3)\n",
    "                                    obj_spec += [(('obj_{}'.format(k), f'Fshape_[{w},{h}]'), 'Attr')]     \n",
    "                                elif distractor_shape == \"Randshape\":\n",
    "                                    max_rect_size = canvas_size // 2\n",
    "                                    w, h = np.random.randint(2, max_rect_size+1, size=2) # hard-code for the size\n",
    "                                    obj_spec += [(('obj_{}'.format(k), f'randomShape_[{w},{h}]'), 'Attr')]\n",
    "                                elif distractor_shape == \"ARCshape\":\n",
    "                                    max_rect_size = canvas_size // 2\n",
    "                                    w, h = np.random.randint(2, max_rect_size+1, size=2) # hard-code for the size\n",
    "                                    obj_spec += [(('obj_{}'.format(k), f'arcShape_[{w},{h}]'), 'Attr')]\n",
    "                        # get all objects include distractors\n",
    "                        all_obj_set = set([])\n",
    "                        for spec in obj_spec:\n",
    "                            if spec[1] == \"Attr\":\n",
    "                                all_obj_set.add(spec[0][0])\n",
    "                            else:\n",
    "                                all_obj_set.add(spec[0][0])\n",
    "                                all_obj_set.add(spec[0][1])\n",
    "\n",
    "                        repre_dict = dataset.sample_single_canvas_by_core_edges(\n",
    "                            OrderedDict(obj_spec),\n",
    "                            allow_connect=self.allow_connect,\n",
    "                            rainbow_prob=rainbow_prob,\n",
    "                            is_plot=False,\n",
    "                            color_avail=self.color_avail,\n",
    "                        )\n",
    "                        if repre_dict == -1:\n",
    "                            continue\n",
    "                        in_canvas = Canvas(repre_dict=repre_dict)\n",
    "\n",
    "                        # Operate on the input:\n",
    "                        chosen_obj_id = np.random.choice(len(operatable_obj_set))\n",
    "                        chosen_obj_name = operatable_obj_set[chosen_obj_id]\n",
    "                        chosen_op = np.random.choice(self.concept_collection)\n",
    "                        if chosen_op in [\"Identity\"]:\n",
    "                            inplace = True if random.random() < 0.5 else False\n",
    "                            out_canvas_list, concept = OperatorEngine().operator_identity(\n",
    "                                [in_canvas],\n",
    "                                [[chosen_obj_name]],\n",
    "                                inplace=inplace,\n",
    "                            )\n",
    "                            if out_canvas_list == -1:\n",
    "                                continue\n",
    "                        elif chosen_op in [\n",
    "                            \"RotateA\", \"RotateB\", \"RotateC\", \n",
    "                            \"hFlip\", \"vFlip\", \"DiagFlipA\", \"DiagFlipB\"\n",
    "                        ]:\n",
    "                            out_canvas_list, concept = OperatorEngine().operate_rotate(\n",
    "                                [in_canvas],\n",
    "                                [[chosen_obj_name]],\n",
    "                                operator_tag=f\"#{chosen_op}\",\n",
    "                                allow_connect=self.allow_connect,\n",
    "                                allow_shape_break=False,\n",
    "                            )\n",
    "                            if out_canvas_list == -1:\n",
    "                                continue\n",
    "                        elif chosen_op in [\"Move\"]:\n",
    "                            # create operator spec as move is a complex operator\n",
    "                            move_spec = OperatorMoveSpec(\n",
    "                                            autonomous=False,\n",
    "                                            direction=random.randint(0,3), \n",
    "                                            distance=-1, \n",
    "                                            hit_type=None, # either wall, agent or None\n",
    "                                            linkage_move=False, \n",
    "                                            linkage_move_distance_ratio=None,\n",
    "                                        )\n",
    "                            out_canvas_list, concept = OperatorEngine().operator_move(\n",
    "                                [in_canvas],\n",
    "                                [[chosen_obj_name]],\n",
    "                                [[move_spec]], \n",
    "                                allow_overlap=False, \n",
    "                                allow_shape_break=False,\n",
    "                                allow_connect=self.allow_connect,\n",
    "                                allow_stay=False,\n",
    "                            )\n",
    "                            if out_canvas_list == -1:\n",
    "                                continue\n",
    "                        else:\n",
    "                            raise Exception(f\"operator={chosen_op} is not supported!\")\n",
    "                        \n",
    "                        if n_operators > 1:\n",
    "                            # operator distractor can act on all objects\n",
    "                            addition_operators = min(len(all_obj_set)-1,n_operators-1) # we need to have minimum number of objs\n",
    "                            operated_obj_name = set([])\n",
    "                            operated_obj_name.add(chosen_obj_name)\n",
    "                            \n",
    "                            exclude_ops = set([chosen_op])\n",
    "                            \n",
    "                            # we need to operate on other objects.\n",
    "                            for _ in range(n_operators-1):\n",
    "                                addition_obj_set = set(all_obj_set) - operated_obj_name\n",
    "                                addition_obj_name = np.random.choice(list(addition_obj_set))\n",
    "                                \n",
    "                                addition_ops = set(self.concept_collection) - exclude_ops\n",
    "                                addition_op = np.random.choice(list(addition_ops))\n",
    "                                exclude_ops.add(addition_op)\n",
    "                                \n",
    "                                # operate the the previous ouput canvas\n",
    "                                if addition_op in [\"Identity\"]:\n",
    "                                    inplace = True if random.random() < 0.5 else False\n",
    "                                    out_canvas_list, concept = OperatorEngine().operator_identity(\n",
    "                                        [out_canvas_list[0]],\n",
    "                                        [[addition_obj_name]],\n",
    "                                        inplace=inplace,\n",
    "                                    )\n",
    "                                    if out_canvas_list == -1:\n",
    "                                        break\n",
    "                                elif addition_op in [\"RotateA\", \"RotateB\", \"RotateC\", \"hFlip\", \"vFlip\", \"DiagFlipA\", \"DiagFlipB\"]:\n",
    "                                    out_canvas_list, concept = OperatorEngine().operate_rotate(\n",
    "                                        [out_canvas_list[0]],\n",
    "                                        [[addition_obj_name]],\n",
    "                                        operator_tag=f\"#{addition_op}\",\n",
    "                                        allow_connect=self.allow_connect,\n",
    "                                        allow_shape_break=False,\n",
    "                                    )\n",
    "                                    if out_canvas_list == -1:\n",
    "                                        break\n",
    "                                elif addition_op in [\"Move\"]:\n",
    "                                    # create operator spec as move is a complex operator\n",
    "                                    move_spec = OperatorMoveSpec(\n",
    "                                                    autonomous=False,\n",
    "                                                    direction=random.randint(0,3), \n",
    "                                                    distance=-1, \n",
    "                                                    hit_type=None, # either wall, agent or None\n",
    "                                                    linkage_move=False, \n",
    "                                                    linkage_move_distance_ratio=None,\n",
    "                                                )\n",
    "                                    out_canvas_list, concept = OperatorEngine().operator_move(\n",
    "                                        [out_canvas_list[0]],\n",
    "                                        [[addition_obj_name]],\n",
    "                                        [[move_spec]], \n",
    "                                        allow_overlap=False, \n",
    "                                        allow_shape_break=False,\n",
    "                                        allow_connect=self.allow_connect,\n",
    "                                        allow_stay=False,\n",
    "                                    )\n",
    "                                    if out_canvas_list == -1:\n",
    "                                        break\n",
    "                                else:\n",
    "                                    raise Exception(f\"operator={addition_op} is not supported!\")\n",
    "                                operated_obj_name.add(addition_obj_name)\n",
    "                        if out_canvas_list == -1:\n",
    "                            continue\n",
    "                        # Add to self.data:\n",
    "                        in_canvas_dict = in_canvas.repr_as_dict()\n",
    "                        out_canvas_dict = out_canvas_list[0].repr_as_dict()\n",
    "                        \n",
    "                        in_mask = in_canvas_dict[\"id_object_mask\"][in_canvas_dict[\"node_id_map\"][chosen_obj_name]][None]\n",
    "                        out_mask = out_canvas_dict[\"id_object_mask\"][in_canvas_dict[\"node_id_map\"][chosen_obj_name]][None]\n",
    "                        # TODO: remove deprecated codes.\n",
    "                        # in_mask = in_canvas_dict[\"id_object_mask\"][chosen_obj_id][None]\n",
    "                        # out_mask = out_canvas_dict[\"id_object_mask\"][chosen_obj_id][None]\n",
    "                        info = {\"obj_spec\": obj_spec}\n",
    "                        self.data.append(\n",
    "                            ((to_one_hot(in_canvas_dict[\"image_t\"]), to_one_hot(out_canvas_dict[\"image_t\"])),\n",
    "                             (in_mask, out_mask),\n",
    "                             chosen_op,\n",
    "                             Dictionary(info),\n",
    "                            )\n",
    "                        )\n",
    "                        if len(self.data) >= n_examples:\n",
    "                            break\n",
    "                        if i > n_examples * 2 and len(self.data) < n_examples * 0.05:\n",
    "                            raise Exception(\"Sampled {} times and only {} of them satisfies the specified condition. Try relaxing the condition!\".format(i, len(self.data)))\n",
    "                else:\n",
    "                    raise Exception(\"concept_collection {} is out of scope!\".format(self.concept_collection))\n",
    "            if \"obj\" in self.w_type and \"mask\" not in self.w_type:\n",
    "                self.data = mask_to_obj(self.data)\n",
    "            self.idx_list = list(range(len(self.data)))\n",
    "            if len(self.idx_list) < n_examples:\n",
    "                p.print(\"Dataset created with {} examples, less than {} specified.\".format(len(self.idx_list), n_examples))\n",
    "            else:\n",
    "                p.print(\"Dataset for {} created.\".format(mode))\n",
    "        else:\n",
    "            self.data = data\n",
    "            self.idx_list = idx_list\n",
    "            self.concept_collection = concept_collection\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.idx_list)\n",
    "\n",
    "    def __repr__(self):\n",
    "        return \"ConceptDataset({})\".format(len(self))\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"Get data instance, where idx can be a number or a slice.\"\"\"\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "\n",
    "        elif isinstance(idx, slice):\n",
    "            return self.__class__(\n",
    "                mode=self.mode,\n",
    "                canvas_size=self.canvas_size,\n",
    "                n_examples=self.n_examples,\n",
    "                rainbow_prob=self.rainbow_prob,\n",
    "                data=self.data,\n",
    "                idx_list=self.idx_list[idx],\n",
    "                concept_collection=self.concept_collection,\n",
    "                w_type=self.w_type,\n",
    "                color_avail=self.color_avail,\n",
    "                max_n_distractors=self.max_n_distractors,\n",
    "                n_operators=self.n_operators,\n",
    "                transform=self.transform,\n",
    "            )\n",
    "        sample = self.data[self.idx_list[idx]]\n",
    "        if self.transform:\n",
    "            sample = self.transform(sample)\n",
    "        return sample\n",
    "\n",
    "    def draw(self, idx):\n",
    "        \"\"\"Draw one of multiple data instances.\"\"\"\n",
    "        if not isinstance(idx, Iterable):\n",
    "            idx = [idx]\n",
    "        for index in idx:\n",
    "            sample = self[index]\n",
    "            if len(sample) == 4:\n",
    "                p.print(\"example {}, {}:\".format(index, sample[2]))\n",
    "                if isinstance(sample[0], tuple):\n",
    "                    visualize_matrices([sample[0][0].argmax(0), sample[0][1].argmax(0)])\n",
    "                else:\n",
    "                    visualize_matrices([sample[0].argmax(0)])\n",
    "                plot_matrices([sample[1][i].squeeze() for i in range(len(sample[1]))], images_per_row=6)\n",
    "\n",
    "\n",
    "class ConceptDataset3D(MineDataset):\n",
    "    def draw(self, idx):\n",
    "        \"\"\"Draw one of multiple data instances.\"\"\"\n",
    "        if not isinstance(idx, Iterable):\n",
    "            idx = [idx]\n",
    "        for index in idx:\n",
    "            sample = self[index]\n",
    "            if len(sample) == 4:\n",
    "                p.print(\"example {}, {}:\".format(index, sample[2]))\n",
    "                if isinstance(sample[0], tuple):\n",
    "                    visualize_matrices([sample[0][0].argmax(0), sample[0][1].argmax(0)], use_color_dict=False)\n",
    "                else:\n",
    "                    visualize_matrices([sample[0]], use_color_dict=False)\n",
    "                visualize_matrices([sample[1][i].squeeze() for i in range(len(sample[1]))])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 ConceptFewshotDataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_fewshot_dataset(args, concept_mode=\"standard\", n_shot=1, n_queries_per_class=15):\n",
    "    \"\"\"For parsing + classify.\n",
    "    The args.dataset has the format:\n",
    "        f\"pc-{concept_1}+{concept_2}+...+{concept_n}\"\n",
    "    \"\"\"\n",
    "    assert args.dataset.startswith(\"pc\") or args.dataset.startswith(\"pg\") or args.dataset.startswith(\"yc\") \n",
    "    str_split = args.dataset.split(\"-\")[1].split(\"^\")\n",
    "    concept_collection = str_split[-1].split(\"+\") \n",
    "    concept_dict = {}\n",
    "    assert n_shot == 1\n",
    "    # Generate samples for each concept:\n",
    "    for concept in concept_collection:\n",
    "        concept_args = init_args({\n",
    "            \"dataset\": \"c-{}\".format(concept),\n",
    "            \"seed\": args.seed,\n",
    "            \"n_examples\": args.n_examples * n_shot,\n",
    "            \"canvas_size\": args.canvas_size,\n",
    "            \"rainbow_prob\": 0.,\n",
    "            \"w_type\": \"image+mask\",\n",
    "            \"color_avail\": args.color_avail,\n",
    "            \"max_n_distractors\": 0,\n",
    "            \"min_n_distractors\": 0,\n",
    "            \"allow_connect\": True, # No effect\n",
    "            \"parsing_check\": False,\n",
    "        })\n",
    "        concept_dict[concept] = get_dataset(concept_args, verbose=False)[0]\n",
    "    if args.dataset.startswith(\"pc\"):\n",
    "        example_dict = {}\n",
    "        for concept in concept_collection:\n",
    "            example_args = init_args({\n",
    "                \"dataset\": \"c-{}\".format(concept),\n",
    "                \"seed\": args.seed + 1,\n",
    "                \"n_examples\": args.n_examples * n_queries_per_class,\n",
    "                \"canvas_size\": args.canvas_size,\n",
    "                \"rainbow_prob\": 0.,\n",
    "                \"w_type\": \"image+mask\",\n",
    "                \"color_avail\": args.color_avail,\n",
    "                \"max_n_distractors\": 0,\n",
    "                \"min_n_distractors\": 0,\n",
    "                \"allow_connect\": True,\n",
    "                \"parsing_check\": False,\n",
    "            })\n",
    "            example_dict[concept] = get_dataset(example_args, verbose=False)[0]\n",
    "    elif args.dataset.startswith(\"yc\"):\n",
    "        example_dict = {}\n",
    "        for concept in concept_collection:\n",
    "            example_args = init_args({\n",
    "                \"dataset\": \"y-{}\".format(concept),\n",
    "                \"seed_3d\": args.seed_3d + args.num_processes_3d,\n",
    "                \"num_processes_3d\": args.num_processes_3d,\n",
    "                \"color_map_3d\": args.color_map_3d,\n",
    "                \"add_thick_surf\": args.add_thick_surf,\n",
    "                \"add_thick_depth\": args.add_thick_depth,\n",
    "                \"image_size_3d\": args.image_size_3d,\n",
    "                \"n_examples\": args.n_examples * n_queries_per_class,\n",
    "                # 2D examples\n",
    "                \"seed\": args.seed + 1,\n",
    "                \"use_seed_2d\": args.use_seed_2d,\n",
    "                \"canvas_size\": args.canvas_size,\n",
    "                \"rainbow_prob\": 0.,\n",
    "                \"w_type\": \"image+mask\",\n",
    "                \"color_avail\": args.color_avail,\n",
    "                \"max_n_distractors\": 0,\n",
    "                \"min_n_distractors\": 0,\n",
    "                \"allow_connect\": True,\n",
    "                \"parsing_check\": False,\n",
    "            })\n",
    "            example_dict[concept] = get_dataset(example_args, verbose=False, is_load=True)[0]\n",
    "    elif args.dataset.startswith(\"pg\"):\n",
    "        example_dict = {}\n",
    "        examples_collection = str_split[0].split(\"+\")\n",
    "        # The tasks are split evenly among the possible concepts for demonstration\n",
    "        n_tasks = [args.n_examples // len(concept_collection)] * len(concept_collection)\n",
    "        n_tasks[-1] += args.n_examples % len(concept_collection)\n",
    "        for idx, concept in enumerate(concept_collection):\n",
    "            example_args = init_args({\n",
    "                \"dataset\": \"c-{}+{}^{}\".format(concept, str_split[0], concept),\n",
    "                \"seed\": args.seed + 1,\n",
    "                \"n_examples\": n_tasks[idx] * n_queries_per_class,\n",
    "                \"canvas_size\": args.canvas_size,\n",
    "                \"rainbow_prob\": 0.,\n",
    "                \"w_type\": \"image+mask\",\n",
    "                \"color_avail\": args.color_avail,\n",
    "                \"max_n_distractors\": 2, # Important: There can be distractors in examples\n",
    "                \"parsing_check\": False,\n",
    "            })\n",
    "            example_dict[concept] = get_dataset(example_args, verbose=False)[0]\n",
    "    else:\n",
    "        raise\n",
    "    if args.dataset[1] == \"c\":\n",
    "        if concept_mode == \"standard\":\n",
    "            dataset = generate_fewshot_dataset_standard(\n",
    "                concept_dict=concept_dict,\n",
    "                example_dict=example_dict,\n",
    "                concept_collection=concept_collection,\n",
    "                n_examples=args.n_examples,\n",
    "                n_shot=n_shot,\n",
    "                n_queries_per_class=n_queries_per_class,\n",
    "            )\n",
    "        elif concept_mode == \"random\":\n",
    "            dataset = generate_fewshot_dataset_random(\n",
    "                concept_dict=concept_dict,\n",
    "                example_dict=example_dict,\n",
    "                concept_collection=concept_collection,\n",
    "                n_examples=args.n_examples,\n",
    "                n_shot=n_shot,\n",
    "            )\n",
    "        else:\n",
    "            raise Exception(\"concept_mode '{}' is not valid!\".format(concept_mode))\n",
    "    elif args.dataset[1] == \"g\":\n",
    "        if concept_mode == \"standard\":\n",
    "            dataset = generate_fewshot_grounding_dataset(\n",
    "                concept_dict=concept_dict,\n",
    "                example_dict=example_dict,\n",
    "                concept_collection=concept_collection,\n",
    "                n_tasks=n_tasks,\n",
    "                n_shot=n_shot,\n",
    "                n_queries_per_class=n_queries_per_class,\n",
    "            )\n",
    "        else:\n",
    "            raise Exception(\"concept_mode '{}' is not valid!\".format(concept_mode))\n",
    "    else:\n",
    "        raise\n",
    "    dataset.concept_collection = concept_collection\n",
    "    return dataset\n",
    "\n",
    "\n",
    "def generate_fewshot_dataset_standard(\n",
    "    concept_dict,\n",
    "    example_dict,\n",
    "    concept_collection,\n",
    "    n_examples,\n",
    "    n_shot=1,\n",
    "    n_queries_per_class=15,\n",
    "):\n",
    "    \"\"\"\n",
    "    Format:\n",
    "        Each data in data_list has the format:\n",
    "            (data_concept, data_concept_mask, data_concept_id), (data_examples, data_examples_mask, data_examples_id), info\n",
    "            where data_examples, data_examples_mask, data_examples_id are all lists of all concept examples \n",
    "            according to concept_collection.\n",
    "    \"\"\"\n",
    "    # Generate dataset:\n",
    "    data_list = []\n",
    "    concept_id_dict = {}\n",
    "    example_id_dict = {}\n",
    "    data_list = []\n",
    "    info = {}\n",
    "    for concept in concept_collection:\n",
    "        concept_id_dict[concept] = np.random.choice(n_examples * n_shot, size=n_examples * n_shot, replace=False)\n",
    "        example_id_dict[concept] = np.random.choice(n_examples * n_queries_per_class, size=n_examples * n_queries_per_class, replace=False)\n",
    "\n",
    "    for i in range(n_examples):\n",
    "        # Obtaining the concept examples:\n",
    "        concept_collection_permute = np.random.choice(concept_collection, size=len(concept_collection), replace=False)\n",
    "        data_concepts = []\n",
    "        data_concepts_mask = []\n",
    "        data_concepts_id = []\n",
    "        info[\"concept_info\"] = []\n",
    "        for j, concept in enumerate(concept_collection_permute):\n",
    "            concept_id = concept_id_dict[concept][i]\n",
    "            data_ele = concept_dict[concept][concept_id]\n",
    "            data_concepts.append(data_ele[0])\n",
    "            data_concepts_mask.append(data_ele[1][0])\n",
    "            data_concepts_id.append(data_ele[2])\n",
    "            info[\"concept_info\"].append(data_ele[3])\n",
    "\n",
    "        # Obtaining the query set:\n",
    "        data_examples = []\n",
    "        data_examples_mask = []\n",
    "        data_examples_id = []\n",
    "        info[\"example_info\"] = []\n",
    "        for example_id in concept_collection:\n",
    "            for kk in range(n_queries_per_class):\n",
    "                data_ele = example_dict[example_id][i * n_queries_per_class + kk]\n",
    "                assert data_ele[2] in example_id\n",
    "                data_examples.append(data_ele[0])\n",
    "                data_examples_mask.append(data_ele[1][0])\n",
    "                data_examples_id.append(data_ele[2])\n",
    "                info[\"example_info\"].append(data_ele[3])\n",
    "        example_num = n_queries_per_class * len(concept_collection)\n",
    "        assert len(data_examples) == example_num\n",
    "        permute_ids_example = np.random.choice(example_num, size=example_num, replace=False)\n",
    "        data_examples = [data_examples[id] for id in permute_ids_example]\n",
    "        data_examples_mask = [data_examples_mask[id] for id in permute_ids_example]\n",
    "        data_examples_id = [data_examples_id[id] for id in permute_ids_example]\n",
    "        data = ((tuple(data_concepts), tuple(data_concepts_mask), tuple(data_concepts_id)),\n",
    "                (tuple(data_examples), tuple(data_examples_mask), tuple(data_examples_id)), info)\n",
    "        data_list.append(data)\n",
    "    permuted_ids = np.random.choice(n_examples, size=n_examples, replace=False)\n",
    "    data_list_final = [data_list[i] for i in permuted_ids]\n",
    "    data_list_final = ConceptFewshotDataset(data=data_list_final, concept_mode=\"standard\")\n",
    "    return data_list_final\n",
    "\n",
    "\n",
    "def generate_fewshot_grounding_dataset(\n",
    "    concept_dict,\n",
    "    example_dict,\n",
    "    concept_collection,\n",
    "    n_tasks,\n",
    "    n_shot=1,\n",
    "    n_queries_per_class=15,\n",
    "):\n",
    "    task_id = []\n",
    "    for idx, num_tasks in enumerate(n_tasks):\n",
    "        # Keep track of the demonstrated concept and the task index for that concept\n",
    "        task_id.extend([(i, concept_collection[idx]) for i in range(num_tasks)])\n",
    "    n_examples = len(task_id)\n",
    "    assert n_examples == sum(n_tasks)\n",
    "    \n",
    "    data_list = []\n",
    "    concept_id_dict = {}\n",
    "    example_id_dict = {}\n",
    "    info = {}\n",
    "    # Randomize the order of concept demonstrations and examples\n",
    "    for idx, concept in enumerate(concept_collection):\n",
    "        concept_id_dict[concept] = np.random.choice(n_examples * n_shot, size=n_examples * n_shot, replace=False)\n",
    "        example_id_dict[concept] = np.random.choice(n_tasks[idx] * n_queries_per_class, size=n_tasks[idx] * n_queries_per_class, replace=False)\n",
    "    \n",
    "    # Each task has a single concept that is demonstrated\n",
    "    for i, concept_id in task_id:\n",
    "        # Get concept demonstrations\n",
    "        concept_data_ele = [concept_dict[concept_id][concept_id_dict[concept_id][i * n_shot + j]] for j in range(n_shot)]\n",
    "        data_concepts = []\n",
    "        data_concepts_mask = []\n",
    "        data_concepts_id = []\n",
    "        info[\"concept_info\"] = []\n",
    "        for data_ele in concept_data_ele:\n",
    "            data_concepts.append(data_ele[0])\n",
    "            data_concepts_mask.append(data_ele[1][0])\n",
    "            data_concepts_id.append(data_ele[2])\n",
    "            info[\"concept_info\"].append(data_ele[3])\n",
    "        \n",
    "        # Get query set\n",
    "        data_examples = []\n",
    "        data_examples_mask = []\n",
    "        data_examples_id = []\n",
    "        info[\"example_info\"] = []\n",
    "        for kk in range(n_queries_per_class):\n",
    "            data_ele = example_dict[concept_id][example_id_dict[concept_id][i * n_queries_per_class + kk]]\n",
    "            assert data_ele[2] == concept_id\n",
    "            data_examples.append(data_ele[0])\n",
    "            data_examples_mask.append(data_ele[1][0])\n",
    "            data_examples_id.append(data_ele[2])\n",
    "            info[\"example_info\"].append(data_ele[3])\n",
    "        example_num = n_queries_per_class\n",
    "        assert len(data_examples) == example_num\n",
    "        data = ((tuple(data_concepts), tuple(data_concepts_mask), tuple(data_concepts_id)),\n",
    "                (tuple(data_examples), tuple(data_examples_mask), tuple(data_examples_id)), info)\n",
    "        data_list.append(data)\n",
    "    permuted_ids = np.random.choice(n_examples, size=n_examples, replace=False)\n",
    "    data_list_final = [data_list[i] for i in permuted_ids]\n",
    "    data_list_final = ConceptFewshotDataset(data=data_list_final, concept_mode=\"standard\")\n",
    "    return data_list_final\n",
    "\n",
    "\n",
    "def generate_fewshot_dataset_random(concept_dict, example_dict, concept_collection, n_examples, n_shot=1):\n",
    "    \"\"\"\n",
    "    Format:\n",
    "        Each data in data_list has the format:\n",
    "            (data_concept,), (data_examples, data_examples_mask, data_examples_id), info\n",
    "            where data_examples, data_examples_mask, data_examples_id are all lists of all concept examples \n",
    "            according to concept_collection.\n",
    "    \"\"\"\n",
    "    # Generate dataset:\n",
    "    data_list = []\n",
    "    query_id_dict = {}\n",
    "    example_id_dict = {}\n",
    "    data_list = []\n",
    "    num = int(np.ceil(n_examples / len(concept_collection)))\n",
    "    for concept in concept_collection:\n",
    "        query_id_dict[concept] = np.random.choice(len(concept_dict[concept]), size=num*n_shot, replace=False)\n",
    "        example_id_dict[concept] = np.random.choice(n_examples * 2, size=num, replace=False)\n",
    "    for concept in concept_collection:\n",
    "        for i, id in enumerate(query_id_dict[concept]):\n",
    "            info = {}\n",
    "            data_concept = (concept_dict[concept][id][0],)\n",
    "            info[\"concept_mask\"] = (concept_dict[concept][id][1][0],)\n",
    "            info[\"concept_id\"] = concept_dict[concept][id][2]\n",
    "            assert concept == concept_dict[concept][id][2]\n",
    "            info[\"concept_info\"] = (concept_dict[concept][id][3],)\n",
    "            concept_collection_permute = np.random.choice(concept_collection, size=len(concept_collection), replace=False)\n",
    "            data_examples = []\n",
    "            data_examples_mask = []\n",
    "            data_examples_id = []\n",
    "            info[\"example_info\"] = []\n",
    "            for j, concept_example in enumerate(concept_collection_permute):\n",
    "                example_id = example_id_dict[concept_example][j]\n",
    "                data_examples.append(example_dict[concept_example][example_id][0])\n",
    "                data_examples_mask.append(example_dict[concept_example][example_id][1][0])\n",
    "                data_examples_id.append(example_dict[concept_example][example_id][2])\n",
    "                info[\"example_info\"].append(example_dict[concept_example][example_id][3])\n",
    "            data = (data_concept, (tuple(data_examples), tuple(data_examples_mask), tuple(data_examples_id)), info)\n",
    "            data_list.append(data)\n",
    "    permuted_ids = np.random.choice(num * len(concept_collection), n_examples, replace=False)\n",
    "    data_list_final = [data_list[i] for i in permuted_ids]\n",
    "    data_list_final = ConceptFewshotDataset(data=data_list_final, concept_mode=\"random\")\n",
    "    return data_list_final\n",
    "\n",
    "\n",
    "class ConceptFewshotDataset(MineDataset):\n",
    "    \"\"\"\n",
    "    Format:\n",
    "        Each data in data_list has the format:\n",
    "            (data_concept,), (data_examples, data_examples_mask, data_examples_id), info\n",
    "            where data_examples, data_examples_mask, data_examples_id are all lists of all concept examples \n",
    "            according to concept_collection.\n",
    "    \"\"\"\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        filtered_kwargs = deepcopy(kwargs)\n",
    "        filtered_kwargs.pop(\"concept_mode\")\n",
    "        super().__init__(*args, **filtered_kwargs)\n",
    "        self.concept_mode = kwargs[\"concept_mode\"]\n",
    "\n",
    "    def draw(self, idx):\n",
    "        \"\"\"Draw one of multiple data instances.\"\"\"\n",
    "        if not isinstance(idx, Iterable):\n",
    "            idx = [idx]\n",
    "        print(\"Drawing dataset with concept_collection of {}:\".format(self.concept_collection))\n",
    "        for index in idx:\n",
    "            sample = self[index]\n",
    "            num_concepts = len(sample[0][0])\n",
    "            num_queries = len(sample[1][0])\n",
    "            num_ex_channels = sample[1][0][0].shape[0]\n",
    "            is_ex_3d = (num_ex_channels == 3)\n",
    "            if self.concept_mode == \"standard\":\n",
    "                print(f\"Concept demonstration for task {index}:\")\n",
    "                visualize_matrices([concept_image.argmax(0) for concept_image in sample[0][0]], subtitles=sample[0][2], images_per_row=num_concepts)\n",
    "                print(\"Instances and masks:\")\n",
    "                if is_ex_3d:\n",
    "                    visualize_matrices([F.interpolate(concept_instance[None], (32,32), mode=\"nearest\")[0] for concept_instance in sample[1][0]], \n",
    "                               use_color_dict=False, subtitles=sample[1][2], images_per_row=min(6, num_concepts))\n",
    "                    plot_matrices([concept_instance_mask.squeeze(0) for concept_instance_mask in sample[1][1]], images_per_row=min(6, num_concepts))\n",
    "                else:\n",
    "                    visualize_matrices([concept_instance.argmax(0) for concept_instance in sample[1][0]], \n",
    "                                       subtitles=sample[1][2], images_per_row=min(6, num_concepts))\n",
    "                    plot_matrices([concept_instance_mask.squeeze(0) for concept_instance_mask in sample[1][1]], images_per_row=min(6, num_concepts))\n",
    "            elif self.concept_mode == \"random\":\n",
    "                label = sample[2][\"concept_id\"]\n",
    "                print(\"example {}, concept '{}':\".format(index, label))\n",
    "                print(f\"Concept demonstration for task {index}:\")\n",
    "                visualize_matrices([concept_image.argmax(0) for concept_image in sample[0]], subtitles=[label], images_per_row=6)\n",
    "                print(\"Instances and masks:\")\n",
    "                if is_ex_3d:\n",
    "                    visualize_matrices([concept_instance for concept_instance in sample[1][0]], \n",
    "                               use_color_dict=False, subtitles=[\"[{}]\".format(ele) if ele == label else ele for ele in sample[1][2]])\n",
    "                    visualize_matrices([concept_instance_mask.squeeze(0) for concept_instance_mask in sample[1][1]])\n",
    "                else:\n",
    "                    visualize_matrices([concept_instance.argmax(0) for concept_instance in sample[1][0]], \n",
    "                                       subtitles=[\"[{}]\".format(ele) if ele == label else ele for ele in sample[1][2]])\n",
    "                    visualize_matrices([concept_instance_mask.squeeze(0) for concept_instance_mask in sample[1][1]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 ConceptCompositionDataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConceptCompositionDataset(Dataset):\n",
    "    \"\"\"Concept Composition dataset for learning to compose concepts from elementary concepts.\n",
    "\n",
    "    mode:\n",
    "        Concepts:  E(x; a; c)\n",
    "            \"Pixel\": one or many pixels\n",
    "            \"Line\": one or many lines\n",
    "            \"Rect\": hollow rectangles\n",
    "            \"{}+{}+...\": each \"{}\" can be a concept.\n",
    "\n",
    "        Relations: E(x; a1, a2; c)\n",
    "            \"Vertical\": lines where some of them are vertical\n",
    "            \"Parallel\": lines where some of them are parallel\n",
    "            \"Vertical+Parallel\": lines where some of them are vertical or parallel\n",
    "            \"IsInside\": obj_1 is inside obj_2\n",
    "            \"SameRow\": obj_1 and obj_2 are at the same row\n",
    "            \"SameCol\": obj_1 and obj_2 are at the same column\n",
    "\n",
    "        Operations: E(x1,x2; a1,a2; c1,c2)\n",
    "            \"RotateA+vFlip(Line+Rect)\": two images where some object1 in image1 is rotated or vertically-flipped w.r.t. some object2 in image2, and the objects are chosen from Line or Rect.\n",
    "            \"hFlip(Lshape)\", \"vFlip(Lshape+Line)\": two images where some object1 in image1 is flipped w.r.t. some object2 in image2.\n",
    "\n",
    "        ARC+:\n",
    "            \"arc^{}\": ARC images with property \"{}\" masked as above.\n",
    "        \"\"\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        canvas_size=8,\n",
    "        n_examples=10000,\n",
    "        concept_avail=None,\n",
    "        relation_avail=None,\n",
    "        additional_concepts=None,\n",
    "        n_concepts_range=(2,3),\n",
    "        relation_structure=\"None\",\n",
    "        rainbow_prob=0.,\n",
    "        data=None,\n",
    "        idx_list=None,\n",
    "        color_avail=\"-1\",\n",
    "        min_n_distractors=0,\n",
    "        max_n_distractors=0,\n",
    "        n_examples_per_task=5,\n",
    "        transform=None,\n",
    "    ):\n",
    "        self.canvas_size = canvas_size\n",
    "        self.n_examples = n_examples\n",
    "        self.concept_avail = concept_avail\n",
    "        self.relation_avail = relation_avail\n",
    "        self.additional_concepts = additional_concepts\n",
    "        self.n_concepts_range = n_concepts_range\n",
    "        self.relation_structure = relation_structure\n",
    "        self.rainbow_prob = rainbow_prob\n",
    "        self.min_n_distractors = min_n_distractors\n",
    "        self.max_n_distractors = max_n_distractors\n",
    "        self.n_examples_per_task = n_examples_per_task\n",
    "\n",
    "        if isinstance(color_avail, str):\n",
    "            if color_avail == \"-1\":\n",
    "                self.color_avail = None\n",
    "            else:\n",
    "                self.color_avail = [int(c) for c in color_avail.split(\",\")]\n",
    "                for c in self.color_avail:\n",
    "                    assert c >= 1 and c <= 9\n",
    "        else:\n",
    "            self.color_avail = color_avail\n",
    "\n",
    "        if idx_list is None:\n",
    "            assert data is None\n",
    "            dataset_engine = BabyARCDataset(\n",
    "                pretrained_obj_cache=os.path.join(get_root_dir(), 'datasets/arc_objs.pt'),\n",
    "                save_directory=get_root_dir() + \"/datasets/BabyARCDataset/\",\n",
    "                object_limit=None,\n",
    "                noise_level=0,\n",
    "                canvas_size=canvas_size,\n",
    "            )\n",
    "            self.data = []\n",
    "            for i in range(self.n_examples * 3):\n",
    "                task = sample_selector(\n",
    "                    dataset_engine=dataset_engine,\n",
    "                    concept_avail=concept_avail,\n",
    "                    relation_avail=relation_avail,\n",
    "                    additional_concepts=self.additional_concepts,\n",
    "                    n_concepts_range=n_concepts_range,\n",
    "                    relation_structure=relation_structure,\n",
    "                    min_n_distractors=min_n_distractors,\n",
    "                    max_n_distractors=max_n_distractors,\n",
    "                    canvas_size=canvas_size,\n",
    "                    color_avail=self.color_avail,\n",
    "                    n_examples_per_task=n_examples_per_task,\n",
    "                    max_n_trials=5,\n",
    "                    isplot=False,\n",
    "                )\n",
    "                if len(task) == n_examples_per_task:\n",
    "                    self.data.append(task)\n",
    "                if len(self.data) % 100 == 0:\n",
    "                    p.print(\"Number of tasks generated: {}\".format(len(self.data)))\n",
    "                if len(self.data) >= self.n_examples:\n",
    "                    break\n",
    "\n",
    "            self.idx_list = list(range(len(self.data)))\n",
    "            if len(self.idx_list) < n_examples:\n",
    "                p.print(\"Dataset created with {} examples, less than {} specified.\".format(len(self.idx_list), n_examples))\n",
    "            else:\n",
    "                p.print(\"Dataset created with {} examples.\".format(len(self.idx_list)))\n",
    "        else:\n",
    "            self.data = data\n",
    "            self.idx_list = idx_list\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.idx_list)\n",
    "\n",
    "    def __repr__(self):\n",
    "        return \"ConceptDataset({})\".format(len(self))\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"Get data instance, where idx can be a number or a slice.\"\"\"\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "\n",
    "        elif isinstance(idx, slice):\n",
    "            return self.__class__(\n",
    "                canvas_size=self.canvas_size,\n",
    "                n_examples=self.n_examples,\n",
    "                concept_avail=self.concept_avail,\n",
    "                relation_avail=self.relation_avail,\n",
    "                additional_concepts=self.additional_concepts,\n",
    "                n_concepts_range=self.n_concepts_range,\n",
    "                relation_structure=self.relation_structure,\n",
    "                rainbow_prob=self.rainbow_prob,\n",
    "                data=self.data,\n",
    "                idx_list=self.idx_list[idx],\n",
    "                color_avail=self.color_avail,\n",
    "                min_n_distractors=self.min_n_distractors,\n",
    "                max_n_distractors=self.max_n_distractors,\n",
    "                n_examples_per_task=self.n_examples_per_task,\n",
    "                transform=self.transform,\n",
    "            )\n",
    "\n",
    "        sample = self.data[self.idx_list[idx]]\n",
    "        if self.transform:\n",
    "            sample = self.transform(sample)\n",
    "        return sample\n",
    "\n",
    "    def to_dict(self):\n",
    "        Dict = {}\n",
    "        for id in self.idx_list:\n",
    "            Dict[str(id)] = self.data[id]\n",
    "        return Dict\n",
    "\n",
    "    def draw(self, idx):\n",
    "        \"\"\"Draw one of multiple data instances.\"\"\"\n",
    "        if not isinstance(idx, Iterable):\n",
    "            idx = [idx]\n",
    "        for index in idx:\n",
    "            task = self[index]\n",
    "            if len(task) == self.n_examples_per_task:\n",
    "                info = task[0][3]\n",
    "                p.print(\"structure: {}\".format(info[\"structure\"]))\n",
    "                p.print(\"obj_spec_core:\")\n",
    "                pp.pprint(info[\"obj_spec_core\"])\n",
    "                for k, example in enumerate(task):\n",
    "                    p.print(\"task {}, example {}:\".format(index, k))\n",
    "                    if isinstance(example[0], tuple):\n",
    "                        visualize_matrices([example[0][0].argmax(0), example[0][1].argmax(0) if example[0][1].shape[0] == 10 else example[0][1].squeeze(0)])\n",
    "                    else:\n",
    "                        visualize_matrices([example[0].argmax(0)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.4 ConceptClevrDataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConceptClevrDataset(MineDataset):\n",
    "    def draw(self, idx, filename=None):\n",
    "        \"\"\"Draw one of multiple data instances.\"\"\"\n",
    "        if not isinstance(idx, Iterable):\n",
    "            idx = [idx]\n",
    "        for index in idx:\n",
    "            sample = self[index]\n",
    "            if len(sample) == 4:\n",
    "                p.print(\"example {}, {}:\".format(index, sample[2]))\n",
    "                if isinstance(sample[0], tuple):\n",
    "                    visualize_matrices([sample[0][0].argmax(0), sample[0][1].argmax(0)], use_color_dict=False, filename=filename)\n",
    "                else:\n",
    "                    visualize_matrices([sample[0]], use_color_dict=False, filename=filename)\n",
    "                plot_matrices([sample[1][i].squeeze() for i in range(len(sample[1]))], images_per_row=6, filename=filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.4 Dataset helper functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_selector(\n",
    "    dataset_engine,\n",
    "    concept_avail=None,\n",
    "    relation_avail=None,\n",
    "    additional_concepts=None,\n",
    "    n_concepts_range=None,\n",
    "    relation_structure=None,\n",
    "    max_n_distractors=0,\n",
    "    min_n_distractors=0,\n",
    "    canvas_size=8,\n",
    "    color_avail=None,\n",
    "    n_examples_per_task=5,\n",
    "    max_n_trials=5,\n",
    "    isplot=False,\n",
    "):\n",
    "    def is_exist_all(obj_full_relations, key):\n",
    "        is_all = True\n",
    "        for tuple_ele, lst in obj_full_relations.items():\n",
    "            assert isinstance(tuple_ele, tuple)\n",
    "            if key not in lst and (tuple_ele[0].startswith(\"obj_\") and tuple_ele[1].startswith(\"obj_\")):\n",
    "                is_all = False\n",
    "                break\n",
    "        return is_all\n",
    "    assert max_n_distractors >= min_n_distractors\n",
    "    if concept_avail is None:\n",
    "        concept_avail = [\n",
    "            \"Line\", \"Rect\", \"RectSolid\", \"Randshape\", \"ARCshape\",\n",
    "            \"Lshape\", \"Tshape\", \"Eshape\",\n",
    "            \"Hshape\", \"Cshape\", \"Ashape\", \"Fshape\",\n",
    "        ]\n",
    "    if relation_avail is None:\n",
    "        relation_avail = [\n",
    "            \"SameAll\", \"SameShape\", \"SameColor\",\n",
    "            \"SameRow\", \"SameCol\", \n",
    "            \"SubsetOf\", \"IsInside\", \"IsTouch\",\n",
    "        ]\n",
    "    allowed_shape_concept = concept_avail\n",
    "    if additional_concepts is not None:\n",
    "        allowed_shape_concept = allowed_shape_concept + additional_concepts\n",
    "    concept_str_reverse_mapping = {\n",
    "        \"Line\": \"line\", \n",
    "        \"Rect\": \"rectangle\", \n",
    "        \"RectSolid\": \"rectangleSolid\", \n",
    "        \"Lshape\": \"Lshape\", \n",
    "        \"Tshape\": \"Tshape\", \n",
    "        \"Eshape\": \"Eshape\", \n",
    "        \"Hshape\": \"Hshape\", \n",
    "        \"Cshape\": \"Cshape\", \n",
    "        \"Ashape\": \"Ashape\", \n",
    "        \"Fshape\": \"Fshape\",\n",
    "        \"Randshape\": \"randomShape\",\n",
    "        \"ARCshape\": \"arcShape\",\n",
    "    }\n",
    "\n",
    "    if relation_structure == \"None\":\n",
    "        # Only concept dataset:\n",
    "        n_concepts_range = (n_concepts_range, n_concepts_range) if (not isinstance(n_concepts_range, tuple)) and n_concepts_range > 0 else n_concepts_range\n",
    "        assert isinstance(n_concepts_range, tuple)\n",
    "        # Sample concepts:\n",
    "        obj_spec_core = obj_spec_fun(\n",
    "            concept_collection=concept_avail,\n",
    "            min_n_objs=n_concepts_range[0],\n",
    "            max_n_objs=n_concepts_range[1],\n",
    "            canvas_size=canvas_size,\n",
    "            color_avail=color_avail,\n",
    "        )\n",
    "        obj_id = len(obj_spec_core)\n",
    "        refer_node_id = None\n",
    "        structure = None\n",
    "    else:\n",
    "        structure_dict = {\n",
    "            \"2a\": [\"pivot\", (0,1), \"(refer)\"],\n",
    "            \"2ai\":[\"pivot:Rect\", (1,0,\"IsInside\"), \"(refer)\"],\n",
    "            \"3a\": [\"pivot\", (0,1), \"(concept)\", (1,2), \"(refer)\"],\n",
    "            \"3ai\":[\"pivot:Rect\", (1,0,\"IsInside\"), \"(concept)\", (1,2), \"(refer)\"],\n",
    "            \"3b\": [\"pivot\", \"pivot\", (0,2), (1,2), \"(refer)\"],\n",
    "            \"4a\": [\"pivot\", (0,1), \"concept\", (1,2), \"(concept)\", (2,3), \"(refer)\"],\n",
    "            \"4ai\":[\"pivot:Rect\", (1,0,\"IsInside\"), \"(concept)\", (1,2), \"(concept)\", (2,3), \"(refer)\"],\n",
    "            \"4b\": [\"pivot\", \"pivot\", (0,2), (1,2), \"(concept)\", (2,3), \"(refer)\"],\n",
    "        }\n",
    "\n",
    "        # Sample pivot concept:\n",
    "        structure_key = np.random.choice(relation_structure.split(\"+\"))\n",
    "        structure = structure_dict[structure_key]\n",
    "        is_valid = False\n",
    "        for i in range(3):\n",
    "            obj_id = 0\n",
    "            obj_spec_core = []\n",
    "            refer_node_id = None\n",
    "            relations_all = []\n",
    "            for j, element in enumerate(structure):\n",
    "                if isinstance(element, tuple):\n",
    "                    if len(element) == 2:\n",
    "                        if structure_key in [\"2a\", \"3a\"] and j == 1:\n",
    "                            relation = np.random.choice(remove_elements(relation_avail, [\"SameShape\", \"SameAll\"]))\n",
    "                        else:\n",
    "                            relation = np.random.choice(remove_elements(relation_avail, [\"SameAll\"]))\n",
    "                    elif len(element) == 3:\n",
    "                        relation = np.random.choice(element[2].split(\"+\"))\n",
    "                    obj_spec_core.append(\n",
    "                        [(\"obj_{}\".format(element[0]),\n",
    "                          \"obj_{}\".format(element[1])),\n",
    "                          relation,\n",
    "                    ])\n",
    "                    relations_all.append(relation)\n",
    "                elif element.startswith(\"pivot\"):\n",
    "                    if \":\" in element:\n",
    "                        concept_avail_core = element.split(\":\")[1].split(\"+\")\n",
    "                    else:\n",
    "                        concept_avail_core = concept_avail\n",
    "                    obj_spec = obj_spec_fun(\n",
    "                        concept_collection=concept_avail_core,\n",
    "                        min_n_objs=1,\n",
    "                        max_n_objs=1,\n",
    "                        canvas_size=canvas_size,\n",
    "                        color_avail=color_avail,\n",
    "                        idx_start=obj_id,\n",
    "                    )[0]\n",
    "                    obj_spec_core.append(obj_spec)\n",
    "                    obj_id += 1\n",
    "                elif element in [\"concept\", \"refer\", \"(concept)\", \"(refer)\"]:\n",
    "                    if not element.startswith(\"(\"):\n",
    "                        obj_spec = obj_spec_fun(\n",
    "                            concept_collection=allowed_shape_concept,\n",
    "                            min_n_objs=1,\n",
    "                            max_n_objs=1,\n",
    "                            canvas_size=canvas_size,\n",
    "                            color_avail=color_avail,\n",
    "                            idx_start=obj_id,\n",
    "                        )[0]\n",
    "                        obj_spec_core.append(obj_spec)\n",
    "                    if element in [\"refer\", \"(refer)\"]:\n",
    "                        assert refer_node_id is None\n",
    "                        refer_node_id = obj_id\n",
    "                    obj_id += 1\n",
    "                else:\n",
    "                    raise\n",
    "            relations_unique = np.unique(relations_all)\n",
    "            if len(relations_unique) == 1:\n",
    "                if structure_key.startswith(\"3\") and relations_unique[0] in [\"SameColor\"]:\n",
    "                    pass\n",
    "                elif relations_unique[0] in [\"SameShape\"] or structure_key.startswith(\"3\"):\n",
    "                    continue\n",
    "            is_valid = True\n",
    "            break\n",
    "        if not is_valid:\n",
    "            return []\n",
    "\n",
    "    task = []\n",
    "    for k in range(n_examples_per_task * 4):\n",
    "        selector_dict = OrderedDict()\n",
    "        if max_n_distractors > 0:\n",
    "            n_distractors = np.random.choice(range(min_n_distractors, max_n_distractors + 1))\n",
    "            obj_spec_distractor = obj_spec_fun(\n",
    "                concept_collection=additional_concepts,\n",
    "                min_n_objs=n_distractors,\n",
    "                max_n_objs=n_distractors,\n",
    "                canvas_size=canvas_size,\n",
    "                color_avail=color_avail,\n",
    "                idx_start=obj_id,\n",
    "            )\n",
    "        else:\n",
    "            obj_spec_distractor = []\n",
    "\n",
    "        obj_spec_all = obj_spec_core + obj_spec_distractor\n",
    "        selector_dict = OrderedDict(obj_spec_all)\n",
    "\n",
    "        is_valid = False\n",
    "        for j in range(max_n_trials):\n",
    "            canvas_dict = dataset_engine.sample_single_canvas_by_core_edges(\n",
    "                selector_dict,\n",
    "                allow_connect=True, is_plot=False, rainbow_prob=0.0,\n",
    "                concept_collection=[concept_str_reverse_mapping[s] for s in concept_avail],\n",
    "                parsing_check=True,\n",
    "                color_avail=color_avail,\n",
    "            )\n",
    "            if canvas_dict == -1:\n",
    "                continue\n",
    "            else:\n",
    "                is_valid = True\n",
    "                if isplot:\n",
    "                    canvas = Canvas(repre_dict=canvas_dict)\n",
    "                    canvas.render()\n",
    "                    plt.show()\n",
    "                break\n",
    "\n",
    "        if is_valid:\n",
    "            image = to_one_hot(canvas_dict[\"image_t\"])\n",
    "            info = Dictionary({\"obj_masks\" : {}})\n",
    "            for k, v in canvas_dict[\"node_id_map\"].items():\n",
    "                info[\"obj_masks\"][k] = canvas_dict[\"id_object_mask\"][v]\n",
    "            info[\"obj_full_relations\"] = canvas_dict[\"partial_relation_edges\"]\n",
    "            # Make sure that if the structure has \"i\" (IsInside), only the obj_1 is inside obj_0 and no other objects:\n",
    "            if relation_structure != \"None\":\n",
    "                n_objs = len(info[\"obj_masks\"])\n",
    "                for i in range(2, n_objs):\n",
    "                    if (f\"obj_{i}\", \"obj_0\") in info[\"obj_full_relations\"] and \"IsInside\" in info[\"obj_full_relations\"][(f\"obj_{i}\", \"obj_0\")]:\n",
    "                        p.print(f\"obj_{i} is also inside the Rect!\")\n",
    "                        is_valid = False\n",
    "                        break\n",
    "                if len(info['obj_masks']) == 3 and is_exist_all(info['obj_full_relations'], key=\"SameColor\"):\n",
    "                    # is_valid = False\n",
    "                    pass\n",
    "                if len(info['obj_masks']) == 2 and is_exist_all(info['obj_full_relations'], key=\"SameShape\"):\n",
    "                    is_valid = False\n",
    "            if not is_valid:\n",
    "                continue\n",
    "\n",
    "            info[\"obj_spec_core\"] = obj_spec_core\n",
    "            info[\"obj_spec_all\"] = obj_spec_all\n",
    "            info[\"obj_spec_distractor\"] = obj_spec_distractor\n",
    "            info[\"refer_node_id\"] = \"obj_{}\".format(refer_node_id)\n",
    "            info[\"structure\"] = structure\n",
    "            masks = None\n",
    "            chosen_concepts = None\n",
    "            if relation_structure != \"None\":\n",
    "                target = info[\"obj_masks\"][info[\"refer_node_id\"]][None]\n",
    "                example = ((image, target), masks, chosen_concepts, info)\n",
    "            else:\n",
    "                example = (image, masks, chosen_concepts, info)\n",
    "            task.append(example)\n",
    "        else:\n",
    "            continue\n",
    "\n",
    "        if len(task) >= n_examples_per_task:\n",
    "            break\n",
    "    return task\n",
    "\n",
    "\n",
    "def get_c_core(c):\n",
    "    if isinstance(c, list):\n",
    "        return [get_c_core(ele) for ele in c]\n",
    "    else:\n",
    "        assert isinstance(c, str)\n",
    "        return c.split(\"[\")[0]\n",
    "\n",
    "\n",
    "def get_c_size(c):\n",
    "    assert isinstance(c, str)\n",
    "    if \"[\" in c:\n",
    "        string = \"[{}]\".format(c.split(\"[\")[1][:-1])\n",
    "        min_size, max_size = eval(string)\n",
    "    else:\n",
    "        min_size, max_size = None, None\n",
    "    return min_size, max_size\n",
    "\n",
    "\n",
    "def get_masks(concept_dict, allowed_concepts, canvas_size):\n",
    "    canvas_all = []\n",
    "    concepts_all = []\n",
    "    for key, item in concept_dict.items():\n",
    "        if key in allowed_concepts:\n",
    "            canvas = torch.zeros(len(item), 1, canvas_size, canvas_size)\n",
    "            for j, pos in enumerate(item):\n",
    "                canvas[j, :, pos[0]: pos[0]+pos[2], pos[1]: pos[1]+pos[3]] = 1\n",
    "            canvas_all.append(canvas)\n",
    "            concepts_all += [key] * len(item)\n",
    "    if len(canvas_all) > 0:\n",
    "        canvas_all = torch.cat(canvas_all)\n",
    "        return canvas_all, concepts_all\n",
    "    else:\n",
    "        return None, None\n",
    "\n",
    "\n",
    "def obj_spec_fun(\n",
    "    concept_collection, \n",
    "    min_n_objs, max_n_objs, \n",
    "    canvas_size, \n",
    "    allowed_shape_concept=None,\n",
    "    is_conjuncture=True,\n",
    "    color_avail=None,\n",
    "    idx_start=0,\n",
    "    focus_type=None,\n",
    "):\n",
    "    \"\"\"Generate specs for several objects for BabyARC.\n",
    "\n",
    "    Args:\n",
    "        idx_start: obj id to start with.\n",
    "    \"\"\"\n",
    "    n_objs = np.random.randint(min_n_objs, max_n_objs+1)\n",
    "    obj_spec = []\n",
    "    if focus_type is not None:\n",
    "        assert focus_type in concept_collection\n",
    "    if set(get_c_core(concept_collection)).issubset({\n",
    "        \"Line\", \"Rect\", \"RectSolid\", \n",
    "        \"Lshape\", \"Randshape\", \"ARCshape\", \n",
    "        \"Tshape\", \"Eshape\", \n",
    "        \"Hshape\", \"Cshape\", \"Ashape\", \"Fshape\",\n",
    "        \"RectE1a\", \"RectE1b\", \"RectE1c\", \n",
    "        \"RectE2a\", \"RectE2b\", \"RectE2c\",\n",
    "        \"RectE3a\", \"RectE3b\", \"RectE3c\", \n",
    "        \"RectF1a\", \"RectF1b\", \"RectF1c\", \n",
    "        \"RectF2a\", \"RectF2b\", \"RectF2c\",\n",
    "        \"RectF3a\", \"RectF3b\", \"RectF3c\",\n",
    "    }):\n",
    "        if focus_type is None:\n",
    "            partition = np.sort(np.random.choice(n_objs+1, len(concept_collection)-1, replace=True))\n",
    "            max_rect_size = canvas_size // 2\n",
    "            for k in range(idx_start, n_objs + idx_start):\n",
    "                if len(concept_collection) == 1:\n",
    "                    chosen_concept = concept_collection[0]\n",
    "                else:\n",
    "                    gt = k-idx_start >= partition  # gt: greater_than_vector\n",
    "                    if gt.any():\n",
    "                        id = np.where(gt)[0][-1] + 1\n",
    "                    else:\n",
    "                        id = 0\n",
    "                    chosen_concept = concept_collection[id]\n",
    "                chosen_concept_core = get_c_core(chosen_concept)\n",
    "                min_size, max_size = get_c_size(chosen_concept)\n",
    "                if chosen_concept_core == \"Line\":\n",
    "                    if min_size is None:\n",
    "                        obj_spec.append((('obj_{}'.format(k), 'line_[-1,1,-1]'), 'Attr'))\n",
    "                    else:\n",
    "                        h = np.random.randint(min_size, max_size+1)\n",
    "                        obj_spec.append((('obj_{}'.format(k), f'line_[{h},1,-1]'), 'Attr'))\n",
    "                elif chosen_concept_core == \"Rect\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(3, max_rect_size+1, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    obj_spec.append((('obj_{}'.format(k), 'rectangle_[{},{}]'.format(w,h)), 'Attr'))\n",
    "                elif chosen_concept_core == \"RectSolid\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(2, max_rect_size+1, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    obj_spec.append((('obj_{}'.format(k), 'rectangleSolid_[{},{}]'.format(w,h)), 'Attr'))\n",
    "                elif chosen_concept_core == \"Lshape\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(3, max_rect_size+1, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    direction = np.random.randint(4)\n",
    "                    obj_spec.append((('obj_{}'.format(k), 'Lshape_[{},{},{}]'.format(w,h,direction)), 'Attr'))\n",
    "                elif chosen_concept_core == \"Tshape\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(3, max_rect_size+2, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    obj_spec += [(('obj_{}'.format(k), f'Tshape_[{w},{h}]'), 'Attr')]\n",
    "                elif chosen_concept_core == \"Eshape\":\n",
    "                    if min_size is None:\n",
    "                        w = np.random.randint(3, max_rect_size+1)\n",
    "                        h = np.random.randint(5, max_rect_size+3)\n",
    "                    else:\n",
    "                        w = np.random.randint(min_size, max_size-1)\n",
    "                        h = np.random.randint(min_size, max_size+1)\n",
    "                    obj_spec += [(('obj_{}'.format(k), f'Eshape_[{w},{h}]'), 'Attr')]\n",
    "                elif chosen_concept_core == \"Hshape\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(3, max_rect_size+2, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    obj_spec += [(('obj_{}'.format(k), f'Hshape_[{w},{h}]'), 'Attr')]\n",
    "                elif chosen_concept_core == \"Cshape\":\n",
    "                    if min_size is None:\n",
    "                        w = np.random.randint(3, max_rect_size+1)\n",
    "                        h = np.random.randint(3, max_rect_size+2)\n",
    "                    else:\n",
    "                        w = np.random.randint(min_size, max_size)\n",
    "                        h = np.random.randint(min_size, max_size+1)\n",
    "                    obj_spec += [(('obj_{}'.format(k), f'Cshape_[{w},{h}]'), 'Attr')]\n",
    "                elif chosen_concept_core == \"Ashape\":\n",
    "                    if min_size is None:\n",
    "                        w = np.random.randint(3, max_rect_size+2)\n",
    "                        h = np.random.randint(4, max_rect_size+3)\n",
    "                    else:\n",
    "                        w = np.random.randint(min_size, max_size)\n",
    "                        h = np.random.randint(min_size, max_size+1)\n",
    "                    obj_spec += [(('obj_{}'.format(k), f'Ashape_[{w},{h}]'), 'Attr')]\n",
    "                elif chosen_concept_core == \"Fshape\":\n",
    "                    if min_size is None:\n",
    "                        w = np.random.randint(3, max_rect_size+1)\n",
    "                        h = np.random.randint(4, max_rect_size+3)\n",
    "                    else:\n",
    "                        w = np.random.randint(min_size, max_size-1)\n",
    "                        h = np.random.randint(min_size, max_size+1)\n",
    "                    obj_spec += [(('obj_{}'.format(k), f'Fshape_[{w},{h}]'), 'Attr')]   \n",
    "                elif chosen_concept_core == \"Randshape\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(2, max_rect_size+1, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    obj_spec.append((('obj_{}'.format(k), 'randomShape_[{},{}]'.format(w,h)), 'Attr'))\n",
    "                elif chosen_concept_core == \"ARCshape\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(2, max_rect_size+1, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    obj_spec.append((('obj_{}'.format(k), 'arcShape_[{},{}]'.format(w,h)), 'Attr'))\n",
    "                elif chosen_concept_core in [\n",
    "                    \"RectE1a\", \"RectE1b\", \"RectE1c\", \n",
    "                    \"RectE2a\", \"RectE2b\", \"RectE2c\",\n",
    "                    \"RectE3a\", \"RectE3b\", \"RectE3c\", \n",
    "                    \"RectF1a\", \"RectF1b\", \"RectF1c\", \n",
    "                    \"RectF2a\", \"RectF2b\", \"RectF2c\",\n",
    "                    \"RectF3a\", \"RectF3b\", \"RectF3c\",\n",
    "                ]:\n",
    "                    w, h = -1, -1 # let the canvas size drives here!\n",
    "                    obj_spec.append((('obj_{}'.format(k), '{}_[{},{}]'.format(chosen_concept_core, w,h)), 'Attr'))\n",
    "                else:\n",
    "                    raise\n",
    "            obj_spec = np.random.permutation(obj_spec).tolist()\n",
    "        else:\n",
    "            concept_collection = deepcopy(concept_collection)\n",
    "            concept_collection.remove(focus_type)\n",
    "            partition = np.sort(np.random.choice(n_objs, len(concept_collection)-1, replace=True))\n",
    "            max_rect_size = canvas_size // 2\n",
    "            for k in range(idx_start, n_objs + idx_start):\n",
    "                if k == n_objs + idx_start - 1:\n",
    "                    chosen_concept = focus_type\n",
    "                else:\n",
    "                    if len(concept_collection) == 1:\n",
    "                        chosen_concept = concept_collection[0]\n",
    "                    else:\n",
    "                        gt = k-idx_start >= partition  # gt: greater_than_vector\n",
    "                        if gt.any():\n",
    "                            id = np.where(gt)[0][-1] + 1\n",
    "                        else:\n",
    "                            id = 0\n",
    "                        chosen_concept = concept_collection[id]\n",
    "                chosen_concept_core = get_c_core(chosen_concept)\n",
    "                min_size, max_size = get_c_size(chosen_concept)\n",
    "                if chosen_concept_core == \"Line\":\n",
    "                    if min_size is None:\n",
    "                        obj_spec.append((('obj_{}'.format(k), 'line_[-1,1,-1]'), 'Attr'))\n",
    "                    else:\n",
    "                        h = np.random.randint(min_size, max_size+1)\n",
    "                        obj_spec.append((('obj_{}'.format(k), f'line_[{h},1,-1]'), 'Attr'))\n",
    "                elif chosen_concept_core == \"Rect\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(3, max_rect_size+1, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    obj_spec.append((('obj_{}'.format(k), 'rectangle_[{},{}]'.format(w,h)), 'Attr'))\n",
    "                elif chosen_concept_core == \"RectSolid\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(2, max_rect_size+1, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    obj_spec.append((('obj_{}'.format(k), 'rectangleSolid_[{},{}]'.format(w,h)), 'Attr'))\n",
    "                elif chosen_concept_core == \"Lshape\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(3, max_rect_size+1, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    direction = np.random.randint(4)\n",
    "                    obj_spec.append((('obj_{}'.format(k), 'Lshape_[{},{},{}]'.format(w,h,direction)), 'Attr'))\n",
    "                elif chosen_concept_core == \"Tshape\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(3, max_rect_size+2, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    obj_spec += [(('obj_{}'.format(k), f'Tshape_[{w},{h}]'), 'Attr')]\n",
    "                elif chosen_concept_core == \"Eshape\":\n",
    "                    if min_size is None:\n",
    "                        w = np.random.randint(3, max_rect_size+1)\n",
    "                        h = np.random.randint(5, max_rect_size+3)\n",
    "                    else:\n",
    "                        w = np.random.randint(min_size, max_size-1)\n",
    "                        h = np.random.randint(min_size, max_size+1)\n",
    "                    obj_spec += [(('obj_{}'.format(k), f'Eshape_[{w},{h}]'), 'Attr')]\n",
    "                elif chosen_concept_core == \"Hshape\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(3, max_rect_size+2, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    obj_spec += [(('obj_{}'.format(k), f'Hshape_[{w},{h}]'), 'Attr')]\n",
    "                elif chosen_concept_core == \"Cshape\":\n",
    "                    if min_size is None:\n",
    "                        w = np.random.randint(3, max_rect_size+1)\n",
    "                        h = np.random.randint(3, max_rect_size+2)\n",
    "                    else:\n",
    "                        w = np.random.randint(min_size, max_size)\n",
    "                        h = np.random.randint(min_size, max_size+1)\n",
    "                    obj_spec += [(('obj_{}'.format(k), f'Cshape_[{w},{h}]'), 'Attr')]\n",
    "                elif chosen_concept_core == \"Ashape\":\n",
    "                    if min_size is None:\n",
    "                        w = np.random.randint(3, max_rect_size+2)\n",
    "                        h = np.random.randint(4, max_rect_size+3)\n",
    "                    else:\n",
    "                        w = np.random.randint(min_size, max_size)\n",
    "                        h = np.random.randint(min_size, max_size+1)\n",
    "                    obj_spec += [(('obj_{}'.format(k), f'Ashape_[{w},{h}]'), 'Attr')]\n",
    "                elif chosen_concept_core == \"Fshape\":\n",
    "                    if min_size is None:\n",
    "                        w = np.random.randint(3, max_rect_size+1)\n",
    "                        h = np.random.randint(4, max_rect_size+3)\n",
    "                    else:\n",
    "                        w = np.random.randint(min_size, max_size-1)\n",
    "                        h = np.random.randint(min_size, max_size+1)\n",
    "                    obj_spec += [(('obj_{}'.format(k), f'Fshape_[{w},{h}]'), 'Attr')]   \n",
    "                elif chosen_concept_core == \"Randshape\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(2, max_rect_size+1, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    obj_spec.append((('obj_{}'.format(k), 'randomShape_[{},{}]'.format(w,h)), 'Attr'))\n",
    "                elif chosen_concept_core == \"ARCshape\":\n",
    "                    if min_size is None:\n",
    "                        w, h = np.random.randint(2, max_rect_size+1, size=2)\n",
    "                    else:\n",
    "                        w, h = np.random.randint(min_size, max_size+1, size=2)\n",
    "                    obj_spec.append((('obj_{}'.format(k), 'arcShape_[{},{}]'.format(w,h)), 'Attr'))\n",
    "                elif chosen_concept_core in [\n",
    "                    \"RectE1a\", \"RectE1b\", \"RectE1c\", \n",
    "                    \"RectE2a\", \"RectE2b\", \"RectE2c\",\n",
    "                    \"RectE3a\", \"RectE3b\", \"RectE3c\", \n",
    "                    \"RectF1a\", \"RectF1b\", \"RectF1c\", \n",
    "                    \"RectF2a\", \"RectF2b\", \"RectF2c\",\n",
    "                    \"RectF3a\", \"RectF3b\", \"RectF3c\",\n",
    "                ]:\n",
    "                    w, h = -1, -1 # let the canvas size drives here!\n",
    "                    obj_spec.append((('obj_{}'.format(k), '{}_[{},{}]'.format(chosen_concept_core, w,h)), 'Attr'))\n",
    "                else:\n",
    "                    raise\n",
    "            obj_spec = obj_spec[-1:] + np.random.permutation(obj_spec[:-1]).tolist()\n",
    "    elif set(concept_collection).issubset({\"SameColor\", \"IsTouch\"}):\n",
    "        if len(concept_collection) > 1:\n",
    "            if is_conjuncture:\n",
    "                # Hard code probability\n",
    "                if color_avail == None:\n",
    "                    random_color = np.random.randint(1, 10)\n",
    "                else:\n",
    "                    random_color = random.choice(color_avail)\n",
    "                obj_spec.append((('obj_{}'.format(idx_start), 'obj_{}'.format(idx_start+1)), 'IsTouch'))\n",
    "                obj_spec.append((('obj_{}'.format(idx_start), f'color_[{random_color}]'), 'Attr'))\n",
    "                obj_spec.append((('obj_{}'.format(idx_start+1), f'color_[{random_color}]'), 'Attr'))\n",
    "            else:\n",
    "                pass # TODO: not implemented\n",
    "        else:\n",
    "            if len(concept_collection) == 1:\n",
    "                chosen_concept = concept_collection[0]\n",
    "                if chosen_concept == \"SameColor\":\n",
    "                    obj_spec.append((('obj_{}'.format(idx_start), 'obj_{}'.format(idx_start+1)), 'SameColor'))\n",
    "                else:\n",
    "                    obj_spec.append((('obj_{}'.format(idx_start), 'obj_{}'.format(idx_start+1)), 'IsTouch'))\n",
    "    # complex shape.\n",
    "    elif set(concept_collection).issubset({\"RectE1a\", \"RectE1b\", \"RectE1c\", \n",
    "                                           \"RectE2a\", \"RectE2b\", \"RectE2c\",\n",
    "                                           \"RectE3a\", \"RectE3b\", \"RectE3c\", \n",
    "                                           \"RectF1a\", \"RectF1b\", \"RectF1c\", \n",
    "                                           \"RectF2a\", \"RectF2b\", \"RectF2c\",\n",
    "                                           \"RectF3a\", \"RectF3b\", \"RectF3c\",}):\n",
    "        chosen_concept = random.choice(concept_collection)\n",
    "        if chosen_concept == \"RectE1a\" or chosen_concept == \"RectF1a\" or chosen_concept == \"RectE1b\" or chosen_concept == \"RectF1b\":\n",
    "            char_shape = \"E\" if \"E\" in chosen_concept else \"F\"\n",
    "            out_w = np.random.randint(5, 17)\n",
    "            out_h = np.random.randint(5, 17)\n",
    "            in_w = np.random.randint(4, out_w)\n",
    "            in_h = np.random.randint(4, out_h)\n",
    "            char_w = np.random.randint(3, 9)\n",
    "            char_h = np.random.randint(5, 9)\n",
    "            obj_spec = [(('obj_0', f'rectangle_[{out_w},{out_h}]'), 'Attr'), \n",
    "                         (('obj_1', f'rectangle_[{in_w},{in_h}]'), 'Attr'), \n",
    "                         (('obj_0', 'obj_1'), 'IsOutside'),\n",
    "                         (('obj_2', f'{char_shape}shape_[{char_w},{char_h}]'), 'Attr')]\n",
    "        elif chosen_concept == \"RectE1c\" or chosen_concept == \"RectF1c\":\n",
    "            char_shape = \"E\" if \"E\" in chosen_concept else \"F\"\n",
    "            out_w = np.random.randint(5, 17)\n",
    "            out_h = np.random.randint(5, 17)\n",
    "            in_w = np.random.randint(4, out_w)\n",
    "            in_h = np.random.randint(4, out_h)\n",
    "            char_w = np.random.randint(3, 9)\n",
    "            char_h = np.random.randint(5, 9)\n",
    "            obj_spec = [(('obj_0', f'rectangle_[{out_w},{out_h}]'), 'Attr'), \n",
    "                         (('obj_1', f'rectangle_[{in_w},{in_h}]'), 'Attr'), \n",
    "                         (('obj_0', 'obj_1'), 'IsOutside'),\n",
    "                         (('obj_2', f'{char_shape}shape_[{char_w},{char_h}]'), 'Attr'), \n",
    "                         (('obj_1', 'obj_2'), 'SameColor')]\n",
    "        elif chosen_concept == \"RectE2a\" or chosen_concept == \"RectF2a\" or chosen_concept == \"RectE2b\" or chosen_concept == \"RectF2b\":\n",
    "            char_shape = \"E\" if \"E\" in chosen_concept else \"F\"\n",
    "            out_w = np.random.randint(5, 17)\n",
    "            out_h = np.random.randint(5, 17)\n",
    "            in_w = np.random.randint(4, 8)\n",
    "            in_h = np.random.randint(4, 8)\n",
    "            char_w = np.random.randint(3, 8)\n",
    "            char_h = np.random.randint(5, 8)\n",
    "            obj_spec = [(('obj_0', f'rectangle_[{out_w},{out_h}]'), 'Attr'), \n",
    "                         (('obj_1', f'rectangle_[{in_w},{in_h}]'), 'Attr'), \n",
    "                         (('obj_0', 'obj_1'), 'IsOutside'),\n",
    "                         (('obj_2', f'{char_shape}shape_[{char_w},{char_h}]'), 'Attr'), \n",
    "                         (('obj_0', 'obj_2'), 'IsOutside')]\n",
    "        elif chosen_concept == \"RectE2c\" or chosen_concept == \"RectF2c\":\n",
    "            char_shape = \"E\" if \"E\" in chosen_concept else \"F\"\n",
    "            out_w = np.random.randint(5, 17)\n",
    "            out_h = np.random.randint(5, 17)\n",
    "            in_w = np.random.randint(4, 8)\n",
    "            in_h = np.random.randint(4, 8)\n",
    "            char_w = np.random.randint(3, 8)\n",
    "            char_h = np.random.randint(5, 8)\n",
    "            obj_spec = [(('obj_0', f'rectangle_[{out_w},{out_h}]'), 'Attr'), \n",
    "                         (('obj_1', f'rectangle_[{in_w},{in_h}]'), 'Attr'), \n",
    "                         (('obj_0', 'obj_1'), 'IsOutside'),\n",
    "                         (('obj_2', f'{char_shape}shape_[{char_w},{char_h}]'), 'Attr'), \n",
    "                         (('obj_0', 'obj_2'), 'IsOutside'), \n",
    "                         (('obj_1', 'obj_2'), 'SameColor')]\n",
    "        elif chosen_concept == \"RectE3a\" or chosen_concept == \"RectF3a\" or chosen_concept == \"RectE3b\" or chosen_concept == \"RectF3b\":\n",
    "            char_shape = \"E\" if \"E\" in chosen_concept else \"F\"\n",
    "            out_w = np.random.randint(9, 17)\n",
    "            out_h = np.random.randint(9, 17)\n",
    "            in_w = np.random.randint(7, out_w-1)\n",
    "            in_h = np.random.randint(7, out_h-1)\n",
    "            char_w = np.random.randint(3, in_w-1)\n",
    "            char_h = np.random.randint(5, in_h-1)\n",
    "            obj_spec = [(('obj_0', f'rectangle_[{out_w},{out_h}]'), 'Attr'), \n",
    "                         (('obj_1', f'rectangle_[{in_w},{in_h}]'), 'Attr'), \n",
    "                         (('obj_0', 'obj_1'), 'IsOutside'),\n",
    "                         (('obj_2', f'Eshape_[{char_w},{char_h}]'), 'Attr'), \n",
    "                         (('obj_1', 'obj_2'), 'IsOutside')]\n",
    "        elif chosen_concept == \"RectE3c\" or chosen_concept == \"RectF3c\":\n",
    "            char_shape = \"E\" if \"E\" in chosen_concept else \"F\"\n",
    "            out_w = np.random.randint(9, 17)\n",
    "            out_h = np.random.randint(9, 17)\n",
    "            in_w = np.random.randint(7, out_w-1)\n",
    "            in_h = np.random.randint(7, out_h-1)\n",
    "            char_w = np.random.randint(3, in_w-1)\n",
    "            char_h = np.random.randint(5, in_h-1)\n",
    "            obj_spec = [(('obj_0', f'rectangle_[{out_w},{out_h}]'), 'Attr'), \n",
    "                         (('obj_1', f'rectangle_[{in_w},{in_h}]'), 'Attr'), \n",
    "                         (('obj_0', 'obj_1'), 'IsOutside'),\n",
    "                         (('obj_2', f'Eshape_[{char_w},{char_h}]'), 'Attr'), \n",
    "                         (('obj_1', 'obj_2'), 'IsOutside'), \n",
    "                         (('obj_1', 'obj_2'), 'SameColor')]\n",
    "    else:\n",
    "        raise Exception(\"concept_collection {} must be a subset of 'Line', 'Rect', 'Lshape', 'Randshape'!\".format(concept_collection))\n",
    "    return obj_spec\n",
    "\n",
    "\n",
    "def generate_samples(\n",
    "    dataset,\n",
    "    obj_spec_fun,\n",
    "    n_examples,\n",
    "    mode,\n",
    "    concept_collection,\n",
    "    min_n_objs,\n",
    "    max_n_objs,\n",
    "    canvas_size,\n",
    "    rainbow_prob=0.,\n",
    "    allow_connect=True,\n",
    "    parsing_check=False,\n",
    "    focus_type=None,\n",
    "    inspect_interval=\"auto\",\n",
    "    save_interval=-1,\n",
    "    save_filename=None,\n",
    "    **kwargs\n",
    "):\n",
    "    data = []\n",
    "    if inspect_interval == \"auto\":\n",
    "        inspect_interval = max(1, n_examples // 100)\n",
    "    for i in range(int(n_examples * 150)):\n",
    "        info = {}\n",
    "        obj_spec = obj_spec_fun(\n",
    "            concept_collection=concept_collection,\n",
    "            min_n_objs=min_n_objs,\n",
    "            max_n_objs=max_n_objs,\n",
    "            canvas_size=canvas_size,\n",
    "            allowed_shape_concept=kwargs[\"allowed_shape_concept\"],\n",
    "            color_avail=kwargs[\"color_avail\"],\n",
    "            focus_type=focus_type,\n",
    "        )\n",
    "        info[\"obj_spec\"] = obj_spec\n",
    "\n",
    "        if mode == \"relation\":\n",
    "            concept_limits = {kwargs[\"concept_str_reverse_mapping\"][get_c_core(c)]: get_c_size(c) for c in kwargs[\"allowed_shape_concept\"] if get_c_size(c)[0] is not None}\n",
    "            canvas_dict = dataset.sample_single_canvas_by_core_edges(\n",
    "                OrderedDict(obj_spec),\n",
    "                allow_connect=allow_connect,\n",
    "                rainbow_prob=rainbow_prob,\n",
    "                is_plot=False,\n",
    "                concept_collection=[kwargs[\"concept_str_reverse_mapping\"][s] for s in get_c_core(kwargs[\"allowed_shape_concept\"])],\n",
    "                parsing_check=parsing_check,\n",
    "                color_avail=kwargs[\"color_avail\"],\n",
    "                concept_limits=concept_limits,\n",
    "            )\n",
    "        else:\n",
    "            canvas_dict = dataset.sample_single_canvas_by_core_edges(\n",
    "                OrderedDict(obj_spec),\n",
    "                allow_connect=allow_connect,\n",
    "                rainbow_prob=rainbow_prob,\n",
    "                is_plot=False,\n",
    "                parsing_check=parsing_check,\n",
    "                color_avail=kwargs[\"color_avail\"],\n",
    "            )\n",
    "        if canvas_dict != -1:\n",
    "            info[\"node_id_map\"] = canvas_dict[\"node_id_map\"]\n",
    "            info[\"id_object_mask\"] = canvas_dict[\"id_object_mask\"]\n",
    "            n_sampled_objs = len(canvas_dict['id_object_mask'])\n",
    "            if mode == \"concept\":\n",
    "                if focus_type is None:\n",
    "                    for k in range(n_sampled_objs):\n",
    "                        # The order of id is the same as its first appearance in the obj_spec:\n",
    "                        data.append((\n",
    "                            to_one_hot(canvas_dict[\"image_t\"]),\n",
    "                            (canvas_dict['id_object_mask'][k][None],),\n",
    "                            kwargs[\"concept_str_mapping\"][obj_spec[k][0][1].split(\"_\")[0]],\n",
    "                            Dictionary(info),\n",
    "                        ))\n",
    "                        if len(data) >= n_examples:\n",
    "                            break\n",
    "                else:\n",
    "                    data.append((\n",
    "                        to_one_hot(canvas_dict[\"image_t\"]),\n",
    "                        (canvas_dict['id_object_mask'][0][None],),\n",
    "                        kwargs[\"concept_str_mapping\"][obj_spec[0][0][1].split(\"_\")[0]],\n",
    "                        Dictionary(info),\n",
    "                    ))\n",
    "            elif mode == \"concept-image\":\n",
    "                for k in range(n_sampled_objs):\n",
    "                    # The order of id is the same as its first appearance in the obj_spec:\n",
    "                    data.append((\n",
    "                        to_one_hot(canvas_dict[\"image_t\"]),\n",
    "                        (canvas_dict['id_object_mask'][k][None],),\n",
    "                        \"Image\",\n",
    "                        Dictionary(info),\n",
    "                    ))\n",
    "                    if len(data) >= n_examples:\n",
    "                        break\n",
    "            elif mode == \"relation\":\n",
    "                if set(concept_collection).issubset({\"Parallel\", \"Vertical\"}):\n",
    "                    def get_chosen_rel(canvas_dict, obj_ids):\n",
    "                        chosen_direction = []\n",
    "                        for id in obj_ids:\n",
    "                            shape = canvas_dict['id_object_map'][id].shape\n",
    "                            if shape[0] > shape[1]:\n",
    "                                chosen_direction.append(\"0\")\n",
    "                            elif shape[0] < shape[1]:\n",
    "                                chosen_direction.append(\"1\")\n",
    "                            else:\n",
    "                                raise Exception(\"Line must have unequal height and width!\")\n",
    "                        if len(set(chosen_direction)) == 1:\n",
    "                            chosen_concept = \"Parallel\"\n",
    "                        else:\n",
    "                            assert len(set(chosen_direction)) == 2\n",
    "                            chosen_concept = \"Vertical\"\n",
    "                        return chosen_concept\n",
    "\n",
    "                    chosen_obj_ids = np.random.choice(n_sampled_objs, size=2, replace=False)\n",
    "                    masks = list(canvas_dict['id_object_mask'].values())\n",
    "                    chosen_obj_types = [obj_spec[id][0][1].split(\"_\")[0] for id in chosen_obj_ids]\n",
    "                    assert set(np.unique(chosen_obj_types)) == {\"line\"}\n",
    "                    chosen_masks = [masks[id][None] for id in chosen_obj_ids]  # after: each mask has shape [1, H, W]\n",
    "                    chosen_concept = get_chosen_rel(canvas_dict, chosen_obj_ids)\n",
    "                    if chosen_concept not in concept_collection:\n",
    "                        continue\n",
    "                    # Consider all relations \n",
    "                    relations = []\n",
    "                    for id1 in range(n_sampled_objs):\n",
    "                        for id2 in range(id1 + 1, n_sampled_objs):\n",
    "                            relation = get_chosen_rel(canvas_dict, [id1, id2])\n",
    "                            relations.append((id1, id2, relation))\n",
    "                    info[\"relations\"] = relations\n",
    "                    data.append((\n",
    "                        to_one_hot(canvas_dict[\"image_t\"]),\n",
    "                        tuple(chosen_masks),\n",
    "                        chosen_concept,\n",
    "                        Dictionary(info),\n",
    "                    ))\n",
    "                if set(concept_collection).issubset({\n",
    "                    \"SameAll\", \"SameShape\", \"SameColor\", \n",
    "                    \"SameRow\", \"SameCol\", \"IsInside\", \"IsTouch\",\n",
    "                    \"IsNonOverlapXY\", \"IsEnclosed\",\n",
    "                }):\n",
    "                    masks = list(canvas_dict['id_object_mask'].values())\n",
    "                    chosen_concept = obj_spec[0][1]\n",
    "                    if chosen_concept not in concept_collection:\n",
    "                        continue\n",
    "                    chosen_obj_ids = [0,1] # we assum it is the first two objs have the relation type always!\n",
    "                    if chosen_concept == \"IsInside\":\n",
    "                        chosen_obj_ids = [1,0] # reverse it.\n",
    "                        chosen_masks = [masks[id][None] for id in chosen_obj_ids]\n",
    "                    else:\n",
    "                        chosen_masks = [masks[id][None] for id in chosen_obj_ids]\n",
    "                    data.append((\n",
    "                        to_one_hot(canvas_dict[\"image_t\"]),\n",
    "                        tuple(chosen_masks),\n",
    "                        chosen_concept,\n",
    "                        Dictionary(info),\n",
    "                    ))\n",
    "            elif mode == \"compositional-concept\":\n",
    "                # we need to filter.\n",
    "                assert len(concept_collection) == 1 # we only allow 1.\n",
    "                chosen_concept = concept_collection[0]\n",
    "                failed = False\n",
    "                if \"1\" in chosen_concept:\n",
    "                    mid_ax_y_l = canvas_dict[\"id_position_map\"][1][0].tolist() + (canvas_dict[\"id_object_map\"][1].shape[0])//2\n",
    "                    mid_ax_x_l = canvas_dict[\"id_position_map\"][1][1].tolist() + (canvas_dict[\"id_object_map\"][1].shape[1])//2\n",
    "\n",
    "                    mid_ax_y_r = canvas_dict[\"id_position_map\"][2][0].tolist() + (canvas_dict[\"id_object_map\"][2].shape[0])//2\n",
    "                    mid_ax_x_r = canvas_dict[\"id_position_map\"][2][1].tolist() + (canvas_dict[\"id_object_map\"][2].shape[1])//2\n",
    "\n",
    "                    if abs(mid_ax_y_l-mid_ax_y_r) <= 2 or abs(mid_ax_x_l-mid_ax_x_r) <= 2:\n",
    "                        if ('obj_2', 'obj_0') in canvas_dict[\"partial_relation_edges\"]:\n",
    "                            if \"IsInside\" in canvas_dict[\"partial_relation_edges\"][('obj_2', 'obj_0')]:\n",
    "                                failed = True\n",
    "                        if ('obj_2', 'obj_1') in canvas_dict[\"partial_relation_edges\"]:\n",
    "                            if \"IsInside\" in canvas_dict[\"partial_relation_edges\"][('obj_2', 'obj_1')]:\n",
    "                                failed = True\n",
    "                    else:\n",
    "                        failed = True\n",
    "                if \"2\" in chosen_concept:\n",
    "                    if ('obj_2', 'obj_1') in canvas_dict[\"partial_relation_edges\"]:\n",
    "                        if \"IsInside\" in canvas_dict[\"partial_relation_edges\"][('obj_2', 'obj_1')]:\n",
    "                            failed = True\n",
    "                if \"3\" in chosen_concept:\n",
    "                    pass\n",
    "                if \"b\" in chosen_concept:\n",
    "                    if ('obj_1', 'obj_2') in canvas_dict[\"partial_relation_edges\"]:\n",
    "                        if \"SameColor\" in canvas_dict[\"partial_relation_edges\"][('obj_1', 'obj_2')]:\n",
    "                            failed = True\n",
    "                if not failed:\n",
    "                    data.append((\n",
    "                        to_one_hot(canvas_dict[\"image_t\"]),\n",
    "                        (canvas_dict[\"image_t\"][None].bool().float(),),\n",
    "                        \"Compositional-Image\",\n",
    "                        Dictionary(info),\n",
    "                    ))\n",
    "                    if inspect_interval != \"None\" and len(data) % inspect_interval == 0:\n",
    "                        p.print(\"Finished generating {} out of {} examples.\".format(len(data), n_examples))\n",
    "\n",
    "        if inspect_interval != \"None\" and len(data) % inspect_interval == 0:\n",
    "            p.print(\"Finished generating {} out of {} examples.\".format(len(data), n_examples))\n",
    "        if save_filename is not None and save_interval != -1 and len(data) % save_interval == 0:\n",
    "            try_call(pdump, args=[data, save_filename])\n",
    "            p.print(\"Save intermediate file at {}, with {} examples.\".format(save_filename, len(data)))\n",
    "        if len(data) >= n_examples:\n",
    "            break\n",
    "        if i > n_examples * 2 and len(data) < n_examples * 0.005:\n",
    "            raise Exception(\"Sampled {} times and only {} of them satisfies the specified condition. Try relaxing the condition!\".format(i, len(data)))\n",
    "    return data\n",
    "\n",
    "\n",
    "def get_chosen_line_rel(pos1, pos2):\n",
    "    chosen_direction = []\n",
    "    # \"0\" corresponds to upright, \"1\" corresponds to horizontal\n",
    "    chosen_direction.append(\"0\" if pos1[2] > pos1[3] else \"1\")\n",
    "    chosen_direction.append(\"0\" if pos2[2] > pos2[3] else \"1\")\n",
    "    if len(set(chosen_direction)) == 1:\n",
    "        chosen_concept = \"Parallel\"\n",
    "    else:\n",
    "        assert len(set(chosen_direction)) == 2\n",
    "        # Set the horizontal and upright \n",
    "        if pos1[2] > pos1[3]:\n",
    "            hori = pos2\n",
    "            upr = pos1\n",
    "        else:\n",
    "            hori = pos1\n",
    "            upr = pos2\n",
    "        # Determine whether the intersection is more like a T shape or a rotated T shape\n",
    "        # by checking if the upright line falls between horizontal line's left and right edges\n",
    "        isT = upr[1] >= hori[1] and upr[1] < hori[1] + hori[3]\n",
    "        isRotT= hori[0] >= upr[0] and hori[0] < upr[0] + upr[2]\n",
    "        if isT:\n",
    "            # First check for separate by comparing top edge of vertical with hori\n",
    "            dist1 = abs(upr[0] - hori[0])\n",
    "            dist2 = abs(upr[0] + upr[2] - 1 - hori[0])\n",
    "            if min(dist1, dist2) > 1:\n",
    "                chosen_concept = \"VerticalSepa\"\n",
    "            else:\n",
    "                # Check where the lines are touching w.r.t. horizontal line\n",
    "                if upr[1] == hori[1] or upr[1] == hori[1] + hori[3] - 1:\n",
    "                    chosen_concept = \"VerticalEdge\"\n",
    "                else:\n",
    "                    chosen_concept = \"VerticalMid\"\n",
    "        elif isRotT:\n",
    "            dist1 = abs(hori[1] - upr[1])\n",
    "            dist2 = abs(hori[1] + hori[3] - 1 - upr[1])\n",
    "            if min(dist1, dist2) > 1:\n",
    "                chosen_concept = \"VerticalSepa\"\n",
    "            else:\n",
    "                # Check where the lines are touching w.r.t. vertical line\n",
    "                if hori[0] == upr[0] or hori[0] == upr[0] + upr[2] - 1:\n",
    "                    chosen_concept = \"VerticalEdge\"\n",
    "                else:\n",
    "                    chosen_concept = \"VerticalMid\"\n",
    "        else:\n",
    "            chosen_concept = \"VerticalSepa\"\n",
    "    return chosen_concept\n",
    "\n",
    "\n",
    "# Line with \"VerticalMid\", \"VerticalEdge\", \"VerticalNot\", \"Parallel\":\n",
    "def generate_lines_full_vertical_parallel(\n",
    "    n_examples,\n",
    "    concept_collection=[\"VerticalMid\", \"VerticalEdge\", \"VerticalNot\", \"Parallel\"],\n",
    "    min_n_objs=2,\n",
    "    max_n_objs=4,\n",
    "    canvas_size=16,\n",
    "    min_size=2,\n",
    "    max_size=None,\n",
    "    color_avail=None,\n",
    "    isplot=False,\n",
    "):                  \n",
    "    if color_avail is None:\n",
    "        color_avail = [1,2,3,4,5,6,7,8,9]\n",
    "    data = []\n",
    "    if max_size is None:\n",
    "        max_size = canvas_size\n",
    "    for i in range(int(n_examples*1.5)):\n",
    "        if i % 1000 == 0:\n",
    "            p.print(i)\n",
    "        image = torch.zeros(1, canvas_size, canvas_size)\n",
    "        # Sample relation from concept_collection:\n",
    "        relation = np.random.choice(concept_collection)\n",
    "        if relation == \"Parallel\":\n",
    "            direction = np.random.choice([0,1])\n",
    "            image, mask1, pos1, _ = get_line(image, direction=direction, pos=None, min_size=min_size, max_size=max_size, color_avail=color_avail)\n",
    "            image, mask2, pos2, _ = get_line(image, direction=direction, pos=None, min_size=min_size, max_size=max_size, color_avail=color_avail)\n",
    "            if (direction == 0 and pos1[0] == pos2[0]) or (direction == 1 and pos1[1] == pos2[1]):\n",
    "                # The two lines cannot be on the same straight line:\n",
    "                continue\n",
    "        elif relation == \"VerticalMid\":\n",
    "            direction = np.random.choice([0,1])\n",
    "            image, mask1, pos1, _ = get_line(image, direction=direction, pos=None, min_size=max(3, min_size), max_size=max_size, color_avail=color_avail)\n",
    "            pos_new = get_pos_new_mid(pos1, direction, min_size=min_size, canvas_size=canvas_size)\n",
    "            if pos_new is None:\n",
    "                continue\n",
    "            image, mask2, pos2, _ = get_line(image, direction=None, pos=pos_new, min_size=min_size, max_size=max_size, color_avail=color_avail)\n",
    "        elif relation == \"VerticalEdge\":\n",
    "            direction = np.random.choice([0,1])\n",
    "            image, mask1, pos1, _ = get_line(image, direction=direction, pos=None, min_size=max(3, min_size), max_size=max_size, color_avail=color_avail)\n",
    "            pos_new = get_pos_new_edge(pos1, direction, min_size=min_size, canvas_size=canvas_size)\n",
    "            if pos_new is None:\n",
    "                continue\n",
    "            image, mask2, pos2, _ = get_line(image, direction=None, pos=pos_new, min_size=min_size, max_size=max_size, color_avail=color_avail)\n",
    "        elif relation == \"VerticalSepa\":\n",
    "            direction = np.random.choice([0,1])\n",
    "            image, mask1, pos1, _ = get_line(image, direction=direction, pos=None, min_size=min_size, max_size=max_size, color_avail=color_avail)\n",
    "            pos_new = get_pos_new_not_touching(mask1, direction=1-direction, min_size=min_size, max_size=max_size, canvas_size=canvas_size, pos=pos1, min_distance=1)\n",
    "            if pos_new is None:\n",
    "                continue\n",
    "            image, mask2, pos2, _ = get_line(image, direction=None, pos=pos_new, min_size=min_size, max_size=max_size, color_avail=color_avail)\n",
    "        \n",
    "        # Randomly permute the mask\n",
    "        if np.random.choice([0,1]) == 1:\n",
    "            mask1, mask2 = mask2, mask1\n",
    "            pos1, pos2 = pos2, pos1\n",
    "\n",
    "        info = {}\n",
    "        info[\"id_object_mask\"] = {0: mask1, 1: mask2}\n",
    "        info[\"id_object_pos\"] = {0: pos1, 1: pos2}\n",
    "        info[\"obj_spec\"] = [((\"obj_0\", \"line\"), \"Attr\"), ((\"obj_1\", \"line\"), \"Attr\")]\n",
    "        info[\"node_id_map\"] = {\"obj_0\": 0, \"obj_1\": 1}\n",
    "        obj_idx = 2\n",
    "\n",
    "        # Add distractors with position and color:\n",
    "        max_n_distractors = max_n_objs - min_n_objs\n",
    "        if max_n_distractors > 0:\n",
    "            n_distractors = np.random.randint(1, max_n_distractors+1)\n",
    "            mask = (image != 0).float()\n",
    "            for k in range(n_distractors):\n",
    "                mask = (image != 0).float()\n",
    "                pos_new = get_pos_new_not_touching(mask, direction=np.random.choice([0,1]), min_size=min_size, max_size=max_size, canvas_size=canvas_size, min_distance=0)\n",
    "                image, obj_mask, obj_pos, _ = get_line(image, direction=None, pos=pos_new, color_avail=color_avail)\n",
    "                # Update info\n",
    "                info[\"id_object_mask\"][obj_idx] = obj_mask\n",
    "                info[\"id_object_pos\"][obj_idx] = obj_pos\n",
    "                info[\"obj_spec\"].append(((f\"obj_{obj_idx}\", \"line\"), \"Attr\"))\n",
    "                info[\"node_id_map\"][f\"obj_{obj_idx}\"] = obj_idx\n",
    "                obj_idx += 1\n",
    "        \n",
    "        # Get all relations\n",
    "        relations = []\n",
    "        n_objs = len(info[\"id_object_mask\"])\n",
    "        for id1 in range(n_objs):\n",
    "            for id2 in range(id1 + 1, n_objs):\n",
    "                rel = get_chosen_line_rel(info[\"id_object_pos\"][id1], info[\"id_object_pos\"][id2])\n",
    "                relations.append((id1, id2, rel))\n",
    "        info[\"relations\"] = relations\n",
    "\n",
    "        data.append(\n",
    "            (to_one_hot(image)[0],\n",
    "             (mask1, mask2),\n",
    "             relation,\n",
    "             Dictionary(info),\n",
    "            )\n",
    "        )\n",
    "        if len(data) >= n_examples:\n",
    "            break\n",
    "        if isplot:\n",
    "            visualize_matrices(image)\n",
    "            p.print(relation)\n",
    "            plot_matrices([mask1.squeeze(), mask2.squeeze()])\n",
    "            print('\\n\\n')\n",
    "    return data\n",
    "\n",
    "def get_line(image, direction=None, pos=None, min_size=2, max_size=10, color_avail=[1,2]):\n",
    "    \"\"\"\n",
    "    Direction: 0: horizontal; 1: vertical.\n",
    "    \"\"\"\n",
    "    canvas_size = image.shape[-1]\n",
    "    mask = torch.zeros(image.shape)\n",
    "    if pos is None:\n",
    "        assert direction is not None\n",
    "        for i in range(10):\n",
    "            if direction == -1:\n",
    "                direction_core = np.random.choice([0,1])\n",
    "            else:\n",
    "                direction_core = direction\n",
    "            if direction_core == 0:\n",
    "                # horizontal:\n",
    "                h = 1\n",
    "                w = np.random.randint(min_size, max_size+1)\n",
    "            elif direction_core == 1:\n",
    "                h = np.random.randint(min_size, max_size+1)\n",
    "                w = 1\n",
    "            row_start = 0\n",
    "            row_end = canvas_size - h\n",
    "            if row_end <= row_start:\n",
    "                continue\n",
    "            row = np.random.randint(row_start, row_end+1)\n",
    "\n",
    "            col_start = 0\n",
    "            col_end = canvas_size - w\n",
    "            if col_end <= col_start:\n",
    "                continue\n",
    "            col = np.random.randint(col_start, col_end+1)\n",
    "            pos = (row, col, h, w)\n",
    "    else:\n",
    "        row, col, h, w = pos\n",
    "\n",
    "    color = np.random.choice(color_avail)\n",
    "    image[..., row: row+h, col: col+w] = color\n",
    "    mask[..., row: row+h, col: col+w] = 1\n",
    "    return image, mask, pos, direction\n",
    "\n",
    "\n",
    "def get_pos_new_mid(pos, direction, min_size, canvas_size):\n",
    "    pos_mid = (pos[0] + pos[2]//2, pos[1] + pos[3]//2)\n",
    "    pos_new = None\n",
    "    for k in range(10):\n",
    "        orientation = np.random.choice([0,1])\n",
    "        if direction == 1:\n",
    "            # second line is horizontal:\n",
    "            if orientation == 0:\n",
    "                # second line is on the right:\n",
    "                if canvas_size-pos_mid[1] <= min_size:\n",
    "                    continue\n",
    "                pos_new = (pos_mid[0], pos_mid[1]+1, 1, np.random.randint(min_size, canvas_size-pos_mid[1]+1))\n",
    "            elif orientation == 1:\n",
    "                # second line is on the left:\n",
    "                if pos[1] <= min_size:\n",
    "                    continue\n",
    "                w_mid = np.random.randint(min_size, pos[1]+1)\n",
    "                pos_new = (pos_mid[0], pos_mid[1] - w_mid, 1, w_mid)\n",
    "        elif direction == 0:\n",
    "            # second line is vertical:\n",
    "            if orientation == 0:\n",
    "                # second line is on the bottom:\n",
    "                if canvas_size-pos_mid[0] <= min_size:\n",
    "                    continue\n",
    "                pos_new = (pos_mid[0]+1, pos_mid[1], np.random.randint(min_size, canvas_size-pos_mid[0]+1), 1)\n",
    "            elif orientation == 1:\n",
    "                # second line is on the top:\n",
    "                if pos[0] <= min_size:\n",
    "                    continue\n",
    "                h_mid = np.random.randint(min_size, pos[0]+1)\n",
    "                pos_new = (pos_mid[0] - h_mid, pos_mid[1], h_mid, 1)\n",
    "        if pos_new is not None:\n",
    "            break\n",
    "    return pos_new\n",
    "\n",
    "\n",
    "def get_pos_new_edge(pos, direction, min_size, canvas_size):\n",
    "    pos_new = None\n",
    "    for k in range(10):\n",
    "        orientation = np.random.choice([0,1])\n",
    "        edge_side = np.random.choice([0,1])\n",
    "        if direction == 1:\n",
    "            # second line is horizontal:\n",
    "            if orientation == 0:\n",
    "                # second line is on the right:\n",
    "                if canvas_size-pos[1] <= min_size:\n",
    "                    continue\n",
    "                if edge_side == 0:\n",
    "                    pos_new = (pos[0], pos[1]+1, 1, np.random.randint(min_size, canvas_size-pos[1]+1))\n",
    "                elif edge_side == 1:\n",
    "                    pos_new = (pos[0]+pos[2]-1, pos[1]+1, 1, np.random.randint(min_size, canvas_size-pos[1]+1))\n",
    "            elif orientation == 1:\n",
    "                # second line is on the left:\n",
    "                if pos[1] <= min_size:\n",
    "                    continue\n",
    "                w_mid = np.random.randint(min_size, pos[1]+1)\n",
    "                if edge_side == 0:\n",
    "                    pos_new = (pos[0], pos[1] - w_mid, 1, w_mid)\n",
    "                elif edge_side == 1:\n",
    "                    pos_new = (pos[0]+pos[2]-1, pos[1] - w_mid, 1, w_mid)\n",
    "        elif direction == 0:\n",
    "            # second line is vertical:\n",
    "            if orientation == 0:\n",
    "                # second line is on the bottom:\n",
    "                if canvas_size-pos[0] <= min_size:\n",
    "                    continue\n",
    "                if edge_side == 0:\n",
    "                    pos_new = (pos[0]+1, pos[1], np.random.randint(min_size, canvas_size-pos[0]+1), 1)\n",
    "                elif edge_side == 1:\n",
    "                    pos_new = (pos[0]+1, pos[1]+pos[3]-1, np.random.randint(min_size, canvas_size-pos[0]+1), 1)\n",
    "            elif orientation == 1:\n",
    "                # second line is on the top:\n",
    "                if pos[0] <= min_size:\n",
    "                    continue\n",
    "                h_mid = np.random.randint(min_size, pos[0]+1)\n",
    "                if edge_side == 0:\n",
    "                    pos_new = (pos[0] - h_mid, pos[1], h_mid, 1)\n",
    "                elif edge_side == 1:\n",
    "                    pos_new = (pos[0] - h_mid, pos[1]+pos[3]-1, h_mid, 1)\n",
    "        if pos_new is not None:\n",
    "            break\n",
    "    return pos_new\n",
    "\n",
    "\n",
    "def get_pos_new_not_touching(mask1, direction, min_size, max_size, canvas_size, pos=None, min_distance=0):\n",
    "    mask = deepcopy(mask1)\n",
    "    if min_distance > 0:\n",
    "        mask[...,max(0,pos[0]-min_distance):pos[0]+pos[2]+min_distance, max(0, pos[1]-min_distance):pos[1]+pos[3]+min_distance] = 1\n",
    "    for k in range(30):\n",
    "        if direction == 0:\n",
    "            # horizontal:\n",
    "            h = 1\n",
    "            w = np.random.randint(min_size, max_size+1)\n",
    "        elif direction == 1:\n",
    "            h = np.random.randint(min_size, max_size+1)\n",
    "            w = 1\n",
    "        row_start = 0\n",
    "        row_end = canvas_size - h\n",
    "        if row_end <= row_start:\n",
    "            continue\n",
    "        row = np.random.randint(row_start, row_end)\n",
    "\n",
    "        col_start = 0\n",
    "        col_end = canvas_size - w\n",
    "        if col_end <= col_start:\n",
    "            continue\n",
    "        col = np.random.randint(col_start, col_end)\n",
    "\n",
    "        mask2 = torch.zeros(mask1.shape)\n",
    "        mask2[..., row: row+h, col: col+w] = 1\n",
    "        if (mask + mask2 == 2).any():\n",
    "            continue\n",
    "        else:\n",
    "            break\n",
    "    pos_new = (row, col, h, w)\n",
    "    return pos_new\n",
    "\n",
    "\n",
    "\n",
    "def mask_to_obj(data):\n",
    "    \"\"\"Transform the data with format 'image+mask' to format 'image+obj'.\"\"\"\n",
    "    def get_obj_from_mask(img, mask_ele):\n",
    "        \"\"\"\n",
    "        image: [C, H, W], where the first dimension in C is for 0.\n",
    "        mask:  [1, H, W]\n",
    "        \"\"\"\n",
    "        assert len(img.shape) == len(mask_ele.shape) == 3\n",
    "        obj_ele = torch.cat([img[:1] * (1 - mask_ele), img[1:] * mask_ele], 0)\n",
    "        return obj_ele\n",
    "    if data[0][1] is not None:\n",
    "        data_new = []\n",
    "        for data_item in data:\n",
    "            image, mask, c_repr, info = data_item\n",
    "            if isinstance(image, tuple):\n",
    "                assert len(image) == len(mask) == 2\n",
    "                obj = (get_obj_from_mask(image[0], mask[0]), get_obj_from_mask(image[1], mask[1]))\n",
    "            else:\n",
    "                # Concept:\n",
    "                assert len(mask) == 1\n",
    "                obj = (get_obj_from_mask(image, mask[0]),)\n",
    "            data_new.append((image, obj, c_repr, info))\n",
    "        return data_new\n",
    "    else:\n",
    "        return data\n",
    "\n",
    "\n",
    "def get_simple_1D_dataset(\n",
    "    n_tasks=1000,\n",
    "    noise_std=0.1,\n",
    "    image_size=(64,),\n",
    "    in_channels=2,\n",
    "    shape_types=[\"rectangle\", \"triangle\"],\n",
    "    max_n_shapes=4,\n",
    "    max_shape_height=2,\n",
    "    max_shape_width=5,\n",
    "    is_expand_2d=True,\n",
    "    isplot=False,\n",
    "):\n",
    "    \"\"\"Generate a dataset on a 1D image, where there can be multiple different shapes and colors.\"\"\"\n",
    "    import matplotlib.pylab as plt\n",
    "    def is_no_collision(shape_info_example, min_distance):\n",
    "        locs = np.array([ele[\"loc\"] for ele in shape_info_example])\n",
    "        distance_matrix = np.abs(locs[None] - locs[:,None])\n",
    "        rows, cols = np.triu_indices(len(distance_matrix), 1)\n",
    "        all_distance = distance_matrix[rows, cols]\n",
    "        return (all_distance > min_distance).all()\n",
    "\n",
    "    def render_shape(shape_info_example, image_size):\n",
    "        x = torch.zeros(in_channels, *image_size)\n",
    "        for shape_info in shape_info_example:\n",
    "            x_ele = torch.zeros(in_channels, *image_size)\n",
    "            if shape_info[\"shape_type\"] == \"rectangle\":\n",
    "                x_ele[:, shape_info[\"loc\"]-shape_info[\"width\"]: shape_info[\"loc\"]+shape_info[\"width\"]+1] = torch.FloatTensor(shape_info[\"color\"][:, None]) * shape_info[\"height\"]\n",
    "            elif shape_info[\"shape_type\"] == \"triangle\":\n",
    "                x_ele[:, shape_info[\"loc\"]-shape_info[\"width\"]: shape_info[\"loc\"]+shape_info[\"width\"]+1] = torch.cat(\n",
    "                    [torch.linspace(0, shape_info[\"height\"], shape_info[\"width\"]+1),\n",
    "                     torch.linspace(shape_info[\"height\"], 0, shape_info[\"width\"]+1)[1:]])[None] * shape_info[\"color\"][:,None]\n",
    "            else:\n",
    "                pass\n",
    "            x = x + x_ele\n",
    "        return x\n",
    "\n",
    "\n",
    "    # First, sample number of shapes:\n",
    "    n_shapes_list = np.random.randint(1, max_n_shapes+1, size=n_tasks*3)\n",
    "    min_shape_height = max_shape_height * 0.25\n",
    "    min_shape_width = np.maximum(2, max_shape_width * 0.25)\n",
    "    max_shape_width = np.round(max_shape_width)\n",
    "\n",
    "    task_dict = {}\n",
    "    jj = 0\n",
    "    for i in range(n_tasks * 3):\n",
    "        shape_info_example = []\n",
    "\n",
    "        n_shapes = n_shapes_list[i]\n",
    "        for k in range(n_shapes):\n",
    "            shape_info = {}\n",
    "            # Second, sample the shape types:\n",
    "            shape_info[\"shape_type\"] = np.random.choice(shape_types)\n",
    "\n",
    "            # Third, sample the features of each shape except position:\n",
    "            shape_info[\"height\"] = np.random.rand() *(max_shape_height - min_shape_height) + min_shape_height\n",
    "            shape_info[\"width\"] = np.random.randint(min_shape_width, max_shape_width+1)\n",
    "            shape_info[\"color\"] = np.random.rand(2) * 0.5 + 1\n",
    "\n",
    "            # Fourth, sample the position of each shape:\n",
    "            shape_info[\"loc\"] = np.random.randint(np.round(max_shape_width), np.round(image_size[0] - max_shape_width))\n",
    "\n",
    "            # Append to shape_info_example:\n",
    "            shape_info_example.append(shape_info)\n",
    "\n",
    "        if is_no_collision(shape_info_example, min_distance=max_shape_width*2):\n",
    "            x_example = render_shape(shape_info_example, image_size)\n",
    "            if noise_std > 0:\n",
    "                x_example = x_example + torch.randn(x_example.shape) * noise_std\n",
    "            if is_expand_2d:\n",
    "                task_dict[str(jj)] = [(x_example[...,None], None, None, {\"shape_info\": shape_info_example})]\n",
    "            jj += 1\n",
    "        else:\n",
    "            continue\n",
    "        if len(task_dict) >= n_tasks:\n",
    "            break\n",
    "    if len(task_dict) < n_tasks:\n",
    "        p.print(\"Generated {} tasks, less than the required {} tasks\".format(len(task_dict), n_tasks))\n",
    "    if isplot:\n",
    "        for i in range(10):\n",
    "            plt.plot(task_dict[str(i)][0][0].T)\n",
    "            plt.show()\n",
    "    return task_dict\n",
    "\n",
    "\n",
    "def get_dataset(args, n_examples=None, isplot=False, is_load=False, is_rewrite=False, verbose=True, is_save_inter=False):\n",
    "    \"\"\"Generate the dataset according to specifications.\n",
    "\n",
    "    Args:\n",
    "        is_load: if True, will load the previously-saved file if there is one, and write to file if there isn't.\n",
    "        is_rewrite: if False, will obey the behavior of is_load. If True, will re-generate the data, and save to file.\n",
    "    \"\"\"\n",
    "    if n_examples is None:\n",
    "        n_examples = args.n_examples\n",
    "    args.is_mask = False\n",
    "    seed = args.seed\n",
    "    if args.dataset.startswith(\"y-\") and not args.use_seed_2d:\n",
    "        seed = args.seed_3d\n",
    "    if ((args.dataset.startswith(\"c-\") or args.dataset.startswith(\"y-\")) and args.max_n_distractors != 2) or args.dataset.startswith(\"h-\"):\n",
    "        dataset_2d_param = \"ex_{}_seed_{}_cav_{}_rain_{}_color_{}_distr_{}\".format(\n",
    "            n_examples, seed,\n",
    "            args.canvas_size if hasattr(args, \"canvas_size\") else None,\n",
    "            args.rainbow_prob if hasattr(args, \"rainbow_prob\") else None,\n",
    "            args.color_avail if hasattr(args, \"color_avail\") else None,\n",
    "            args.max_n_distractors)\n",
    "    else:\n",
    "        dataset_2d_param = \"ex_{}_seed_{}_cav_{}_rain_{}_color_{}\".format(\n",
    "            n_examples, seed,\n",
    "            args.canvas_size if hasattr(args, \"canvas_size\") else None,\n",
    "            args.rainbow_prob if hasattr(args, \"rainbow_prob\") else None,\n",
    "            args.color_avail if hasattr(args, \"color_avail\") else None)\n",
    "    if args.min_n_distractors != 0:\n",
    "        dataset_2d_param += \"_mindistr_{}\".format(args.min_n_distractors)\n",
    "    if args.allow_connect is False:\n",
    "        dataset_2d_param += \"_connect_{}\".format(args.allow_connect)\n",
    "\n",
    "    dataset_filename = REA_PATH + \"/data/{}-{}.p\".format(args.dataset if not args.dataset.startswith(\"y-\") else \"c\" + args.dataset[1:], dataset_2d_param)\n",
    "    dataset = None\n",
    "    if is_load and os.path.isfile(dataset_filename) and not is_rewrite:\n",
    "        dataset, args_load = pload(dataset_filename)\n",
    "        if (len(dataset[0]) > 3 and \"id_object_mask\" in dataset[0][3]) or not args.dataset.startswith(\"c-\"):\n",
    "            if verbose and not args.dataset.startswith(\"y-\"):\n",
    "                p.print(\"Dataset loaded from {}.\".format(dataset_filename))\n",
    "            args.n_classes = args_load.n_classes\n",
    "            args.image_size = args_load.image_size\n",
    "            args.in_channels = args_load.in_channels\n",
    "            args.is_mask = args_load.is_mask\n",
    "            if args.dataset.startswith(\"c-\"):\n",
    "                args.concept_collection = dataset.concept_collection\n",
    "            if args.dataset.startswith(\"u-\"):\n",
    "                args.concept_collection = args_load.concept_collection\n",
    "            if not args.dataset.startswith(\"y-\"):\n",
    "                return dataset, args\n",
    "        else:\n",
    "            if verbose:\n",
    "                p.print(\"Dataset's info is an old version, regenerate.\")\n",
    "    if is_load and not is_rewrite and not os.path.isfile(dataset_filename):\n",
    "        p.print(\"Do not find dataset {}.\".format(dataset_filename))\n",
    "    set_seed(seed)\n",
    "    if args.dataset in [\"cifar10\"]:\n",
    "        \"\"\"Standard datasets\"\"\"\n",
    "        dataset = datasets.CIFAR10(get_root_dir() + '/datasets/', download=True, transform=transforms.ToTensor())\n",
    "        args.in_channels = dataset[0][0].shape[0]\n",
    "        args.image_size = dataset[0][0].shape[-2:]\n",
    "        args.concept_collection = None\n",
    "        args.n_classes = len(dataset.classes)\n",
    "    elif args.dataset.startswith(\"c-\"):\n",
    "        \"\"\"BabyARC single concept/relation/operator dataset with ground-truth masks.\"\"\"\n",
    "        mode = args.dataset.split(\"-\")[1]\n",
    "        if \"^\" in mode:\n",
    "            mode, focus_type = mode.split(\"^\")\n",
    "        else:\n",
    "            focus_type = None\n",
    "        dataset = ConceptDataset(\n",
    "            mode=mode,\n",
    "            n_examples=n_examples,\n",
    "            canvas_size=args.canvas_size,\n",
    "            rainbow_prob=args.rainbow_prob,\n",
    "            w_type=args.w_type,\n",
    "            max_n_distractors=args.max_n_distractors,\n",
    "            min_n_distractors=args.min_n_distractors,\n",
    "            color_avail=args.color_avail,\n",
    "            allow_connect=args.allow_connect,\n",
    "            parsing_check=args.parsing_check if hasattr(args, \"parsing_check\") else False,\n",
    "            focus_type=focus_type,\n",
    "            save_filename=dataset_filename[:-2] + \"_inter.p\" if is_save_inter else None,\n",
    "        )\n",
    "        args.in_channels = 10\n",
    "        args.n_classes = 1\n",
    "        if isinstance(dataset[0][0], tuple):\n",
    "            args.image_size = dataset[0][0][0].shape[-2:]\n",
    "        else:\n",
    "            args.image_size = dataset[0][0].shape[-2:]\n",
    "        args.concept_collection = dataset.concept_collection\n",
    "        args.is_mask = True\n",
    "    elif args.dataset.startswith(\"u-\"):\n",
    "        \"\"\"\n",
    "        args.dataset = \"u-concept-Red+Green+Blue+Cube+Cylinder+Large+Small\"\n",
    "        \"\"\"\n",
    "        mode = args.dataset[2:]\n",
    "        if mode.startswith(\"concept\"):\n",
    "            if set(mode.split(\"concept-\")[1].split(\"-\")) == set(\"Red+Green+Blue+Cube+Cylinder+Large+Small\".split(\"-\")):\n",
    "                data_list_1 = pload(\"/dfs/user/tailin/.results/CLEVR_relation/clevr-concept-relation-saved/data_list_canvas_64_ex_25000_1.p\")\n",
    "                data_list_2 = pload(\"/dfs/user/tailin/.results/CLEVR_relation/clevr-concept-relation-saved/data_list_canvas_64_ex_30000_2.p\")\n",
    "                data_list_3 = pload(\"/dfs/user/tailin/.results/CLEVR_relation/clevr-concept-relation-saved/data_list_canvas_64_ex_11000_3.p\")\n",
    "                data_list = data_list_1 + data_list_2 + data_list_3 \n",
    "            elif set(mode.split(\"concept-\")[1].split(\"-\")) == set(\"Red+Cube+Large\".split(\"-\")):\n",
    "                data_list_1 = pload(\"/dfs/user/tailin/.results/CLEVR_relation/clevr-concept-relation-saved/data_list_canvas_Red+Cube+Large_64_ex_20000_1.p\")\n",
    "                data_list_2 = pload(\"/dfs/user/tailin/.results/CLEVR_relation/clevr-concept-relation-saved/data_list_canvas_Red+Cube+Large_64_ex_20000_2.p\")\n",
    "                data_list_3 = pload(\"/dfs/user/tailin/.results/CLEVR_relation/clevr-concept-relation-saved/data_list_canvas_Red+Cube+Large_64_ex_4000_3.p\")\n",
    "                data_list = data_list_1 + data_list_2 + data_list_3 \n",
    "            else:\n",
    "                raise\n",
    "            # data_list = pload(\"/dfs/user/tailin/.results/CLEVR_relation/clevr-concept-relation-saved/data_list_canvas_64_ex_440.p\")\n",
    "        elif mode.startswith(\"relation\"):\n",
    "            if set(mode.split(\"relation-\")[1].split(\"-\")) == set(\"SameColor+SameShape+SameSize\".split(\"-\")):\n",
    "                data_list_1 = pload(\"/dfs/user/tailin/.results/CLEVR_relation/clevr-concept-relation-saved/data_list_canvas_relation_64_ex_25000_1.p\")\n",
    "                data_list_2 = pload(\"/dfs/user/tailin/.results/CLEVR_relation/clevr-concept-relation-saved/data_list_canvas_relation_64_ex_30000_2.p\")\n",
    "                data_list_3 = pload(\"/dfs/user/tailin/.results/CLEVR_relation/clevr-concept-relation-saved/data_list_canvas_relation_64_ex_11000_3.p\")\n",
    "                data_list = data_list_1 + data_list_2 + data_list_3 \n",
    "            else:\n",
    "                raise\n",
    "            # data_list = pload(\"/dfs/user/tailin/.results/CLEVR_relation/clevr-concept-relation-saved/data_list_canvas_relation_64_ex_440.p\")\n",
    "        elif mode.startswith(\"graph\"):\n",
    "            if set(mode.split(\"graph-\")[1].split(\"-\")) == set(\"Graph1+Graph2+Graph3\".split(\"-\")):\n",
    "                data_list = pload(\"/dfs/user/tailin/.results/CLEVR_relation/clevr-concept-relation-saved/data_list_canvas_graph_Graph1+Graph2+Graph3_64_ex_200_4.p\") # Need to fix\n",
    "            else:\n",
    "                raise\n",
    "        else:\n",
    "            raise\n",
    "        data_list = data_list[:n_examples]\n",
    "        dataset = ConceptClevrDataset(\n",
    "            data=data_list,\n",
    "        )\n",
    "        args.concept_collection = mode.split(\"-\")[1].split(\"+\")\n",
    "        args.in_channels = 3\n",
    "        args.image_size = dataset[0][0].shape[-2:]\n",
    "        args.n_classes = 1\n",
    "        args.is_mask = True\n",
    "    elif args.dataset.startswith(\"y-\"):\n",
    "        dataset_3d_param = \"seed_{}_proc_{}_sz_{}_{}_color_{}_thxy_{}_{}_thz_{}_{}\".format(\n",
    "                            args.seed_3d, args.num_processes_3d, args.image_size_3d[0], args.image_size_3d[1],\n",
    "                            args.color_map_3d[:4], args.add_thick_surf[0],  args.add_thick_surf[1],\n",
    "                            args.add_thick_depth[0], args.add_thick_depth[1])\n",
    "        dataset_filename = REA_PATH + \"/data/{}-3d_{}_2d_{}.p\".format(args.dataset, dataset_3d_param, dataset_2d_param)\n",
    "        if os.path.isfile(dataset_filename) and not is_rewrite:\n",
    "            dataset, args_load = pload(dataset_filename)\n",
    "            if verbose:\n",
    "                p.print(\"Dataset loaded from {}.\".format(dataset_filename))\n",
    "            args.n_classes = args_load.n_classes\n",
    "            args.image_size = args_load.image_size\n",
    "            args.in_channels = args_load.in_channels\n",
    "            args.is_mask = args_load.is_mask\n",
    "            args.concept_collection = args_load.concept_collection\n",
    "            return dataset, args\n",
    "        if dataset is None:\n",
    "            args_2d = deepcopy(args)\n",
    "            args_2d.seed = seed\n",
    "            args_2d.dataset = \"c\" + args.dataset[1:]\n",
    "            dataset, _ = get_dataset(args_2d, n_examples=n_examples, isplot=isplot, is_load=is_load, \n",
    "                                     is_rewrite=is_rewrite, verbose=verbose)\n",
    "        args.concept_collection = dataset.concept_collection\n",
    "        dataset = ConceptDataset3D(data=convert_babyarc(dataset, args))\n",
    "        args.in_channels = 3\n",
    "        args.image_size = args.image_size_3d\n",
    "        args.n_classes = 1\n",
    "        args.is_mask = True\n",
    "    elif args.dataset.startswith(\"pc-\") or args.dataset.startswith(\"pg-\"):\n",
    "        dataset = generate_fewshot_dataset(args, n_shot=1, n_queries_per_class=args.n_queries_per_class)\n",
    "        args.in_channels = 10\n",
    "        args.image_size = dataset[0][0][0][0].shape[-2:]\n",
    "        args.n_classes = 1\n",
    "        args.is_mask = True\n",
    "    elif args.dataset.startswith(\"yc-\"):\n",
    "        dataset = generate_fewshot_dataset(args, n_shot=1, n_queries_per_class=args.n_queries_per_class)\n",
    "        args.in_channels = 3\n",
    "        args.image_size = dataset[0][1][0][0].shape[-2:]\n",
    "        args.n_classes = 1\n",
    "        args.is_mask = True\n",
    "    elif args.dataset.startswith(\"h-\"):\n",
    "        \"\"\"\n",
    "        BabyARC composition dataset.\n",
    "\n",
    "        args.dataset:\n",
    "            h-c^(1,2):Eshape+Ashape-d^1:RandShape :\n",
    "                only concept with concept_avail=[\"Eshape\",\"Ashape\"], n_concepts_range=(1,2), max_n_distractors=1, additional_shape=\"Randshape\"\n",
    "            h-r^2a+2ai+3a+3b+3bi:SameShape+SameColor(Line+Rect+RectSolid+Lshape)-d^1:Randshape :\n",
    "                relation_structure=\"2a+2ai+3a+3b+3bi\",\n",
    "                relation_avail=[\"SameShape\",\"SameColor\"],\n",
    "                concept_avail=[\"Line\",\"Rect\",\"RectSolid\",\"Lshape\"]\n",
    "                additional_concepts=[\"Randshape\"]\n",
    "                max_n_distractors=1\n",
    "        \"\"\"\n",
    "        mode_split = args.dataset[2:].split(\"-\")\n",
    "        settings = {}\n",
    "        for content in mode_split:\n",
    "            if content.startswith(\"c^\"):\n",
    "                settings[\"n_concepts_range\"] = eval(content.split(\":\")[0][2:])\n",
    "                settings[\"concept_avail\"] = content.split(\":\")[1].split(\"+\")\n",
    "            elif content.startswith(\"r^\"):\n",
    "                settings[\"relation_structure\"] = content.split(\":\")[0][2:]\n",
    "                settings[\"relation_avail\"] = content.split(\":\")[1].split(\"(\")[0].split(\"+\")\n",
    "                settings[\"concept_avail\"] = content.split(\":\")[1][:-1].split(\"(\")[1].split(\"+\")\n",
    "            elif content.startswith(\"d^\"):\n",
    "                settings[\"max_n_distractors\"] = eval(content.split(\":\")[0][2:])\n",
    "                settings[\"additional_concepts\"] = content.split(\":\")[1].split(\"+\") if len(content.split(\":\")) > 1 else None\n",
    "            else:\n",
    "                raise\n",
    "        if hasattr(args, \"max_n_distractors\"):\n",
    "            settings[\"max_n_distractors\"] = args.max_n_distractors\n",
    "        if hasattr(args, \"min_n_distractors\"):\n",
    "            settings[\"min_n_distractors\"] = args.min_n_distractors\n",
    "        dataset = ConceptCompositionDataset(\n",
    "            canvas_size=args.canvas_size,\n",
    "            n_examples=n_examples,\n",
    "            rainbow_prob=args.rainbow_prob,\n",
    "            color_avail=args.color_avail,\n",
    "            concept_avail=settings[\"concept_avail\"] if \"concept_avail\" in settings else None,\n",
    "            relation_avail=settings[\"relation_avail\"] if \"relation_avail\" in settings else None,\n",
    "            relation_structure=settings[\"relation_structure\"] if \"relation_structure\" in settings else \"None\",\n",
    "            additional_concepts=settings[\"additional_concepts\"] if \"additional_concepts\" in settings else None,\n",
    "            n_concepts_range=settings[\"n_concepts_range\"] if \"n_concepts_range\" in settings else 2,\n",
    "            min_n_distractors=settings[\"min_n_distractors\"] if \"min_n_distractors\" in settings else 0,\n",
    "            max_n_distractors=settings[\"max_n_distractors\"] if \"max_n_distractors\" in settings else 0,\n",
    "            n_examples_per_task=5 if \"c^\" in args.dataset else 6,\n",
    "        )\n",
    "        if isinstance(dataset[0][0][0], tuple):\n",
    "            args.image_size = dataset[0][0][0][0].shape[-2:]\n",
    "        else:\n",
    "            args.image_size = dataset[0][0][0].shape[-2:]\n",
    "        args.in_channels = 10\n",
    "        args.n_classes = 1\n",
    "        args.is_mask = True\n",
    "    elif args.dataset.startswith(\"arc\"):\n",
    "        \"\"\"ARC image dataset\"\"\"\n",
    "        dataset = ARCDataset(n_examples=n_examples, output_mode=\"energy\")\n",
    "        args.in_channels = 10\n",
    "        args.n_classes = 1\n",
    "        args.image_size = dataset[0][0].shape[-2:]\n",
    "    elif args.dataset.startswith(\"v-\"):\n",
    "        \"\"\"CLEVR dataset:\"\"\"\n",
    "        parts = args.dataset.split(\"-\")\n",
    "        split_name = parts[1]\n",
    "        start_idx = 0\n",
    "        if len(parts) > 2:\n",
    "            start_idx = int(parts[2])\n",
    "        if \"interpolate_mode\" not in args.__dict__:\n",
    "            interpolate_mode = \"bilinear\" if \"train\" in split_name else \"nearest\"\n",
    "        else:\n",
    "            interpolate_mode = args.interpolate_mode\n",
    "        # Process the images the same way as slot attention. This \n",
    "        # rescales RGB values to the passed in range\n",
    "        min_val, max_val = args.image_value_range.split(\",\") if \"image_value_range\" in args.__dict__ else args.selector_image_value_range.split(\",\")\n",
    "        min_val = float(min_val)\n",
    "        max_val = float(max_val)\n",
    "        rgb_std = 1.0 / (max_val - min_val)\n",
    "        rgb_mean = -(rgb_std * min_val)\n",
    "\n",
    "        if split_name.startswith(\"CLEVR6\") or split_name.startswith(\"CLEVR10\"):\n",
    "            processor = ClevrImagePreprocessor((128, 128), crop=(29, 221, 64, 256), rgb_mean=rgb_mean, rgb_std=rgb_std)\n",
    "            dataset = CLEVR(root=\"/dfs/user/tailin/.results/CLEVR_concept/CLEVR_with_masks\", \n",
    "                            split_name=split_name)\n",
    "            i = start_idx\n",
    "            n_tasks = n_examples if n_examples != None and n_examples > 0 else len(dataset)\n",
    "            end_idx = min(start_idx + n_tasks, len(dataset))\n",
    "            task_dict = {}\n",
    "            while i < end_idx:\n",
    "                img, info = dataset[i]\n",
    "                img_name = info['image_name']\n",
    "                obj_masks = info['mask']\n",
    "                # Make sure to process image the same as in slot attention\n",
    "                true_img = processor(img.unsqueeze(0), interpolate_mode=interpolate_mode).squeeze(0)\n",
    "                task_dict[img_name] = [(true_img, None, None, {'obj_masks': obj_masks})]\n",
    "                i += 1\n",
    "            args.image_size = true_img.shape[-2:]\n",
    "        elif split_name.startswith(\"CLEVRRelation\"):\n",
    "            resolution = (88, 128)\n",
    "            crop = (32, 208, 32, 288)\n",
    "            processor = ClevrImagePreprocessor(resolution, crop=crop, rgb_mean=rgb_mean, rgb_std=rgb_std)\n",
    "\n",
    "            \"split_name is 'CLEVRRelation:{dataset_split}' where dataset_split is either train, val, or test\"\n",
    "            dataset_split = split_name.split(\":\")[-1]\n",
    "            train_set, val_set, test_set = create_easy_dataset(output_type=\"mask-only\")\n",
    "            dataset_dct = {\"train\": train_set, \"val\": val_set, \"test\": test_set}\n",
    "            dataset = dataset_dct[dataset_split]\n",
    "            \n",
    "            i = start_idx\n",
    "            n_tasks = n_examples if n_examples != None and n_examples > 0 else len(dataset)\n",
    "            end_idx = min(start_idx + n_tasks, len(dataset))\n",
    "            task_dict = {}\n",
    "            while i < end_idx:\n",
    "                task = dataset[i]\n",
    "                task_dict[i] = []\n",
    "                # Iterate through the task example pairs.\n",
    "                for idx in range(5):\n",
    "                    # Take the RGB channels\n",
    "                    input_img = processor(task[\"inputs\"][idx][\"image\"][:3].unsqueeze(0)).squeeze(0)\n",
    "                    target_mask = task[\"outputs\"][idx].unsqueeze(0)\n",
    "                    target_mask = target_mask[..., crop[0]:crop[1], crop[2]:crop[3]] if crop else target_mask\n",
    "                    target_mask = F.interpolate(target_mask.float(), resolution, mode=\"nearest\").round().squeeze(0)\n",
    "                    task_dict[i].append(((input_img, target_mask), None, None, {'obj_masks': task[\"outputs_mask_only\"][idx]}))\n",
    "                # Test example:\n",
    "                test_img = processor(task[\"test_input\"][\"image\"][:3].unsqueeze(0)).squeeze(0)\n",
    "                test_mask = task[\"test_output\"].unsqueeze(0)\n",
    "                test_mask = test_mask[..., crop[0]:crop[1], crop[2]:crop[3]] if crop else test_mask\n",
    "                test_mask = F.interpolate(test_mask.float(), resolution, mode=\"nearest\").round().squeeze(0)\n",
    "                task_dict[i].append(((test_img, test_mask), None, None, {}))\n",
    "\n",
    "                i += 1\n",
    "            args.image_size = target_mask.shape[-2:]\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "        args.in_channels = 3\n",
    "        args.n_classes = 1\n",
    "        dataset = task_dict\n",
    "    elif args.dataset.startswith(\"m-\") or args.dataset.startswith(\"t-\"):\n",
    "        \"\"\"Multi-Dsprites and Tetrominoes datasets\"\"\"\n",
    "        \n",
    "        parts = args.dataset.split(\"-\")\n",
    "        split_name = parts[1]\n",
    "        if args.dataset.startswith(\"m-\"):\n",
    "            dataset = MultiDsprites(root=\"/dfs/user/tailin/.results/CLEVR_concept/multi_dsprites\", split_name=split_name)\n",
    "        else:\n",
    "            dataset = Tetrominoes(root=\"/dfs/user/tailin/.results/CLEVR_concept/tetrominoes\", split_name=split_name)\n",
    "        # Process the images the same way as slot attention. This \n",
    "        # rescales RGB values to the passed in range\n",
    "        min_val, max_val = args.image_value_range.split(\",\") if \"image_value_range\" in args.__dict__ else args.selector_image_value_range.split(\",\")\n",
    "        min_val = float(min_val)\n",
    "        max_val = float(max_val)\n",
    "        rgb_std = 1.0 / (max_val - min_val)\n",
    "        rgb_mean = -(rgb_std * min_val)\n",
    "        args.image_size = dataset[0][0].shape[-2:]\n",
    "        # Perform normalization, but not cropping / resizing\n",
    "        processor = ClevrImagePreprocessor(args.image_size, rgb_mean=rgb_mean, rgb_std=rgb_std)    \n",
    "\n",
    "        start_idx = 0\n",
    "        if len(parts) > 2:\n",
    "            start_idx = int(parts[2])\n",
    "        n_tasks = n_examples if n_examples != None and n_examples > 0 else len(dataset)\n",
    "        end_idx = min(start_idx + n_tasks, len(dataset))\n",
    "        task_dict = {} \n",
    "        i = start_idx\n",
    "        while i < end_idx:\n",
    "            img, info = dataset[i]\n",
    "            img_name = info['image_name']\n",
    "            obj_masks = info['mask']\n",
    "            # Make sure to process image the same as in slot attention\n",
    "            true_img = processor(img.unsqueeze(0)).squeeze(0)\n",
    "            task_dict[img_name] = [(true_img, None, None, {'obj_masks': obj_masks})]\n",
    "            i += 1\n",
    "        args.in_channels = 3\n",
    "        args.n_classes = 1\n",
    "        dataset = task_dict\n",
    "    elif args.dataset.startswith(\"s-\"):\n",
    "        \"\"\"Simple 1D dataset:\"\"\"\n",
    "        n_tasks = n_examples if n_examples is not None else 1000\n",
    "        max_n_shapes_str, max_n_shapes, width_str, width, noise_std_str, noise_std = args.dataset.split(\"-\")[1:]\n",
    "        assert max_n_shapes_str == \"nsh\"\n",
    "        assert width_str == \"width\"\n",
    "        assert noise_std_str == \"noise\"\n",
    "        dataset = get_simple_1D_dataset(\n",
    "            n_tasks=n_tasks,\n",
    "            noise_std=eval(noise_std),\n",
    "            image_size=(args.canvas_size,),\n",
    "            in_channels=2,\n",
    "            shape_types=[\"rectangle\", \"triangle\"],\n",
    "            max_n_shapes=eval(max_n_shapes),\n",
    "            max_shape_height=2,\n",
    "            max_shape_width=eval(width),\n",
    "            isplot=isplot,\n",
    "        )\n",
    "        args.in_channels = 2\n",
    "        args.n_classes = 1\n",
    "        args.image_size = dataset[\"0\"][0][0].shape[2:]\n",
    "    else:\n",
    "        raise Exception(\"dataset '{}' is not recognized!\".format(args.dataset))\n",
    "    if is_rewrite or is_load:\n",
    "        pdump((dataset, args), dataset_filename)\n",
    "        if verbose:\n",
    "            p.print(\"dataset saved at {}.\".format(dataset_filename))\n",
    "    return dataset, args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_discriminative_task(\n",
    "    concepts,\n",
    "    n_examples_per_concept,\n",
    "    n_examples_per_task,\n",
    "    n_tasks,\n",
    "    canvas_size=8,\n",
    "    isplot=False,\n",
    "):\n",
    "    data_dict = {}\n",
    "    for key in concepts:\n",
    "        args = init_args({\n",
    "            \"dataset\": \"c-{}\".format(key),\n",
    "            \"seed\": 1,\n",
    "            \"n_examples\": n_examples_per_concept,\n",
    "            \"canvas_size\": canvas_size,\n",
    "            \"rainbow_prob\": 0.,\n",
    "            \"color_avail\": \"-1\",\n",
    "            \"w_type\": \"mask\",\n",
    "            \"max_n_distractors\": 0,\n",
    "        })\n",
    "        dataset, args = get_dataset(args, is_load=True)\n",
    "        data_dict[key] = dataset.data\n",
    "\n",
    "    # Generate task:\n",
    "    task_list = []\n",
    "    for i in range(n_tasks):\n",
    "        # randomly sample a concept:\n",
    "        concept_chosen = np.random.choice(concepts)\n",
    "        example_ids = np.random.choice(len(data_dict[concept_chosen]), replace=False, size=n_examples_per_task+1)\n",
    "        concept_examples = torch.stack([data_dict[concept_chosen][i][0] for i in example_ids[:-1]])\n",
    "\n",
    "        concept_others = deepcopy(concepts)\n",
    "        concept_others.remove(concept_chosen)\n",
    "        other_ids_all = list(itertools.product(concept_others, range(n_examples_per_concept)))\n",
    "        other_ids = np.random.choice(len(other_ids_all), replace=False, size=n_examples_per_task+1)\n",
    "        other_keys = [other_ids_all[id] for id in other_ids]\n",
    "        other_examples = torch.stack([data_dict[other_key][id][0] for other_key, id in other_keys[:-1]])\n",
    "\n",
    "        is_concept = np.random.choice([True, False])\n",
    "        if is_concept:\n",
    "            test_example = data_dict[concept_chosen][example_ids[-1]][0][None]\n",
    "        else:\n",
    "            test_example = data_dict[other_keys[-1][0]][other_keys[-1][1]][0][None]\n",
    "        task_list.append(((concept_examples, other_examples), (test_example, is_concept)))\n",
    "\n",
    "    if isplot:\n",
    "        p.print(\"concept: {}\".format(concept_chosen))\n",
    "        visualize_matrices(concept_examples.argmax(1), use_color_dict=True)\n",
    "        visualize_matrices(other_examples.argmax(1), use_color_dict=True)\n",
    "        p.print(\"is_concept: {}\".format(is_concept))\n",
    "        visualize_matrices(test_example.argmax(1), use_color_dict=True)\n",
    "    return task_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.5 Datasets for paper:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "is_load_example_dataset = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.5.0 CLEVR-Concept:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### CLEVR-Concept, training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"u-concept-Red+Cube+Large\",\n",
    "        \"seed\": 3,\n",
    "        \"n_examples\": 66000,\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"canvas_size\": 64,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"allow_connect\": True,\n",
    "        \"rainbow_prob\": 0.0,\n",
    "    })\n",
    "    concept_dataset, _ = get_dataset(args, is_load=True)\n",
    "    # concept_dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# concept_dataset.draw(range(40,80))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Red\n",
    "# data_c1 = concept_dataset[5]\n",
    "# visualize_matrices([data_c1[0]], use_color_dict=False, filename=\"clevr_img/c1_img.pdf\")\n",
    "# plot_matrices(data_c1[1][0], images_per_row=6, no_xlabel=True, filename=\"clevr_img/c1_mask.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # Cube\n",
    "# data_c2 = concept_dataset[54]\n",
    "# visualize_matrices([data_c2[0]], use_color_dict=False, filename=\"clevr_img/c2_img.pdf\")\n",
    "# plot_matrices(data_c2[1][0], images_per_row=6, no_xlabel=True, filename=\"clevr_img/c2_mask.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # Large\n",
    "# data_c3 = concept_dataset[59]\n",
    "# visualize_matrices([data_c3[0]], use_color_dict=False, filename=\"clevr_img/c3_img.pdf\")\n",
    "# plot_matrices(data_c3[1][0], images_per_row=6, no_xlabel=True, filename=\"clevr_img/c3_mask.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# relation_dataset.draw(range(40))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# relation_dataset.draw(range(40, 80))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # SameColor\n",
    "# data_r1 = relation_dataset[30]\n",
    "# visualize_matrices([data_r1[0]], use_color_dict=False, filename=\"clevr_img/r1_img.pdf\")\n",
    "# plot_matrices(data_r1[1][0], images_per_row=6, no_xlabel=True, filename=\"clevr_img/r1_mask_1.pdf\")\n",
    "# plot_matrices(data_r1[1][1], images_per_row=6, no_xlabel=True, filename=\"clevr_img/r1_mask_2.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # SameShape:\n",
    "# data_r2 = relation_dataset[24]\n",
    "# visualize_matrices([data_r2[0]], use_color_dict=False, filename=\"clevr_img/r2_img.pdf\")\n",
    "# plot_matrices(data_r2[1][0], images_per_row=6, no_xlabel=True, filename=\"clevr_img/r2_mask_1.pdf\")\n",
    "# plot_matrices(data_r2[1][1], images_per_row=6, no_xlabel=True, filename=\"clevr_img/r2_mask_2.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # SameSize:\n",
    "# data_r3 = relation_dataset[49]\n",
    "# visualize_matrices([data_r3[0]], use_color_dict=False, filename=\"clevr_img/r3_img.pdf\")\n",
    "# plot_matrices(data_r3[1][0], images_per_row=6, no_xlabel=True, filename=\"clevr_img/r3_mask_1.pdf\")\n",
    "# plot_matrices(data_r3[1][1], images_per_row=6, no_xlabel=True, filename=\"clevr_img/r3_mask_2.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Graph2:\n",
    "# data_g1 = graph_dataset[33]\n",
    "# visualize_matrices([data_g1[0]], use_color_dict=False, filename=\"clevr_img/g1_img.pdf\")\n",
    "# # plot_matrices(data_g2[1][0], images_per_row=6, no_xlabel=True, filename=\"clevr_img/g2_mask_1.pdf\")\n",
    "# # plot_matrices(data_g2[1][1], images_per_row=6, no_xlabel=True, filename=\"clevr_img/g2_mask_2.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Graph2:\n",
    "# data_g2 = graph_dataset[12]\n",
    "# visualize_matrices([data_g2[0]], use_color_dict=False, filename=\"clevr_img/g2_img.pdf\")\n",
    "# # plot_matrices(data_g2[1][0], images_per_row=6, no_xlabel=True, filename=\"clevr_img/g2_mask_1.pdf\")\n",
    "# # plot_matrices(data_g2[1][1], images_per_row=6, no_xlabel=True, filename=\"clevr_img/g2_mask_2.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Graph3:\n",
    "# data_g3 = graph_dataset[77]\n",
    "# visualize_matrices([data_g3[0]], use_color_dict=False, filename=\"clevr_img/g3_img.pdf\")\n",
    "# # plot_matrices(data_g2[1][0], images_per_row=6, no_xlabel=True, filename=\"clevr_img/g2_mask_1.pdf\")\n",
    "# # plot_matrices(data_g2[1][1], images_per_row=6, no_xlabel=True, filename=\"clevr_img/g2_mask_2.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### CLEVR-Relation, training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"u-relation-SameColor+SameShape+SameSize\",\n",
    "        \"seed\": 3,\n",
    "        \"n_examples\": 66000,\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"canvas_size\": 64,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"allow_connect\": True,\n",
    "        \"rainbow_prob\": 0.0,\n",
    "    })\n",
    "    relation_dataset, _ = get_dataset(args, is_load=True)\n",
    "    # relation_dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### CLEVR-graph, inference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"u-graph-Graph1+Graph2+Graph3\",\n",
    "        \"seed\": 2,\n",
    "        \"n_examples\": 400,\n",
    "        \"canvas_size\": 32,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"max_n_distractors\": 0,\n",
    "        \"allow_connect\": True,\n",
    "        \"parsing_check\": False,\n",
    "    })\n",
    "    graph_dataset, _ = get_dataset(args, is_load=True)\n",
    "    # dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.5.1 HD-Letter:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### HD-Letter, training concept:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"c-Line\",\n",
    "        \"seed\": 1,\n",
    "        \"n_examples\": 44000,\n",
    "        \"canvas_size\": 16,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 2,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"allow_connect\": True,\n",
    "        \"parsing_check\": False,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### HD-Letter, training relation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"c-Parallel+VerticalMid+VerticalEdge\",\n",
    "        \"seed\": 1,\n",
    "        \"n_examples\": 44000,\n",
    "        \"canvas_size\": 16,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 3,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"allow_connect\": True,\n",
    "        \"parsing_check\": False,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### HD-Letter, classification at inference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"c-Eshape+Fshape+Ashape\",\n",
    "        \"seed\": 2,\n",
    "        \"n_examples\": 400,\n",
    "        \"canvas_size\": 16,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 0,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"allow_connect\": True,\n",
    "        \"parsing_check\": False,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### HD-Letter, detection at inference:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Eshape:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Eshape:\n",
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"c-Eshape[6,8]+Cshape+Lshape+Tshape+Rect+RectSolid^Eshape[6,8]\",\n",
    "        \"seed\": 2,\n",
    "        \"n_examples\": 400,\n",
    "        \"canvas_size\": 16,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 2,\n",
    "        \"min_n_distractors\": 1,\n",
    "        \"allow_connect\": False,\n",
    "        \"parsing_check\": True,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Fshape:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fshape:\n",
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"c-Fshape[6,8]+Cshape+Lshape+Tshape+Rect+RectSolid^Fshape[6,8]\",\n",
    "        \"seed\": 2,\n",
    "        \"n_examples\": 400,\n",
    "        \"canvas_size\": 16,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 2,\n",
    "        \"min_n_distractors\": 1,\n",
    "        \"allow_connect\": False,\n",
    "        \"parsing_check\": True,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Ashape:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ashape:\n",
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"c-Fshape[6,8]+Cshape+Lshape+Tshape+Rect+RectSolid^Fshape[6,8]\",\n",
    "        \"seed\": 2,\n",
    "        \"n_examples\": 400,\n",
    "        \"canvas_size\": 16,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 2,\n",
    "        \"min_n_distractors\": 1,\n",
    "        \"allow_connect\": False,\n",
    "        \"parsing_check\": True,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.5.2 HD-Concept:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### HD-Concept, training concept:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"c-Rect[4,16]+Eshape[3,10]\",\n",
    "        \"seed\": 1,\n",
    "        \"n_examples\": 44000,\n",
    "        \"canvas_size\": 20,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 2,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"allow_connect\": True,\n",
    "        \"parsing_check\": False,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### HD-Concept, training relation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    # BabyARC-fewshot dataset for classification:\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"c-IsNonOverlapXY+IsInside+IsEnclosed(Rect[4,16]+Randshape[3,8]+Lshape[3,10]+Tshape[3,10])\",\n",
    "        \"seed\": 1,\n",
    "        \"n_examples\": 44000,\n",
    "        \"canvas_size\": 20,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 1,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"allow_connect\": True,\n",
    "        \"parsing_check\": False,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### HD-Concept, classification at inference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"c-RectE1a+RectE2a+RectE3a\",\n",
    "        \"seed\": 2,\n",
    "        \"n_examples\": 200,\n",
    "        \"canvas_size\": 20,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 0,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"allow_connect\": True,\n",
    "        \"parsing_check\": False,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### HD-Concept, detection at inference:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### RectE1a:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if is_load_example_dataset:\n",
    "# RectE1a:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"c-RectE1a+Cshape+Lshape+Tshape+Rect+RectSolid^RectE1a\",\n",
    "        \"seed\": 2,\n",
    "        \"n_examples\": 200,\n",
    "        \"canvas_size\": 20,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 1,\n",
    "        \"min_n_distractors\": 1,\n",
    "        \"allow_connect\": False,\n",
    "        \"parsing_check\": True,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### RectE2a:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# RectE2a:\n",
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"c-RectE2a+Cshape+Lshape+Tshape+Rect+RectSolid^RectE2a\",\n",
    "        \"seed\": 2,\n",
    "        \"n_examples\": 200,\n",
    "        \"canvas_size\": 20,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 1,\n",
    "        \"min_n_distractors\": 1,\n",
    "        \"allow_connect\": False,\n",
    "        \"parsing_check\": True,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### RectE3a:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RectE3a:\n",
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"c-RectE3a+Cshape+Lshape+Tshape+Rect+RectSolid^RectE3a\",\n",
    "        \"seed\": 2,\n",
    "        \"n_examples\": 200,\n",
    "        \"canvas_size\": 20,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 1,\n",
    "        \"min_n_distractors\": 1,\n",
    "        \"allow_connect\": True,\n",
    "        \"parsing_check\": True,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.5.3 Acquiring concepts between domains:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# RectE3a:\n",
    "if is_load_example_dataset:\n",
    "    import sys, os\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "    sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "    from zeroc.train import get_dataset, ConceptDataset, ConceptClevrDataset\n",
    "    from zeroc.concept_library.util import init_args\n",
    "\n",
    "    args = init_args({\n",
    "        \"dataset\": \"yc-Eshape[5,9]+Fshape[5,9]+Ashape[5,9]\",\n",
    "        \"seed_3d\": 42,\n",
    "        \"n_examples\": 200,\n",
    "        \"num_processes_3d\": 5,\n",
    "        \"n_queries_per_class\": 1,\n",
    "        # 2D examples\n",
    "        \"seed\": 102,\n",
    "        \"use_seed_2d\": True,\n",
    "        \"canvas_size\": 16,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 0,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"parsing_check\": False,\n",
    "        \"allow_connect\": True,\n",
    "    })\n",
    "    dataset, _ = get_dataset(args, is_load=True)\n",
    "    dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.6 Dataset test:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test get_chosen_rel\n",
    "def test_line_rel(pos1, pos2):\n",
    "    image = torch.zeros(1, 16, 16)\n",
    "    image, mask1, pos1, _ = get_line(image, direction=None, pos=pos1, min_size=3, max_size=None, color_avail=[1])\n",
    "    image, mask2, pos2, _ = get_line(image, direction=None, pos=pos2, min_size=3, max_size=None, color_avail=[2])\n",
    "    print(image)\n",
    "    print(get_chosen_line_rel(pos1, pos2)) # Vertical mid\n",
    "    print()\n",
    "\n",
    "# test_line_rel((9, 4, 1, 7), (0, 3, 9, 1))\n",
    "\n",
    "# # Rotated T-shape tests\n",
    "# test_line_rel((9, 4, 1, 7), (1, 3, 9, 1)) # Test bottom edge\n",
    "# test_line_rel((8, 4, 1, 7), (1, 3, 9, 1))\n",
    "# test_line_rel((1, 4, 1, 7), (1, 3, 9, 1)) # Test upper edge\n",
    "# test_line_rel((2, 4, 1, 7), (1, 3, 9, 1))\n",
    "\n",
    "# test_line_rel((9, 1, 1, 5), (1, 6, 9, 1)) # Test bottom edge\n",
    "# test_line_rel((8, 1, 1, 5), (1, 6, 9, 1)) \n",
    "# test_line_rel((1, 1, 1, 5), (1, 6, 9, 1)) \n",
    "# test_line_rel((2, 1, 1, 5), (1, 6, 9, 1)) \n",
    "\n",
    "\n",
    "# # T-shape tests\n",
    "# test_line_rel((2, 9, 6, 1), (1, 2, 1, 8)) # Test right edge\n",
    "# test_line_rel((2, 8, 6, 1), (1, 2, 1, 8))\n",
    "# test_line_rel((2, 2, 6, 1), (1, 2, 1, 8)) # Test left edge\n",
    "# test_line_rel((2, 3, 6, 1), (1, 2, 1, 8))\n",
    "# test_line_rel((2, 5, 6, 1), (1, 2, 1, 8)) # Test middle\n",
    "\n",
    "# test_line_rel((1, 9, 6, 1), (7, 2, 1, 8)) # Test right edge\n",
    "# test_line_rel((1, 8, 6, 1), (7, 2, 1, 8)) \n",
    "# test_line_rel((1, 2, 6, 1), (7, 2, 1, 8)) \n",
    "# test_line_rel((1, 3, 6, 1), (7, 2, 1, 8)) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"c-RectE3a+Eshape+Rect+Tshape+Fshape+Ashape^RectE3a\",\n",
    "#     \"seed\": 2,\n",
    "#     \"n_examples\": 200,\n",
    "#     \"canvas_size\": 20,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 1,\n",
    "#     \"min_n_distractors\": 1,\n",
    "#     \"allow_connect\": True,\n",
    "#     \"parsing_check\": False,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True, is_save_inter=True)\n",
    "# dataset.draw(range(40))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"c-RectE3a+Cshape[2,5]+Lshape[2,5]+Tshape[2,5]+Rect[2,5]+RectSolid[2,5]^RectE3a\",\n",
    "#     \"seed\": 2,\n",
    "#     \"n_examples\": 200,\n",
    "#     \"canvas_size\": 20,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 1,\n",
    "#     \"min_n_distractors\": 1,\n",
    "#     \"allow_connect\": True,\n",
    "#     \"parsing_check\": False,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True, is_save_inter=True)\n",
    "# dataset.draw(range(40))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"c-Cshape[2,5]+Lshape[2,5]+Tshape[2,5]+Rect[2,5]+RectSolid[2,5]^Cshape[2,5]\",\n",
    "#     \"seed\": 2,\n",
    "#     \"n_examples\": 200,\n",
    "#     \"canvas_size\": 20,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 1,\n",
    "#     \"min_n_distractors\": 1,\n",
    "#     \"allow_connect\": False,\n",
    "#     \"parsing_check\": True,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True, is_rewrite=True, is_save_inter=True)\n",
    "# dataset.draw(range(40))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # BabyARC-relation dataset:\n",
    "# relation_args = init_args({\n",
    "#     \"dataset\": \"c-IsNonOverlapXY+IsInside+IsEnclosed(Rect[4,16]+Randshape[3,8]+Lshape[3,10]+Tshape[3,10])\",\n",
    "#     \"seed\": 1,\n",
    "#     \"n_examples\": 44000,\n",
    "#     \"canvas_size\": 20,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"max_n_distractors\": 1,\n",
    "#     \"min_n_distractors\": 0,\n",
    "#     \"allow_connect\": True,\n",
    "#     \"parsing_check\": True,\n",
    "# })\n",
    "# relation_dataset, args = get_dataset(relation_args, is_load=True)\n",
    "# relation_dataset.draw(range(40))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# from zeroc.concepts_shapes import SameShape, SameColor, SameAll, SameRow, SameCol, SubsetOf, IsInside, IsNonOverlapXY\n",
    "# for data in relation_dataset:\n",
    "#     pos_id = data[2]\n",
    "#     info = data[3]\n",
    "#     relations = get_arc_relations(data[0], info)\n",
    "#     relations = [ele[2] for ele in relations]\n",
    "#     print(pos_id, relations)\n",
    "#     print(data[2])\n",
    "#     plot_matrices([data[1][0].squeeze(), data[1][1].squeeze()])\n",
    "#         # raise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# concept_args = init_args({\n",
    "#     \"dataset\": \"c-Parallel+VerticalMid+VerticalEdge\",\n",
    "#     \"seed\": 1,\n",
    "#     \"n_examples\": 44000,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"max_n_distractors\": 3,\n",
    "#     \"min_n_distractors\": 0,\n",
    "#     \"allow_connect\": True,\n",
    "# })\n",
    "# concept_dataset, args = get_dataset(concept_args, is_load=True, is_rewrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"c-Eshape+Ashape^Eshape\",\n",
    "#     \"seed\": 2,\n",
    "#     \"n_examples\": 40,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 1,\n",
    "#     \"min_n_distractors\": 1,\n",
    "#     \"allow_connect\": True,\n",
    "#     \"parsing_check\": True,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True, is_rewrite=True)\n",
    "# dataset.draw(range(40))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"c-IsNonOverlapXY+IsInside+IsEnclosed(Rect[4,16]+Randshape[3,8]+Lshape[3,10]+Tshape[3,10])\",\n",
    "#     \"seed\": 2,\n",
    "#     \"n_examples\": 400,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"min_n_distractors\": 0,\n",
    "#     \"max_n_distractors\": 1,\n",
    "#     \"parsing_check\": True,\n",
    "#     \"allow_connect\": True,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True)\n",
    "# dataset.draw(range(40))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"c-RectE1a+Lshape^RectE1a\",\n",
    "#     \"seed\": 2,\n",
    "#     \"n_examples\": 40,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"min_n_distractors\": 1,\n",
    "#     \"max_n_distractors\": 1,\n",
    "#     \"parsing_check\": True,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=False)\n",
    "# dataset.draw(range(40))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"pc-RectE1b+RectE1c+RectE2b+RectE2c+RectE3b+RectE3c\",\n",
    "#     \"seed\": 2,\n",
    "#     \"n_examples\": 100,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 0,\n",
    "#     \"parsing_check\": True,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True)\n",
    "# dataset.draw(range(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"pc-RectE1a+RectE2a+RectE3a\",\n",
    "#     \"seed\": 2,\n",
    "#     \"n_examples\": 400,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 0,\n",
    "#     \"parsing_check\": True,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True, is_rewrite=True)\n",
    "# dataset.draw(range(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # import sys, os\n",
    "# # sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# # sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"c-Rect[4,16]+Eshape[3,10]\",\n",
    "#     \"seed\": 1,\n",
    "#     \"n_examples\": 44000,\n",
    "#     \"canvas_size\": 20,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 2,\n",
    "#     \"min_n_distractors\": 0,\n",
    "#     \"allow_connect\": True,\n",
    "#     \"parsing_check\": False,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True)\n",
    "# dataset.draw(range(40))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # import sys, os\n",
    "# # sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# # sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"c-IsNonOverlapXY+IsInside+IsEnclosed(Rect[4,16]+Randshape[3,8]+Lshape[3,10]+Tshape[3,10])\",\n",
    "#     \"seed\": 1,\n",
    "#     \"n_examples\": 44000,\n",
    "#     \"canvas_size\": 20,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 1,\n",
    "#     \"min_n_distractors\": 0,\n",
    "#     \"allow_connect\": True,\n",
    "#     \"parsing_check\": False,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True)\n",
    "# dataset.draw(range(40))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"c-Rect[4,15]+Eshape[5,12]\",\n",
    "#     \"seed\": 2,\n",
    "#     \"n_examples\": 40,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 1,\n",
    "#     \"parsing_check\": True,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True)\n",
    "# dataset.draw(range(40))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"pc-Cshape+Eshape+Fshape+Ashape+Hshape+Rect\",\n",
    "#     \"seed\": 102,\n",
    "#     \"n_examples\": 400,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 0,\n",
    "#     \"parsing_check\": False,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True, is_rewrite=True)\n",
    "# dataset.draw(range(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"pg-Cshape+Lshape+Tshape+Rect^Eshape+Fshape+Ashape\",\n",
    "#     \"seed\": 12,\n",
    "#     \"n_examples\": 400,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 0,\n",
    "#     \"parsing_check\": False,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True, is_rewrite=False)\n",
    "# dataset.draw(range(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-fewshot dataset for classification:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "\n",
    "# args = init_args({\n",
    "#     \"dataset\": \"yc-Eshape[5,9]+Fshape[5,9]+Ashape[5,9]\",\n",
    "#     \"seed_3d\": 42,\n",
    "#     \"n_examples\": 200,\n",
    "#     \"num_processes_3d\": 5,\n",
    "#     \"n_queries_per_class\": 1,\n",
    "#     # 2D examples\n",
    "#     \"seed\": 102,\n",
    "#     \"use_seed_2d\": True,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 0,\n",
    "#     \"min_n_distractors\": 0,\n",
    "#     \"parsing_check\": False,\n",
    "#     \"allow_connect\": True,\n",
    "# })\n",
    "# dataset, _ = get_dataset(args, is_load=True)\n",
    "# dataset.draw(range(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# composite_args = init_args({\n",
    "#     \"dataset\": \"c-Eshape+Cshape+Lshape+Tshape+Rect+RectSolid^Eshape\",\n",
    "#     \"seed\": 2,\n",
    "#     \"n_examples\": 400,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"max_n_distractors\": 4,\n",
    "#     \"parsing_check\": True,\n",
    "# })\n",
    "# dataset, composite_args = get_dataset(composite_args, is_rewrite=True, is_load=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # BabyARC-relation dataset, 3D:\n",
    "# relation_args = init_args({\n",
    "#     \"dataset\": \"y-Parallel+VerticalMid+VerticalEdge\",\n",
    "#     \"seed\": 2,\n",
    "#     \"n_examples\": 44000,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"color_map_3d\": \"same\",\n",
    "#     \"add_thick_surf\": (0, 0.5),\n",
    "#     \"add_thick_depth\": (0, 0.5),\n",
    "#     \"max_n_distractors\": 2,\n",
    "#     \"seed_3d\": 42,\n",
    "#     \"num_processes_3d\": 20,\n",
    "#     \"image_size_3d\": (256,256),\n",
    "# })\n",
    "# relation_dataset, args = get_dataset(relation_args, is_rewrite=False, is_load=True)\n",
    "# relation_dataset.draw(range(40))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # import sys, os\n",
    "# # sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# # sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # # BabyARC-concept dataset:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset, ConceptFewshotDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "# concept_args = init_args({\n",
    "#     \"dataset\": \"c-Line\",\n",
    "#     \"seed\": 1,\n",
    "#     \"n_examples\": 44000,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"max_n_distractors\": 0,\n",
    "# })\n",
    "# concept_dataset, args = get_dataset(concept_args, is_load=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-concept dataset:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "# concept_args = init_args({\n",
    "#     \"dataset\": \"c-Line\",\n",
    "#     \"seed\": 1,\n",
    "#     \"n_examples\": 44000,\n",
    "#     \"canvas_size\": 8,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"max_n_distractors\": 2,\n",
    "# })\n",
    "# concept_dataset, args = get_dataset(concept_args, is_load=True)\n",
    "\n",
    "# BabyARC-relation dataset:\n",
    "# relation_args = init_args({\n",
    "#     \"dataset\": \"c-Parallel+Vertical\",\n",
    "#     \"seed\": 1,\n",
    "#     \"n_examples\": 44000,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"max_n_distractors\": 2,\n",
    "# })\n",
    "# relation_dataset, args = get_dataset(relation_args, is_rewrite=True, is_load=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import sys, os\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "# sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "# # BabyARC-concept dataset:\n",
    "# from zeroc.concept_library.train import get_dataset, ConceptDataset\n",
    "# from zeroc.concept_library.util import init_args\n",
    "# concept_args = init_args({\n",
    "#     \"dataset\": \"c-Line+Lshape+Rect+RectSolid\",\n",
    "#     \"seed\": 1,\n",
    "#     \"n_examples\": 44000,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"max_n_distractors\": 2,\n",
    "# })\n",
    "# concept_dataset, args = get_dataset(concept_args, is_load=False)\n",
    "\n",
    "# # BabyARC-relation dataset:\n",
    "# relation_args = init_args({\n",
    "#     \"dataset\": \"c-SameShape+SameColor+IsInside(Line+Lshape+Rect+RectSolid)\",\n",
    "#     \"seed\": 1,\n",
    "#     \"n_examples\": 44000,\n",
    "#     \"canvas_size\": 16,\n",
    "#     \"rainbow_prob\": 0.,\n",
    "#     \"color_avail\": \"1,2\",\n",
    "#     \"w_type\": \"image+mask\",\n",
    "#     \"max_n_distractors\": 2,\n",
    "# })\n",
    "# relation_dataset, args = get_dataset(relation_args, is_load=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     args = init_args({\n",
    "#         \"dataset\": \"c-Eshape+Fshape+Ashape\",\n",
    "#         \"seed\": 1,\n",
    "#         \"n_examples\":100,\n",
    "#         \"canvas_size\": 16,\n",
    "#         \"rainbow_prob\": 0.,\n",
    "#         \"color_avail\": \"1,2\",\n",
    "#         \"w_type\": \"mask\",\n",
    "#         \"max_n_distractors\": 0,\n",
    "#     })\n",
    "#     dataset, args = get_dataset(args, is_load=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     args = init_args({\n",
    "#         \"dataset\": \"c-Line+Lshape+Rect+RectSolid\",\n",
    "#         \"seed\": 1,\n",
    "#         \"n_examples\": 20,\n",
    "#         \"canvas_size\": 8,\n",
    "#         \"rainbow_prob\": 0.,\n",
    "#         \"color_avail\": \"-1\",\n",
    "#         \"w_type\": \"mask\",\n",
    "#         \"max_n_distractors\": 0,\n",
    "#     })\n",
    "#     dataset, args = get_dataset(args, is_load=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # Get dataset:\n",
    "# concepts = [\"Lshape\", \"Tshape\", \"Eshape\", \"Hshape\", \"Cshape\", \"Ashape\", \"Fshape\"]\n",
    "# n_examples_per_concept = 10\n",
    "# n_examples_per_task = 5\n",
    "# n_tasks = 4\n",
    "\n",
    "# task_list = generate_discriminative_task(\n",
    "#     concepts,\n",
    "#     n_examples_per_concept,\n",
    "#     n_examples_per_task,\n",
    "#     n_tasks,\n",
    "#     isplot=True,\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     # CLEVR-relation dataset easy:\n",
    "#     args = init_args({\n",
    "#         \"dataset\": \"v-CLEVRRelation:train\",\n",
    "#         \"seed\": 1,\n",
    "#         \"n_examples\": 3000,\n",
    "#         \"canvas_size\": 8,\n",
    "#         \"rainbow_prob\": 0.,\n",
    "#         \"color_avail\": \"-1\",\n",
    "#         \"selector_image_value_range\": \"-1,1\",\n",
    "#     })\n",
    "#     dataset, args = get_dataset(args, is_load=False, is_rewrite=True)\n",
    "#     for i in range(50):\n",
    "#         visualize_matrices([(dataset[i][0][0][0] + 1)/2], use_color_dict=False)\n",
    "#         plot_matrices(dataset[i][0][0][1], images_per_row=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     # BabyARC-relation dataset easy:\n",
    "#     args = init_args({\n",
    "#         \"dataset\": \"h-r^2ai+2a+3ai+3a+3b:SameShape+SameColor(Line+Rect+RectSolid+Lshape)-d^1:Line+Randshape\",\n",
    "#         \"seed\": 1,\n",
    "#         \"n_examples\": 40,\n",
    "#         \"canvas_size\": 16,\n",
    "#         \"rainbow_prob\": 0.,\n",
    "#         \"color_avail\": \"-1\",\n",
    "#         \"max_n_distractors\": 2,\n",
    "#         \"min_n_distractors\": 1,\n",
    "#         \"allow_connect\": False,\n",
    "#     })\n",
    "#     dataset, args = get_dataset(args, is_load=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     # BabyARC-relation dataset easy validation:\n",
    "#     args = init_args({\n",
    "#         \"dataset\": \"h-r^2ai+2a+3ai+3a+3b:SameShape+SameColor(Line+Rect+RectSolid+Lshape)\",\n",
    "#         \"seed\": 1,\n",
    "#         \"n_examples\": 3000,\n",
    "#         \"canvas_size\": 16,\n",
    "#         \"rainbow_prob\": 0.,\n",
    "#         \"color_avail\": \"1,2\",\n",
    "#     })\n",
    "#     dataset, args = get_dataset(args, is_load=False, is_rewrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     # BabyARC-relation dataset easy validation:\n",
    "#     args = init_args({\n",
    "#         \"dataset\": \"h-r^2ai+2a+3ai+3a+3b+4ai+4a+4b:SameShape+SameColor(Line+Rect+RectSolid+Lshape)\",\n",
    "#         \"seed\": 1,\n",
    "#         \"n_examples\": 3000,\n",
    "#         \"canvas_size\": 16,\n",
    "#         \"rainbow_prob\": 0.,\n",
    "#         \"color_avail\": \"1,2\",\n",
    "#     })\n",
    "#     dataset, args = get_dataset(args, is_load=False, is_rewrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     args = init_args({\n",
    "#         \"seed\": 1,\n",
    "#         \"dataset\": \"s-nsh-4-noise-0.01\",\n",
    "#         \"n_examples\": 1000,\n",
    "#     })\n",
    "#     dataset = get_dataset(args, isplot=True, is_load=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     args = init_args({\n",
    "#         \"dataset\": \"h-c^(1,2):Eshape+Ashape-d^1:Randshape\",\n",
    "#         \"seed\": 1,\n",
    "#         \"n_examples\": 10,\n",
    "#         \"canvas_size\": 8,\n",
    "#         \"rainbow_prob\": 0.,\n",
    "#         \"color_avail\": \"-1\",\n",
    "#     })\n",
    "#     dataset, args = get_dataset(args, is_load=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     # BabyARC-concept dataset:\n",
    "#     from zeroc.concept_library.train import get_dataset, ConceptDataset\n",
    "#     from zeroc.concept_library.util import init_args\n",
    "#     args = init_args({\n",
    "#         \"dataset\": \"h-c^(1,2):Line+Lshape+Rect+RectSolid\",\n",
    "#         \"seed\": 1,\n",
    "#         \"n_examples\": 40000,\n",
    "#         \"canvas_size\": 8,\n",
    "#         \"rainbow_prob\": 0.,\n",
    "#         \"color_avail\": \"1,2\",\n",
    "#     })\n",
    "#     dataset, args = get_dataset(args, is_load=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     # BabyARC-concept dataset:\n",
    "#     args = init_args({\n",
    "#         \"dataset\": \"h-c^(1,2):Line+Lshape+Rect+RectSolid\",\n",
    "#         \"seed\": 1,\n",
    "#         \"n_examples\": 40000,\n",
    "#         \"canvas_size\": 8,\n",
    "#         \"rainbow_prob\": 0.,\n",
    "#         \"color_avail\": \"1,2\",\n",
    "#     })\n",
    "#     dataset, args = get_dataset(args, is_load=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     args = get_pdict()(\n",
    "#         # dataset=\"c-Image\",\n",
    "# #         dataset=\"c-Line+RectSolid\",\n",
    "# #         dataset=\"c-IsInside(ARCshape)\",\n",
    "#         dataset=\"c-SameAll+SameShape+SameColor+SameRow+SameCol+IsInside(Line+Rect+RectSolid+Lshape+Randshape+ARCshape)\",\n",
    "#         # dataset=\"c-arc^Line+RectSolid+Rect\",\n",
    "#         canvas_size=8, n_examples=100, rainbow_prob=0,\n",
    "#         max_n_distractors=-1,\n",
    "#         color_avail=\"-1\",\n",
    "#     )\n",
    "#     dataset, args = get_dataset(args)\n",
    "# #     dataset.draw(range(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     args = get_pdict()(\n",
    "#         dataset=\"c-RotateA+RotateB+RotateC+hFlip+vFlip+DiagFlipA+DiagFlipB+Move(Line+Rect+RectSolid+Lshape+Randshape+ARCshape)\",\n",
    "#         canvas_size=8, n_examples=100, rainbow_prob=0,\n",
    "#         max_n_distractors=-1,\n",
    "#         color_avail=\"2\",\n",
    "#     )\n",
    "#     dataset, args = get_dataset(args)\n",
    "#     dataset.draw(range(100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Helper functions:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Data Transformations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_bboxes(mask):\n",
    "    \"\"\"Compute bounding boxes from masks.\n",
    "    mask: [height, width, num_instances]. Mask pixels are either 1 or 0.\n",
    "    Returns: bbox array [num_instances, (y1, x1, y2, x2)].\n",
    "    \"\"\"\n",
    "    boxes = np.zeros([mask.shape[-1], 4])\n",
    "    for i in range(mask.shape[-1]):\n",
    "        m = mask[:, :, i]\n",
    "        # Bounding box.\n",
    "        horizontal_indices = np.where(m.any(dim=0))[0]\n",
    "        vertical_indices = np.where(m.any(dim=1))[0]\n",
    "        if horizontal_indices.shape[0]:\n",
    "            x1, x2 = horizontal_indices[[0, -1]]\n",
    "            y1, y2 = vertical_indices[[0, -1]]\n",
    "            x2 += 1\n",
    "            y2 += 1\n",
    "        else:\n",
    "            # No mask for this instance. Might happen due to\n",
    "            # resizing or cropping. Set bbox to zeros\n",
    "            x1, x2, y1, y2 = 0, 0, 0, 0\n",
    "        boxes[i] = np.array([x1, y1, x2, y2])\n",
    "    return boxes.astype(np.float32)\n",
    "\n",
    "\n",
    "class BabyARC2DPad(nn.Module):\n",
    "    \"\"\" \n",
    "    Pad 10 channel BabyARC images by padding the 0 channel with one's and the rest\n",
    "    with zero's\n",
    "    \"\"\"\n",
    "    def __init__(self, pad_size):\n",
    "        super(BabyARC2DPad, self).__init__()\n",
    "        self.pad_size = pad_size\n",
    "\n",
    "    def forward(self, img):\n",
    "        assert img.shape[-3] == 10\n",
    "        one_pad = nn.ConstantPad2d(self.pad_size, 1)\n",
    "        zero_pad = nn.ConstantPad2d(self.pad_size, 0)\n",
    "        zero_channel = one_pad(img[0:1, :, :])\n",
    "        remainder = zero_pad(img[1:, :, :])\n",
    "        new_img = torch.cat((zero_channel, remainder))\n",
    "        return new_img\n",
    "\n",
    "\n",
    "class PermuteChannels(nn.Module):\n",
    "    \"\"\"\n",
    "    Permutes all channels following some starting channel\n",
    "    \"\"\"\n",
    "    def __init__(self, start_channel, num_channels):\n",
    "        super(PermuteChannels, self).__init__()\n",
    "        self.start_channel = start_channel\n",
    "        self.num_channels = num_channels\n",
    "    \n",
    "    def forward(self, sample):\n",
    "        # Assume sample is a tuple of (negative data, positive mask)\n",
    "        perm = torch.randperm(self.num_channels - self.start_channel)\n",
    "        perm = perm + self.start_channel\n",
    "        indices = torch.cat((torch.tensor(list(range(0, self.start_channel))) , perm))\n",
    "        permuted = sample[0][0][indices, ...]\n",
    "        new_sample = ((permuted, sample[0][1], sample[0][2], sample[0][3]), sample[1])\n",
    "        return new_sample\n",
    "\n",
    "\n",
    "class BabyARC2DColorTransform(PermuteChannels):\n",
    "    \"\"\"\n",
    "    Permutes all channels following some starting channel\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        super().__init__(start_channel=1, num_channels=10)\n",
    "        \n",
    "        \n",
    "class BabyARC3DColorTransform(nn.Module):\n",
    "    \"\"\"\n",
    "    Performs random color jitter and random grayscale\n",
    "    \"\"\"\n",
    "    def __init__(self, s=0.5):\n",
    "        # s is the strength of color distortion.\n",
    "        super().__init__()\n",
    "        self.s = s\n",
    "        \n",
    "    def get_color_distortion(self):\n",
    "        color_jitter = transforms.ColorJitter(0.8*self.s, 0.8*self.s, 0.8*self.s, 0.4*self.s)\n",
    "        rnd_color_jitter = transforms.RandomApply([color_jitter], p=0.8)\n",
    "        rnd_gray = transforms.RandomGrayscale(p=0.2)\n",
    "        color_distort = transforms.Compose([\n",
    "            rnd_color_jitter,\n",
    "            rnd_gray])\n",
    "        return color_distort\n",
    "    \n",
    "    def forward(self, sample):\n",
    "        color_transform = self.get_color_distortion()\n",
    "        new_img = color_transform(sample[0][0])\n",
    "        new_sample = ((new_img, sample[0][1], sample[0][2], sample[0][3]), sample[1])\n",
    "        return new_sample\n",
    "\n",
    "\n",
    "class GenericRandomResizedCrop(nn.Module):\n",
    "    \"\"\"\n",
    "    Crops an area with fixed aspect ratio that fully includes \n",
    "    ground truth mask\n",
    "    \"\"\"\n",
    "    def __init__(self, size, pad, interpolation=InterpolationMode.NEAREST):\n",
    "        super().__init__()\n",
    "        self.size = size # Size of the result, which is equivalent to the canvas_size\n",
    "        self.pad = pad\n",
    "        self.mask_pad = nn.ConstantPad2d(size // 2, 0)\n",
    "        self.interpolation = interpolation\n",
    "        self.MAX_ITER = 2\n",
    "    \n",
    "    def __call__(self, sample):\n",
    "        # Assume sample is a tuple of (negative data, positive mask)\n",
    "        neg_data, pos_mask = sample\n",
    "        half_canvas = int(neg_data[0].shape[-2] / 2)\n",
    "        assert self.size == max(*neg_data[0].shape[-2:])\n",
    "        # Important: get the bounding box of the masks in the image\n",
    "        total_mask = pos_mask[0] if len(pos_mask) == 1 else pos_mask[0] + pos_mask[1]\n",
    "        # Top left, bottom right corners\n",
    "        y1, x1, y2, x2 = extract_bboxes(self.mask_pad(total_mask.squeeze()).unsqueeze(-1))[0].astype(np.int32)\n",
    "        box_h = x2 - x1\n",
    "        box_w = y2 - y1\n",
    "        min_crop_size = max(box_h, box_w)\n",
    "        mask_sums = torch.tensor([0] * len(pos_mask))\n",
    "        i = 0\n",
    "        # Make sure both positive masks are non-empty\n",
    "        while torch.any(mask_sums == 0) and i < self.MAX_ITER:\n",
    "            # Randomly sample crop sizes\n",
    "            crop_size = torch.randint(low=min_crop_size, high=min_crop_size + half_canvas, size=()).item()\n",
    "            # Sample the top left corner. Make sure that the crop area is within the padded canvas\n",
    "            bottom = neg_data[0].shape[-2] + 2 * half_canvas\n",
    "            right = neg_data[0].shape[-1] + 2 * half_canvas\n",
    "            x1_sample = torch.randint(low=max(x1 - (crop_size - box_h ), 0), high=min(x1+1, bottom - crop_size + 1), size=()).item()\n",
    "            y1_sample = torch.randint(low=max(y1 - (crop_size - box_w ), 0), high=min(y1+1, right - crop_size + 1), size=()).item()\n",
    "            assert x1_sample >= 0 and x1_sample + crop_size <= bottom \\\n",
    "                and y1_sample >= 0 and y1_sample + crop_size <= right\n",
    "            # Performing cropping \n",
    "            new_img = F_tr.resized_crop(self.pad(neg_data[0]), x1_sample, y1_sample, crop_size, crop_size, self.size, self.interpolation)\n",
    "            assert new_img.shape[-2:] == neg_data[0].shape[-2:]\n",
    "            new_neg_masks = [F_tr.resized_crop(self.mask_pad(mask), x1_sample, y1_sample, crop_size, crop_size, self.size, self.interpolation) for mask in neg_data[1]]\n",
    "            new_pos_mask = [F_tr.resized_crop(self.mask_pad(mask), x1_sample, y1_sample, crop_size, crop_size, self.size, self.interpolation) for mask in pos_mask] \n",
    "            new_sample = ((new_img, tuple(new_neg_masks), neg_data[2], neg_data[3]), new_pos_mask)\n",
    "            mask_sums = torch.tensor([torch.sum(mask) for mask in new_pos_mask])\n",
    "            i += 1\n",
    "        if torch.any(mask_sums == 0):\n",
    "            return sample\n",
    "        return new_sample\n",
    "\n",
    "\n",
    "class BabyARC2DRandomResizedCrop(GenericRandomResizedCrop):\n",
    "    \"\"\"\n",
    "    Crops an area with fixed aspect ratio that fully includes \n",
    "    ground truth mask\n",
    "    \"\"\"\n",
    "    def __init__(self, size):\n",
    "        if not isinstance(size, int):\n",
    "            size = size[-1]\n",
    "        super().__init__(size=size, pad=BabyARC2DPad(size // 2))\n",
    "\n",
    "\n",
    "class BabyARC3DRandomResizedCrop(GenericRandomResizedCrop):\n",
    "    \"\"\"\n",
    "    Crops an area with fixed aspect ratio that fully includes \n",
    "    ground truth mask\n",
    "    \"\"\"\n",
    "    def __init__(self, size):\n",
    "        super().__init__(size=size, pad=nn.ConstantPad2d(size // 2, 0))\n",
    "\n",
    "\n",
    "class GenericRandomFlip(nn.Module):\n",
    "    \"\"\"flip the given image randomly (horizontally or vertically) with a given probability.\n",
    "    If the image is torch Tensor, it is expected\n",
    "    to have [..., H, W] shape, where ... means an arbitrary number of leading\n",
    "    dimensions\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, sample):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            img (PIL Image or Tensor): Image to be flipped (Identity, hFlip, vFlip).\n",
    "\n",
    "        Returns:\n",
    "            PIL Image or Tensor: Randomly flipped image.\n",
    "        \"\"\"\n",
    "        neg_data, pos_mask = sample\n",
    "        funcs = [\n",
    "            nn.Identity(),\n",
    "            F_tr.hflip,\n",
    "            F_tr.vflip,\n",
    "        ]\n",
    "        func = np.random.choice(funcs)\n",
    "        new_neg_data = (func(neg_data[0]),\n",
    "                        tuple([func(mask) for mask in neg_data[1]]), neg_data[2], neg_data[3])\n",
    "        new_pos_mask = tuple([func(mask) for mask in pos_mask])\n",
    "        return (new_neg_data, new_pos_mask)\n",
    "\n",
    "\n",
    "class GenericRandomRotate(nn.Module):\n",
    "    \"\"\"Randomly rotate with a given probability.\n",
    "    If the image is torch Tensor, it is expected\n",
    "    to have [..., H, W] shape, where ... means an arbitrary number of leading\n",
    "    dimensions\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, sample):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            img (PIL Image or Tensor): Image to be rotated (Identity, RotateA, RotateB, RotateC).\n",
    "\n",
    "        Returns:\n",
    "            PIL Image or Tensor: Randomly flipped image.\n",
    "        \"\"\"\n",
    "        neg_data, pos_mask = sample\n",
    "        funcs = [\n",
    "            nn.Identity(),\n",
    "            partial(F_tr.rotate, angle=90),\n",
    "            partial(F_tr.rotate, angle=180),\n",
    "            partial(F_tr.rotate, angle=270),\n",
    "        ]\n",
    "        func = np.random.choice(funcs)\n",
    "        new_neg_data = (func(neg_data[0]),\n",
    "                        tuple([func(mask) for mask in neg_data[1]]), neg_data[2], neg_data[3])\n",
    "        new_pos_mask = tuple([func(mask) for mask in pos_mask])\n",
    "        return (new_neg_data, new_pos_mask)\n",
    "\n",
    "\n",
    "class AddRandomImgPixels(nn.Module):\n",
    "    def __init__(self, size):\n",
    "        super().__init__()\n",
    "        self.size = size\n",
    "\n",
    "    def forward(self, sample):\n",
    "        neg_data, pos_mask = sample\n",
    "        neg_img, neg_mask, neg_id, neg_info = neg_data\n",
    "        assert not isinstance(neg_img, tuple) and not isinstance(neg_img, list)\n",
    "        assert neg_img.shape[-3] == 10\n",
    "        neg_img_argmax = neg_img.argmax(-3, keepdims=True)\n",
    "        remainder = (neg_img_argmax == 0)  # [1, H, W]\n",
    "        p = (neg_img_argmax!=0).float().mean().item()  # non_empty probability\n",
    "        p = torch.rand(1)[0] * 1.5 * p + 0.5 * p\n",
    "        size = (self.size, self.size) if isinstance(self.size, Number) else self.size \n",
    "        new_neg_img = torch.randint(10, size=(1, *size))  # [B, H, W]\n",
    "        mask_pixels = ((torch.rand(1, *size) < p) & remainder).float()\n",
    "        pixels = new_neg_img * mask_pixels\n",
    "        new_neg_img = neg_img_argmax + pixels\n",
    "        new_neg_img = to_one_hot(new_neg_img[0], n_channels=10)\n",
    "        new_neg_data = (new_neg_img, neg_mask, neg_id, neg_info)\n",
    "        return (new_neg_data, pos_mask)\n",
    "\n",
    "\n",
    "class AddRandomImgPatch(nn.Module):\n",
    "    def __init__(self, size, color_avail):\n",
    "        \"\"\"Add random image patch to the place where there is no mask.\"\"\"\n",
    "        super().__init__()\n",
    "        self.size = size\n",
    "        self.color_avail = color_avail\n",
    "\n",
    "    def forward(self, sample):\n",
    "        neg_data, pos_mask = sample\n",
    "        neg_img, neg_mask, neg_id, neg_info = neg_data\n",
    "        neg_img = deepcopy(neg_img)\n",
    "        assert not isinstance(neg_img, tuple) and not isinstance(neg_img, list)\n",
    "        assert neg_img.shape[-3] == 10\n",
    "        \"\"\"\n",
    "        1. First, sample one or two patches\n",
    "        2. Randomly permute the color, and rotate the patch\n",
    "        3. Randomly put into a place with no mask.\n",
    "        \"\"\"\n",
    "        neg_img_argmax = neg_img.argmax(-3)\n",
    "        canvas_h, canvas_w = neg_img.shape[-2:]\n",
    "        num = np.random.choice([1,2])\n",
    "        for _ in range(num):\n",
    "            patch0, pos0 = shrink(neg_img_argmax.float())\n",
    "            row0, col0, h0, w0 = pos0\n",
    "            # Sample a section of the patch:\n",
    "            h = torch.randint(low=1, high=max(2, h0+1), size=(1,))[0]\n",
    "            w = torch.randint(low=1, high=max(2, w0+1), size=(1,))[0]\n",
    "            row = torch.randint(low=0, high=max(1,h0-h+1), size=(1,))[0]\n",
    "            col = torch.randint(low=0, high=max(1,w0-w+1), size=(1,))[0]\n",
    "            patch = patch0[row:row+h, col:col+w]\n",
    "            # Sample direction:\n",
    "            rand_dir = np.random.choice([0,1])\n",
    "            if rand_dir == 1:\n",
    "                patch = torch.rot90(patch, dims=[0,1])\n",
    "            # Permute color:\n",
    "            patch = to_one_hot(patch)\n",
    "            if self.color_avail == \"-1\":\n",
    "                color_avail = \"1,2,3,4,5,6,7,8,9\"\n",
    "            else:\n",
    "                color_avail = self.color_avail\n",
    "            n_colors = len(color_avail.split(\",\"))\n",
    "            permute_idx = torch.randperm(n_colors) + 1\n",
    "            permute_idx = torch.cat([torch.tensor([0]), permute_idx, torch.arange(n_colors+1, 10)])\n",
    "            patch = patch[permute_idx]\n",
    "            # Randomly elongate the image:\n",
    "            h1, w1 = patch.shape[-2:]\n",
    "            is_elongate = np.random.choice([0,1])\n",
    "            if is_elongate == 1:\n",
    "                elongate_ratio = np.random.rand() + 1\n",
    "                if h1 > w1:\n",
    "                    h1_new = min(canvas_h, int(np.round(h1 * elongate_ratio)))\n",
    "                    patch = F.interpolate(patch[None], size=(h1_new, w1), mode=\"nearest\")[0]\n",
    "                elif w1 > h1:\n",
    "                    w1_new = min(canvas_w, int(np.round(w1 * elongate_ratio)))\n",
    "                    patch = F.interpolate(patch[None], size=(h1, w1_new), mode=\"nearest\")[0]\n",
    "\n",
    "            # Add to the current image:\n",
    "            h1, w1 = patch.shape[-2:]\n",
    "            for j in range(10):\n",
    "                row1 = torch.randint(low=0, high=max(1, canvas_h-h1+1), size=(1,))[0]\n",
    "                col1 = torch.randint(low=0, high=max(1, canvas_w-w1+1), size=(1,))[0]\n",
    "                if neg_img_argmax[row1:row1+h1, col1:col1+w1].sum() == 0:\n",
    "                    neg_img[:, row1:row1+h1, col1:col1+w1] = patch\n",
    "                    break\n",
    "                else:\n",
    "                    continue\n",
    "            neg_img_argmax = neg_img.argmax(-3)\n",
    "        new_neg_data = (neg_img, neg_mask, neg_id, neg_info)\n",
    "        return new_neg_data, pos_mask\n",
    "\n",
    "\n",
    "def get_augment(transforms_str, canvas_size, is_rgb=False, color_avail=None):\n",
    "    transforms_lst = []\n",
    "    assert \":\" not in transforms_str\n",
    "    if \"color\" in transforms_str.split(\"+\"):\n",
    "        transforms_lst.append(BabyARC2DColorTransform() if not is_rgb else BabyARC3DColorTransform())\n",
    "    if \"flip\" in transforms_str.split(\"+\"):\n",
    "        transforms_lst.append(GenericRandomFlip())\n",
    "    if \"rotate\" in transforms_str.split(\"+\"):\n",
    "        transforms_lst.append(GenericRandomRotate())\n",
    "    if \"resize\" in transforms_str.split(\"+\"):\n",
    "        transforms_lst.append(BabyARC2DRandomResizedCrop(canvas_size) if not is_rgb else BabyARC3DRandomResizedCrop(canvas_size))\n",
    "    if \"rand\" in transforms_str.split(\"+\"):\n",
    "        transforms_lst.append(AddRandomImgPixels(canvas_size))\n",
    "    if \"randpatch\" in transforms_str.split(\"+\"):\n",
    "        transforms_lst.append(AddRandomImgPatch(canvas_size, color_avail=color_avail))\n",
    "    return transforms.Compose(transforms_lst)\n",
    "\n",
    "\n",
    "def transform_pos_data(pos_data, transforms_str, color_avail):\n",
    "    \"\"\"Transforms a batch of positive data, a tuple of (img, masks, label, info)\n",
    "    Args:\n",
    "        pos_data: (img, tuple(masks), pos_ids, pos_infos)\n",
    "            img: has shape of [B, C, H, W]\n",
    "            tuple(masks) is a tuple of masks, each of which has shape of [B, 1, H, W]\n",
    "            pos_ids and pos_infos will remain as they are.\n",
    "        transforms_str: has format of e.g. \"color+flip+rotate+resize:0.7\", which means that\n",
    "            the img and masks will have 0.7 probability of making the transform \n",
    "            (and 0.3 probability of remaining at the current status), and when transform\n",
    "            happens, will peform transforms of color, flip, rotate, resize.\n",
    "\n",
    "    Returns:\n",
    "        transformed_pos_data: (transformed_img, tuple(transformed_masks), pos_ids, pos_infos)\n",
    "    \"\"\"\n",
    "    transforms_str_split = transforms_str.split(\":\")\n",
    "    if len(transforms_str_split) == 1:\n",
    "        p_transforms = 1\n",
    "        transforms_str_core = transforms_str\n",
    "    else:\n",
    "        p_transforms = eval(transforms_str_split[1])\n",
    "        transforms_str_core = transforms_str_split[0]\n",
    "    if np.random.rand() > p_transforms:\n",
    "        return pos_data\n",
    "    pos_imgs, pos_masks, pos_ids, pos_infos = pos_data\n",
    "    is_rgb = True if pos_imgs.shape[1] == 3 else False\n",
    "    canvas_size = pos_imgs.shape[-2:]\n",
    "    augment = get_augment(transforms_str_core, canvas_size, is_rgb=is_rgb, color_avail=color_avail)\n",
    "    transformed_img = torch.zeros(pos_imgs.shape)\n",
    "    transformed_masks = [torch.zeros(mask.shape) for mask in pos_masks] \n",
    "    for idx in range(pos_imgs.shape[0]):\n",
    "        ex_mask = tuple(mask[idx] for mask in pos_masks)\n",
    "        transformed_tup, new_pos_masks = augment(((pos_imgs[idx], ex_mask, None, None), ex_mask))\n",
    "        transformed_img[idx, ...] = transformed_tup[0]\n",
    "        for mask_idx in range(len(transformed_masks)):\n",
    "            transformed_masks[mask_idx][idx, ...] = new_pos_masks[mask_idx]\n",
    "    return (transformed_img, tuple(transformed_masks), pos_ids, pos_infos)\n",
    "\n",
    "\n",
    "def rescale_data(pos_data, rescaled_size, rescale_mode=\"nearest\"):\n",
    "    \"\"\"Rescale the data to some given size.\n",
    "\n",
    "    Args:\n",
    "        rescaled_size: Choose from \"None\", or e.g. \"16,16\" (size of (16,16))\n",
    "    \"\"\"\n",
    "    if rescaled_size == \"None\":\n",
    "        return pos_data\n",
    "    assert len(rescaled_size.split(\",\")) == 2\n",
    "    rescaled_size = eval(rescaled_size)\n",
    "    pos_imgs, pos_masks, pos_ids, pos_infos = pos_data\n",
    "    if isinstance(pos_imgs, tuple):\n",
    "        assert len(pos_imgs) == 2\n",
    "        pos_imgs = (F.interpolate(pos_imgs[0], size=rescaled_size, mode=rescale_mode),\n",
    "                    F.interpolate(pos_imgs[1], size=rescaled_size, mode=rescale_mode))\n",
    "    else:\n",
    "        pos_imgs = F.interpolate(pos_imgs, size=rescaled_size, mode=rescale_mode)\n",
    "    pos_masks = tuple(F.interpolate(mask_ele, size=rescaled_size, mode=\"nearest\") for mask_ele in pos_masks)\n",
    "    return pos_imgs, pos_masks, pos_ids, pos_infos\n",
    "\n",
    "\n",
    "def rescale_tensor(tensor, rescaled_size):\n",
    "    \"\"\"Rescale the tensor to some given size.\n",
    "\n",
    "    Args:\n",
    "        rescaled_size: Choose from \"None\", or e.g. \"16,16\" (size of (16,16))\n",
    "    \"\"\"\n",
    "    if rescaled_size == \"None\":\n",
    "        return tensor\n",
    "    assert len(rescaled_size.split(\",\")) == 2\n",
    "    rescaled_size = eval(rescaled_size)\n",
    "    if len(tensor.shape) == 3:\n",
    "        tensor = tensor[None]\n",
    "        is_size_3 = True\n",
    "    else:\n",
    "        is_size_3 = False\n",
    "    tensor = F.interpolate(tensor, size=rescaled_size, mode=\"nearest\")\n",
    "    if is_size_3:\n",
    "        tensor = tensor[0]\n",
    "    return tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     args = init_args({\n",
    "#         \"dataset\": \"c-Eshape+Cshape+Fshape+Tshape^Fshape\",\n",
    "#         \"seed\": 1,\n",
    "#         \"n_examples\":40,\n",
    "#         \"canvas_size\": 16,\n",
    "#         \"rainbow_prob\": 0.,\n",
    "#         \"color_avail\": \"1,2\",\n",
    "#         \"w_type\": \"mask\",\n",
    "#         \"max_n_distractors\": 2,\n",
    "#     })\n",
    "#     dataset, args = get_dataset(args, is_load=True)\n",
    "#     # Test color permutation\n",
    "#     permute_transform = BabyARC2DColorTransform()\n",
    "#     example = (dataset[0][0], (torch.zeros(1, 16, 16),), dataset[0][2], dataset[0][3])\n",
    "#     new_example, pos_mask = permute_transform((example, dataset[0][1]))\n",
    "#     visualize_matrices([dataset[0][0].argmax(0)])\n",
    "#     visualize_matrices([new_example[0].argmax(0)])\n",
    "#     visualize_matrices(torch.cat(pos_mask))\n",
    "    \n",
    "#     # Test Hflip\n",
    "#     flip_transform = GenericRandomFlip()\n",
    "#     example = (dataset[0][0], (torch.zeros(1, 16, 16),), dataset[0][2], dataset[0][3])\n",
    "#     new_example, pos_mask = flip_transform((example, dataset[0][1]))\n",
    "#     visualize_matrices([dataset[0][0].argmax(0)])\n",
    "#     visualize_matrices([new_example[0].argmax(0)])\n",
    "#     visualize_matrices(new_example[1][0])\n",
    "#     visualize_matrices(torch.cat(pos_mask))\n",
    "\n",
    "    \n",
    "#     # Test ResizeCrop\n",
    "#     crop_transform = BabyARC2DRandomResizedCrop(size=dataset[0][0].shape[-1])\n",
    "#     example = (dataset[0][0], (torch.zeros(1, 16, 16),), dataset[0][2], dataset[0][3])\n",
    "#     new_example, pos_mask = crop_transform((example, dataset[0][1]))\n",
    "#     visualize_matrices([dataset[0][0].argmax(0)])\n",
    "#     visualize_matrices([new_example[0].argmax(0)])\n",
    "#     visualize_matrices(new_example[1][0])\n",
    "#     visualize_matrices(torch.cat(pos_mask))\n",
    "    \n",
    "#     # Test composition\n",
    "#     composed = transforms.Compose([BabyARC2DColorTransform(), GenericRandomFlip(), BabyARC2DRandomResizedCrop(size=dataset[0][0].shape[-1])])\n",
    "#     example = (dataset[0][0], (torch.zeros(1, 16, 16),), dataset[0][2], dataset[0][3])\n",
    "#     new_example, pos_mask = composed((example, dataset[0][1]))\n",
    "#     visualize_matrices([dataset[0][0].argmax(0)])\n",
    "#     visualize_matrices([new_example[0].argmax(0)])\n",
    "#     visualize_matrices(new_example[1][0])\n",
    "#     visualize_matrices(torch.cat(pos_mask))\n",
    "    \n",
    "#     composed = transforms.Compose([BabyARC2DColorTransform(), GenericRandomFlip(), BabyARC2DRandomResizedCrop(size=dataset[0][0].shape[-1])])\n",
    "#     example = (dataset[1][0], (torch.zeros(1, 16, 16),), dataset[1][2], dataset[1][3])\n",
    "#     new_example, pos_mask = composed((example, dataset[1][1]))\n",
    "#     visualize_matrices([dataset[1][0].argmax(0)])\n",
    "#     visualize_matrices([new_example[0].argmax(0)])\n",
    "#     visualize_matrices(new_example[1][0])\n",
    "#     visualize_matrices(torch.cat(pos_mask))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     relation_args = init_args({\n",
    "#         \"dataset\": \"y-Parallel+VerticalMid+VerticalEdge\",\n",
    "#         \"seed\": 2,\n",
    "#         \"n_examples\": 40,\n",
    "#         \"canvas_size\": 16,\n",
    "#         \"rainbow_prob\": 0.,\n",
    "#         \"color_avail\": \"1,2\",\n",
    "#         \"w_type\": \"image+mask\",\n",
    "#         \"color_map_3d\": \"same\",\n",
    "#         \"add_thick_surf\": (0, 0.5),\n",
    "#         \"add_thick_depth\": (0, 0.5),\n",
    "#         \"max_n_distractors\": 2,\n",
    "#         \"seed_3d\": 42,\n",
    "#         \"num_processes_3d\": 10,\n",
    "#         \"image_size_3d\": (256,256),\n",
    "#         \"use_seed_2d\": False,\n",
    "#     })\n",
    "#     relation_dataset, args = get_dataset(relation_args, is_rewrite=False, is_load=True)\n",
    "#     # Test color permutation\n",
    "#     permute_transform = BabyARC3DColorTransform(s=0.5)\n",
    "#     example = (relation_dataset[0][0], (relation_dataset[0][1][1],relation_dataset[0][1][0]), relation_dataset[0][2], relation_dataset[0][3])\n",
    "#     new_example, pos_mask = permute_transform((example, relation_dataset[0][1]))\n",
    "#     visualize_matrices([relation_dataset[0][0]], use_color_dict=False)\n",
    "#     visualize_matrices([new_example[0]], use_color_dict=False)\n",
    "#     visualize_matrices(torch.cat(new_example[1]))\n",
    "#     visualize_matrices(torch.cat(pos_mask))\n",
    "    \n",
    "#     new_example, pos_mask = permute_transform((example, relation_dataset[0][1]))\n",
    "#     visualize_matrices([relation_dataset[0][0]], use_color_dict=False)\n",
    "#     visualize_matrices([new_example[0]], use_color_dict=False)\n",
    "    \n",
    "#     new_example, pos_mask = permute_transform((example, relation_dataset[0][1]))\n",
    "#     visualize_matrices([relation_dataset[0][0]], use_color_dict=False)\n",
    "#     visualize_matrices([new_example[0]], use_color_dict=False)\n",
    "    \n",
    "#     new_example, pos_mask = permute_transform((example, relation_dataset[0][1]))\n",
    "#     visualize_matrices([relation_dataset[0][0]], use_color_dict=False)\n",
    "#     visualize_matrices([new_example[0]], use_color_dict=False)\n",
    "\n",
    "#     # Test Hflip\n",
    "#     flip_transform = GenericRandomFlip()\n",
    "#     example = (relation_dataset[0][0], (relation_dataset[0][1][1],relation_dataset[0][1][0]), relation_dataset[0][2], relation_dataset[0][3])\n",
    "#     new_example, pos_mask = flip_transform((example, relation_dataset[0][1]))\n",
    "#     visualize_matrices([relation_dataset[0][0]], use_color_dict=False)\n",
    "#     visualize_matrices([new_example[0]], use_color_dict=False)\n",
    "#     visualize_matrices(torch.cat(new_example[1]))\n",
    "#     visualize_matrices(torch.cat(pos_mask))\n",
    "    \n",
    "#     # Test ResizeCrop\n",
    "#     crop_transform = BabyARC3DRandomResizedCrop(size=relation_dataset[0][0].shape[-1])\n",
    "#     example = (relation_dataset[0][0], (relation_dataset[0][1][0],torch.zeros(1, 256, 256)), relation_dataset[0][2], relation_dataset[0][3])\n",
    "#     new_example, pos_mask = crop_transform((example, relation_dataset[0][1]))\n",
    "#     visualize_matrices([relation_dataset[0][0]], use_color_dict=False)\n",
    "#     visualize_matrices([new_example[0]], use_color_dict=False)\n",
    "#     visualize_matrices(torch.cat(new_example[1]))\n",
    "#     visualize_matrices(torch.cat(pos_mask))\n",
    "    \n",
    "#     example = (relation_dataset[0][0], (torch.zeros(1, 256, 256),relation_dataset[0][1][0]), relation_dataset[0][2], relation_dataset[0][3])\n",
    "#     new_example, pos_mask = crop_transform((example, relation_dataset[0][1]))\n",
    "#     visualize_matrices([relation_dataset[0][0]], use_color_dict=False)\n",
    "#     visualize_matrices([new_example[0]], use_color_dict=False)\n",
    "#     visualize_matrices(torch.cat(new_example[1]))\n",
    "#     visualize_matrices(torch.cat(pos_mask))\n",
    "    \n",
    "#     new_example, pos_mask = crop_transform((example, relation_dataset[0][1]))\n",
    "#     visualize_matrices([relation_dataset[0][0]], use_color_dict=False)\n",
    "#     visualize_matrices([new_example[0]], use_color_dict=False)\n",
    "#     visualize_matrices(torch.cat(new_example[1]))\n",
    "#     visualize_matrices(torch.cat(pos_mask))\n",
    "\n",
    "#     # Test composed\n",
    "#     composed = transforms.Compose([BabyARC3DColorTransform(s=0.5), GenericRandomFlip(), BabyARC3DRandomResizedCrop(size=relation_dataset[0][0].shape[-1])])\n",
    "#     example = (relation_dataset[0][0], (torch.zeros(1, 256, 256),relation_dataset[0][1][0]), relation_dataset[0][2], relation_dataset[0][3])\n",
    "#     new_example, pos_mask = composed((example, relation_dataset[0][1]))\n",
    "#     visualize_matrices([relation_dataset[0][0]], use_color_dict=False)\n",
    "#     visualize_matrices([new_example[0]], use_color_dict=False)\n",
    "#     visualize_matrices(torch.cat(new_example[1]))\n",
    "#     visualize_matrices(torch.cat(pos_mask))\n",
    "    \n",
    "#     new_example, pos_mask = composed((example, relation_dataset[0][1]))\n",
    "#     visualize_matrices([relation_dataset[0][0]], use_color_dict=False)\n",
    "#     visualize_matrices([new_example[0]], use_color_dict=False)\n",
    "#     visualize_matrices(torch.cat(new_example[1]))\n",
    "#     visualize_matrices(torch.cat(pos_mask))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Sample buffer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class SampleBuffer(object):\n",
    "    def __init__(\n",
    "        self,\n",
    "        is_mask=False,\n",
    "        is_two_branch=False,\n",
    "        max_samples=10000,\n",
    "    ):\n",
    "        self.max_samples = max_samples\n",
    "        self.is_mask = is_mask\n",
    "        self.is_two_branch = is_two_branch\n",
    "        self.mask_arity = None\n",
    "        self.buffer = []\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.buffer)\n",
    "\n",
    "    def push(self, samples, masks=None, class_ids=None, c_reprs=None, infos=None):\n",
    "        if self.is_two_branch:\n",
    "            samples = (samples[0].detach().to('cpu'), samples[1].detach().to('cpu'))\n",
    "        else:\n",
    "            samples = samples.detach().to('cpu')\n",
    "        if self.is_mask:\n",
    "            masks = tuple(masks[k].detach().to('cpu') for k in range(len(masks)))\n",
    "            if self.mask_arity is None:\n",
    "                self.mask_arity = len(masks)\n",
    "            else:\n",
    "                assert self.mask_arity == len(masks)\n",
    "            c_reprs = c_reprs.detach().to('cpu')\n",
    "        else:\n",
    "            if isinstance(class_ids, torch.Tensor):\n",
    "                class_ids = class_ids.detach().to('cpu')\n",
    "\n",
    "        if self.is_mask:\n",
    "            assert class_ids is None\n",
    "            for i in range(len(c_reprs)):\n",
    "                if self.is_two_branch:\n",
    "                    assert len(samples) == 2\n",
    "                    sample = (samples[0][i], samples[1][i])\n",
    "                else:\n",
    "                    sample = samples[i]\n",
    "                mask = tuple(masks[k][i] for k in range(len(masks)))\n",
    "                c_repr = c_reprs[i]\n",
    "                info = infos[i]\n",
    "\n",
    "                self.buffer.append((sample, mask, c_repr, info))\n",
    "\n",
    "                if len(self.buffer) > self.max_samples:\n",
    "                    self.buffer.pop(0)\n",
    "        else:\n",
    "            assert c_reprs is None\n",
    "            for sample, class_id in zip(samples, class_ids):\n",
    "                self.buffer.append((sample.detach(), class_id))\n",
    "\n",
    "                if len(self.buffer) > self.max_samples:\n",
    "                    self.buffer.pop(0)\n",
    "\n",
    "    def get(self, n_samples, device='cuda'):\n",
    "        def stack_general(List, dim=0, device=device):\n",
    "            if isinstance(List[0], torch.Tensor):\n",
    "                return torch.stack(List, dim=dim).to(device)\n",
    "            elif isinstance(List[0], str):\n",
    "                return List\n",
    "            elif isinstance(List[0], Dictionary):\n",
    "                return List\n",
    "            else:\n",
    "                raise\n",
    "        items = random.choices(self.buffer, k=n_samples)\n",
    "        if self.is_mask:\n",
    "            \"\"\"\n",
    "            Zip(*[(('a',2), 1), (('b',3), 2), (('c',3), 3), (('d',2), 4)], function=function)\n",
    "                ==> [[function(['a', 'b', 'c', 'd']), function([2, 3, 3, 2])], function([1, 2, 3, 4])]\n",
    "            \"\"\"\n",
    "            samples, masks, c_reprs, infos = Zip(*items, function=stack_general)\n",
    "            return samples, masks, c_reprs, infos\n",
    "        else:\n",
    "            samples, class_ids = Zip(*items, function=stack_general)\n",
    "            return samples, class_ids\n",
    "\n",
    "\n",
    "def sample_buffer(buffer, in_channels, n_classes, image_size, batch_size=128, p=0.95, is_mask=True, is_two_branch=False, w_type=\"image+mask\", mask_arity=None, device='cuda'):\n",
    "    if isinstance(image_size, Number):\n",
    "        image_size = (image_size, image_size)\n",
    "    if is_mask:\n",
    "        w_dim = 1 if \"mask\" in w_type else in_channels\n",
    "        if len(buffer) < 1:\n",
    "            if buffer.mask_arity is None:\n",
    "                if mask_arity is None:\n",
    "                    buffer.mask_arity = 2 if is_two_branch else 1\n",
    "                else:\n",
    "                    buffer.mask_arity = mask_arity\n",
    "            if is_two_branch:\n",
    "                return (\n",
    "                    (torch.rand(batch_size, in_channels, *image_size, device=device), torch.rand(batch_size, in_channels, *image_size, device=device)),\n",
    "                    tuple(torch.rand(batch_size, w_dim, *image_size, device=device) for _ in range(buffer.mask_arity)),\n",
    "                    torch.rand(batch_size, REPR_DIM, device=device),\n",
    "                    [Dictionary()] * batch_size,\n",
    "                )\n",
    "            else:\n",
    "                return (\n",
    "                    torch.rand(batch_size, in_channels, *image_size, device=device),\n",
    "                    tuple(torch.rand(batch_size, w_dim, *image_size, device=device) for _ in range(buffer.mask_arity)),\n",
    "                    torch.rand(batch_size, REPR_DIM, device=device),\n",
    "                    [Dictionary()] * batch_size,\n",
    "                )\n",
    "\n",
    "        n_replay = (np.random.rand(batch_size) < p).sum()\n",
    "\n",
    "        replay_sample, replay_mask, replay_repr, replay_info = buffer.get(n_replay, device=device)\n",
    "        if is_two_branch:\n",
    "            random_sample = (torch.rand(batch_size - n_replay, in_channels, *image_size, device=device), torch.rand(batch_size - n_replay, in_channels, *image_size, device=device))\n",
    "        else:\n",
    "            random_sample = torch.rand(batch_size - n_replay, in_channels, *image_size, device=device)\n",
    "        random_mask = tuple(torch.rand(batch_size - n_replay, w_dim, *image_size, device=device) for _ in range(buffer.mask_arity))\n",
    "        random_repr = torch.rand(batch_size - n_replay, REPR_DIM, device=device)\n",
    "        random_info = [Dictionary()] * (batch_size - n_replay)\n",
    "\n",
    "        combined_mask = tuple(torch.cat([replay_mask[k], random_mask[k]], 0) for k in range(buffer.mask_arity))\n",
    "\n",
    "        if is_two_branch:\n",
    "            return (\n",
    "                (torch.cat([replay_sample[0], random_sample[0]], 0), torch.cat([replay_sample[1], random_sample[1]], 0)),\n",
    "                combined_mask,\n",
    "                torch.cat([replay_repr, random_repr], 0),\n",
    "                replay_info + random_info,\n",
    "            )\n",
    "        else:\n",
    "            return (\n",
    "                torch.cat([replay_sample, random_sample], 0),\n",
    "                combined_mask,\n",
    "                torch.cat([replay_repr, random_repr], 0),\n",
    "                replay_info + random_info,\n",
    "            )\n",
    "    else:\n",
    "        if len(buffer) < 1:\n",
    "            return (\n",
    "                torch.rand(batch_size, in_channels, *image_size, device=device),\n",
    "                torch.randint(0, n_classes, (batch_size,), device=device),\n",
    "            )\n",
    "\n",
    "        n_replay = (np.random.rand(batch_size) < p).sum()\n",
    "\n",
    "        replay_sample, replay_id = buffer.get(n_replay, device=device)\n",
    "        random_sample = torch.rand(batch_size - n_replay, in_channels, *image_size, device=device)\n",
    "        random_id = torch.randint(0, n_classes, (batch_size - n_replay,), device=device)\n",
    "\n",
    "        return (\n",
    "            torch.cat([replay_sample, random_sample], 0),\n",
    "            torch.cat([replay_id, random_id], 0),\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SampleBuffer_Conditional(object):\n",
    "    def __init__(\n",
    "        self,\n",
    "        is_two_branch=False,\n",
    "        max_samples=10000,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        self.buffer is a dictionary, and has the form of \n",
    "        {\n",
    "            (img_hash, class_id, ebm_target): (img, mask, c_repr, info),\n",
    "            ...\n",
    "        }\n",
    "            When img_hash is not None, it means that img == pos_img, and mask and c_repr have at least \n",
    "                one negative examples, and ebm_target should not have \"image\" (e.g. \"mask\", \"repr\", \"mask+repr\").\n",
    "            When class_id is not None, it means that c_repr == pos_repr, and both img and mask are negative\n",
    "                examples, and the ebm_target == \"image+mask\".\n",
    "        \"\"\"\n",
    "        self.max_samples = max_samples\n",
    "        self.is_two_branch = is_two_branch\n",
    "        self.mask_arity = None\n",
    "        self.buffer = {}\n",
    "\n",
    "    def __len__(self):\n",
    "        return sum([len(self.buffer[key]) for key in self.buffer])\n",
    "\n",
    "    def get_length(self, ebm_target):\n",
    "        return sum([len(value) for key, value in self.buffer.items() if key[2] == ebm_target])\n",
    "\n",
    "    def push(self, imgs=None, masks=None, c_reprs=None, class_ids=None, infos=None, ebm_target=None):\n",
    "        # Check validity:\n",
    "        if isinstance(imgs, tuple) or isinstance(imgs, list):\n",
    "            is_image_tuple = True\n",
    "            assert isinstance(imgs[0], torch.Tensor)\n",
    "        else:\n",
    "            is_image_tuple = False\n",
    "            assert isinstance(imgs, torch.Tensor)\n",
    "        # Detach and move to cpu:\n",
    "        imgs = to_device_recur(imgs, \"cpu\", is_detach=True)\n",
    "        masks = to_device_recur(masks, \"cpu\", is_detach=True)\n",
    "        if self.mask_arity is None:\n",
    "            self.mask_arity = len(masks)\n",
    "        else:\n",
    "            assert self.mask_arity == len(masks)\n",
    "        c_reprs = c_reprs.detach().to(\"cpu\")\n",
    "\n",
    "        for i in range(len(c_reprs)):\n",
    "            # (img_hash, class_id, ebm_target): (img, mask, c_repr)\n",
    "            if is_image_tuple:\n",
    "                img = (imgs[0][i], imgs[1][i])\n",
    "            else:\n",
    "                img = imgs[i]\n",
    "            img_hash = get_image_hashing(img) if ebm_target != \"image+mask\" else None\n",
    "            class_id = class_ids[i] if ebm_target == \"image+mask\" else None\n",
    "            mask = tuple(masks[k][i] for k in range(len(masks)))\n",
    "            c_repr = c_reprs[i]\n",
    "            info = infos[i]\n",
    "            record_data(self.buffer, [(img, mask, c_repr, info)], [(img_hash, class_id, ebm_target)], recent_record=self.max_samples)\n",
    "\n",
    "    def get(\n",
    "        self,\n",
    "        pos_data,\n",
    "        p=1,\n",
    "        ebm_target=None,\n",
    "        transforms=\"None\",\n",
    "        color_avail=None,\n",
    "        device='cuda',\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Given the positive examples, will sample the negative examples according to the ebm_target.\n",
    "            If ebm_target == \"image+mask\", the neg_examples will have the same c_repr as the positive examples.\n",
    "            Otherwise, will first sample from the key that has the same img_hash and ebm_target, and if not available,\n",
    "            sample random instance that has the same ebm_target.\n",
    "\n",
    "        Zip(*[(('a',2), 1), (('b',3), 2), (('c',3), 3), (('d',2), 4)], function=function)\n",
    "            ==> [[function(['a', 'b', 'c', 'd']), function([2, 3, 3, 2])], function([1, 2, 3, 4])]\n",
    "\n",
    "        Args:\n",
    "            p: probability of using the samples from buffer. Otherwise use random samples.\n",
    "\n",
    "        Returns:\n",
    "            neg_data: contains (neg_imgs, neg_masks, neg_reprs, neg_infos).\n",
    "        \"\"\"\n",
    "        def stack_general(List, dim=0, device=device):\n",
    "            if isinstance(List[0], torch.Tensor):\n",
    "                return torch.stack(List, dim=dim).to(device)\n",
    "            elif isinstance(List[0], str):\n",
    "                return List\n",
    "            elif isinstance(List[0], Dictionary):\n",
    "                return List\n",
    "            else:\n",
    "                raise\n",
    "        assert ebm_target is not None\n",
    "        pos_imgs, pos_masks, pos_ids, pos_infos = pos_data\n",
    "        neg_imgs, neg_masks, neg_reprs, neg_infos = [], [], [], []\n",
    "        if isinstance(pos_imgs, tuple) or isinstance(pos_imgs, list):\n",
    "            is_image_tuple = True\n",
    "            assert isinstance(pos_imgs[0], torch.Tensor)\n",
    "        else:\n",
    "            is_image_tuple = False\n",
    "            assert isinstance(pos_imgs, torch.Tensor)\n",
    "        if ebm_target == \"image+mask\":\n",
    "            # The c_repr is the same as that of the pos_data:\n",
    "            for i, pos_id in enumerate(pos_ids):\n",
    "                if np.random.rand() < p:\n",
    "                    # From buffer:\n",
    "                    lst = buffer.buffer[(None, pos_id, ebm_target)]\n",
    "                    id = np.random.choice(len(lst))\n",
    "                    # Apply transformations\n",
    "                    if transforms == \"None\":\n",
    "                        img, mask, c_repr, info = lst[id]\n",
    "                    else:\n",
    "                        transforms_split = transforms.split(\":\")\n",
    "                        if len(transforms_split) == 1:\n",
    "                            p_transforms = 1\n",
    "                            transforms_core = transforms\n",
    "                        else:\n",
    "                            p_transforms = eval(transforms_split[1])\n",
    "                            transforms_core = transforms_split[0]\n",
    "                        if np.random.rand() > p_transforms:\n",
    "                            img, mask, c_repr, info = lst[id]\n",
    "                        else:\n",
    "                            pos_mask = tuple(pos_masks[k][i] for k in range(len(pos_masks)))\n",
    "                            is_rgb = True if pos_imgs[0].shape[0] == 3 else False\n",
    "                            augment = get_augment(transforms_core, pos_imgs[0].shape[-1], is_rgb=is_rgb, color_avail=color_avail)\n",
    "                            img, mask, c_repr, info = augment((lst[id], pos_mask))[0]\n",
    "                    if i == 0:\n",
    "                        pos_repr = id_to_tensor([pos_id], CONCEPTS=CONCEPTS, OPERATORS=OPERATORS)[0]\n",
    "                        assert (c_repr == pos_repr).all()\n",
    "                    info[\"src\"] = \"buffer-same\"\n",
    "                else:\n",
    "                    # From random:\n",
    "                    img = (torch.rand_like(pos_imgs[0][0]), torch.rand_like(pos_imgs[1][0])) if is_image_tuple else torch.rand_like(pos_imgs[0])\n",
    "                    mask = tuple(torch.rand_like(pos_masks[0][0]) for _ in range(len(pos_masks)))\n",
    "                    c_repr = id_to_tensor([pos_id], CONCEPTS=CONCEPTS, OPERATORS=OPERATORS)[0]\n",
    "                    info = Dictionary({\"src\": \"random\"})\n",
    "                neg_imgs.append(img)\n",
    "                neg_masks.append(mask)\n",
    "                neg_reprs.append(c_repr)\n",
    "                neg_infos.append(info)\n",
    "        else:\n",
    "            # The neg_img is the same as that of the pos_data:\n",
    "            for i in range(len(pos_ids)):\n",
    "                if is_image_tuple:\n",
    "                    img = (pos_imgs[0][i], pos_imgs[1][i])\n",
    "                else:\n",
    "                    img = pos_imgs[i]\n",
    "                if np.random.rand() < p:\n",
    "                    # From buffer:\n",
    "                    img_hash = get_image_hashing(img)\n",
    "                    key = (img_hash, None, ebm_target)\n",
    "                    if key in buffer.buffer:\n",
    "                        # Find the items that has the same img_hash and ebm_target as the pos_data:\n",
    "                        lst = buffer.buffer[key]\n",
    "                    else:\n",
    "                        # Find the items that has the same ebm_target as the given ebm_target:\n",
    "                        lst = list(itertools.chain.from_iterable([\n",
    "                            item for key, item in buffer.buffer.items() if key[2] == ebm_target]))\n",
    "                    id = np.random.choice(len(lst))\n",
    "                    # Apply transformations\n",
    "                    if transforms == \"None\":\n",
    "                        img_load, mask, c_repr, info = lst[id]\n",
    "                        if i == 0 and key in buffer.buffer:\n",
    "                            assert img_hash == get_image_hashing(img_load)\n",
    "                    else:\n",
    "                        transforms_split = transforms.split(\":\")\n",
    "                        if len(transforms_split) == 1:\n",
    "                            p_transforms = 1\n",
    "                            transforms_core = transforms\n",
    "                        else:\n",
    "                            p_transforms = eval(transforms_split[1])\n",
    "                            transforms_core = transforms_split[0]\n",
    "                        if np.random.rand() > p_transforms:\n",
    "                            img_load, mask, c_repr, info = lst[id]\n",
    "                        else:\n",
    "                            pos_mask = tuple(pos_masks[k][i] for k in range(len(pos_masks)))\n",
    "                            is_rgb = True if pos_imgs[0].shape[0] == 3 else False\n",
    "                            augment = get_augment(transforms_core, pos_imgs[0].shape[-1], is_rgb=is_rgb, color_avail=color_avail)\n",
    "                            img, mask, c_repr, info = augment((lst[id], pos_mask))[0]\n",
    "                    if key in buffer.buffer:\n",
    "                        info[\"src\"] = \"buffer-same\"\n",
    "                    else:\n",
    "                        info[\"src\"] = \"buffer-diff\"\n",
    "                else:\n",
    "                    # From random:\n",
    "                    if \"mask\" in ebm_target:\n",
    "                        mask = tuple(torch.rand_like(pos_masks[0][0]) for _ in range(len(pos_masks)))\n",
    "                    else:\n",
    "                        # If \"mask\" is not in ebm_target, then use the ground-truth:\n",
    "                        mask = tuple(pos_masks[k][i] for k in range(len(pos_masks)))\n",
    "                    if \"repr\" in ebm_target:\n",
    "                        c_repr = torch.rand(REPR_DIM)\n",
    "                    else:\n",
    "                        # If \"repr\" is not in ebm_target, then use the ground-truth:\n",
    "                        c_repr = id_to_tensor([pos_ids[i]], CONCEPTS=CONCEPTS, OPERATORS=OPERATORS)[0]\n",
    "                    info = Dictionary({\"src\": \"random\"})\n",
    "                neg_imgs.append(img)\n",
    "                neg_masks.append(mask)\n",
    "                neg_reprs.append(c_repr)\n",
    "                neg_infos.append(info)\n",
    "        if is_image_tuple:\n",
    "            neg_imgs = Zip(*neg_imgs, function=stack_general)\n",
    "        else:\n",
    "            neg_imgs = torch.stack(neg_imgs).to(device)\n",
    "        neg_masks = Zip(*neg_masks, function=stack_general)\n",
    "        neg_reprs = torch.stack(neg_reprs).to(device)\n",
    "        neg_data = (neg_imgs, neg_masks, neg_reprs, neg_infos)\n",
    "        return neg_data\n",
    "\n",
    "    def __repr__(self):\n",
    "        return \"SampleBuffer_Conditional(keys={}, examples={})\".format(len(self.buffer), len(self))\n",
    "\n",
    "\n",
    "def sample_buffer_conditional(\n",
    "    buffer,\n",
    "    pos_data,\n",
    "    ebm_target,\n",
    "    in_channels,\n",
    "    image_size,\n",
    "    batch_size=128,\n",
    "    p=0.95,\n",
    "    is_two_branch=False,\n",
    "    is_image_tuple=False,\n",
    "    w_type=\"image+mask\",\n",
    "    transforms=\"None\",\n",
    "    color_avail=\"-1\",\n",
    "    device='cuda',\n",
    "):\n",
    "    \"\"\"\n",
    "    Given the positive examples, will sample the negative examples according to the ebm_target.\n",
    "    E.g. if ebm_target == \"image+mask\", the neg_examples will have the same c_repr as the positive examples.\n",
    "    \"\"\"\n",
    "    if isinstance(image_size, Number):\n",
    "        image_size = (image_size, image_size)\n",
    "    w_dim = 1 if \"mask\" in w_type else in_channels\n",
    "    if buffer.get_length(ebm_target) <= len(pos_data[2]):\n",
    "        if buffer.mask_arity is None:\n",
    "            buffer.mask_arity = 2 if is_two_branch else 1\n",
    "        random_img = (torch.rand(batch_size, in_channels, *image_size, device=device),\n",
    "                      torch.rand(batch_size, in_channels, *image_size, device=device)) if is_image_tuple else \\\n",
    "                    torch.rand(batch_size, in_channels, *image_size, device=device)\n",
    "        return (\n",
    "            random_img,\n",
    "            tuple(torch.rand(batch_size, w_dim, *image_size, device=device) for _ in range(buffer.mask_arity)),\n",
    "            torch.rand(batch_size, REPR_DIM, device=device),\n",
    "            [Dictionary()] * batch_size,\n",
    "        )\n",
    "    else:\n",
    "        return buffer.get(pos_data, p=p, ebm_target=ebm_target, transforms=transforms, color_avail=color_avail, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if __name__ == \"__main__\":\n",
    "#     imgs=neg_img\n",
    "#     masks=neg_mask\n",
    "#     class_ids=[\"Hshape\"] * 128\n",
    "#     c_reprs=neg_repr\n",
    "#     infos=neg_info\n",
    "#     ebm_target=\"image+mask\"\n",
    "\n",
    "#     buffer = SampleBuffer_Conditional()\n",
    "#     buffer.push(imgs, masks, c_reprs, class_ids, infos, ebm_target)\n",
    "\n",
    "#     sample = buffer.get(pos_data, p=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Negative examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_neg_examples(pos_img, pos_mask, pos_repr, neg_mode):\n",
    "    \"\"\"Generate negative examples according to the neg_mode.\n",
    "\n",
    "    Args:\n",
    "        neg_mode: choose from\n",
    "                    \"addrand\" (add random pixels from the image),\n",
    "                    \"permlabel\" (randomly permute the pos_repr),\n",
    "                    \"addallrand\" (add random pixels on the empty place of the positive masks),\n",
    "                    or using \"+\" to combine any subset of them.\n",
    "\n",
    "    Returns:\n",
    "        neg_mask_gen: generated negative masks\n",
    "        neg_repr_gen: generated (permuted c_repr)\n",
    "        neg_gen_valid: valid negative examples\n",
    "    \"\"\"\n",
    "    if \"+\" in neg_mode:\n",
    "        # multiple modes, randomly assign one mode for each example:\n",
    "        neg_mode_split = neg_mode.split(\"+\")\n",
    "        is_two_branch = isinstance(pos_img, tuple) or isinstance(pos_img, list)\n",
    "        device = pos_repr.device\n",
    "        length = len(pos_img[0]) if is_two_branch else len(pos_img)\n",
    "        id_assign = torch.randint(len(neg_mode_split), size=(length,))\n",
    "        neg_mask_gen = tuple(-torch.ones_like(pos_mask[k]).to(device)\n",
    "                             for k in range(len(pos_mask)))\n",
    "        neg_repr_gen = -torch.ones_like(pos_repr).to(device)\n",
    "        neg_gen_valid = -torch.ones(length, 1).to(device)\n",
    "\n",
    "        for i in range(len(neg_mode_split)):\n",
    "            id_assign_i = id_assign == i\n",
    "            if is_two_branch:\n",
    "                pos_img_ele = (pos_img[0][id_assign_i],\n",
    "                               pos_img[1][id_assign_i])\n",
    "            else:\n",
    "                pos_img_ele = pos_img[id_assign_i]\n",
    "            pos_mask_ele = tuple(pos_mask[k][id_assign_i]\n",
    "                                 for k in range(len(pos_mask)))\n",
    "            pos_repr_ele = pos_repr[id_assign_i]\n",
    "            neg_mask_gen_i, neg_repr_gen_i, neg_gen_valid_i = generate_neg_examples(\n",
    "                pos_img_ele,\n",
    "                pos_mask_ele,\n",
    "                pos_repr_ele,\n",
    "                neg_mode=neg_mode_split[i],  # use the assigned mode\n",
    "            )\n",
    "            for k in range(len(neg_mask_gen)):\n",
    "                neg_mask_gen[k][id_assign_i] = neg_mask_gen_i[k]\n",
    "            neg_repr_gen[id_assign_i] = neg_repr_gen_i\n",
    "            neg_gen_valid[id_assign_i] = neg_gen_valid_i\n",
    "\n",
    "        # Make sure that all examples are operated on:\n",
    "        for k in range(len(neg_mask_gen)):\n",
    "            assert (neg_mask_gen[k] > -1).all()\n",
    "        if \"softmax\" not in args.c_repr_mode and args.is_pos_repr_learnable is False:\n",
    "            assert (neg_repr_gen > -1).all()\n",
    "        assert (neg_gen_valid > -1).all()\n",
    "        # Returns:\n",
    "        return neg_mask_gen, neg_repr_gen, neg_gen_valid\n",
    "\n",
    "    else:\n",
    "        # Single mode:\n",
    "        if neg_mode == \"addrand\":\n",
    "            # Randomly add pixels to the other places of the image where there is object:\n",
    "            if len(pos_mask) > 1:\n",
    "                if not (isinstance(pos_img, tuple) or isinstance(pos_img, list)):\n",
    "                    pos_img = (pos_img, pos_img)\n",
    "                assert pos_img[0].shape[1] == 10, \"if channel_size is 10, then cannot use neg_mask_mode of 'addrand'.\"\n",
    "                # [B, 1, H, W]; pos_mask[0]: [B, 1, H, W]\n",
    "                pos_non_zero = (pos_img[0][:, :1] != 1, pos_img[1][:, :1] != 1)\n",
    "                remainder = (pos_non_zero[0] & (\n",
    "                    ~pos_mask[0].bool()), pos_non_zero[1] & (~pos_mask[1].bool()))\n",
    "                # valid if the remainder is not all-zero:\n",
    "                neg_gen_valid = ((remainder[0].sum((1, 2, 3)) > 0) | (\n",
    "                    remainder[1].sum((1, 2, 3)) > 0)).unsqueeze(1).float()\n",
    "                remainder_random = (remainder[0]*torch.randint(2, size=remainder[0].shape).to(\n",
    "                    remainder[0].device), remainder[1]*torch.randint(2, size=remainder[0].shape).to(remainder[1].device))\n",
    "                neg_mask_gen = (\n",
    "                    pos_mask[0]+remainder_random[0], pos_mask[1]+remainder_random[1])\n",
    "                neg_repr_gen = pos_repr\n",
    "            else:\n",
    "                assert pos_img.shape[1] == 10, \"if channel_size is 10, then cannot use neg_mask_mode of 'addrand'.\"\n",
    "                assert len(pos_mask) == 1\n",
    "                # [B, 1, H, W]; pos_mask: [B, 1, H, W]\n",
    "                pos_non_zero = pos_img[:, :1] != 1\n",
    "                remainder = pos_non_zero & (~pos_mask[0].bool())\n",
    "                # valid if the remainder is not all-zero:\n",
    "                neg_gen_valid = (remainder.sum((1, 2, 3))\n",
    "                                 > 0).unsqueeze(1).float()\n",
    "                remainder_random = remainder * \\\n",
    "                    torch.randint(2, size=remainder.shape).to(remainder.device)\n",
    "                neg_mask_gen = (pos_mask[0] + remainder_random,)\n",
    "                neg_repr_gen = pos_repr\n",
    "        elif neg_mode == \"addallrand\":\n",
    "            # Randomly add pixels to the other places of the image where there is no object:\n",
    "            if len(pos_mask) > 1:\n",
    "                if not (isinstance(pos_img, tuple) or isinstance(pos_img, list)):\n",
    "                    pos_img = (pos_img, pos_img)\n",
    "                # [B, 1, H, W]; pos_mask[0]: [B, 1, H, W]\n",
    "                remainder = (~pos_mask[0].bool(), ~pos_mask[1].bool())\n",
    "                p0, p1 = pos_mask[0].mean().item(), pos_mask[1].mean().item()\n",
    "                remainder_random = ((torch.rand(pos_mask[0].shape, device=pos_mask[0].device) < p0).float(),\n",
    "                                    (torch.rand(pos_mask[1].shape, device=pos_mask[1].device) < p1).float())\n",
    "                remainder_random = (remainder[0]*remainder_random[0], remainder[1]*remainder_random[1])\n",
    "\n",
    "                # valid if the remainder is not all-zero:\n",
    "                neg_gen_valid = ((remainder_random[0].sum((1, 2, 3)) > 0) | (\n",
    "                    remainder_random[1].sum((1, 2, 3)) > 0)).unsqueeze(1).float()\n",
    "                neg_mask_gen = (\n",
    "                    pos_mask[0]+remainder_random[0], pos_mask[1]+remainder_random[1])\n",
    "                neg_repr_gen = pos_repr\n",
    "            else:\n",
    "                assert len(pos_mask) == 1\n",
    "                # [B, 1, H, W]; pos_mask: [B, 1, H, W]\n",
    "                remainder = ~pos_mask[0].bool()\n",
    "                p0 = pos_mask[0].mean().item()\n",
    "                # valid if the remainder is not all-zero:\n",
    "                remainder_random = (torch.rand(pos_mask[0].shape, device=pos_mask[0].device) < p0).float() * remainder\n",
    "                neg_gen_valid = (remainder_random.sum((1, 2, 3))\n",
    "                                 > 0).unsqueeze(1).float()\n",
    "                neg_mask_gen = (pos_mask[0] + remainder_random,)\n",
    "                neg_repr_gen = pos_repr\n",
    "        elif neg_mode == \"delrand\":\n",
    "            if len(pos_mask) > 1:\n",
    "                remove_random = (pos_mask[0]*torch.randint(2, size=pos_mask[0].shape).to(pos_mask[0].device),\n",
    "                                 pos_mask[1]*torch.randint(2, size=pos_mask[0].shape).to(pos_mask[1].device))\n",
    "                # valid if the remove_random is not all-zero:\n",
    "                neg_gen_valid = ((remove_random[0].sum((1, 2, 3)) > 0) | (\n",
    "                    remove_random[1].sum((1, 2, 3)) > 0)).unsqueeze(1).float()\n",
    "                neg_mask_gen = (\n",
    "                    pos_mask[0]-remove_random[0], pos_mask[1]-remove_random[1])\n",
    "                neg_repr_gen = pos_repr\n",
    "            else:\n",
    "                assert len(pos_mask) == 1\n",
    "                # pos_mask: [B, 1, H, W]\n",
    "                remove_random = pos_mask[0]*torch.randint(2, size=pos_mask[0].shape).to(pos_mask[0].device)\n",
    "\n",
    "                # valid if the remainder is not all-zero:\n",
    "                neg_gen_valid = (remove_random.sum((1, 2, 3)) > 0).unsqueeze(1).float()\n",
    "                neg_mask_gen = (pos_mask[0] - remove_random,)\n",
    "                neg_repr_gen = pos_repr\n",
    "        elif neg_mode == \"permlabel\":\n",
    "            # Permute pos_repr:\n",
    "            id_rand = torch.randperm(len(pos_repr))\n",
    "            neg_repr_gen = pos_repr[id_rand]\n",
    "            # Valid if the c_repr is different:\n",
    "            neg_gen_valid = (neg_repr_gen != pos_repr).any(\n",
    "                1, keepdims=True).float()\n",
    "            neg_mask_gen = pos_mask\n",
    "        else:\n",
    "            raise Exception(\"neg_mode '{}' is not supported!\".format(neg_mode))\n",
    "        return neg_mask_gen, neg_repr_gen, neg_gen_valid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Loss function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_loss_core(args, pos_out, neg_out, emp_out=None, neg_out_gen=None):\n",
    "    \"\"\"\n",
    "    Compute the core loss from pos_out, neg_out, emp_out, based on the energy_mode.\n",
    "    If the loss involves emp_out, will only be valid if is_emp_loss(args) is True.\n",
    "\n",
    "    Args:\n",
    "        pos_out, neg_out, emp_out, neg_out_gen: tensors of the same shape, typically [B, 1].\n",
    "        energy_mode:\n",
    "            \"standard:0.3\": (E_pos - E_neg) * 0.3\n",
    "            \"margin^0.2:0.3\": max(0, 0.3 + E_pos - E_neg) * 0.2\n",
    "            \"mid^0.2:0.3\": (max(0, 0.2 + E_pos - E_empty) + max(0, 0.2 + E_empty - E_neg)) * 0.3\n",
    "            \"mid^0.2^adapt:0.3\":\n",
    "                (max(0, gamma + E_pos - E_empty) + max(0, gamma + E_empty - E_neg)) * 0.3\n",
    "                    where gamma = max(0, StopGrad(E_neg - E_pos)/2) + 0.2\n",
    "            \"standard:0.5+mid^0.2^adapt:0.3\":\n",
    "                (E_pos - E_neg) * 0.5 + (max(0, gamma + E_pos - E_empty) + max(0, gamma + E_empty - E_neg)) * 0.3,\n",
    "                    where gamma = max(0, StopGrad(E_neg - E_pos)/2) + 0.2\n",
    "            \"standard+center^stop\": (E_pos - E_neg) * 1 + ((E_pos+E_neg).detach()/2 - E_empty).abs()\n",
    "                \"stop\": stop gradient, and each empty loss is computed per example\n",
    "                \"stopgen\": similar to \"stop\", but the negative energy is the mean of neg_out and neg_out_gen, per example.\n",
    "                \"stopmean\": stop gradient, and each empty loss is computed per minibatch\n",
    "                \"stopgenmean\": similar to \"stopmean\", but the negative energy is the mean of neg_out and neg_out_gen, per minibatch.\n",
    "    Returns:\n",
    "        loss_core: total loss, in the same shape as pos_out, etc.\n",
    "        loss_core_info: dictionary of each of the average loss for the component of the energy_mode.\n",
    "    \"\"\"\n",
    "    loss_core = 0\n",
    "    loss_core_info = {}\n",
    "    device = pos_out.device\n",
    "    for energy_mode_ele in args.energy_mode.split(\"+\"):\n",
    "        if len(energy_mode_ele.split(\":\")) > 1:\n",
    "            coef = eval(energy_mode_ele.split(\":\")[-1])\n",
    "        else:\n",
    "            coef = 1\n",
    "        energy_mode_ele = energy_mode_ele.split(\":\")[0]\n",
    "        if energy_mode_ele.startswith(\"standard\"):\n",
    "            loss_ele = (pos_out - neg_out) * coef\n",
    "        elif energy_mode_ele.startswith(\"margin\"):\n",
    "            gamma = eval(energy_mode_ele.split(\"^\")[1])\n",
    "            loss_ele = torch.maximum(torch.tensor(0).to(device), pos_out - neg_out + gamma) * coef\n",
    "        elif energy_mode_ele.startswith(\"mid\"):\n",
    "            if not is_emp_loss(args):\n",
    "                continue\n",
    "            gamma = eval(energy_mode_ele.split(\"^\")[1])\n",
    "            is_adapt = \"adapt\" in energy_mode_ele\n",
    "            if is_adapt:\n",
    "                gamma = gamma + max(0, (neg_out - pos_out).mean().item() / 2)\n",
    "            loss_ele = torch.maximum(torch.tensor(0).to(device), pos_out - emp_out + gamma) * coef + \\\n",
    "                       torch.maximum(torch.tensor(0).to(device), emp_out - neg_out + gamma) * coef\n",
    "        elif energy_mode_ele.startswith(\"center\"):\n",
    "            if not is_emp_loss(args):\n",
    "                continue\n",
    "            grad_type = energy_mode_ele.split(\"^\")[1] if len(energy_mode_ele) > 1 else \"None\"\n",
    "            if grad_type == \"None\":\n",
    "                loss_ele = ((pos_out + neg_out) / 2 - emp_out).abs() * coef\n",
    "            elif grad_type == \"stop\":\n",
    "                loss_ele = ((pos_out + neg_out).detach() / 2 - emp_out).abs() * coef\n",
    "            elif grad_type == \"stopsq\":\n",
    "                loss_ele = ((pos_out + neg_out).detach() / 2 - emp_out).square() * coef\n",
    "            elif grad_type == \"stopgen\":\n",
    "                if neg_out_gen is not None:\n",
    "                    loss_ele = ((pos_out + (neg_out + neg_out_gen)/2).mean() / 2 - emp_out).abs() * coef\n",
    "                else:\n",
    "                    loss_ele = ((pos_out + neg_out).mean() / 2 - emp_out).abs() * coef\n",
    "            elif grad_type == \"stopmean\":\n",
    "                loss_ele = ((pos_out + neg_out).detach().mean() / 2 - emp_out).abs() * coef\n",
    "            elif grad_type == \"stopsqmean\":\n",
    "                loss_ele = ((pos_out + neg_out).detach().mean() / 2 - emp_out).square() * coef\n",
    "            elif grad_type == \"stopgenmean\":\n",
    "                if neg_out_gen is not None:\n",
    "                    loss_ele = ((pos_out + (neg_out + neg_out_gen)/2).detach().mean() / 2 - emp_out).abs() * coef\n",
    "                else:\n",
    "                    loss_ele = ((pos_out + neg_out).detach().mean() / 2 - emp_out).abs() * coef\n",
    "            else:\n",
    "                raise\n",
    "        else:\n",
    "            raise\n",
    "        loss_core = loss_core + loss_ele\n",
    "        loss_core_info[energy_mode_ele.split(\"^\")[0]] = loss_ele.mean().item()\n",
    "    return loss_core, loss_core_info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.5 Helper functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_image_hashing(img):\n",
    "    \"\"\"Get the hashing of an image.\"\"\"\n",
    "    if isinstance(img, tuple) or isinstance(img, list):\n",
    "        hashing = get_hashing(str(torch.stack(img).long().view(-1)), length=10)\n",
    "    else:\n",
    "        hashing = get_hashing(str(img.long().view(-1)), length=10)\n",
    "    return hashing\n",
    "\n",
    "\n",
    "def get_concept_embeddings(CONCEPTS, OPERATORS):\n",
    "    concept_embeddings = {key: to_np_array(CONCEPTS[key].get_node_repr()) for key in CONCEPTS}\n",
    "    operator_embeddings = {key: to_np_array(OPERATORS[key].get_node_repr()) for key in OPERATORS}\n",
    "    concept_embeddings.update(operator_embeddings)\n",
    "    return concept_embeddings\n",
    "\n",
    "\n",
    "def unittest(model, args, device, CONCEPTS, OPERATORS):\n",
    "    \"\"\"Make sure that the loaded model is the same as the original model.\"\"\"\n",
    "    if isinstance(model, nn.parallel.DataParallel):\n",
    "        model = model.module\n",
    "    model.eval()\n",
    "    buffer = SampleBuffer()\n",
    "    neg_data = sample_buffer(\n",
    "        buffer,\n",
    "        in_channels=args.in_channels,\n",
    "        n_classes=args.n_classes,\n",
    "        image_size=args.image_size if args.rescaled_size == \"None\" else eval(args.rescaled_size),\n",
    "        batch_size=args.batch_size,\n",
    "        is_mask=args.is_mask,\n",
    "        is_two_branch=args.is_two_branch,\n",
    "        w_type=args.w_type,\n",
    "        p=args.p_buffer,\n",
    "        device=device,\n",
    "    )\n",
    "    model2 = load_model_energy(model.model_dict, device=device)\n",
    "    model2.eval()\n",
    "    if args.is_mask:\n",
    "        neg_img, neg_mask, neg_repr, neg_info = neg_data\n",
    "        neg_out = model(neg_img, mask=neg_mask, c_repr=neg_repr)\n",
    "        neg_out2 = model2(neg_img, mask=neg_mask, c_repr=neg_repr)\n",
    "    else:\n",
    "        neg_img, neg_id = neg_data\n",
    "        neg_out = model(neg_img, neg_id)\n",
    "        neg_out2 = model2(neg_img, neg_id)\n",
    "    diff_max = (neg_out2 - neg_out).abs().max().item()\n",
    "\n",
    "    if diff_max < 8e-6:\n",
    "        p.print(\"The largest diff for sample neg_img is {}, smaller than 8e-6. Unittest passed!\".format(diff_max))\n",
    "    else:\n",
    "        raise Exception(\"The largest diff for sample neg_img is {}, greater than 8e-6. Check model loading and saving!\".format(diff_max))\n",
    "\n",
    "\n",
    "def get_filename(args, short_str_dict, is_local_path):\n",
    "    filename_short = get_filename_short(\n",
    "        short_str_dict.keys(),\n",
    "        short_str_dict,\n",
    "        args_dict=args.__dict__,\n",
    "    )\n",
    "    dirname = EXP_PATH + \"/{}_{}/\".format(args.exp_id, args.date_time)\n",
    "    if args.exp_name != \"None\":\n",
    "        # If args.exp_name != \"None\", the experiments are saved under \"{exp_id}_{date_time}/{exp_name}/\"\n",
    "        dirname += \"{}/\".format(args.exp_name)\n",
    "    filename = dirname + filename_short[:-2] + \"_{}.p\".format(get_machine_name())\n",
    "    make_dir(filename)\n",
    "    p.print(filename, banner_size=100)\n",
    "    return dirname, filename\n",
    "\n",
    "\n",
    "def get_ebm_target(args):\n",
    "    if args.ebm_target_mode.startswith(\"r-\"):\n",
    "        # randomly sample mode:\n",
    "        ebm_str = args.ebm_target_mode.split(\"-\")[1]\n",
    "        ebm_str_dict = {\"r\": \"repr\", \"m\": \"mask\", \"b\": \"mask+repr\", \"x\": \"image+mask\"}\n",
    "        ebm_target_collection = [ebm_str_dict[key] for key in ebm_str]\n",
    "        args.ebm_target = np.random.choice(ebm_target_collection)\n",
    "    elif args.ebm_target_mode == \"None\":\n",
    "        pass\n",
    "    else:\n",
    "        raise\n",
    "    return args\n",
    "\n",
    "\n",
    "def is_emp_loss(args):\n",
    "    \"\"\"Return True if the current args.ebm_target belongs to the args.emp_target.\"\"\"\n",
    "    if args.emp_target_mode == \"all\":\n",
    "        emp_target_mode = args.ebm_target_mode.split(\"-\")[1]\n",
    "    else:\n",
    "        assert args.emp_target_mode.startswith(\"r-\")\n",
    "        emp_target_mode = args.emp_target_mode.split(\"-\")[1]\n",
    "    ebm_str_reverse_dict = {\"repr\": \"r\", \"mask\": \"m\", \"mask+repr\": \"b\", \"image+mask\": \"x\"}\n",
    "    current_str = ebm_str_reverse_dict[args.ebm_target]\n",
    "    return current_str in emp_target_mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_concepts(n_concepts, n_relations):\n",
    "    from zeroc.concept_library.concepts import Concept, Graph, Placeholder, Tensor\n",
    "    CONCEPTS = OrderedDict()\n",
    "    OPERATORS = OrderedDict()\n",
    "    IS_CUDA = False\n",
    "    num_colors = 10\n",
    "    CONCEPTS[\"Image\"] = Concept(name=\"Image\",\n",
    "        repr=to_Variable(torch.rand(REPR_DIM), is_cuda=IS_CUDA),\n",
    "        inherit_to=[f\"c{i}\" for i in range(n_concepts)],\n",
    "        value=Placeholder(Tensor(dtype=\"cat\", range=range(num_colors))),\n",
    "    )\n",
    "\n",
    "    CONCEPTS[\"Bool\"] = Concept(name=\"Bool\",\n",
    "        repr=to_Variable(torch.rand(REPR_DIM), is_cuda=IS_CUDA),\n",
    "        value=Placeholder(Tensor(dtype=\"bool\", shape=(1,), range=[True, False])),\n",
    "    )\n",
    "\n",
    "    for i in range(n_concepts):\n",
    "        CONCEPTS[f\"c{i}\"] = Concept(name=f\"c{i}\",\n",
    "            repr=to_Variable(torch.rand(REPR_DIM), is_cuda=IS_CUDA),\n",
    "            inherit_from=[\"Image\"],\n",
    "            value=Placeholder(Tensor(dtype=\"cat\", range=range(num_colors))),\n",
    "        )\n",
    "\n",
    "    for i in range(n_relations):\n",
    "        OPERATORS[f\"r{i}\"] = Graph(name=f\"r{i}\",\n",
    "            repr=to_Variable(torch.rand(REPR_DIM), is_cuda=IS_CUDA),\n",
    "            forward={\"args\": [Placeholder(\"Image\"), Placeholder(\"Image\")],\n",
    "                     \"output\": Placeholder(\"Bool\"),\n",
    "                     \"fun\": identity_fun,\n",
    "                    })\n",
    "    return CONCEPTS, OPERATORS\n",
    "\n",
    "\n",
    "def init_concepts_with_repr(concept_repr_dict=None, relation_repr_dict=None):\n",
    "    from zeroc.concept_library.concepts import Concept, Graph, Placeholder, Tensor\n",
    "    CONCEPTS = OrderedDict()\n",
    "    OPERATORS = OrderedDict()\n",
    "    IS_CUDA = False\n",
    "    num_colors = 10\n",
    "    CONCEPTS[\"Image\"] = Concept(name=\"Image\",\n",
    "        repr=to_Variable(torch.rand(REPR_DIM), is_cuda=IS_CUDA),\n",
    "        inherit_to=list(concept_repr_dict.keys()),\n",
    "        value=Placeholder(Tensor(dtype=\"cat\", range=range(num_colors))),\n",
    "    )\n",
    "\n",
    "    CONCEPTS[\"Bool\"] = Concept(name=\"Bool\",\n",
    "        repr=to_Variable(torch.rand(REPR_DIM), is_cuda=IS_CUDA),\n",
    "        value=Placeholder(Tensor(dtype=\"bool\", shape=(1,), range=[True, False])),\n",
    "    )\n",
    "\n",
    "    if concept_repr_dict is not None:\n",
    "        for c_str, c_repr in concept_repr_dict.items():\n",
    "            CONCEPTS[c_str] = Concept(name=c_str,\n",
    "                repr=c_repr,\n",
    "                inherit_from=[\"Image\"],\n",
    "                value=Placeholder(Tensor(dtype=\"cat\", range=range(num_colors))),\n",
    "            )\n",
    "    if relation_repr_dict is not None:\n",
    "        for c_str, c_repr in relation_repr_dict.items():\n",
    "            OPERATORS[c_str] = Graph(name=c_str,\n",
    "                repr=c_repr,\n",
    "                forward={\"args\": [Placeholder(\"Image\"), Placeholder(\"Image\")],\n",
    "                         \"output\": Placeholder(\"Bool\"),\n",
    "                         \"fun\": identity_fun,\n",
    "                        }\n",
    "            )\n",
    "    return CONCEPTS, OPERATORS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Test_acc:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_arc_relations(inp_img, info):\n",
    "    \"\"\"Returns a list of relations in the format of (id1, id2, relation_type),\n",
    "    where relation_type is SameShape, SameColor, IsInside, etc.\n",
    "    \"\"\"\n",
    "    # Convert masks into objects\n",
    "    obj_lst = []\n",
    "    for i in range(2):\n",
    "        obj_id = f\"obj_{i}\"\n",
    "        mask = info['id_object_mask'][info['node_id_map'][obj_id]]\n",
    "        obj_value, obj_pos = shrink(get_obj_from_mask(inp_img.argmax(0), mask))\n",
    "        obj_lst.append((i, CONCEPTS[DEFAULT_OBJ_TYPE].copy().set_node_value(obj_value).set_node_value(obj_pos, \"pos\")))\n",
    "\n",
    "    relations = []\n",
    "    for idx1 in range(len(obj_lst)):\n",
    "        for idx2 in range(idx1 + 1, len(obj_lst)):\n",
    "            # Use functions defined in concepts_ARC2\n",
    "            rel_types = []\n",
    "            id1, obj1 = obj_lst[idx1]\n",
    "            id2, obj2 = obj_lst[idx2]\n",
    "            if SameShape(obj1, obj2):\n",
    "                relations.append((id1, id2, \"SameShape\"))\n",
    "            if SameColor(obj1, obj2):\n",
    "                relations.append((id1, id2, \"SameColor\"))\n",
    "            if SameAll(obj1, obj2):\n",
    "                relations.append((id1, id2, \"SameAll\"))\n",
    "            if SameRow(obj1, obj2):\n",
    "                relations.append((id1, id2, \"SameRow\"))\n",
    "            if SameCol(obj1, obj2):\n",
    "                relations.append((id1, id2, \"SameCol\"))\n",
    "            if IsNonOverlapXY(obj1, obj2):\n",
    "                relations.append((id1, id2, \"IsNonOverlapXY\"))\n",
    "            # Uncommutative relations\n",
    "            if SubsetOf(obj1, obj2):\n",
    "                relations.append((id1, id2, \"SubsetOf\"))\n",
    "            if SubsetOf(obj2, obj1):\n",
    "                relations.append((id2, id1, \"SubsetOf\"))\n",
    "            if IsInside(obj1, obj2):\n",
    "                relations.append((id1, id2, \"IsInside\"))\n",
    "            if IsInside(obj2, obj1):\n",
    "                relations.append((id1, id2, \"IsEnclosed\"))\n",
    "    return relations\n",
    "\n",
    "\n",
    "def get_all_masks(pos_img, pos_id, pos_info, args):\n",
    "    # For each example in the batch, get all masks that match the pos_id\n",
    "    # max_num_occur is the maximum number of occurrences of a concept type in an example\n",
    "    concept_str_mapping = {\n",
    "        \"line\": \"Line\",\n",
    "        \"rectangle\": \"Rect\",\n",
    "        \"rectangleSolid\": \"RectSolid\",\n",
    "        \"Lshape\": \"Lshape\",\n",
    "        \"Tshape\": \"Tshape\",\n",
    "        \"Eshape\": \"Eshape\",\n",
    "        \"Hshape\": \"Hshape\",\n",
    "        \"Cshape\": \"Cshape\",\n",
    "        \"Ashape\": \"Ashape\",\n",
    "        \"Fshape\": \"Fshape\",\n",
    "        \"randomShape\": \"Randshape\",\n",
    "        \"arcShape\": \"ARCshape\",\n",
    "        \"Red\": \"Red\",\n",
    "        \"Blue\": \"Blue\",\n",
    "        \"Green\": \"Green\",\n",
    "        \"Cube\": \"Cube\",\n",
    "        \"Cylinder\": \"Cylinder\",\n",
    "        \"Large\": \"Large\",\n",
    "        \"Small\": \"Small\",\n",
    "    }\n",
    "    concept_str_reverse_mapping = {item: key for key, item in concept_str_mapping.items()}\n",
    "    batch_masks  = []\n",
    "    # Go through the batch\n",
    "    for idx, info in enumerate(pos_info):\n",
    "        ex_masks = torch.zeros(args.max_num_occur, 2 if args.is_two_branch else 1, *(args.image_size if args.rescaled_size == \"None\" else eval(args.rescaled_size)))\n",
    "        if args.is_two_branch:\n",
    "            found_relation = False\n",
    "            ex_ind = 0\n",
    "            relations = info[\"relations\"] if \"relations\" in info else get_arc_relations(pos_img[idx], info)\n",
    "            for tup in relations:\n",
    "                if tup[2] == pos_id[idx]:\n",
    "                    # Create the mask mask dimension\n",
    "                    if args.rescaled_size == \"None\":\n",
    "                        if tup[2] != \"IsInside\":\n",
    "                            ex_masks[ex_ind,] = torch.stack((info['id_object_mask'][tup[0]].squeeze(),\n",
    "                                                             info['id_object_mask'][tup[1]].squeeze()))\n",
    "                        else:\n",
    "                            ex_masks[ex_ind,] = torch.stack((info['id_object_mask'][tup[1]].squeeze(),\n",
    "                                                             info['id_object_mask'][tup[0]].squeeze()))\n",
    "                    else:\n",
    "                        if tup[2] != \"IsInside\":\n",
    "                            ex_masks[ex_ind] = torch.cat((rescale_tensor(info['id_object_mask'][tup[0]], rescaled_size=args.rescaled_size),\n",
    "                                                          rescale_tensor(info['id_object_mask'][tup[1]], rescaled_size=args.rescaled_size)))\n",
    "                        else:\n",
    "                            ex_masks[ex_ind] = torch.cat((rescale_tensor(info['id_object_mask'][tup[1]], rescaled_size=args.rescaled_size),\n",
    "                                                          rescale_tensor(info['id_object_mask'][tup[0]], rescaled_size=args.rescaled_size)))\n",
    "                    found_relation = True\n",
    "                    ex_ind += 1\n",
    "            assert found_relation, \"Should be at least one relation matching the passed in pos_id\"\n",
    "        else:\n",
    "            found_concept = False\n",
    "            ex_ind = 0\n",
    "            \"\"\"\n",
    "            Example for obj_spec: \n",
    "            [[('obj_0', 'line_[-1,1,-1]'), 'Attr'],\n",
    "             [('obj_1', 'line_[-1,1,-1]'), 'Attr'],\n",
    "            ]\n",
    "            \"\"\"\n",
    "            for obj_lst in info[\"obj_spec\"]:\n",
    "                # Get obj name\n",
    "                tup = obj_lst[0]  # tup: ('obj_0', 'line_[-1,1,-1]')\n",
    "                obj_name, obj_type = tup[0], tup[1].split(\"_\")[0].split(\"+\")  # obj_name: 'obj_0', obj_type: ['line']\n",
    "                if concept_str_reverse_mapping[pos_id[idx]] in obj_type:\n",
    "                    if args.rescaled_size == \"None\":\n",
    "                        ex_masks[ex_ind] = info[\"id_object_mask\"][info[\"node_id_map\"][obj_name]].unsqueeze(0)\n",
    "                    else:\n",
    "                        ex_masks[ex_ind] = rescale_tensor(\n",
    "                            info[\"id_object_mask\"][info[\"node_id_map\"][obj_name]].unsqueeze(0),\n",
    "                            rescaled_size=args.rescaled_size)\n",
    "                    found_concept = True\n",
    "                    ex_ind += 1\n",
    "            assert found_concept, \"Should be at least one concept matching the passed in pos_id\"\n",
    "        ex_masks = ex_masks.unsqueeze(2) # Create the channel dimension\n",
    "        batch_masks.append(ex_masks)\n",
    "    return torch.stack(batch_masks) # Create the batch dimension\n",
    "\n",
    "\n",
    "def test_acc_mask(\n",
    "    pos_data,\n",
    "    args,\n",
    "    neg_mask,\n",
    "    CONCEPTS,\n",
    "    OPERATORS,\n",
    "    reduction=\"mean\",\n",
    "    device=\"cpu\",\n",
    "):\n",
    "    uncommunitable_relation = ['IsInside']\n",
    "    result_dict = {}\n",
    "    pos_img, pos_mask, pos_id, pos_info = pos_data\n",
    "    pos_all_masks = get_all_masks(pos_img, pos_id, pos_info, args)\n",
    "    pos_repr = id_to_tensor(pos_id, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS).to(device) # vector of len 4\n",
    "    pos_img, pos_mask, pos_all_masks = to_device_recur([pos_img, pos_mask, pos_all_masks], device)\n",
    "\n",
    "    # Now we have pos_mask and neg_mask, given c_repr. We calculate the accuracy for neg_mask:\n",
    "    if args.is_two_branch:\n",
    "        pos_img = (pos_img, pos_img) if not args.is_image_tuple else pos_img\n",
    "        is_commutable = torch.Tensor([0 if ele in uncommunitable_relation else 1 for ele in pos_id]).unsqueeze(1).to(device) # [B*1], True or False. \n",
    "        is_uncommutable = 1 - is_commutable # [B*1], True or False.\n",
    "\n",
    "        # Use mask_iou:\n",
    "        neg_repeated = repeat(torch.stack(neg_mask, 1), \"b m c h w -> b p m c h w\", p=args.max_num_occur)\n",
    "        neg_repeated = rearrange(neg_repeated, \"b p m c h w -> (b p) m c h w\")\n",
    "        pos_all_masks = rearrange(pos_all_masks, \"b p m c h w -> (b p) m c h w\")\n",
    "        mask_acc_0_original = mask_iou_score(neg_repeated[:, 0, ...], pos_all_masks[:, 0, ...])\n",
    "        mask_acc_0_commute = mask_iou_score(neg_repeated[:, 0, ...], pos_all_masks[:, 1, ...])\n",
    "        mask_acc_1_original = mask_iou_score(neg_repeated[:, 1, ...], pos_all_masks[:, 1, ...])\n",
    "        mask_acc_1_commute = mask_iou_score(neg_repeated[:, 1, ...], pos_all_masks[:, 0, ...])\n",
    "        orig_sum = (mask_acc_0_original + mask_acc_1_original) / 2\n",
    "        alt_sum = (mask_acc_0_commute + mask_acc_1_commute) / 2\n",
    "        # Rearrange from [B * P, ...] to [B, P, ...]\n",
    "        orig_sum = rearrange(orig_sum, \"(b p) -> b p\", p=args.max_num_occur)\n",
    "        alt_sum = rearrange(alt_sum, \"(b p) -> b p\", p=args.max_num_occur)\n",
    "        # Get the best match among all possible pairs \n",
    "        values1, best_ind1 = orig_sum.max(1)\n",
    "        values2, best_ind2 = alt_sum.max(1)\n",
    "        # For each batch element, whether to use original or commute\n",
    "        bool_val = torch.logical_or(values1 > values2, is_uncommutable.squeeze())\n",
    "        if reduction == \"mean\":\n",
    "            mask_acc = to_np_array(torch.where(bool_val, values1, values2).mean())\n",
    "        elif reduction == \"none\":\n",
    "            mask_acc = to_np_array(torch.where(bool_val, values1, values2))\n",
    "        else:\n",
    "            raise\n",
    "        batch_size = pos_img[0].shape[0]\n",
    "        mask_acc_0_orig_max = rearrange(mask_acc_0_original, \"(b p) -> b p\", p=args.max_num_occur)[torch.arange(batch_size), best_ind1]\n",
    "        mask_acc_1_orig_max = rearrange(mask_acc_1_original, \"(b p) -> b p\", p=args.max_num_occur)[torch.arange(batch_size), best_ind1]\n",
    "        mask_acc_0_com_max = rearrange(mask_acc_0_commute, \"(b p) -> b p\", p=args.max_num_occur)[torch.arange(batch_size), best_ind2]\n",
    "        mask_acc_1_com_max = rearrange(mask_acc_1_commute, \"(b p) -> b p\", p=args.max_num_occur)[torch.arange(batch_size), best_ind2]\n",
    "        mask_acc_0 = torch.where(bool_val, mask_acc_0_orig_max, mask_acc_0_com_max)\n",
    "        mask_acc_1 = torch.where(bool_val, mask_acc_1_orig_max, mask_acc_1_com_max)\n",
    "        result_dict[\"mask_acc_0\"] = to_np_array(mask_acc_0)\n",
    "        result_dict[\"mask_acc_1\"] = to_np_array(mask_acc_1)\n",
    "    else:\n",
    "        assert len(pos_mask) == 1 and len(neg_mask) == 1\n",
    "        # Use mask_iou\n",
    "        neg_repeated = repeat(torch.stack(neg_mask, 1), \"b m c h w -> b p m c h w\", p=args.max_num_occur)\n",
    "        neg_repeated = rearrange(neg_repeated, \"b p m c h w -> (b p) m c h w\")\n",
    "        pos_all_masks = rearrange(pos_all_masks, \"b p m c h w -> (b p) m c h w\")\n",
    "        mask_acc = mask_iou_score(neg_repeated[:, 0, ...], pos_all_masks[:, 0, ...])\n",
    "        mask_acc = rearrange(mask_acc, \"(b p) -> b p\", p=args.max_num_occur)\n",
    "        mask_acc = to_np_array(reduce_tensor(mask_acc.max(1)[0], reduction))\n",
    "    result_dict[\"mask_acc\"] = mask_acc\n",
    "    return result_dict\n",
    "\n",
    "\n",
    "def test_acc(\n",
    "    model,\n",
    "    args,\n",
    "    dataloader,\n",
    "    device,\n",
    "    test_aspects=None,\n",
    "    CONCEPTS=None,\n",
    "    OPERATORS=None,\n",
    "    reduction=\"mean\",\n",
    "    suffix=\"\",\n",
    "):\n",
    "    \"\"\"Test the accuracy of the prediction of mask and c_repr.\n",
    "    Always assume that the pos_img is given.\n",
    "    \"\"\"\n",
    "    is_training = model.training\n",
    "    model.eval()\n",
    "    uncommunitable_relation = ['IsInside']\n",
    "\n",
    "    if test_aspects is None:\n",
    "        test_aspects = [\"mask|c_repr\", \"c_repr|mask\", \"both|c\", \"both\"]\n",
    "    acc_dict = {}\n",
    "\n",
    "    for test_aspect in test_aspects:\n",
    "        if test_aspect == \"mask|c_repr\":\n",
    "            # Given true c_repr, compute acc on mask: per pixel acc, c_repr\n",
    "            args_core = deepcopy(args)\n",
    "            args_core.ebm_target = \"mask\"\n",
    "            mask_acc_list = []\n",
    "            neg_out_list = []\n",
    "            if args.is_two_branch:\n",
    "                mask_acc_list_0 = []\n",
    "                mask_acc_list_1 = []\n",
    "            for i, pos_data in enumerate(dataloader):\n",
    "                if args.rescaled_size != \"None\":\n",
    "                    pos_data = rescale_data(pos_data, rescaled_size=args.rescaled_size, rescale_mode=args.rescale_mode)\n",
    "                if args.transforms_pos != \"None\":\n",
    "                    pos_data = transform_pos_data(pos_data, args.transforms_pos, color_avail=args.color_avail)\n",
    "                pos_img, pos_mask, pos_id, pos_info = pos_data\n",
    "                pos_img = to_device_recur(pos_img, device)\n",
    "                pos_repr = id_to_tensor(pos_id, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS).to(device)\n",
    "                (_, neg_mask, _, _, _, _), info = neg_mask_sgd(model, pos_img, c_repr=pos_repr, args=args_core)\n",
    "                neg_out = info[\"neg_out_list\"][-1]  # [B]\n",
    "                neg_out_list.append(reduce_tensor(neg_out, reduction))\n",
    "\n",
    "                result_dict = test_acc_mask(\n",
    "                    pos_data=pos_data,\n",
    "                    args=args_core,\n",
    "                    neg_mask=neg_mask,\n",
    "                    reduction=reduction,\n",
    "                    CONCEPTS=CONCEPTS,\n",
    "                    OPERATORS=OPERATORS,\n",
    "                    device=device,\n",
    "                )\n",
    "                mask_acc_list.append(result_dict[\"mask_acc\"])\n",
    "                if args.is_two_branch:\n",
    "                    mask_acc_list_0.append(result_dict[\"mask_acc_0\"])\n",
    "                    mask_acc_list_1.append(result_dict[\"mask_acc_1\"])\n",
    "            if reduction == \"mean\":\n",
    "                acc_dict[\"iou:mask|c_repr\"+suffix] = np.mean(mask_acc_list)\n",
    "                acc_dict[\"E:neg|c_repr\"+suffix] = np.mean(neg_out_list)\n",
    "            elif reduction == \"none\":\n",
    "                acc_dict[\"iou:mask|c_repr\"+suffix] = np.concatenate(mask_acc_list)\n",
    "                acc_dict[\"E:neg|c_repr\"+suffix] = np.concatenate(neg_out_list)\n",
    "            if args.is_two_branch:\n",
    "                if reduction == \"mean\":\n",
    "                    acc_dict[\"iou:mask_0|c_repr\"+suffix] = np.mean(mask_acc_list_0)\n",
    "                    acc_dict[\"iou:mask_1|c_repr\"+suffix] = np.mean(mask_acc_list_1)\n",
    "                elif reduction == \"none\":\n",
    "                    acc_dict[\"iou:mask_0|c_repr\"+suffix] = np.concatenate(mask_acc_list_0)\n",
    "                    acc_dict[\"iou:mask_1|c_repr\"+suffix] = np.concatenate(mask_acc_list_1)\n",
    "            del neg_mask\n",
    "\n",
    "        elif test_aspect == \"c_repr|mask\":\n",
    "            # Given true mask, compute acc on c_repr:\n",
    "            repr_acc_list = []\n",
    "            pos_out_list = []\n",
    "            emp_out_list = []\n",
    "            for i, pos_data in enumerate(dataloader):\n",
    "                if args.rescaled_size != \"None\":\n",
    "                    pos_data = rescale_data(pos_data, rescaled_size=args.rescaled_size, rescale_mode=args.rescale_mode)\n",
    "                if args.transforms_pos != \"None\":\n",
    "                    pos_data = transform_pos_data(pos_data, args.transforms_pos, color_avail=args.color_avail)\n",
    "                pos_img, pos_mask, pos_id, pos_info = pos_data\n",
    "                length = len(pos_id)\n",
    "                pos_img, pos_mask = to_device_recur([pos_img, pos_mask], device)\n",
    "                pos_repr = id_to_tensor(pos_id, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS).to(device)\n",
    "                pos_out = model(pos_img, mask=pos_mask, c_repr=pos_repr)\n",
    "                pos_out_list.append(to_np_array(pos_out.squeeze()))\n",
    "                emp_mask = tuple(torch.zeros_like(mask_ele) for mask_ele in pos_mask)\n",
    "                emp_out = model(pos_img, mask=emp_mask, c_repr=pos_repr)\n",
    "                emp_out_list.append(to_np_array(emp_out.squeeze()))\n",
    "                c_repr_energy = []\n",
    "                for j in range(len(args.concept_collection)):\n",
    "                    # args.concept_collection is the collection of concept id that is parsed from input\n",
    "                    c_repr = id_to_tensor([get_c_core(args.concept_collection[j])] * length, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS).to(device) # len 4\n",
    "                    neg_energy = model(pos_img, mask=pos_mask, c_repr=c_repr)\n",
    "                    c_repr_energy.append(neg_energy)\n",
    "                c_repr_energy = torch.cat(c_repr_energy, 1)\n",
    "                c_repr_argmin = to_np_array(c_repr_energy.argmin(1))\n",
    "\n",
    "                pos_repr_int = np.array([get_c_core(args.concept_collection).index(pos_id[k]) for k in range(length)])\n",
    "                repr_acc = reduce_tensor((c_repr_argmin == pos_repr_int), reduction)\n",
    "                repr_acc_list.append(repr_acc)\n",
    "\n",
    "            if reduction == \"mean\":\n",
    "                acc_dict[\"acc:c_repr|mask\"+suffix] = np.mean(repr_acc_list)\n",
    "                acc_dict[\"E:pos\"+suffix] = np.mean(pos_out_list)\n",
    "                acc_dict[\"E:emp\"+suffix] = np.mean(emp_out_list)\n",
    "            elif reduction == \"none\":\n",
    "                acc_dict[\"acc:c_repr|mask\"+suffix] = np.concatenate(repr_acc_list)\n",
    "                acc_dict[\"E:pos\"+suffix] = np.concatenate(pos_out_list)\n",
    "                acc_dict[\"E:emp\"+suffix] = np.concatenate(emp_out_list)\n",
    "            acc_dict[\"E:pos_std\"+suffix] = np.std(pos_out_list)\n",
    "            acc_dict[\"E:emp_std\"+suffix] = np.std(emp_out_list)\n",
    "\n",
    "        elif test_aspect == \"both|c\":\n",
    "            # For each c_repr, sgd -> mask and energy, find the lowest mask, compute acc on mask and c_repr\n",
    "            args_core = deepcopy(args)\n",
    "            args_core.ebm_target = \"mask\"\n",
    "            c_repr_min_list = []\n",
    "            repr_acc_list = []\n",
    "            mask_acc_list = []\n",
    "            if args.is_two_branch:\n",
    "                mask_acc_list_0 = []\n",
    "                mask_acc_list_1 = []\n",
    "\n",
    "            for i, pos_data in enumerate(dataloader):\n",
    "                if args.rescaled_size != \"None\":\n",
    "                    pos_data = rescale_data(pos_data, rescaled_size=args.rescaled_size, rescale_mode=args.rescale_mode)\n",
    "                if args.transforms_pos != \"None\":\n",
    "                    pos_data = transform_pos_data(pos_data, args.transforms_pos, color_avail=args.color_avail)\n",
    "                pos_img, pos_mask, pos_id, pos_info = pos_data\n",
    "                pos_all_masks = get_all_masks(pos_img, pos_id, pos_info, args)\n",
    "                pos_img, pos_mask, pos_all_masks = to_device_recur([pos_img, pos_mask, pos_all_masks], device)\n",
    "                length = len(pos_id)\n",
    "\n",
    "                c_repr_energy = []\n",
    "                c_repr_masks = []\n",
    "                for j in range(len(args.concept_collection)):\n",
    "                    # args.concept_collection is the collection of concept id that is parsed from input\n",
    "                    c_repr = id_to_tensor([get_c_core(args.concept_collection[j])] * length, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS).to(device) # len 4\n",
    "                    (_, c_repr_mask, _, _, _, _), info = neg_mask_sgd(model, pos_img, c_repr=c_repr, args=args_core)\n",
    "                    neg_energy = info[\"neg_out_list\"][-1]\n",
    "                    c_repr_masks.append(c_repr_mask)\n",
    "                    c_repr_energy.append(neg_energy)\n",
    "                c_repr_energy = np.stack(c_repr_energy, 1)  # [batch_size, n_concepts]\n",
    "                c_repr_min_list.append(reduce_tensor(c_repr_energy.min(1), reduction))\n",
    "                c_repr_argmin = c_repr_energy.argmin(1)  # [batch_size]\n",
    "                c_repr_argmin_onehot = torch.LongTensor(np.eye(len(args.concept_collection))[c_repr_argmin]).bool()\n",
    "                \n",
    "                if args.is_two_branch:\n",
    "                    c_repr_masks_core = Zip(*c_repr_masks, function=lambda x: torch.stack(x, 1))  # ( (b, n_cs, c, h, w), (b, n_cs, c, h, w))\n",
    "                    c_repr_masks_best = (c_repr_masks_core[0][c_repr_argmin_onehot], c_repr_masks_core[1][c_repr_argmin_onehot])  # ((b, c, h, w), (b, c, h, w))\n",
    "                else:\n",
    "                    c_repr_masks_core = Zip(*c_repr_masks, function=lambda x: torch.stack(x, 1))\n",
    "                    c_repr_masks_best = (c_repr_masks_core[0][c_repr_argmin_onehot],)\n",
    "\n",
    "                result_dict = test_acc_mask(\n",
    "                    pos_data=pos_data,\n",
    "                    args=args_core,\n",
    "                    neg_mask=c_repr_masks_best,\n",
    "                    reduction=reduction,\n",
    "                    CONCEPTS=CONCEPTS,\n",
    "                    OPERATORS=OPERATORS,\n",
    "                    device=device,\n",
    "                )\n",
    "                mask_acc_list.append(result_dict[\"mask_acc\"])\n",
    "                if args.is_two_branch:\n",
    "                    mask_acc_list_0.append(result_dict[\"mask_acc_0\"])\n",
    "                    mask_acc_list_1.append(result_dict[\"mask_acc_1\"])\n",
    "\n",
    "                # Calculate the accuracy for repr:\n",
    "                pos_repr_int = np.array([get_c_core(args.concept_collection).index(pos_id[k]) for k in range(length)])\n",
    "                repr_acc = reduce_tensor(c_repr_argmin == pos_repr_int, reduction)\n",
    "                repr_acc_list.append(repr_acc)\n",
    "\n",
    "            if reduction == \"mean\":\n",
    "                acc_dict[\"acc:c_repr|c\"+suffix] = np.mean(repr_acc_list)\n",
    "                acc_dict[\"iou:mask|c\"+suffix] = np.mean(mask_acc_list)\n",
    "                acc_dict[\"E:neg|c\"+suffix] = np.mean(c_repr_min_list)\n",
    "            elif reduction == \"none\":\n",
    "                acc_dict[\"acc:c_repr|c\"+suffix] = np.concatenate(repr_acc_list)\n",
    "                acc_dict[\"iou:mask|c\"+suffix] = np.concatenate(mask_acc_list)\n",
    "                acc_dict[\"E:neg|c\"+suffix] = np.concatenate(c_repr_min_list)\n",
    "            if args.is_two_branch:\n",
    "                if reduction == \"mean\":\n",
    "                    acc_dict[\"iou:mask_0|c\"+suffix] = np.mean(mask_acc_list_0)\n",
    "                    acc_dict[\"iou:mask_1|c\"+suffix] = np.mean(mask_acc_list_1)\n",
    "                elif reduction == \"none\":\n",
    "                    acc_dict[\"iou:mask_0|c\"+suffix] = np.concatenate(mask_acc_list_0)\n",
    "                    acc_dict[\"iou:mask_1|c\"+suffix] = np.concatenate(mask_acc_list_1)\n",
    "            del c_repr_masks_best\n",
    "            del c_repr_min_list\n",
    "        elif test_aspect == \"both\":\n",
    "            # Perform SGD on both mask and c_repr, and compute the acc on both:\n",
    "            test_aspect_core = \"mask\"\n",
    "            args_core = deepcopy(args)\n",
    "            args_core.ebm_target = \"mask+repr\"\n",
    "            mask_acc_list = []\n",
    "            repr_acc_list = []\n",
    "            neg_out_list = []\n",
    "            if args.is_two_branch:\n",
    "                mask_acc_list_0 = []\n",
    "                mask_acc_list_1 = []\n",
    "            c_reprs = id_to_tensor(get_c_core(args.concept_collection), CONCEPTS=CONCEPTS, OPERATORS=OPERATORS).to(device)  # [n_concepts, REPR_DIM]\n",
    "            for i, pos_data in enumerate(dataloader):\n",
    "                if args.rescaled_size != \"None\":\n",
    "                    pos_data = rescale_data(pos_data, rescaled_size=args.rescaled_size, rescale_mode=args.rescale_mode)\n",
    "                if args.transforms_pos != \"None\":\n",
    "                    pos_data = transform_pos_data(pos_data, args.transforms_pos, color_avail=args.color_avail)\n",
    "                pos_img, pos_mask, pos_id, pos_info = pos_data\n",
    "                if i == 0:\n",
    "                    c_reprs = c_reprs[None].expand(len(pos_id), *c_reprs.shape)  # [batch_size, n_concepts, REPR_DIM]\n",
    "                pos_repr = id_to_tensor(pos_id, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS).to(device) # vector of len 4\n",
    "                pos_all_masks = get_all_masks(pos_img, pos_id, pos_info, args)\n",
    "                pos_img, pos_mask, pos_all_masks = to_device_recur([pos_img, pos_mask, pos_all_masks], device)\n",
    "\n",
    "                # Now we have pos_mask and neg_mask, given c_repr. We calculate the accuracy for neg_mask:\n",
    "                (_, neg_mask, neg_repr, _, _, _), info = neg_mask_sgd(model, pos_img, args=args_core)\n",
    "                neg_out_list.append(to_np_array(info[\"neg_out_list\"][-1].squeeze()))\n",
    "                \n",
    "                result_dict = test_acc_mask(\n",
    "                    pos_data=pos_data,\n",
    "                    args=args_core,\n",
    "                    neg_mask=neg_mask,\n",
    "                    reduction=reduction,\n",
    "                    CONCEPTS=CONCEPTS,\n",
    "                    OPERATORS=OPERATORS,\n",
    "                    device=device,\n",
    "                )\n",
    "                mask_acc_list.append(result_dict[\"mask_acc\"])\n",
    "                if args.is_two_branch:\n",
    "                    mask_acc_list_0.append(result_dict[\"mask_acc_0\"])\n",
    "                    mask_acc_list_1.append(result_dict[\"mask_acc_1\"])\n",
    "\n",
    "                # The chosen repr is the one with the smallest distance to the ground-truth:\n",
    "                c_repr_argmin = to_np_array((c_reprs - neg_repr[:, None]).square().sum(-1).argmin(-1))\n",
    "                pos_repr_int = np.array([get_c_core(args.concept_collection).index(id) for id in pos_id])\n",
    "                repr_acc_list.append(reduce_tensor(pos_repr_int == c_repr_argmin, reduction))\n",
    "            if reduction == \"mean\":\n",
    "                acc_dict[\"acc:c_repr\"+suffix] = np.mean(repr_acc_list)\n",
    "                acc_dict[\"iou:mask\"+suffix] = np.mean(mask_acc_list)\n",
    "                acc_dict[\"E:neg\"+suffix] = np.mean(neg_out_list)\n",
    "            elif reduction == \"none\":\n",
    "                acc_dict[\"acc:c_repr\"+suffix] = np.concatenate(repr_acc_list)\n",
    "                acc_dict[\"iou:mask\"+suffix] = np.concatenate(mask_acc_list)\n",
    "                acc_dict[\"E:neg\"+suffix] = np.concatenate(neg_out_list)\n",
    "            acc_dict[\"E:neg_std\"+suffix] = np.std(neg_out_list)\n",
    "            if args.is_two_branch:\n",
    "                if reduction == \"mean\":\n",
    "                    acc_dict[\"iou:mask_0\"+suffix] = np.mean(mask_acc_list_0)\n",
    "                    acc_dict[\"iou:mask_1\"+suffix] = np.mean(mask_acc_list_1)\n",
    "                elif reduction == \"none\":\n",
    "                    acc_dict[\"iou:mask_0\"+suffix] = np.concatenate(mask_acc_list_0)\n",
    "                    acc_dict[\"iou:mask_1\"+suffix] = np.concatenate(mask_acc_list_1)\n",
    "            del neg_mask\n",
    "            del neg_repr\n",
    "        else:\n",
    "            raise Exception(\"Oh no! The test aspect '{}' is not valid!\".format(test_aspect))\n",
    "    acc_dict = {key: acc_dict[key] for key in sorted(acc_dict.keys())}\n",
    "    if is_training:\n",
    "        model.train()\n",
    "    return acc_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# #test commutable\n",
    "# if __name__ == \"__main__\":\n",
    "#     dirname = \"/dfs/user/tailin/.results/ebm_4-30/\"\n",
    "#     filenames = sorted(filter_filename(dirname, include=[\".p\"]))\n",
    "#     #print(filenames)\n",
    "    \n",
    "#     set_seed(seed=1)\n",
    "#     # Make sure that the initialization of CONCEPTS and OPERATORS (including their embedding) is after setting the seed\n",
    "#     from zeroc.concepts_shapes import OPERATORS, CONCEPTS, load_task, seperate_concept\n",
    "#     #filename = \"c-RotateA+RotateB+RotateC(Lshape)_cz_8_model_CEBM_alpha_1_lambd_0.005_size_20.0_sams_60_et_mask_pl_False_neg_addrand+permlabel_nco_0.1_mask_mulcat_tbm_concat_p_0.2_id_1_Hash_7iTPXjV8_turing2.p\"\n",
    "#     #filename = \"c-Parallel+Vertical_cz_8_model_CEBM_alpha_1_las_0.1_size_20.0_sams_60_et_mask_pl_False_neg_addrand+permlabel_nco_0.2_mask_concat_tbm_concat_cm_c2_cf_1_p_0.2_id_1_Hash_1iHtZRVo_turing3.p\"\n",
    "#     filename ='c-IsInside+IsTouch_cz_8_model_CEBM_alpha_1_las_0.1_size_20.0_sams_60_et_mask_pl_False_neg_addrand+permlabel_nco_0.2_mask_concat_tbm_concat_cm_c1_cf_0_p_0.2_id_1_Hash_Io3xVIYa_turing1.p'\n",
    "#     device = \"cuda:0\"\n",
    "#     data_record = pickle.load(open(dirname + filename, \"rb\"))\n",
    "#     model = load_model_energy(data_record[\"model_dict\"][5]).to(device)\n",
    "#     set_seed(seed=2)\n",
    "#     args = init_args(update_default_hyperparam(data_record[\"args\"]))\n",
    "#     args.n_examples = 500\n",
    "#     dataset, args = get_dataset(args)\n",
    "#     dataloader = DataLoader(dataset, batch_size=len(dataset), shuffle=False)\n",
    "\n",
    "#     set_seed(seed=1)\n",
    "#     # Make sure that the initialization of CONCEPTS and OPERATORS (including their embedding) is after setting the seed\n",
    "#     zeroc.concepts_shapes import OPERATORS, CONCEPTS, load_task, seperate_concept\n",
    "#     val_acc_dict = test_acc(model, args, dataloader, device, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS)\n",
    "#     print(val_acc_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # test \n",
    "# if __name__ == \"__main__\":\n",
    "\n",
    "#     # Make sure that the initialization of CONCEPTS and OPERATORS (including their embedding) is after setting the seed\n",
    "#     zeroc.concepts_shapes import OPERATORS, CONCEPTS, load_task, seperate_concept\n",
    "#     concept_embeddings = {key: to_np_array(CONCEPTS[key].get_node_repr()) for key in CONCEPTS}\n",
    "#     operator_embeddings = {key: to_np_array(OPERATORS[key].get_node_repr()) for key in OPERATORS}\n",
    "\n",
    "#     concept_embeddings.update(operator_embeddings)    \n",
    "#     #device = get_device(args) \n",
    "#     device = 'cpu'\n",
    "#     # Get dataset and dataloader:\n",
    "#     # if is_two_branch is True, will use for operators, in which the image x and mask a \n",
    "#     #    fed to the EBM E(x; a; c) will actually be a tuple of x=(x_in, x_out) and mask=(mask_in, mask_out):\n",
    "\n",
    "\n",
    "#     dirname = '/dfs/user/xyang23/results/ebm_4-4_model/'\n",
    "#     filename = 'c-Parallel+Vertical_model_CEBM_alpha_1_lambd_0.005_size_50.0_sams_60_mask_mulcat_p_0.5_id_0_Hash_V6a3+T6S_turing3.p'\n",
    "#     #filename = 'c-Lshape_model_CEBM_alpha_1_lambd_0.005_size_50.0_sams_60_mask_concat_p_0.2_id_0_Hash_7UaCu0QQ_turing3.p'\n",
    "\n",
    "#     data_record = pickle.load(open(dirname + filename, \"rb\"))\n",
    "#     args = init_args(update_default_hyperparam(data_record[\"args\"]))\n",
    "#     args.n_examples = 400\n",
    "#     dataset, args = get_dataset(args)\n",
    "#     args.n_examples = 400\n",
    "#     dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.n_workers, drop_last=True)\n",
    "\n",
    "#     args.is_two_branch = isinstance(dataset[0][0], tuple)\n",
    "#     print(args.concept_collection)\n",
    "#     pp.pprint(args.__dict__)\n",
    "#     model = load_model_energy(data_record[\"model_dict\"][20], device=device)\n",
    "#     test = test_acc_operator(model, args, dataloader, test_aspects=None)\n",
    "#     print(test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    # Obtain args and setting specific args for jupyter for easy testing:\n",
    "    args = get_args_EBM()\n",
    "    try:\n",
    "        get_ipython().run_line_magic('matplotlib', 'inline')\n",
    "        %env CUDA_VISIBLE_DEVICES=4,5\n",
    "        # Experiment management:\n",
    "#         args.exp_id = \"ebm_test\"\n",
    "#         args.date_time = \"{}-{}\".format(datetime.now().month, datetime.now().day)\n",
    "#         args.inspect_interval = 5\n",
    "#         args.save_interval = 20\n",
    "#         args.gpuid = \"5\"\n",
    "#         args.seed = 1\n",
    "#         args.id = \"1\"      # For differentiating experiments with different settings\n",
    "#         # args.batch_size = 16\n",
    "\n",
    "#         # Dataset:\n",
    "#         # args.dataset = \"c-Parallel+VerticalMid+VerticalEdge\"\n",
    "#         # args.dataset = \"c-SameShape+SameColor+IsInside(Line+Rect+RectSolid+Lshape)\"\n",
    "# #         args.dataset = \"c-arc^Line+RectSolid\"  # choose from c-{}, cifar10. E.g., c-Line, c-Rect, c-Line+Rect, c-arc^Line, c-arc^Line+Rect\n",
    "# #         args.dataset = \"c-Parallel+Vertical\"\n",
    "# #         args.dataset = \"y-Parallel+Vertical\"\n",
    "#         args.dataset = \"c-Line\"\n",
    "#         args.dataset = \"c-IsNonOverlapXY+IsInside+IsEnclosed(Rect[4,16]+Randshape[3,8]+Lshape[3,10]+Tshape[3,10])\"\n",
    "#         # args.dataset = \"c-Rect[4,16]+Eshape[5,12]\"\n",
    "#         # args.dataset = \"c-Eshape[5,12]+Rect[4,16]\"\n",
    "#         # args.dataset = \"c-IsInside+SameColor(Rect+Lshape+Tshape)\"\n",
    "#         # args.dataset = \"y-Line\"\n",
    "#         args.seed = 1\n",
    "# #         args.dataset = \"c-RotateA+RotateB+RotateC+hFlip+vFlip+DiagFlipA+DiagFlipB(Lshape+Line+Rect)\"  # Full operator\n",
    "#         # args.dataset = \"c-Hshape+Lshape\"\n",
    "# #         args.dataset = \"c-Image\"\n",
    "# #         args.dataset = \"c-arc^RotateA+RotateB()\"\n",
    "# #         args.dataset = \"cifar10\"\n",
    "# #         args.dataset=\"h-c^(1,2):Eshape+Ashape-d^1\"\n",
    "#         args.n_examples = 400   # Use 400 for quick testing, use 10000 for running actual experiments\n",
    "#         args.rainbow_prob = 0.  # Probability of using rainbow color for BabyARC datatset\n",
    "#         args.max_n_distractors = 0 #-1  # The maximum number of distractors besides objects related to the core concept.\n",
    "#         args.color_avail = \"1,2\"\n",
    "#         args.batch_size = 40\n",
    "#         args.canvas_size = 16\n",
    "#         # args.rescaled_size = \"32,32\"\n",
    "        \n",
    "#         # 3D dataset generation\n",
    "#         args.num_processes_3d = 20\n",
    "\n",
    "#         # Model setting:\n",
    "#         # args.model_type = \"CEBMLarge\"  # \"CEBM\" (for all other args.dataset), \"CEBMLarge\", \"IGEBM\" (only for cifar10)\n",
    "#         args.model_type = \"CEBM\"\n",
    "# #         args.model_atom = \"Line^Vertical+Parallel^AdaptRe\"\n",
    "#         args.w_type = \"image+mask\"\n",
    "#         args.mask_mode = \"concat\"\n",
    "#         args.channel_base = 64\n",
    "#         args.two_branch_mode = \"concat\"\n",
    "#         args.c_repr_mode = \"c2\"\n",
    "#         args.c_repr_first = 2\n",
    "#         args.c_repr_base = 2\n",
    "#         args.z_mode = \"None\"\n",
    "#         args.z_first = 2\n",
    "#         args.z_dim = 4\n",
    "#         args.n_workers = 0\n",
    "#         args.aggr_mode = \"max\"\n",
    "#         args.normalization_type = \"gn-2\"\n",
    "#         args.is_spec_norm = \"True\"\n",
    "#         args.dropout = 0\n",
    "#         args.self_attn_mode = \"None\"\n",
    "#         args.last_act_name = \"None\"\n",
    "# #         args.transforms = \"None\"\n",
    "#         args.transforms = \"color+flip+rotate+resize:0.7\"\n",
    "#         args.transforms_pos = \"randpatch\"#\"None\" # \"color+flip+rotate:0.5\"\n",
    "\n",
    "#         # EBM training setting:\n",
    "#         args.train_mode = \"cd\"\n",
    "#         # args.energy_mode = \"standard+center^stopsq:0.1\"\n",
    "#         args.kl_all_step = False\n",
    "#         args.kl_coef = 1\n",
    "#         args.entropy_coef_mask = 0\n",
    "#         args.entropy_coef_repr = 0\n",
    "#         args.entropy_coef_img = 0\n",
    "#         args.pos_consistency_coef = 0.1\n",
    "#         args.neg_consistency_coef = 0.1\n",
    "#         args.SGLD_mutual_exclusive_coef = 0\n",
    "#         args.epsilon_ent = 1e-5\n",
    "#         args.ebm_target_mode = \"r-rmbx\"\n",
    "#         args.ebm_target = \"mask+repr\"\n",
    "#         args.is_pos_repr_learnable = False\n",
    "#         args.neg_mode = \"addallrand+delrand\"  # Choose from \"None\" (default), \"addrand\", \"delrand\", \"addallrand\", \"permlabel\"\n",
    "#         args.neg_mode_coef = 0.2\n",
    "#         args.lambd_start = 0.1\n",
    "#         args.step_size = 20\n",
    "#         args.step_size_repr = 2\n",
    "#         args.step_size_img = -1\n",
    "#         args.sample_step = 60\n",
    "#         args.p_buffer = 0.2\n",
    "#         args.epochs = 500\n",
    "#         args.early_stopping_patience = 5\n",
    "        \n",
    "        \n",
    "        # Using \"u-\"\n",
    "        args.exp_id = \"ebm_test\"\n",
    "        # args.dataset = \"u-concept-Red+Green+Blue+Cube+Cylinder+Large+Small\"\n",
    "        args.dataset = \"u-relation-SameColor+SameShape+SameSize\"\n",
    "        args.color_avail = \"1,2\"\n",
    "        args.n_examples = 400\n",
    "        args.canvas_size = 64\n",
    "        args.train_mode = \"cd\"\n",
    "        args.model_type = \"CEBM\"\n",
    "        args.mask_mode = \"concat\"\n",
    "        args.c_repr_mode = \"c2\"\n",
    "        args.c_repr_first = 2\n",
    "        args.kl_coef = 1\n",
    "        args.entropy_coef_mask = 0\n",
    "        args.entropy_coef_repr = 0\n",
    "        args.ebm_target = \"mask\"\n",
    "        args.ebm_target_mode = \"r-rmbx\"\n",
    "        args.is_pos_repr_learnable = False\n",
    "        args.sample_step = 60\n",
    "        args.step_size = 30\n",
    "        args.step_size_repr = 2\n",
    "        args.lambd_start = 0.1\n",
    "        args.p_buffer=0.2\n",
    "        args.channel_base = 128\n",
    "        args.two_branch_mode = \"concat\"\n",
    "        args.neg_mode = \"addallrand+delrand+permlabel\"\n",
    "        args.aggr_mode = \"max\"\n",
    "        args.neg_mode_coef = 0.2\n",
    "        args.epochs=200\n",
    "        args.early_stopping_patience = -1\n",
    "        args.n_workers = 0\n",
    "        args.seed = 1\n",
    "        args.self_attn_mode = \"None\"\n",
    "        args.transforms = \"color+flip+rotate+resize:0.5\"\n",
    "        args.pos_consistency_coef = 0.1\n",
    "        args.is_res = True\n",
    "        args.lr = 1e-4\n",
    "        args.gpuid = \"5\"\n",
    "        args.id = \"0\"\n",
    "        args.act_name = \"leakyrelu\"\n",
    "        args.rescaled_size = \"32,32\"\n",
    "        args.rescale_mode = \"nearest\"\n",
    "        args.energy_mode = \"standard+center^stop:0.1\"\n",
    "        args.emp_target_mode = \"r-mb\"\n",
    "        args.batch_size = 40\n",
    "        args.max_num_occur = 20\n",
    "        args.parallel_mode = \"None\"\n",
    "        is_jupyter = True\n",
    "    except:\n",
    "        is_jupyter = False\n",
    "    if args.step_size_img == -1:\n",
    "        args.step_size_img = args.step_size\n",
    "    if args.step_size_repr == -1:\n",
    "        args.step_size_repr = args.step_size\n",
    "    if args.step_size_z == -1:\n",
    "        args.step_size_z = args.step_size\n",
    "\n",
    "    set_seed(args.seed)\n",
    "    # Make sure that the initialization of CONCEPTS and OPERATORS (including their embedding) is after setting the seed\n",
    "    from zeroc.concepts_shapes import OPERATORS, CONCEPTS, load_task, seperate_concept\n",
    "    from zeroc.concepts_shapes import SameShape, SameColor, SameAll, SameRow, SameCol, SubsetOf, IsInside, IsNonOverlapXY\n",
    "    concept_embeddings = get_concept_embeddings(CONCEPTS, OPERATORS)\n",
    "\n",
    "    # short_str_dict specifies the options to appear in the filename in their short name:\n",
    "    short_str_dict = {\n",
    "        \"dataset\": \"\",\n",
    "        \"canvas_size\": \"cz\",\n",
    "        \"model_type\": \"model\",\n",
    "        \"alpha\": \"alpha\",\n",
    "        \"lambd_start\": \"las\",\n",
    "        \"step_size\": \"size\",\n",
    "        \"sample_step\": \"sams\",\n",
    "        \"ebm_target_mode\": \"e\",\n",
    "        \"ebm_target\": \"et\",\n",
    "        \"is_pos_repr_learnable\": \"pl\",\n",
    "        \"train_mode\": \"tm\",\n",
    "        \"w_type\": \"w\",\n",
    "        \"neg_mode_coef\": \"nco\",\n",
    "        \"mask_mode\": \"mask\",\n",
    "        \"two_branch_mode\": \"tbm\",\n",
    "        \"c_repr_mode\": \"cm\",\n",
    "        \"c_repr_first\": \"cf\",\n",
    "        \"p_buffer\": \"p\",\n",
    "        \"id\": \"id\",\n",
    "    }\n",
    "    if len(args.dataset) > 60:\n",
    "        short_str_dict.pop(\"dataset\", None)\n",
    "    _, filename = get_filename(args, short_str_dict, is_local_path=is_jupyter)\n",
    "    # Get dataset and dataloader:\n",
    "    dataset, args = get_dataset(args, n_examples=int(args.n_examples*1.1), is_load=True, is_rewrite=args.is_rewrite)\n",
    "    n_train = int(len(dataset)*10/11)\n",
    "    if is_jupyter:\n",
    "        train_dataset = dataset[:400]\n",
    "        val_dataset = dataset[400:440]\n",
    "    else:\n",
    "        train_dataset = dataset[:n_train]\n",
    "        val_dataset = dataset[n_train:]\n",
    "    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, collate_fn=Batch(is_collate_tuple=True).collate(), num_workers=args.n_workers, drop_last=True)\n",
    "    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, collate_fn=Batch(is_collate_tuple=True).collate(), num_workers=args.n_workers, drop_last=True)\n",
    "    # if is_two_branch is True, will use for operators, in which the image x and mask a \n",
    "    #    fed to the EBM E(x; a; c) will actually be a tuple of x=(x_in, x_out) and mask=(mask_in, mask_out):\n",
    "    data_ex = dataset[0]\n",
    "    args.is_two_branch = len(data_ex[1]) > 1\n",
    "    args.is_image_tuple = isinstance(data_ex[0], tuple) or isinstance(data_ex[0], list)\n",
    "    early_stopping = Early_Stopping(patience=args.early_stopping_patience, mode=\"max\")\n",
    "\n",
    "    # Load model. (IBGBM is original model Implicit Generation and Modeling with Energy-Based Models (Du and Mordatch, 2019),\n",
    "    #   and ConceptEBM is our model that learns masks instead of images):\n",
    "    model = get_model_energy(args)\n",
    "    model, device = model_parallel(model, args)\n",
    "\n",
    "    # Unittest to make sure that the loaded model is exactly the same with the original model:\n",
    "    unittest(model, args, device, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS)\n",
    "    # Buffer contains 10000 examples of negative examples to provide as initial condition to perform gradient step\n",
    "    #   (usually 95% from buffer, 5% from unit Gaussian noise):\n",
    "    buffer = SampleBuffer_Conditional(is_two_branch=args.is_two_branch)\n",
    "\n",
    "    all_params = [{'params': model.parameters()}]\n",
    "    if args.is_pos_repr_learnable:\n",
    "        all_params.append({\"params\": [CONCEPTS[key].get_node_repr() for key in CONCEPTS] + [OPERATORS[key].get_node_repr() for key in OPERATORS]})\n",
    "    optimizer = torch.optim.Adam(all_params, lr=args.lr, betas=(0.0, 0.999))\n",
    "\n",
    "    data_record = {\"args\": args.__dict__, \"concept_embeddings\": [concept_embeddings], \"repr_epoch\": [-1], \"acc\": {}}\n",
    "    pp.pprint(args.__dict__)\n",
    "    for epoch in range(args.epochs+1):\n",
    "        # Initialize the recording list:\n",
    "        loss_list = []\n",
    "        if args.train_mode == \"cd\":\n",
    "            pos_out_list = []\n",
    "            neg_out_list = []\n",
    "            loss_core_info_dict = {\"L2\": []}\n",
    "            if \"mid\" in args.energy_mode:\n",
    "                emp_out_list = []\n",
    "                loss_core_info_dict[\"mid\"] = []\n",
    "            if \"center\" in args.energy_mode:\n",
    "                emp_out_list = []\n",
    "                loss_core_info_dict[\"center\"] = []\n",
    "            if \"standard\" in args.energy_mode:\n",
    "                loss_core_info_dict[\"standard\"] = []\n",
    "            if \"margin\" in args.energy_mode:\n",
    "                loss_core_info_dict[\"margin\"] = []\n",
    "            if args.neg_mode_coef > 0 and args.neg_mode != \"None\":\n",
    "                neg_out_gen_list = []\n",
    "                neg_out_gen_valid_list = []\n",
    "            if args.pos_consistency_coef > 0:\n",
    "                loss_pos_consistency_list = []\n",
    "            if args.neg_consistency_coef > 0:\n",
    "                loss_neg_consistency_list = []\n",
    "            if args.emp_consistency_coef > 0:\n",
    "                loss_emp_consistency_list = []\n",
    "        elif args.train_mode == \"sl\":\n",
    "            loss_img_list = []\n",
    "            loss_mask_list = []\n",
    "            loss_repr_list = []\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        for i, pos_data_ori in enumerate(train_loader):\n",
    "            if is_diagnose(loc=\"ebm:0\", filename=filename):\n",
    "                pdb.set_trace()\n",
    "            args = get_ebm_target(args)  # get ebm_target\n",
    "\n",
    "            # Sample negative examples (negative image for IGEBM and negative mask for CEBM,\n",
    "            #    args.p_buffer fraction is from buffer, rest from unit Gaussian):\n",
    "            if args.rescaled_size != \"None\":\n",
    "                pos_data_ori = rescale_data(pos_data_ori, rescaled_size=args.rescaled_size, rescale_mode=args.rescale_mode)\n",
    "            if args.transforms_pos != \"None\":\n",
    "                pos_data = transform_pos_data(pos_data_ori, args.transforms_pos, color_avail=args.color_avail)\n",
    "            else:\n",
    "                pos_data = pos_data_ori\n",
    "            neg_data = sample_buffer_conditional(\n",
    "                buffer,\n",
    "                pos_data=pos_data,\n",
    "                ebm_target=args.ebm_target,\n",
    "                in_channels=args.in_channels,\n",
    "                image_size=args.image_size if args.rescaled_size == \"None\" else eval(args.rescaled_size),\n",
    "                batch_size=args.batch_size,\n",
    "                is_two_branch=args.is_two_branch,\n",
    "                w_type=args.w_type,\n",
    "                p=args.p_buffer,\n",
    "                transforms=args.transforms,\n",
    "                color_avail=args.color_avail,\n",
    "                device=device,\n",
    "            )\n",
    "\n",
    "            if args.is_mask:\n",
    "                pos_img, pos_mask, pos_id, pos_info = pos_data\n",
    "                neg_img, neg_mask, neg_repr, neg_info = neg_data\n",
    "                # if args.ebm_target == \"mask\":\n",
    "                #     visualize_matrices([neg_img[0].argmax(0)])\n",
    "                #     plot_matrices([neg_mask[0][0][0]])\n",
    "                pos_repr = id_to_tensor(pos_id, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS,\n",
    "                                        requires_grad=args.is_pos_repr_learnable).to(device)\n",
    "                # Make sure that neg_mask can be perform gradient descent on:\n",
    "                pos_img, pos_mask = to_device_recur([pos_img, pos_mask], device)\n",
    "\n",
    "                if args.ebm_target in [\"image+mask\"]:\n",
    "                    if args.is_image_tuple:\n",
    "                        # Operator:\n",
    "                        neg_img[1].requires_grad = True\n",
    "                    else:\n",
    "                        # Concept or Relation:\n",
    "                        neg_img.requires_grad = True\n",
    "                else:\n",
    "                    neg_img = deepcopy(pos_img)\n",
    "                if args.ebm_target in [\"mask\", \"mask+repr\", \"image+mask\"]:\n",
    "                    for k in range(model.mask_arity):\n",
    "                        neg_mask[k].requires_grad = True\n",
    "                else:\n",
    "                    neg_mask = deepcopy(tuple(ele.detach() for ele in pos_mask))\n",
    "                if args.ebm_target in [\"repr\", \"mask+repr\"]:\n",
    "                    neg_repr = neg_repr.to(device)\n",
    "                    neg_repr.requires_grad = True\n",
    "                else:\n",
    "                    neg_repr = deepcopy(pos_repr.detach())\n",
    "\n",
    "                if args.train_mode == \"cd\":\n",
    "                    # The for loop below obtains the negative examples (image/mask) using current EBM,\n",
    "                    #   using gradient descent on fix model E(x; a; c) w.r.t. mask a:\n",
    "                    if is_diagnose(loc=\"ebm:1\", filename=filename):\n",
    "                        pdb.set_trace()\n",
    "                    if args.kl_coef == 0:\n",
    "                        (neg_img, neg_mask, neg_repr, _, _, _), info = neg_mask_sgd(\n",
    "                            model, neg_img, neg_mask=neg_mask, c_repr=neg_repr, args=args)\n",
    "                    else:\n",
    "                        (neg_img, neg_mask, neg_repr, _, _, _), (neg_img_kl, neg_mask_kl, neg_repr_kl, _, _, _), info = neg_mask_sgd_with_kl(\n",
    "                            model, neg_img, neg_mask=neg_mask, c_repr=neg_repr, args=args)\n",
    "\n",
    "                    # After obtaining negative examples, now perform gradient descent on fixed pos and neg examples w.r.t model parameter of EBM:\n",
    "                    requires_grad(model.parameters(), True)\n",
    "                    model.train()\n",
    "                    optimizer.zero_grad()\n",
    "\n",
    "                    pos_out = model(pos_img, mask=pos_mask, c_repr=pos_repr)\n",
    "                    neg_out = model(neg_img, mask=neg_mask, c_repr=neg_repr)\n",
    "                    if \"mid\" in args.energy_mode or \"center\" in args.energy_mode:\n",
    "                        emp_mask = tuple(torch.zeros_like(mask_ele) for mask_ele in pos_mask)\n",
    "                        emp_out = model(pos_img, mask=emp_mask, c_repr=pos_repr)\n",
    "                    else:\n",
    "                        emp_out = None\n",
    "                    if args.neg_mode_coef > 0 and args.neg_mode != \"None\" and args.ebm_target not in [\"image+mask\"]:\n",
    "                        neg_mask_gen, neg_repr_gen, neg_gen_valid = generate_neg_examples(pos_img, pos_mask, pos_repr, neg_mode=args.neg_mode)\n",
    "                        neg_out_gen = model(pos_img, mask=neg_mask_gen, c_repr=neg_repr_gen)\n",
    "                    else:\n",
    "                        neg_out_gen = None\n",
    "\n",
    "                    # The loss consists of L2 regularization on both pos and neg examples (to make the energy landscape less steep) and pos_out - neg_out:\n",
    "                    loss = args.alpha * (pos_out ** 2 + neg_out ** 2)\n",
    "                    loss_core, loss_core_info = get_loss_core(\n",
    "                        args=args,\n",
    "                        pos_out=pos_out,\n",
    "                        neg_out=neg_out,\n",
    "                        emp_out=emp_out,\n",
    "                        neg_out_gen=neg_out_gen,\n",
    "                    )\n",
    "                    loss_core_info[\"L2\"] = loss.mean().item()\n",
    "                    loss = loss + loss_core\n",
    "                    if args.neg_mode_coef > 0 and args.neg_mode != \"None\" and args.ebm_target not in [\"image+mask\"]:\n",
    "                        loss = loss + args.neg_mode_coef * neg_gen_valid * (args.alpha * neg_out_gen ** 2 - neg_out_gen)\n",
    "\n",
    "                    if args.kl_coef > 0:\n",
    "                        # KL loss:\n",
    "                        requires_grad(model.parameters(), False)\n",
    "                        loss_kl = model(neg_img_kl if neg_img_kl is not None else neg_img,\n",
    "                                        mask=neg_mask_kl if neg_mask_kl is not None else neg_mask,\n",
    "                                        c_repr=neg_repr_kl if neg_repr_kl is not None else neg_repr)\n",
    "                        requires_grad(model.parameters(), True)\n",
    "                        loss = loss + loss_kl * args.kl_coef\n",
    "\n",
    "                        if (args.entropy_coef_mask > 0 or args.entropy_coef_repr > 0 or args.entropy_coef_img > 0) and args.ebm_target in [\"mask\", \"repr\", \"mask+repr\", \"image+mask\"] and len(buffer) > 1000:\n",
    "                            neg_img_compare, neg_mask_compare, neg_repr_compare, _ = sample_buffer_conditional(\n",
    "                                buffer,\n",
    "                                pos_data=pos_data,\n",
    "                                ebm_target=args.ebm_target,\n",
    "                                in_channels=args.in_channels,\n",
    "                                image_size=args.image_size if args.rescaled_size == \"None\" else eval(args.rescaled_size),\n",
    "                                batch_size=100,\n",
    "                                is_two_branch=args.is_two_branch,\n",
    "                                w_type=args.w_type,\n",
    "                                p=1.,\n",
    "                                device=device,\n",
    "                            )\n",
    "\n",
    "                        if args.entropy_coef_img > 0 and args.ebm_target in [\"image+mask\"] and len(buffer) > 1000:\n",
    "                            neg_img_flat = torch.clamp(neg_img_kl.view(args.batch_size, -1), 0, 1) if not args.is_image_tuple else torch.clamp(neg_img_kl[1].view(args.batch_size, -1), 0, 1)\n",
    "                            neg_img_compare_flat = neg_img_compare.view(100, -1) if not args.is_image_tuple else neg_img_compare[1].view(100, -1)\n",
    "                            dist_matrix_img = torch.norm(neg_img_flat[:,None,:] - neg_img_compare_flat[None,:,:], p=2, dim=-1)\n",
    "                            loss_entropy_img = -torch.log(dist_matrix_img.min(dim=1, keepdims=True)[0] + args.epsilon_ent)\n",
    "                            loss = loss + loss_entropy_img * args.entropy_coef_img\n",
    "                        else:\n",
    "                            loss_entropy_img = torch.zeros(args.batch_size, 1)\n",
    "\n",
    "                        if args.entropy_coef_mask > 0 and args.ebm_target in [\"mask\", \"mask+repr\", \"image+mask\"] and len(buffer) > 1000:\n",
    "                            neg_mask_flat = tuple(torch.clamp(neg_mask_kl[k].view(args.batch_size, -1), 0, 1) for k in range(model.mask_arity))\n",
    "                            neg_mask_compare_flat = tuple(neg_mask_compare[k].view(100, -1) for k in range(model.mask_arity))\n",
    "                            dist_matrix_mask = tuple(torch.norm(neg_mask_flat[k][:,None,:] - neg_mask_compare_flat[k][None,:,:], p=2, dim=-1) for k in range(model.mask_arity))\n",
    "                            loss_entropy_mask = torch.cat([-torch.log(dist_matrix_mask[k].min(dim=1, keepdims=True)[0] + args.epsilon_ent) for k in range(model.mask_arity)], -1).mean(-1, keepdims=True)\n",
    "                            loss = loss + loss_entropy_mask * args.entropy_coef_mask\n",
    "                        else:\n",
    "                            loss_entropy_mask = torch.zeros(args.batch_size, 1)\n",
    "\n",
    "                        if args.entropy_coef_repr > 0 and args.ebm_target in [\"repr\", \"mask+repr\"] and len(buffer) > 1000:\n",
    "                            if \"softmax\" not in args.c_repr_mode:\n",
    "                                neg_repr_flat = torch.clamp(neg_repr_kl, 0, 1)\n",
    "                            else:\n",
    "                                neg_repr_flat = neg_repr_kl\n",
    "                            neg_repr_compare_flat = neg_repr_compare\n",
    "                            dist_matrix_repr = torch.norm(neg_repr_flat[:,None,:] - neg_repr_compare_flat[None,:,:], p=2, dim=-1)\n",
    "                            loss_entropy_repr = -torch.log(dist_matrix_repr.min(dim=1, keepdims=True)[0] + args.epsilon_ent)\n",
    "                            loss = loss + loss_entropy_repr * args.entropy_coef_repr\n",
    "                        else:\n",
    "                            loss_entropy_repr = torch.zeros(args.batch_size, 1)\n",
    "\n",
    "                    loss = loss.mean()\n",
    "                    if args.pos_consistency_coef > 0:\n",
    "                        loss_pos_consistency = pos_out.std() * args.pos_consistency_coef\n",
    "                        loss = loss + loss_pos_consistency\n",
    "                        loss_pos_consistency_list.append(loss_pos_consistency.item())\n",
    "                    if args.neg_consistency_coef > 0:\n",
    "                        loss_neg_consistency = neg_out.std() * args.neg_consistency_coef\n",
    "                        loss = loss + loss_neg_consistency\n",
    "                        loss_neg_consistency_list.append(loss_neg_consistency.item())\n",
    "                    if args.emp_consistency_coef > 0:\n",
    "                        loss_emp_consistency = emp_out.std() * args.emp_consistency_coef\n",
    "                        loss = loss + loss_emp_consistency\n",
    "                        loss_emp_consistency_list.append(loss_emp_consistency.item())\n",
    "\n",
    "                    # Record:\n",
    "                    pos_out_list.append(to_np_array(pos_out.mean()))\n",
    "                    neg_out_list.append(to_np_array(neg_out.mean()))\n",
    "                    if \"mid\" in args.energy_mode or \"center\" in args.energy_mode:\n",
    "                        emp_out_list.append(to_np_array(emp_out.mean()))\n",
    "                    for key in loss_core_info:\n",
    "                        loss_core_info_dict[key].append(loss_core_info[key])\n",
    "                    loss_list.append(to_np_array(loss.mean()))\n",
    "                    if args.neg_mode_coef > 0 and args.neg_mode != \"None\" and args.ebm_target not in [\"image+mask\"]:\n",
    "                        neg_out_gen_valid_list.append(to_np_array(neg_gen_valid).mean())\n",
    "                        neg_out_gen_list.append(to_np_array(neg_gen_valid * neg_out_gen).mean())\n",
    "\n",
    "                elif args.train_mode == \"sl\":\n",
    "                    _, (neg_img_kl, neg_mask_kl, neg_repr_kl, _, _, _), info = neg_mask_sgd_with_kl(\n",
    "                        model, neg_img, neg_mask, c_repr=neg_repr, args=args)\n",
    "\n",
    "                    # After obtaining negative examples, now perform gradient descent on fixed pos and neg examples w.r.t model parameter of EBM:\n",
    "                    requires_grad(model.parameters(), True)\n",
    "                    model.train()\n",
    "                    optimizer.zero_grad()\n",
    "\n",
    "                    loss = 0\n",
    "                    if args.ebm_target in [\"image+mask\"]:\n",
    "                        if args.is_image_tuple:\n",
    "                            loss_img = loss_op_core(neg_img_kl[1], pos_img[1], loss_type=args.supervised_loss_type)\n",
    "                        else:\n",
    "                            loss_img = loss_op_core(neg_img_kl, pos_img, loss_type=args.supervised_loss_type)\n",
    "                        loss = loss + loss_img\n",
    "                        loss_img_list.append(to_np_array(loss_img))\n",
    "                    if args.ebm_target in [\"mask\", \"mask+repr\", \"image+mask\"]:\n",
    "                        loss_mask = torch.stack([loss_op_core(neg_mask_kl[kk], pos_mask[kk], loss_type=args.supervised_loss_type) for kk in range(len(pos_mask))]).mean()\n",
    "                        loss = loss + loss_mask\n",
    "                        loss_mask_list.append(to_np_array(loss_mask))\n",
    "                    if args.ebm_target in [\"repr\", \"mask+repr\"]:\n",
    "                        loss_repr = loss_op_core(neg_repr_kl, pos_repr, loss_type=args.supervised_loss_type)\n",
    "                        loss = loss + loss_repr\n",
    "                        loss_repr_list.append(to_np_array(loss_repr))\n",
    "\n",
    "                    # Record:\n",
    "                    loss_list.append(to_np_array(loss.mean()))\n",
    "                else:\n",
    "                    raise Exception(\"train_mode {} is not valid!\".format(args.train_mode))\n",
    "\n",
    "                # Optimization:\n",
    "                loss.backward()\n",
    "                clip_grad(optimizer)\n",
    "                optimizer.step()\n",
    "\n",
    "                # Put the negative example (pos_img, neg_mask, pos_id) into the buffer:\n",
    "                buffer.push(imgs=neg_img, masks=neg_mask, c_reprs=neg_repr, class_ids=pos_id, infos=pos_info, ebm_target=args.ebm_target)\n",
    "            else:\n",
    "                pos_img, pos_id = pos_data_ori\n",
    "                pos_img, pos_id = pos_img.to(device), pos_id.to(device)\n",
    "                neg_img, neg_id = neg_data\n",
    "                neg_img.requires_grad = True\n",
    "                requires_grad(model.parameters(), False)\n",
    "                model.eval()\n",
    "\n",
    "                if args.lambd_start == -1:\n",
    "                    args.lambd_start = args.lambd\n",
    "                lambd_list = args.lambd + 1/2 * (args.lambd_start - args.lambd) * (1 + torch.cos(torch.arange(args.sample_step)/args.sample_step * np.pi))\n",
    "                if args.step_size_start == -1:\n",
    "                    args.step_size_start = args.step_size\n",
    "                step_size_list = args.step_size + 1/2 * (args.step_size_start - args.step_size) * (1 + torch.cos(torch.arange(args.sample_step)/args.sample_step * np.pi))\n",
    "\n",
    "                for k in range(args.sample_step):\n",
    "                    if noise.shape[0] != neg_img.shape[0]:\n",
    "                        noise = torch.randn(neg_img.shape[0], args.in_channels, *(args.image_size if args.rescaled_size == \"None\" else eval(args.rescaled_size)), device=device)\n",
    "\n",
    "                    noise.normal_(0, lambd_list[k])\n",
    "                    neg_img.data.add_(noise.data)\n",
    "\n",
    "                    neg_out = model(neg_img, neg_id)\n",
    "                    neg_out.sum().backward()\n",
    "                    neg_img.grad.data.clamp_(-0.01, 0.01)\n",
    "\n",
    "                    neg_img.data.add_(neg_img.grad.data, alpha=-step_size_list[k])\n",
    "\n",
    "                    neg_img.grad.detach_()\n",
    "                    neg_img.grad.zero_()\n",
    "\n",
    "                    neg_img.data.clamp_(0, 1)\n",
    "\n",
    "                neg_img = neg_img.detach()\n",
    "\n",
    "                requires_grad(model.parameters(), True)\n",
    "                model.train()\n",
    "\n",
    "                model.zero_grad()\n",
    "\n",
    "                pos_out = model(pos_img, pos_id)\n",
    "                neg_out = model(neg_img, neg_id)\n",
    "\n",
    "                loss = args.alpha * (pos_out ** 2 + neg_out ** 2)\n",
    "                loss = loss + (pos_out - neg_out)\n",
    "                loss = loss.mean()\n",
    "                loss.backward()\n",
    "\n",
    "                clip_grad(optimizer)\n",
    "\n",
    "                optimizer.step()\n",
    "\n",
    "                buffer.push(neg_img, class_ids=neg_id)\n",
    "\n",
    "                pos_out_list.append(to_np_array(pos_out.mean()))\n",
    "                neg_out_list.append(to_np_array(neg_out.mean()))\n",
    "                loss_list.append(to_np_array(loss.mean()))\n",
    "\n",
    "        # Epoch ends, record:\n",
    "        if args.train_mode == \"cd\":\n",
    "            record_data(data_record, [epoch, np.mean(pos_out_list), np.mean(neg_out_list), np.mean(loss_list)],\n",
    "                        [\"epoch\", \"E:pos:train\", \"E:neg|c_repr:train\", \"E:train\"])\n",
    "            if \"mid\" in args.energy_mode or \"center\" in args.energy_mode:\n",
    "                record_data(data_record, [np.mean(emp_out_list)], [\"E:emp:train\"])\n",
    "            loss_core_info_dict = transform_dict(loss_core_info_dict, \"mean\")\n",
    "            record_data(data_record, list(loss_core_info_dict.values()), [\"{}:train\".format(key) for key in loss_core_info_dict])\n",
    "            p.print(\"epoch {}: pos: {:.6f}     neg: {:.6f}\".format(\n",
    "                    str(epoch).zfill(3), np.mean(pos_out_list), np.mean(neg_out_list)), end=\"\")\n",
    "            if \"mid\" in args.energy_mode or \"center\" in args.energy_mode:\n",
    "                p.print(\"     emp: {:.6f}\".format(np.mean(emp_out_list)), is_datetime=False, end=\"\")\n",
    "            p.print(\"     loss: {:.6f}\".format(np.mean(loss_list)), is_datetime=False, end=\"\")\n",
    "            for key in loss_core_info_dict:\n",
    "                p.print(\"     {}: {:.6f}\".format(key, loss_core_info_dict[key]), is_datetime=False, end=\"\")\n",
    "            if args.neg_mode_coef > 0 and args.neg_mode != \"None\":\n",
    "                if args.ebm_target not in [\"image+mask\"]:\n",
    "                    record_data(data_record, [np.mean(neg_out_gen_valid_list), np.mean(neg_out_gen_list)], [\"E:neg_gen_valid:train\", \"E:neg_gen:train\"])\n",
    "                    p.print(\"        neg_gen: {:.6f}     neg_gen_valid: {:.3f}\".format(np.mean(neg_out_gen_list), np.mean(neg_out_gen_valid_list)), is_datetime=False, end=\"\")\n",
    "                else:\n",
    "                    record_data(data_record, [np.NaN, np.NaN], [\"E:neg_gen_valid:train\", \"E:neg_gen:train\"])\n",
    "            if args.kl_coef > 0:\n",
    "                loss_kl_mean = loss_kl.mean().item() * args.kl_coef\n",
    "                loss_entropy_mask_mean = loss_entropy_mask.mean().item() * args.entropy_coef_mask\n",
    "                loss_entropy_repr_mean = loss_entropy_repr.mean().item() * args.entropy_coef_repr\n",
    "                record_data(data_record, [loss_kl_mean, loss_entropy_mask_mean, loss_entropy_repr_mean, args.ebm_target], [\"loss_kl\", \"loss_entropy_mask_mean\", \"loss_entropy_repr_mean\", \"ebm_target\"])\n",
    "                print(\"         loss_kl: {:.6f}     loss_ent_mask: {:.6f}     loss_ent_repr: {:.6f}    ebm_target: {}\".format(loss_kl_mean, loss_entropy_mask_mean, loss_entropy_repr_mean, args.ebm_target), end=\"\")\n",
    "            if \"neg_mask_exceed_energy\" in info:\n",
    "                record_data(data_record, [info[\"neg_mask_exceed_energy\"] / args.batch_size], [\"neg_mask_exceed_energy\"])\n",
    "                print(\"         neg_mask_exceed: {:.6f}\".format(info[\"neg_mask_exceed_energy\"] / args.batch_size), end=\"\")\n",
    "            if args.pos_consistency_coef > 0:\n",
    "                record_data(data_record, [np.mean(loss_pos_consistency_list)], [\"E:pos_consistency:train\"])\n",
    "                print(\"         loss_pos_consistency: {:.6f}\".format(np.mean(loss_pos_consistency_list)), end=\"\")\n",
    "            if args.neg_consistency_coef > 0:\n",
    "                record_data(data_record, [np.mean(loss_neg_consistency_list)], [\"E:neg_consistency:train\"])\n",
    "                print(\"         loss_neg_consistency: {:.6f}\".format(np.mean(loss_neg_consistency_list)), end=\"\")\n",
    "            if args.emp_consistency_coef > 0:\n",
    "                record_data(data_record, [np.mean(loss_emp_consistency_list)], [\"E:emp_consistency:train\"])\n",
    "                print(\"         loss_emp_consistency: {:.6f}\".format(np.mean(loss_emp_consistency_list)), end=\"\")\n",
    "        elif args.train_mode == \"sl\":\n",
    "            record_data(data_record, [epoch, np.mean(loss_mask_list), np.mean(loss_repr_list), np.mean(loss_list)],\n",
    "                        [\"epoch\", \"loss_mask:train\", \"loss_repr:train\", \"loss:train\"])\n",
    "            p.print(\"epoch {}: loss_mask: {:.6f}     loss_repr: {:.6f}     loss: {:.6f}\".format(\n",
    "                    str(epoch).zfill(3), np.mean(loss_mask_list), np.mean(loss_repr_list), np.mean(loss_list)), end=\"\")\n",
    "        else:\n",
    "            raise\n",
    "        print()\n",
    "\n",
    "        if epoch % args.inspect_interval == 0 or epoch == args.epochs - 1:\n",
    "            # Inspection, plotting and recording:\n",
    "            torch.cuda.empty_cache()\n",
    "            image_dir = filename[:-2] + \"_samples/\"\n",
    "            make_dir(image_dir)\n",
    "            concept_embeddings = get_concept_embeddings(CONCEPTS, OPERATORS)\n",
    "            record_data(data_record, [concept_embeddings, epoch], [\"concept_embeddings\", \"repr_epoch\"])\n",
    "            if args.is_mask:\n",
    "                val_acc_dict = test_acc(model, args, val_loader, device, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS, suffix=\":val\")\n",
    "                val_acc_dict[\"acc+iou:mean:val\"] = np.mean([val_acc_dict[key] for key in val_acc_dict if (key.startswith(\"acc:\") or key.startswith(\"iou:\")) and \"_0\" not in key and \"_1\" not in key])\n",
    "                to_stop = early_stopping.monitor(val_acc_dict[\"acc+iou:mean:val\"])\n",
    "                record_data(data_record[\"acc\"], [epoch], [\"epoch:val\"])\n",
    "                record_data(data_record[\"acc\"], list(val_acc_dict.values()), list(val_acc_dict.keys()))\n",
    "                print()\n",
    "                print(filename)\n",
    "                pp.pprint(val_acc_dict)\n",
    "                print()\n",
    "                if args.is_two_branch:\n",
    "                    pos_img_aug = (pos_img, pos_img) if not args.is_image_tuple else pos_img\n",
    "                    if pos_img_aug[0].shape[1] == 3:\n",
    "                        pos_img_core = torch.cat([pos_img_aug[0].detach().to('cpu'), pos_img_aug[1].detach().to('cpu')])\n",
    "                    else:\n",
    "                        pos_img_core = torch.cat([onehot_to_RGB(pos_img_aug[0]), onehot_to_RGB(pos_img_aug[1])])\n",
    "\n",
    "                    if \"mask\" in args.w_type:\n",
    "                        pos_mask_core = torch.cat([pos_mask[0].detach().to('cpu').round(), pos_mask[1].detach().to('cpu').round()])\n",
    "                        neg_mask_core = torch.cat([neg_mask[0].detach().to('cpu').round(), neg_mask[1].detach().to('cpu').round()])\n",
    "                    else:\n",
    "                        pos_mask_core = torch.cat([onehot_to_RGB(to_one_hot(pos_mask[0].detach().to('cpu').argmax(1))), onehot_to_RGB(to_one_hot(pos_mask[1].detach().to('cpu').argmax(1)))])\n",
    "                        neg_mask_core = torch.cat([onehot_to_RGB(to_one_hot(neg_mask[0].detach().to('cpu').argmax(1))), onehot_to_RGB(to_one_hot(neg_mask[1].detach().to('cpu').argmax(1)))])\n",
    "                    if is_jupyter:\n",
    "                        p.print(\"Pos images:\")\n",
    "                        if args.is_image_tuple:\n",
    "                            if pos_img[0].shape[1] == 3:\n",
    "                                visualize_matrices(pos_img[0][:6], images_per_row=6, use_color_dict=False)\n",
    "                                visualize_matrices(pos_img[1][:6], images_per_row=6, use_color_dict=False)\n",
    "                            else:\n",
    "                                visualize_matrices(pos_img[0][:6].argmax(1), images_per_row=6)\n",
    "                                visualize_matrices(pos_img[1][:6].argmax(1), images_per_row=6)\n",
    "                        else:\n",
    "                            if pos_img.shape[1] == 3:\n",
    "                                visualize_matrices(pos_img[:6], images_per_row=6, use_color_dict=False)\n",
    "                            else:\n",
    "                                visualize_matrices(pos_img[:6].argmax(1), images_per_row=6)\n",
    "                        if args.ebm_target in [\"image+mask\"]:\n",
    "                            p.print(\"Neg images from SGLD:\")\n",
    "                            if args.is_image_tuple:\n",
    "                                if neg_img[0].shape[1] == 3:\n",
    "                                    visualize_matrices(neg_img[0][:6], images_per_row=6, use_color_dict=False)\n",
    "                                    visualize_matrices(neg_img[1][:6], images_per_row=6, use_color_dict=False)\n",
    "                                else:\n",
    "                                    visualize_matrices(neg_img[0][:6].argmax(1), images_per_row=6)\n",
    "                                    visualize_matrices(neg_img[1][:6].argmax(1), images_per_row=6)\n",
    "                            else:\n",
    "                                if neg_img.shape[1] == 3:\n",
    "                                    visualize_matrices(neg_img[:6], images_per_row=6, use_color_dict=False)\n",
    "                                else:\n",
    "                                    visualize_matrices(neg_img[:6].argmax(1), images_per_row=6)\n",
    "                        p.print(\"Pos masks:\")\n",
    "                        visualize_matrices(pos_mask[0][:6,0].round() if \"mask\" in args.w_type else pos_mask[0][:6].argmax(1), images_per_row=6, subtitles=pos_id[:6])\n",
    "                        visualize_matrices(pos_mask[1][:6,0].round() if \"mask\" in args.w_type else pos_mask[1][:6].argmax(1), images_per_row=6, subtitles=pos_id[:6])\n",
    "                        if args.train_mode == \"cd\" and args.neg_mode_coef > 0 and args.neg_mode != \"None\" and args.ebm_target not in [\"image+mask\"]:\n",
    "                            p.print(\"Neg masks generated:\")\n",
    "                            visualize_matrices(neg_mask_gen[0][:6,0].round() if \"mask\" in args.w_type else neg_mask_gen[0][:6].argmax(1), images_per_row=6, subtitles=[\"valid\" if is_valid else \"invalid\" for is_valid in neg_gen_valid.squeeze()[:6]])\n",
    "                            visualize_matrices(neg_mask_gen[1][:6,0].round() if \"mask\" in args.w_type else neg_mask_gen[1][:6].argmax(1), images_per_row=6)\n",
    "                        p.print(\"Neg masks from SGLD:\")\n",
    "                        visualize_matrices(neg_mask[0][:6,0].round() if \"mask\" in args.w_type else neg_mask[0][:6].argmax(1), images_per_row=6)\n",
    "                        visualize_matrices(neg_mask[1][:6,0].round() if \"mask\" in args.w_type else neg_mask[1][:6].argmax(1), images_per_row=6)\n",
    "                        print()\n",
    "                    # Save to png files:\n",
    "                    utils.save_image(\n",
    "                        pos_img_core,\n",
    "                        f'{image_dir}{str(epoch).zfill(5)}_image.png',\n",
    "                        nrow=args.batch_size,\n",
    "                        normalize=False,\n",
    "                        value_range=(0, 1),\n",
    "                    )\n",
    "                    utils.save_image(\n",
    "                        pos_mask_core,\n",
    "                        f'{image_dir}{str(epoch).zfill(5)}_pos_mask.png',\n",
    "                        nrow=args.batch_size,\n",
    "                        normalize=False,\n",
    "                        value_range=(0, 1),\n",
    "                    )\n",
    "                    utils.save_image(\n",
    "                        neg_mask_core,\n",
    "                        f'{image_dir}{str(epoch).zfill(5)}_neg_mask.png',\n",
    "                        nrow=args.batch_size,\n",
    "                        normalize=False,\n",
    "                        value_range=(0, 1),\n",
    "                    )\n",
    "                else:\n",
    "                    if pos_img.shape[1] == 3:\n",
    "                        pos_img_core = pos_img.detach().to('cpu')\n",
    "                    else:\n",
    "                        pos_img_core = onehot_to_RGB(pos_img)\n",
    "                    if \"mask\" in args.w_type:\n",
    "                        pos_mask_core = pos_mask[0].detach().to('cpu').round()\n",
    "                        neg_mask_core = neg_mask[0].detach().to('cpu').round()\n",
    "                    else:\n",
    "                        pos_mask_core = onehot_to_RGB(to_one_hot(pos_mask[0].detach().to('cpu').argmax(1)))\n",
    "                        neg_mask_core = onehot_to_RGB(to_one_hot(neg_mask[0].detach().to('cpu').argmax(1)))\n",
    "                    if is_jupyter:\n",
    "                        p.print(\"Pos images:\")\n",
    "                        if pos_img.shape[1] == 3:\n",
    "                            visualize_matrices(pos_img[:36], images_per_row=6, use_color_dict=False)\n",
    "                        else:\n",
    "                            visualize_matrices(pos_img[:36].argmax(1), images_per_row=6)\n",
    "                        if args.ebm_target in [\"image+mask\"]:\n",
    "                            p.print(\"Neg images from SGLD:\")\n",
    "                            if neg_img.shape[1] == 3:\n",
    "                                visualize_matrices(neg_img[:6], images_per_row=6, use_color_dict=False)\n",
    "                            else:\n",
    "                                visualize_matrices(neg_img[:6].argmax(1), images_per_row=6)\n",
    "                        p.print(\"Pos masks:\")\n",
    "                        visualize_matrices(pos_mask[0][:36,0].round() if \"mask\" in args.w_type else pos_mask[0][:6].argmax(1), images_per_row=6, subtitles=pos_id[:36])\n",
    "                        if args.train_mode == \"cd\" and args.neg_mode_coef > 0 and args.neg_mode != \"None\" and args.ebm_target not in [\"image+mask\"]:\n",
    "                            p.print(\"Neg masks generated:\")\n",
    "                            visualize_matrices(neg_mask_gen[0][:6,0].round() if \"mask\" in args.w_type else neg_mask_gen[0][:6].argmax(1), images_per_row=6, subtitles=[\"valid\" if is_valid else \"invalid\" for is_valid in neg_gen_valid.squeeze()[:6]])\n",
    "                        p.print(\"Neg masks from SGLD:\")\n",
    "                        visualize_matrices(neg_mask[0][:6,0].round() if \"mask\" in args.w_type else neg_mask[0][:6].argmax(1), images_per_row=6)\n",
    "                        print()\n",
    "                    utils.save_image(\n",
    "                        pos_img_core,\n",
    "                        f'{image_dir}{str(epoch).zfill(5)}_image.png',\n",
    "                        nrow=16,\n",
    "                        normalize=False,\n",
    "                        value_range=(0, 1),\n",
    "                    )\n",
    "                    utils.save_image(\n",
    "                        pos_mask_core,\n",
    "                        f'{image_dir}{str(epoch).zfill(5)}_pos_mask.png',\n",
    "                        nrow=16,\n",
    "                        normalize=False,\n",
    "                        value_range=(0, 1),\n",
    "                    )\n",
    "                    utils.save_image(\n",
    "                        neg_mask_core,\n",
    "                        f'{image_dir}{str(epoch).zfill(5)}_neg_mask.png',\n",
    "                        nrow=16,\n",
    "                        normalize=False,\n",
    "                        value_range=(0, 1),\n",
    "                    )\n",
    "            else:\n",
    "                if neg_img.shape[1] == 3:\n",
    "                    neg_img_core = neg_img.detach().to('cpu')\n",
    "                else:\n",
    "                    neg_img_core = onehot_to_RGB(neg_img)\n",
    "                if is_jupyter:\n",
    "                    visualize_matrices(neg_img[:6].argmax(1), images_per_row=6)\n",
    "                utils.save_image(\n",
    "                    neg_img_core,\n",
    "                    f'{image_dir}{str(epoch).zfill(5)}.png',\n",
    "                    nrow=16,\n",
    "                    normalize=True,\n",
    "                    value_range=(0, 1),\n",
    "                )\n",
    "        if epoch % args.save_interval == 0 or epoch == args.epochs - 1:\n",
    "            record_data(data_record, [epoch, model.model_dict], [\"save_epoch\", \"model_dict\"])\n",
    "            data_record[\"optimizer_dict\"] = to_cpu_recur(optimizer.state_dict())\n",
    "            try_call(pdump, args=[data_record, filename], max_exp_time=600)\n",
    "        if epoch % args.inspect_interval == 0 and to_stop:\n",
    "            p.print(\"Early-stopping at epoch {}, with acc_mean:val={}\".format(epoch, val_acc_dict[\"acc+iou:mean:val\"]))\n",
    "            record_data(data_record, [epoch, model.model_dict], [\"save_epoch\", \"model_dict\"])\n",
    "            data_record[\"optimizer_dict\"] = to_cpu_recur(optimizer.state_dict())\n",
    "            try_call(pdump, args=[data_record, filename], max_exp_time=600)\n",
    "            break"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  },
  "toc-autonumbering": false,
  "toc-showcode": false
 },
 "nbformat": 4,
 "nbformat_minor": 4
}