Kaggle competition: Abstraction and Reasoning Challenge
<SOS> metric: MeanBestErrorAtK Kaggle data source: abstraction-and-reasoning-challenge<import_modules>
T = torch.Tensor
class CAModel(nn.Module):
def __init__(self, num_states):
super(CAModel, self ).__init__()
self.simple_conv = nn.Sequential(
 nn.Conv2d(num_states, 32, kernel_size=3, padding=1),
nn.ReLU() 
)
 self.conv1 = nn.Conv2d(32, num_states, kernel_size=1)
 self.conv2 = nn.Conv2d(32, num_states+1, kernel_size=1)
 self.conv3 = nn.Conv2d(32, 1, kernel_size=1)
 self.conv4 = nn.Conv2d(32, 1, kernel_size=1)
 self.conv5 = nn.Conv2d(32, 1, kernel_size=1)
 self.soft = nn.Softmax(dim=1)
 def forward(self, x, steps = 1):
for step in range(steps):
x = self.simple_conv(x)
 y1 = self.conv1(x)
 y2 = self.conv2(x)
 y3 = self.conv3(x)
 y3 = torch.flatten(y3, start_dim=1)
 y3 = self.soft(y3)
 y4 = self.conv4(x)
 y4 = torch.flatten(y4, start_dim=1)
 y4 = self.soft(y4)
 y5 = self.conv5(x)
 y5 = torch.flatten(y5, start_dim=1)
 y5 = self.soft(y5)
 x = self.soft(y1)
 return y1, y2, y3, y4, y5<categorify>
SIZE = 1000
EPOCHS = 50
CONV_OUT_1 = 50
CONV_OUT_2 = 100
BATCH_SIZE = 128
TEST_PATH = Path('.. /input/abstraction-and-reasoning-challenge/')
 SUBMISSION_PATH = Path('.. /input/abstraction-and-reasoning-challenge/')
 TEST_PATH = TEST_PATH / 'test'
SUBMISSION_PATH = SUBMISSION_PATH / 'sample_submission.csv'
def get_matching_shapes_helper(img, template):
matches = []
for i in range(img.shape[0]-3):
for j in range(img.shape[1]-3):
match = True
for ti in range(3):
for tj in range(3):
if img[i+ti,j+tj] != template[ti,tj]:
match = False
if match:
matches.append(( i,j))
print(matches)
 def get_matching_shapes(inp):
inp = np.array(inp)
 print(inp.shape)
 for j in range(inp.shape[1]-3):
for i in range(inp.shape[0]-3):
template = np.zeros(( 3,3))
template = inp[i:i+3,j:j+3]
if np.sum(template)!= 0:
plot_pictures([inp, template],['inp', 'template'])
 get_matching_shapes_helper(inp, template)
 def generate_copies_to_right(inp_img, object_coords, anchor_coords, spacing=1):
inp_img = np.array(inp_img)
 img_xlen = inp_img.shape[0]
img_ylen = inp_img.shape[1]
object_left_x = object_coords[0]%img_ylen
object_up_y = object_coords[0]//img_ylen
object_rite_x = object_coords[1]%img_ylen
object_low_y = object_coords[1]//img_ylen
object_width = object_rite_x - object_left_x
object_height = object_up_y - object_low_y
anchor_x = anchor_coords%img_ylen
anchor_y = anchor_coords//img_ylen
object_img = inp_img[object_left_x:object_rite_x,object_low_y:object_up_y]
if object_width <= 0 or object_height <= 0:
return inp_img
elif anchor_x+object_width > img_xlen or anchor_y+object_height > img_ylen:
return inp_img
elif object_left_x < 0 or object_rite_x > img_xlen or object_low_y < 0 or object_up_y > img_ylen:
return inp_img
for i in range(img_xlen):
for j in range(img_ylen):
if i == anchor_x and j == anchor_y:
for ti in range(object_width):
for tj in range(object_height):
copy_value = object_img[ti,tj]
inp_img[i+ti,j+tj]= copy_value
return inp_img
def my_weird_loss(trg_img, copy_img):
trg_img = np.array(trg_img)
 loss = 0
for i in range(trg_img.shape[0]):
for j in range(trg_img.shape[1]):
if trg_img[i,j] != copy_img[i,j]:
loss += -1
return loss<normalization>
test_task_files = sorted(os.listdir(TEST_PATH))
test_tasks = []
for task_file in test_task_files:
with open(str(TEST_PATH / task_file), 'r')as f:
task = json.load(f)
 test_tasks.append(task )
