text_plain_output_1.png
from tqdm import tqdm

import json

import numpy as np # linear algebra

import os

train_path = '/kaggle/input/abstraction-and-reasoning-challenge/training/'
evaluation_path = '/kaggle/input/abstraction-and-reasoning-challenge/evaluation/'
test_path = '/kaggle/input/abstraction-and-reasoning-challenge/test/'

same_shape = []
for ex in tqdm(os.listdir(evaluation_path)):
    with open(evaluation_path + ex, 'r') as train_file:
        all_im = json.load(train_file)
        im_in = np.array(all_im['train'][0]['input'])
        im_out = np.array(all_im['train'][0]['output'])
        if im_in.shape == im_out.shape:
            same_shape.append(ex)

def get_im_with_same_ioshape(file_path, name, show=False, mode='train'):
    train = []
    test = []
    with open(file_path + name, 'r') as train_file:
        all_im = json.load(train_file)
        im_in = np.array(all_im['train'][0]['input'])
        im_out = np.array(all_im['train'][0]['output'])
        if im_in.shape != im_out.shape:
            return None
        for im in all_im['train']:
            im_in = np.array(im['input'])
            im_out = np.array(im['output'])
            mask = np.asarray(np.nan_to_num((im_in - im_out) / (im_in - im_out), 0), 'int8')
            train.append((im_in, im_out, mask))
            if show:
                print('NAME:\n', same_shape[N])
                print('IN:\n', im_in)
                print('OUT:\n', im_out)
                print('MASK:\n', mask)
        if mode == 'train':
            for im in all_im['test']:
                im_in = np.array(im['input'])
                im_out = np.array(im['output'])
                test.append((im_in, im_out))
        if mode == 'test':
            for im in all_im['test']:
                im_in = np.array(im['input'])
                test.append(im_in)
    return (train, test)
train, test = get_im_with_same_ioshape(evaluation_path, same_shape[1], False)