get_matching_shapes(sample_out )<categorify>
Xs_test, Xs_train, ys_train = [], [], []
for task in test_tasks:
X_test, X_train, y_train = [], [], []
for pair in task["test"]:
X_test.append(pair["input"])
 for pair in task["train"]:
X_train.append(pair["input"])
 y_train.append(pair["output"])
 Xs_test.append(X_test)
 Xs_train.append(X_train)
 ys_train.append(y_train )
def inp2img(inp):
inp = np.array(inp)
 img = np.full(( 10, inp.shape[0], inp.shape[1]), 0, dtype=np.uint8)
 for i in range(10):
img[i] =(inp==i)
 return img
def target_maker(samin, samout):
target = np.full(( len(samin), len(samin[0])) , 0, dtype=np.uint8)
 for i in range(len(samin)) :
for j in range(len(samin[0])) :
if samin[i][j]==samout[i][j]:
target[i,j] = 10
else:
target[i,j] = samout[i][j]
return target
def pred_maker(inp, pred, pred_grade):
inp = np.asarray(inp)
 final_pred = np.full(( inp.shape[0], inp.shape[1]), 0, dtype=np.uint8)
 grade_avg = np.mean(pred_grade)
 threshold = 0.66
for i in range(pred.shape[0]):
for j in range(pred.shape[1]):
if pred[i,j] == 10 or pred_grade[i, j]/grade_avg < threshold:
final_pred[i,j] = inp[i,j]
else:
final_pred[i,j] = pred[i,j]
return final_pred
class TaskSolver:
def train(self, task_train, n_epoch=200, max_steps=2):

 self.net = CAModel(10)
 criterion = CrossEntropyLoss()
optimizer = Adam(self.net.parameters() , lr = 0.01)
 for epoch in range(n_epoch):
optimizer.zero_grad()
for sample in task_train:
target = target_maker(sample['input'],sample['output'])
 inputs = FloatTensor(inp2img(sample['input'])).unsqueeze(dim=0)
 labels1 = LongTensor(sample['output'] ).unsqueeze(dim=0)
 labels2 = LongTensor(target ).unsqueeze(dim=0)
 outputs1 = self.net(inputs)[0]
outputs2 = self.net(inputs)[1]
loss = criterion(outputs1, labels1)
 loss.backward()
loss = criterion(outputs2, labels2)
 loss.backward()
object_upl = self.net(inputs)[2]
object_upleft = object_upl.squeeze(dim=0 ).cpu().detach().numpy().argmax(0)
 object_lowr = self.net(inputs)[3]
object_lowrite = object_lowr.squeeze(dim=0 ).cpu().detach().numpy().argmax(0)
 anchor_point = self.net(inputs)[4]
anchor_point = anchor_point.squeeze(dim=0 ).cpu().detach().numpy().argmax(0)
 object_coords =(object_upleft, object_lowrite)
 copied_object_img = generate_copies_to_right(sample['input'], object_coords, anchor_point)
 loss = my_weird_loss(sample['output'],copied_object_img)*torch.mean(object_upl)
 loss.backward()
optimizer.step()
plot_pictures([sample['input'], sample['output'], copied_object_img ], ['Input', 'Label', 'Copies'])
 return self
def predict(self, task_test):
predictions = []
pred_grades = []
with torch.no_grad() :
for sample in task_test:
inputs = FloatTensor(inp2img(sample['input'])).unsqueeze(dim=0)
 outputs = self.net(inputs)[0]
pred = outputs.squeeze(dim=0 ).cpu().numpy().argmax(0)
 final_pred = pred
predictions.append(final_pred)
 return predictions<compute_test_metric>
matrices = []
for X_test in Xs_test:
for X in X_test:
matrices.append(X)
 values = []
for matrix in matrices:
for row in matrix:
for value in row:
values.append(value)
 df = pd.DataFrame(values)
 df.columns = ["values"]
def input_output_shape_is_same(task):
return all([np.array(el['input'] ).shape == np.array(el['output'] ).shape for el in task['train']])
 def calk_score(task_test, predict):
return [int(np.equal(sample['output'], pred ).all())for sample, pred in zip(task_test, predict)]<predict_on_test>
heights = [np.shape(matrix)[0] for matrix in matrices]
widths = [np.shape(matrix)[1] for matrix in matrices]
def evaluate(tasks):
ts = TaskSolver()
result = []
predictions = []
for task in tqdm(tasks):
if input_output_shape_is_same(task):
ts.train(task['train'])
 pred = ts.predict(task['test'])
 score = calk_score(task['test'], pred)
 else:
pred = [el['input'] for el in task['test']]
score = [0]*len(task['test'])
 predictions.append(pred)
 result.append(score)
 return result, predictions<compute_test_metric>
def replace_values(a, d):
return np.array([d.get(i, -1)for i in range(a.min() , a.max() + 1)])[a - a.min() ]
def repeat_matrix(a):
return np.concatenate([a]*(( SIZE // len(a)) + 1)) [:SIZE]
def get_new_matrix(X):
if len(set([np.array(x ).shape for x in X])) > 1:
X = np.array([X[0]])
 return X
def get_outp(outp, dictionary=None, replace=True):
if replace:
outp = replace_values(outp, dictionary)
 outp_matrix_dims = outp.shape
outp_probs_len = outp.shape[0]*outp.shape[1]*10
outp = to_categorical(outp.flatten() ,
num_classes=10 ).flatten()
return outp, outp_probs_len, outp_matrix_dims
train_result, train_predictions = evaluate(train_tasks)
 train_solved = [any(score)for score in train_result]
total = sum([len(score)for score in train_result])
 print(f"solved : {sum(train_solved)} from {total}({sum(train_solved)/total})" )<compute_test_metric>
class ARCDataset(Dataset):
def __init__(self, X, y, stage="train"):
self.X = get_new_matrix(X)
 self.X = repeat_matrix(self.X)
 self.stage = stage
if self.stage == "train":
self.y = get_new_matrix(y)
 self.y = repeat_matrix(self.y)
 def __len__(self):
return SIZE
def __getitem__(self, idx):
inp = self.X[idx]
if self.stage == "train":
outp = self.y[idx]
if idx != 0:
rep = np.arange(10)
 orig = np.arange(10)
 np.random.shuffle(rep)
 dictionary = dict(zip(orig, rep))
inp = replace_values(inp, dictionary)
 if self.stage == "train":
outp, outp_probs_len, outp_matrix_dims = get_outp(outp, dictionary)
 if idx == 0:
if self.stage == "train":
outp, outp_probs_len, outp_matrix_dims = get_outp(outp, None, False)
 return inp, outp, outp_probs_len, outp_matrix_dims, self.y
evaluation_result, evaluation_predictions = evaluate(evaluation_tasks)
 evaluation_solved = [any(score)for score in evaluation_result]
total = sum([len(score)for score in evaluation_result])
 print(f"solved : {sum(evaluation_solved)} from {total}({sum(evaluation_solved)/total})" )<predict_on_test>
class BasicCNNModel(nn.Module):
def __init__(self, inp_dim=(10, 10), outp_dim=(10, 10)) :
super(BasicCNNModel, self ).__init__()
CONV_IN = 3
KERNEL_SIZE = 3
DENSE_IN = CONV_OUT_2
self.relu = nn.ReLU()
self.softmax = nn.Softmax(dim=1)
 self.dense_1 = nn.Linear(DENSE_IN, outp_dim[0]*outp_dim[1]*10)
 if inp_dim[0] < 5 or inp_dim[1] < 5:
KERNEL_SIZE = 1
self.conv2d_1 = nn.Conv2d(CONV_IN, CONV_OUT_1, kernel_size=KERNEL_SIZE)
 self.conv2d_2 = nn.Conv2d(CONV_OUT_1, CONV_OUT_2, kernel_size=KERNEL_SIZE)
 def forward(self, x, outp_dim):
x = torch.cat([x.unsqueeze(0)]*3)
 x = x.permute(( 1, 0, 2, 3)).float()
self.conv2d_1.in_features = x.shape[1]
conv_1_out = self.relu(self.conv2d_1(x))
self.conv2d_2.in_features = conv_1_out.shape[1]
conv_2_out = self.relu(self.conv2d_2(conv_1_out))
self.dense_1.out_features = outp_dim
feature_vector, _ = torch.max(conv_2_out, 2)
 feature_vector, _ = torch.max(feature_vector, 2)
 logit_outputs = self.dense_1(feature_vector)
 out = []
for idx in range(logit_outputs.shape[1]//10):
out.append(self.softmax(logit_outputs[:, idx*10:(idx+1)*10]))
return torch.cat(out, axis=1 )
def flattener(pred):
str_pred = str([row for row in pred])
 str_pred = str_pred.replace(', ', '')
 str_pred = str_pred.replace('[[', '|')
 str_pred = str_pred.replace('][', '|')
 str_pred = str_pred.replace(']]', '|')
 return str_pred
def make_pediction(tasks):
ts = TaskSolver()
result = pd.Series()
for idx, task in tqdm(test_tasks.iteritems()):
if input_output_shape_is_same(task):
ts.train(task['train'])
 pred = ts.predict(task['test'])
 else:
pred = [el['input'] for el in task['test']]
for i, p in enumerate(pred):
result[f'{idx}_{i}'] = flattener(np.array(p ).tolist())
 return result<prepare_output>
def transform_dim(inp_dim, outp_dim, test_dim):
return(test_dim[0]*outp_dim[0]/inp_dim[0],
test_dim[1]*outp_dim[1]/inp_dim[1])
 def resize(x, test_dim, inp_dim):
if inp_dim == test_dim:
return x
else:
return cv2.resize(flt(x), inp_dim,
interpolation=cv2.INTER_AREA)
 def flt(x): return np.float32(x)
 def npy(x): return x.cpu().detach().numpy()
def itg(x): return np.int32(np.round(x))
submission = make_pediction(test_tasks)
 submission.head()<save_to_csv>
idx = 0
start = time.time()
test_predictions = []
for X_train, y_train in zip(Xs_train, ys_train):
print("TASK " + str(idx + 1))
train_set = ARCDataset(X_train, y_train, stage="train")
 train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)
 inp_dim = np.array(X_train[0] ).shape
outp_dim = np.array(y_train[0] ).shape
network = BasicCNNModel(inp_dim, outp_dim ).cuda()
optimizer = Adam(network.parameters() , lr=0.01)
 for epoch in range(EPOCHS):
for train_batch in train_loader:
train_X, train_y, out_d, d, out = train_batch
train_preds = network.forward(train_X.cuda() , out_d.cuda())
 train_loss = nn.MSELoss()(train_preds, train_y.cuda())
 optimizer.zero_grad()
train_loss.backward()
optimizer.step()
end = time.time()
print("Train loss: " + str(np.round(train_loss.item() , 3)) + " " +\
"Total time: " + str(np.round(end - start, 1)) + " s" + "
")
 X_test = np.array([resize(flt(X), np.shape(X), inp_dim)for X in Xs_test[idx-1]])
 for X in X_test:
test_dim = np.array(T(X)).shape
test_preds = npy(network.forward(T(X ).unsqueeze(0 ).cuda() , out_d.cuda()))
test_preds = np.argmax(test_preds.reshape(( 10, *outp_dim)) , axis=0)
 test_predictions.append(itg(resize(test_preds, np.shape(test_preds),
tuple(itg(transform_dim(inp_dim,
outp_dim,
test_dim))))))
idx += 1
submission = submission.reset_index()
submission.columns = ['output_id', 'output']
submission.to_csv('submission.csv', index=False)
 submission<import_modules>
def flattener(pred):
str_pred = str([row for row in pred])
 str_pred = str_pred.replace(', ', '')
 str_pred = str_pred.replace('[[', '|')
 str_pred = str_pred.replace('][', '|')
 str_pred = str_pred.replace(']]', '|')
 return str_pred
import numpy as np
import pandas as pd
import sys
import os
import re
import json
import cv2
import itertools
import traceback
from itertools import chain, product, combinations
from glob import glob
from pathlib import Path
from copy import copy, deepcopy
import matplotlib.pyplot as plt
from matplotlib import colors
from typing import List, Union, Callable
import toolz
import pydash as py
from pydash import py_ as _<categorify>
test_predictions = [[list(pred)for pred in test_pred] for test_pred in test_predictions]
for idx, pred in enumerate(test_predictions):
test_predictions[idx] = flattener(pred)
 submission = pd.read_csv(SUBMISSION_PATH)
 submission["output"] = test_predictions
task_files = glob('.. /input/abstraction-and-reasoning-challenge/**/*.json')
 def load_tasks(task_files):
if isinstance(task_files, str): task_files = glob(task_files)
 tasks = { re.sub(r'^.*/ (.*?/.*$)','\\1',file): json.load(open(file, 'r')) for file in task_files }
for file, task in tasks.items() :
for test_train, specs in task.items() :
for index, spec in enumerate(specs):
for input_output, grid in spec.items() :
tasks[file][test_train][index][input_output] = np.array(grid ).astype('uint8')
 for file, task in tasks.items() :
tasks[file]['file'] = file
return tasks
tasks = load_tasks(task_files)
 task = list(tasks.values())[6]; task<define_variables>
submission.to_csv("submission.csv", index=False )
<feature_engineering><EOS> 
submission.to_csv("submission.csv", index=False )