init

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+color_palette/regressor/colorLoversData.mat filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+/destijl_dataset
+*.jpg
+*.log

app.py

@@ -0,0 +1,397 @@
+from color_palette.regressor.config import config_to_use
+import numpy as np
+from PIL import Image, ImageFont, ImageDraw
+from sklearn.linear_model import LinearRegression
+from color_palette.model.GNN import ColorAttentionClassification
+from color_palette.regressor.model import Color2CubeDataset
+from color_palette.regressor.config import *
+from torch.utils.data import Dataset, DataLoader
+from color_palette.dataset import GraphDestijlDataset
+from color_palette.config import DataConfig
+import random
+import os
+import torch.nn.functional as F
+import torch
+import gradio as gr
+
+config = DataConfig()
+model_name = config.model_name
+dataset_root = config.dataset
+feature_size = config.feature_size
+device = config.device
+image_folder = "img_folder"
+
+if not os.path.exists(image_folder):
+    os.mkdir(image_folder)
+
+def train_regressor(train_loader):
+    X = []
+    y = []
+    for i, (input_data, target) in enumerate(train_loader):
+        input_data = np.squeeze(input_data)
+        target = np.squeeze(target)
+        X.append(input_data)
+        y.append(target)
+
+    X = np.stack(X, axis=0)
+    y = np.squeeze(np.stack(y, axis=0))
+
+    print("Before regressor train!\n")
+
+    reg = LinearRegression().fit(X, y)
+
+    return reg
+
+model_weight_path = "models/" + model_name + "/weights/best.pth"
+
+# palettes = np.load(config_to_use.save_folder+'/new_palettes_purple.npy')
+# original_palettes = np.load(config_to_use.save_folder+'/original_palettes_purple.npy')
+
+graph_test_dataset = GraphDestijlDataset(root=dataset_root, test=True, cube_mapping=True)
+model = ColorAttentionClassification(feature_size).to(device)
+model.load_state_dict(torch.load(model_weight_path)["state_dict"])
+
+dataset = Color2CubeDataset(config=config_to_use)
+train_loader = DataLoader(dataset, batch_size=1, shuffle=False)
+regressor = train_regressor(train_loader=train_loader)
+
+palette_of_the_design = [[0, 0, 0] for i in range(5)]
+all_node_colors = None
+class Demo:
+    def __init__(self, graph_dataset):
+        self.dataset = graph_dataset
+
+        first_sample_idx = random.randint(0, len(self.dataset)-1)
+        self.input_data, self.target_color, node_to_mask, also_normal_values = self.dataset.get(first_sample_idx)
+        global all_node_colors
+        all_node_colors = also_normal_values
+        self.same_indices = None
+        self.generate_img_from_palette([color.detach().numpy()*255 for color in also_normal_values], is_first=True)
+
+    def demo_reset(self):
+        first_sample_idx = random.randint(0, len(self.dataset))
+        self.input_data, self.target_color, node_to_mask, also_normal_values = self.dataset.get(first_sample_idx)
+        global all_node_colors
+        all_node_colors = also_normal_values
+        self.generate_img_from_palette([color.detach().numpy()*255 for color in also_normal_values], is_first=True)
+
+    def generate_img_from_palette(self, palette, canvas_size=512, is_first=False):
+        palette = np.array(palette).astype('int')
+        rgb_bg, rgb_text, rgb_text, rgb_circle, rgb_main_img, rgb_img1, rgb_img2, rgb_img3 = [tuple(color) for color in palette]
+        if is_first:
+            self.same_indices, unique_colors, _ = self.return_all_same_colors(palette=palette)
+        else:
+            _, unique_colors, _ = self.return_all_same_colors(palette=palette)
+
+        # assign the current palette using global keyword 
+        global palette_of_the_design
+        palette_of_the_design = unique_colors
+
+        # Set the background color and create an empty PIL Image to fill with shapes and text
+        image = Image.new("RGB", (canvas_size, canvas_size), color=rgb_bg)
+        # Save background image
+
+        title = "Lorem Ipsum Dolor"
+        undertitle = "Neque porro quisquam est qui dolorem ipsum quia dolor sit amet, \n consectetur, adipisci velit..."
+
+        draw = ImageDraw.Draw(image)
+        # Set settings for the fonts
+        font_title = ImageFont.truetype("Arial.ttf", 32)
+        title_width, title_height = draw.textsize(title, font=font_title)
+        title_x = (canvas_size - title_width) // 2
+        title_y = (canvas_size - title_height) // 2 - 100
+
+        font_undertitle = ImageFont.truetype("Arial.ttf", 15)
+        text_width, text_height = draw.textsize(undertitle, font=font_undertitle)
+        undertitle_x = (canvas_size - text_width) // 2
+        undertitle_y = (canvas_size - text_height) // 2 - 50
+        # Draw titles
+        draw.text((title_x, title_y), title, fill=rgb_text, font=font_title)
+        draw.text((undertitle_x, undertitle_y), undertitle, fill=rgb_text, font=font_undertitle)
+
+        # Draw the circle
+        rad = random.randint(30, 70)
+        x = random.randint(400, 512-(rad+10))
+        y = random.randint(10, title_y-(rad+10))
+
+        draw.ellipse((x, y, x+rad, y+rad), fill=rgb_circle)
+
+        # Draw the image
+        for j, color in enumerate([rgb_main_img, rgb_img1, rgb_img2, rgb_img3]):
+            x = 512-((j+1)*60)
+            y = 512-((j+1)*60)
+            
+            if j == 0:
+                rad = 80
+                draw.rectangle((x, y, x+rad, y+rad), fill=color)
+            else:
+                rad = 40
+                draw.rectangle((x, y, x+rad, y+rad), fill=color)
+
+        image.save(os.path.join("deneme.png"))
+
+    def run_model(self, input_data, target_color, node_to_mask, updated_color):
+
+        global all_node_colors
+        palette = np.array([color.detach().numpy()*255 for color in all_node_colors]).astype('int')
+        same_indices_list, unique_colors, first_indices = self.return_all_same_colors(palette)
+
+        unique_colors = unique_colors/255
+
+        selected_color = torch.Tensor(updated_color)/255
+        map_node_to_mask = -1
+        print("same indices list")
+        print(self.same_indices)
+        print("node to mask")
+        print(node_to_mask)
+        for i, idxs in enumerate(self.same_indices):
+            if node_to_mask in idxs:
+                map_node_to_mask = i
+                print("map node to mask: ", i)
+
+        for i, indices in enumerate(self.same_indices):
+            
+            if i == 0:
+                # update the color [0.15, 0.4908, 0.73]
+                cube_num_of_selected = self.rgb2cube(selected_color*255)
+                one_hot = np.zeros((64,))
+                one_hot[int(cube_num_of_selected)] = 1.0
+                node_to_recommend = map_node_to_mask
+                input_data.x[same_indices_list[map_node_to_mask], 4:] = torch.Tensor(one_hot)
+                unique_colors[map_node_to_mask] = selected_color
+            else:
+                if i == map_node_to_mask:
+                    zeroth_bin = 0
+                    indices = same_indices_list[zeroth_bin]
+                    node_to_recommend = 0
+                    input_data.x[indices[0], 4:] = torch.zeros((input_data.x.shape[1]-4))
+                    node_to_mask = indices[0]
+                else:
+                    node_to_recommend = i
+                    input_data.x[indices[0], 4:] = torch.zeros((input_data.x.shape[1]-4))
+                    node_to_mask = indices[0]
+
+                out = self.forward_pass(model, input_data) # input data has one-hot color features
+                if torch.is_tensor(node_to_mask):
+                    node_to_mask = node_to_mask.item()
+                values, values_indices = torch.topk(F.softmax(out[node_to_mask, :], dim=0), k=3, dim=0) # predict the color cube of the recommendation
+                prediction = values_indices.detach().numpy()[2]
+                # construct a palette using unique RGB palette and one-hot representation of the prediction cube.
+                feature_vector = self.create_rgb_and_one_hot_cube_vector(unique_colors, prediction, node_to_recommend)
+                # map cube to rgb color space using the regressor
+                recommendation = regressor.predict(feature_vector)[0]
+                # we now have the first set of recommendations. Now, we need to update the colors and input_data to propagate information.
+                # update the color in the palette and run the algorithm for rest of the palette.
+                # for that, first map the color to cube and convert to one_hot
+                input_data, unique_colors = self.update_palette(input_data, unique_colors, recommendation, same_indices_list, node_to_recommend)
+                # recursively do this here.
+                # save the results.
+        return np.array(unique_colors*255).astype(int)
+
+    def rgb2cube(self, color):
+        intervals = np.arange(0, 256, 256//4)
+        cube_coordinates = []
+        for channel in color:
+            i = 0
+            for j, value in enumerate(intervals):
+                if value < channel:
+                    i = j
+            cube_coordinates.append(i)
+        
+        cube_num = cube_coordinates[0]*1 + cube_coordinates[1]*4 + cube_coordinates[2]*4*4
+        return cube_num
+    
+    def cube2rgb(self, cube_num):
+        """
+            Return the start of the ranges
+        """
+        cube_num = int(cube_num)
+        intervals = np.arange(0, 256, 256//4)
+        coor2 = cube_num // 16
+        coor1 = (cube_num - coor2*4*4) // 4
+        coor0 = cube_num - coor2*4*4 - coor1*4
+        return [intervals[coor0], intervals[coor1], intervals[coor2]]
+
+    def return_all_same_colors(self, palette):
+
+        indices_list = [[],[],[],[],[]]
+        unique_colors, first_indices = np.unique(palette, axis=0, return_index=True)
+
+        unique_colors = np.array(unique_colors)
+        all_colors = np.array(palette)
+
+        for idx, color in enumerate(unique_colors):
+            for node_num, element in enumerate(all_colors):
+                if np.equal(color, element).all():
+                    indices_list[idx].append(node_num)
+
+        # these palettes and indices also include the masked color
+        return indices_list, unique_colors, first_indices
+
+    def update_palette(self, input_data, unique_rgb_palette, recommendation, indices_list, idx_to_idxs):
+        # convert prediction to one-hot vector
+        cube_num_of_the_changed_color = self.rgb2cube(recommendation*255)
+        one_hot = np.zeros((64,))
+        one_hot[int(cube_num_of_the_changed_color)] = 1.0
+
+        # update the feature vector accordingly for all the same colors
+        for idx in indices_list[idx_to_idxs]:
+            input_data.x[idx, 4:] = torch.Tensor(one_hot)
+
+        # update the unique color vector
+        unique_rgb_palette[idx_to_idxs] = recommendation
+        return input_data, unique_rgb_palette
+
+
+    def create_rgb_and_one_hot_cube_vector(self, rgb_palette, cube_num, node_to_mask):
+        one_hot = np.zeros((64,))
+        one_hot[int(cube_num)] = 1.0
+        removed_palette = np.delete(rgb_palette, node_to_mask, axis=0)
+        feature_vector = np.concatenate((removed_palette.flatten(), one_hot), axis=0)
+        return feature_vector.reshape(1, -1)
+
+    def create_all_one_hot_vector(self, rgb_palette, cube_num, node_to_mask):
+        one_hot = np.zeros((64,))
+        one_hot[int(cube_num)] = 1.0
+        removed_palette = np.delete(rgb_palette, node_to_mask, axis=0)
+        new_input_data = []
+        for color in removed_palette:
+            color_cube_num = self.rgb2cube(color*255)
+            empty_arr = np.zeros((64,))
+            empty_arr[int(color_cube_num)] = 1.0
+            new_input_data.append(empty_arr)
+
+        feature_vector = np.concatenate((np.array(new_input_data).flatten(), one_hot), axis=0)
+        return feature_vector.reshape(1, -1)
+
+    def forward_pass(self, model, data):
+        model.eval()
+        out = model(data.x, data.edge_index.long(), data.edge_weight)
+        return out
+
+    def rearrange_indices_list(self, indices_list, node_to_mask, unique_rgb_palette):
+        # take the node_to_mask indices to the beginning of the list
+        for i in range(len(indices_list)):
+            if node_to_mask in indices_list[i]:
+                index_to_pop = i
+
+        idxs = indices_list.pop(index_to_pop)
+        palette = unique_rgb_palette[index_to_pop]
+        temp_palette = np.delete(unique_rgb_palette, index_to_pop, axis=0)
+        unique_rgb_palette = np.concatenate(([palette], temp_palette), axis=0)
+        return [idxs] + indices_list, unique_rgb_palette
+
+    def update_color(self, updated_color, idx):
+        """
+        Takes a color and assigns it to the palette and the image.
+        """
+        idx = int(idx)
+        color = updated_color[1:-1].split(",")
+        color = [int(num) for num in color]
+
+        index_list = self.same_indices[idx]
+        which_one = random.randint(0, len(index_list)-1)
+        idx_to_update = index_list[which_one]
+
+        unique_colors = self.run_model(self.input_data, self.target_color, idx_to_update, color)
+
+        global palette_of_the_design
+        palette_of_the_design = unique_colors
+        
+        global all_node_colors
+        if torch.is_tensor(all_node_colors):
+            all_node_colors = all_node_colors.detach().numpy()
+        for i, index_list in enumerate(self.same_indices):
+            for index in index_list:
+                all_node_colors[index] = unique_colors[i]
+
+        self.generate_img_from_palette(palette=[color for color in all_node_colors])
+        main_image = Image.open("deneme.png")
+        gradio_elements = []
+        gradio_elements.append(gr.Image(main_image, height=256, width=256))
+        for i in range(len(self.same_indices)):
+            color = unique_colors[i]
+            image = Image.new("RGB", (512, 512), color=tuple(color))
+            gradio_elements.append(gr.Image(image, height=64, width=64))
+            string_version = "["+str(color[0])+", "+ str(color[1])+", " + str(color[2])+"]"
+            gradio_elements.append(gr.Textbox(value=string_version, min_width=64))
+    
+        all_node_colors = torch.Tensor(all_node_colors) / 255
+        return tuple(gradio_elements)
+    
+def perform_reset(button_input):
+    global demo
+    global all_node_colors
+    gradio_elements = []
+
+    demo.demo_reset()
+    main_image = Image.open("deneme.png")
+    gradio_elements = []
+    gradio_elements.append(gr.Image(main_image, height=256, width=256))
+    
+    for color in palette_of_the_design:
+        image = Image.new("RGB", (512, 512), color=tuple(color))
+        gradio_elements.append(gr.Image(image, height=64, width=64))
+        string_version = "["+str(color[0])+", "+ str(color[1])+", " + str(color[2])+"]"
+        gradio_elements.append(gr.Textbox(value=string_version, min_width=64))
+
+    return tuple(gradio_elements)
+
+
+demo = Demo(graph_dataset=graph_test_dataset)
+
+# Form a gradio template to display images and update the colors.
+
+with gr.Blocks() as project_demo:
+    with gr.Row():
+        image = Image.open("deneme.png")
+        design = gr.Image(image, height=256, width=256)
+
+    with gr.Row():
+        with gr.Column(min_width=100):
+            image1 = Image.new("RGB", (512, 512), color=tuple(palette_of_the_design[0]))
+            image1_gr = gr.Image(image1, height=64, width=64)
+            string1 = "["+str(palette_of_the_design[0][0])+", "+ str(palette_of_the_design[0][1])+", " + str(palette_of_the_design[0][2])+"]"
+            color1_update = gr.Textbox(value=string1, min_width=64)
+            color1_button = gr.Button(value="Update Color 1", min_width=64)
+        with gr.Column(min_width=100):
+            image2 = Image.new("RGB", (512, 512), color=tuple(palette_of_the_design[1]))
+            image2_gr = gr.Image(image2, height=64, width=64)
+            string2 = "["+str(palette_of_the_design[1][0])+", "+ str(palette_of_the_design[1][1])+", " + str(palette_of_the_design[1][2])+"]"
+            color2_update = gr.Textbox(value=string2, min_width=64)
+            color2_button = gr.Button(value="Update Color 2", min_width=64)
+        with gr.Column(min_width=100):
+            image3 = Image.new("RGB", (512, 512), color=tuple(palette_of_the_design[2]))
+            image3_gr = gr.Image(image3, height=64, width=64)
+            string3 = "["+str(palette_of_the_design[2][0])+", "+ str(palette_of_the_design[2][1])+", " + str(palette_of_the_design[2][2])+"]"
+            color3_update = gr.Textbox(value=string3, min_width=64)
+            color3_button = gr.Button(value="Update Color 3", min_width=64)
+        with gr.Column(min_width=100):
+            image4 = Image.new("RGB", (512, 512), color=tuple(palette_of_the_design[3]))
+            image4_gr = gr.Image(image4, height=64, width=64)
+            string4 = "["+str(palette_of_the_design[3][0])+", "+ str(palette_of_the_design[3][1])+", " + str(palette_of_the_design[3][2])+"]"
+            color4_update = gr.Textbox(value=string4, min_width=64)
+            color4_button = gr.Button(value="Update Color 4", min_width=64)
+        with gr.Column(min_width=100):
+            image5 = Image.new("RGB", (512, 512), color=tuple(palette_of_the_design[4]))
+            image5_gr = gr.Image(image5, height=64, width=64)
+            string5 = "["+str(palette_of_the_design[4][0])+", "+ str(palette_of_the_design[4][1])+", " + str(palette_of_the_design[4][2])+"]"
+            color5_update = gr.Textbox(value=string5, min_width=64)
+            color5_button = gr.Button(value="Update Color 5", min_width=64)
+        
+    with gr.Row():
+        reset_button = gr.Button(value="Reset the palette", min_width=64)
+
+    zero = gr.Number(value=0, visible=False)
+    one = gr.Number(value=1, visible=False)
+    two = gr.Number(value=2, visible=False)
+    three = gr.Number(value=3, visible=False)
+    four = gr.Number(value=4, visible=False)
+    color1_button.click(fn=demo.update_color, inputs=[color1_update, zero], outputs=[design, image1_gr, color1_update, image2_gr, color2_update, image3_gr, color3_update, image4_gr, color4_update, image5_gr, color5_update])
+    color2_button.click(fn=demo.update_color, inputs=[color2_update, one], outputs=[design, image1_gr, color1_update, image2_gr, color2_update, image3_gr, color3_update, image4_gr, color4_update, image5_gr, color5_update])
+    color3_button.click(fn=demo.update_color, inputs=[color3_update, two], outputs=[design, image1_gr, color1_update, image2_gr, color2_update, image3_gr, color3_update, image4_gr, color4_update, image5_gr, color5_update])
+    color4_button.click(fn=demo.update_color, inputs=[color4_update, three], outputs=[design, image1_gr, color1_update, image2_gr, color2_update, image3_gr, color3_update, image4_gr, color4_update, image5_gr, color5_update])
+    color5_button.click(fn=demo.update_color, inputs=[color5_update, four], outputs=[design, image1_gr, color1_update, image2_gr, color2_update, image3_gr, color3_update, image4_gr, color4_update, image5_gr, color5_update])
+    reset_button.click(fn=perform_reset, inputs=[reset_button], outputs=[design, image1_gr, color1_update, image2_gr, color2_update, image3_gr, color3_update, image4_gr, color4_update, image5_gr, color5_update])
+
+    project_demo.launch()

color_palette/all_one_hot_LR/test_gt.npy

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

color_palette/all_one_hot_LR/test_preds.npy

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

color_palette/all_one_hot_LR/test_preds_graph.npy

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

color_palette/all_one_hot_LR/test_rgb_colors.npy

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

color_palette/all_one_hot_LR_sequential/new_palettes.npy

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

color_palette/all_one_hot_LR_sequential/original_palettes.npy

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

color_palette/all_one_hot_LR_sequential/test_gt.npy

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

color_palette/all_one_hot_LR_sequential/test_preds.npy

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

color_palette/app copy.py

@@ -0,0 +1,326 @@
+from regressor.config import config_to_use
+import numpy as np
+from PIL import Image, ImageFont, ImageDraw
+from sklearn.linear_model import LinearRegression
+from dataset import GraphDestijlDataset
+from config import DataConfig
+import random
+import os
+import torch
+import gradio as gr
+
+config = DataConfig()
+model_name = config.model_name
+dataset_root = config.dataset
+image_folder = "img_folder"
+
+if not os.path.exists(image_folder):
+    os.mkdir(image_folder)
+
+model_weight_path = "../models/" + model_name + "/weights/best.pth"
+
+palettes = np.load(config_to_use.save_folder+'/new_palettes_purple.npy')
+original_palettes = np.load(config_to_use.save_folder+'/original_palettes_purple.npy')
+
+graph_test_dataset = GraphDestijlDataset(root=dataset_root, test=True, cube_mapping=True)
+
+palette_of_the_design = [[0, 0, 0] for i in range(5)]
+
+class Demo:
+    def __init__(self, graph_dataset):
+        self.dataset = graph_dataset
+
+        first_sample_idx = random.randint(0, len(self.dataset))
+        new_data, target_color, node_to_mask, also_normal_values = self.dataset.get(first_sample_idx)
+        self.all_node_colors = also_normal_values
+        self.same_indices = None
+        self.generate_img_from_palette([color.detach().numpy()*255 for color in also_normal_values], is_first=True)
+        
+    def generate_img_from_palette(self, palette, canvas_size=512, is_first=False):
+        palette = np.array(palette).astype('int')
+        rgb_bg, rgb_text, rgb_text, rgb_circle, rgb_main_img, rgb_img1, rgb_img2, rgb_img3 = [tuple(color) for color in palette]
+        
+        if is_first:
+            self.same_indices, unique_colors, _ = self.return_all_same_colors(palette=palette)
+        else:
+            _, unique_colors, _ = self.return_all_same_colors(palette=palette)
+
+        # assign the current palette using global keyword 
+        global palette_of_the_design
+        palette_of_the_design = unique_colors
+
+        # Set the background color and create an empty PIL Image to fill with shapes and text
+        image = Image.new("RGB", (canvas_size, canvas_size), color=rgb_bg)
+        # Save background image
+
+        title = "Lorem Ipsum Dolor"
+        undertitle = "Neque porro quisquam est qui dolorem ipsum quia dolor sit amet, \n consectetur, adipisci velit..."
+
+        draw = ImageDraw.Draw(image)
+        # Set settings for the fonts
+        font_title = ImageFont.truetype("Arial.ttf", 32)
+        title_width, title_height = draw.textsize(title, font=font_title)
+        title_x = (canvas_size - title_width) // 2
+        title_y = (canvas_size - title_height) // 2 - 100
+
+        font_undertitle = ImageFont.truetype("Arial.ttf", 15)
+        text_width, text_height = draw.textsize(undertitle, font=font_undertitle)
+        undertitle_x = (canvas_size - text_width) // 2
+        undertitle_y = (canvas_size - text_height) // 2 - 50
+        # Draw titles
+        draw.text((title_x, title_y), title, fill=rgb_text, font=font_title)
+        draw.text((undertitle_x, undertitle_y), undertitle, fill=rgb_text, font=font_undertitle)
+
+        # Draw the circle
+        rad = random.randint(30, 70)
+        x = random.randint(400, 512-(rad+10))
+        y = random.randint(10, title_y-(rad+10))
+
+        draw.ellipse((x, y, x+rad, y+rad), fill=rgb_circle)
+
+        # Draw the image
+        for j, color in enumerate([rgb_main_img, rgb_img1, rgb_img2, rgb_img3]):
+            x = 512-((j+1)*60)
+            y = 512-((j+1)*60)
+            
+            if j == 0:
+                rad = 80
+                draw.rectangle((x, y, x+rad, y+rad), fill=color)
+            else:
+                rad = 40
+                draw.rectangle((x, y, x+rad, y+rad), fill=color)
+
+        image.save(os.path.join("deneme.png"))
+
+    def run_model(self):
+        also_normal_values = np.squeeze(np.stack(also_normal_values, axis=0)) # rgb
+        same_indices_list, unique_colors, first_indices = return_all_same_colors(also_normal_values)
+        
+        # move node_to_mask to first place in input_data and unique colors
+        # same_indices_list, unique_colors = rearrange_indices_list(same_indices_list, node_to_mask, unique_colors)
+        map_node_to_mask = -1
+        for i, idxs in enumerate(same_indices_list):
+            if node_to_mask in idxs:
+                map_node_to_mask = i
+
+        original_palettes.append(unique_colors.copy())
+        for i, indices in enumerate(same_indices_list):
+            
+            if i == 0:
+                # update the color [0.15, 0.4908, 0.73]
+                selected_color = torch.Tensor([254/255, 254/255, 224/255])
+                cube_num_of_selected = rgb2cube(selected_color*255)
+                one_hot = np.zeros((64,))
+                one_hot[int(cube_num_of_selected)] = 1.0
+                node_to_recommend = map_node_to_mask
+                input_data.x[same_indices_list[map_node_to_mask], 4:] = torch.Tensor(one_hot)
+                unique_colors[map_node_to_mask] = selected_color
+            else:
+                if i == map_node_to_mask:
+                    zeroth_bin = 0
+                    indices = same_indices_list[zeroth_bin]
+                    node_to_recommend = 0
+                    input_data.x[indices[0], 4:] = torch.zeros((input_data.x.shape[1]-4))
+                    node_to_mask = indices[0]
+                else:
+                    node_to_recommend = i
+                    input_data.x[indices[0], 4:] = torch.zeros((input_data.x.shape[1]-4))
+                    node_to_mask = indices[0]
+
+                out = forward_pass(model, input_data) # input data has one-hot color features
+                if torch.is_tensor(node_to_mask):
+                    node_to_mask = node_to_mask.item()
+                values, values_indices = torch.topk(F.softmax(out[node_to_mask, :], dim=0), k=3, dim=0) # predict the color cube of the recommendation
+                prediction = values_indices.detach().numpy()[2]
+                # construct a palette using unique RGB palette and one-hot representation of the prediction cube.
+                feature_vector = create_rgb_and_one_hot_cube_vector(unique_colors, prediction, node_to_recommend)
+                # map cube to rgb color space using the regressor
+                recommendation = regressor.predict(feature_vector)[0]
+                # we now have the first set of recommendations. Now, we need to update the colors and input_data to propagate information.
+                # update the color in the palette and run the algorithm for rest of the palette.
+                # for that, first map the color to cube and convert to one_hot
+                input_data, unique_colors = update_palette(input_data, unique_colors, recommendation, same_indices_list, node_to_recommend)
+                # recursively do this here.
+                # save the results.
+
+    def rgb2cube(color):
+        intervals = np.arange(0, 256, 256//4)
+        cube_coordinates = []
+        for channel in color:
+            i = 0
+            for j, value in enumerate(intervals):
+                if value < channel:
+                    i = j
+            cube_coordinates.append(i)
+        
+        cube_num = cube_coordinates[0]*1 + cube_coordinates[1]*4 + cube_coordinates[2]*4*4
+        return cube_num
+    
+    def cube2rgb(self, cube_num):
+        """
+            Return the start of the ranges
+        """
+        cube_num = int(cube_num)
+        intervals = np.arange(0, 256, 256//4)
+        coor2 = cube_num // 16
+        coor1 = (cube_num - coor2*4*4) // 4
+        coor0 = cube_num - coor2*4*4 - coor1*4
+        return [intervals[coor0], intervals[coor1], intervals[coor2]]
+
+    def return_all_same_colors(self, palette):
+
+        indices_list = [[],[],[],[],[]]
+        unique_colors, first_indices = np.unique(palette, axis=0, return_index=True)
+
+        unique_colors = np.array(unique_colors)
+        all_colors = np.array(palette)
+
+        for idx, color in enumerate(unique_colors):
+            for node_num, element in enumerate(all_colors):
+                if np.equal(color, element).all():
+                    indices_list[idx].append(node_num)
+
+        # these palettes and indices also include the masked color
+        return indices_list, unique_colors, first_indices
+
+    def update_palette(self, input_data, unique_rgb_palette, recommendation, indices_list, idx_to_idxs):
+        # convert prediction to one-hot vector
+        cube_num_of_the_changed_color = self.rgb2cube(recommendation*255)
+        one_hot = np.zeros((64,))
+        one_hot[int(cube_num_of_the_changed_color)] = 1.0
+
+        # update the feature vector accordingly for all the same colors
+        for idx in indices_list[idx_to_idxs]:
+            input_data.x[idx, 4:] = torch.Tensor(one_hot)
+
+        # update the unique color vector
+        unique_rgb_palette[idx_to_idxs] = recommendation
+        return input_data, unique_rgb_palette
+
+
+    def create_rgb_and_one_hot_cube_vector(self, rgb_palette, cube_num, node_to_mask):
+        one_hot = np.zeros((64,))
+        one_hot[int(cube_num)] = 1.0
+        removed_palette = np.delete(rgb_palette, node_to_mask, axis=0)
+        feature_vector = np.concatenate((removed_palette.flatten(), one_hot), axis=0)
+        return feature_vector.reshape(1, -1)
+
+    def create_all_one_hot_vector(self, rgb_palette, cube_num, node_to_mask):
+        one_hot = np.zeros((64,))
+        one_hot[int(cube_num)] = 1.0
+        removed_palette = np.delete(rgb_palette, node_to_mask, axis=0)
+        new_input_data = []
+        for color in removed_palette:
+            color_cube_num = self.rgb2cube(color*255)
+            empty_arr = np.zeros((64,))
+            empty_arr[int(color_cube_num)] = 1.0
+            new_input_data.append(empty_arr)
+
+        feature_vector = np.concatenate((np.array(new_input_data).flatten(), one_hot), axis=0)
+        return feature_vector.reshape(1, -1)
+
+    def forward_pass(self, model, data):
+        model.eval()
+        out = model(data.x, data.edge_index.long(), data.edge_weight)
+        return out
+
+    def train_regressor(self, train_loader):
+        X = []
+        y = []
+        for i, (input_data, target) in enumerate(train_loader):
+            input_data = np.squeeze(input_data)
+            target = np.squeeze(target)
+            X.append(input_data)
+            y.append(target)
+
+        X = np.stack(X, axis=0)
+        y = np.squeeze(np.stack(y, axis=0))
+
+        print("Before regressor train!\n")
+
+        reg = LinearRegression().fit(X, y)
+
+        return reg
+
+    def rearrange_indices_list(self, indices_list, node_to_mask, unique_rgb_palette):
+        # take the node_to_mask indices to the beginning of the list
+        for i in range(len(indices_list)):
+            if node_to_mask in indices_list[i]:
+                index_to_pop = i
+
+        idxs = indices_list.pop(index_to_pop)
+        palette = unique_rgb_palette[index_to_pop]
+        temp_palette = np.delete(unique_rgb_palette, index_to_pop, axis=0)
+        unique_rgb_palette = np.concatenate(([palette], temp_palette), axis=0)
+        return [idxs] + indices_list, unique_rgb_palette
+
+    def update_color(self, updated_color, idx):
+        """
+        Takes a color and assigns it to the palette and the image.
+        """
+        idx = int(idx)
+        color = updated_color[1:-1].split(",")
+        color = [int(num) for num in color]
+
+        global palette_of_the_design
+        palette_of_the_design[idx] = color
+
+        idxs_to_change = self.same_indices[idx]
+        for index in idxs_to_change:
+            self.all_node_colors[index] = torch.Tensor(color)/255
+
+        self.generate_img_from_palette(palette=[color.detach().numpy()*255 for color in self.all_node_colors])
+
+        image = Image.new("RGB", (512, 512), color=tuple(color))
+        main_image = Image.open("deneme.png")
+        return gr.Image(main_image, height=256, width=256), gr.Image(image, height=64, width=64), gr.Textbox(value=str(color), min_width=64)
+    
+if __name__ == "__main__":
+    demo = Demo(graph_dataset=graph_test_dataset)
+
+    # Form a gradio template to display images and update the colors.
+
+    with gr.Blocks() as project_demo:
+        with gr.Row():
+            image = Image.open("deneme.png")
+            design = gr.Image(image, height=256, width=256)
+
+        with gr.Row():
+            with gr.Column(min_width=100):
+                image1 = Image.new("RGB", (512, 512), color=tuple(palette_of_the_design[0]))
+                image1_gr = gr.Image(image1, height=64, width=64)
+                color1_update = gr.Textbox(value=str(palette_of_the_design[0]).replace(" ", ","), min_width=64)
+                color1_button = gr.Button(value="Update Color 1", min_width=64)
+            with gr.Column(min_width=100):
+                image2 = Image.new("RGB", (512, 512), color=tuple(palette_of_the_design[1]))
+                image2_gr = gr.Image(image2, height=64, width=64)
+                color2_update = gr.Textbox(value=str(palette_of_the_design[1]).replace(" ", ","), min_width=64)
+                color2_button = gr.Button(value="Update Color 2", min_width=64)
+            with gr.Column(min_width=100):
+                image3 = Image.new("RGB", (512, 512), color=tuple(palette_of_the_design[2]))
+                image3_gr = gr.Image(image3, height=64, width=64)
+                color3_update = gr.Textbox(value=str(palette_of_the_design[2]).replace(" ", ","), min_width=64)
+                color3_button = gr.Button(value="Update Color 3", min_width=64)
+            with gr.Column(min_width=100):
+                image4 = Image.new("RGB", (512, 512), color=tuple(palette_of_the_design[3]))
+                image4_gr = gr.Image(image4, height=64, width=64)
+                color4_update = gr.Textbox(value=str(palette_of_the_design[3]).replace(" ", ","), min_width=64)
+                color4_button = gr.Button(value="Update Color 4", min_width=64)
+            with gr.Column(min_width=100):
+                image5 = Image.new("RGB", (512, 512), color=tuple(palette_of_the_design[4]))
+                image5_gr = gr.Image(image5, height=64, width=64)
+                color5_update = gr.Textbox(value=str(palette_of_the_design[4]).replace(" ", ","), min_width=64)
+                color5_button = gr.Button(value="Update Color 5", min_width=64)
+
+        zero = gr.Number(value=0, visible=False)
+        one = gr.Number(value=1, visible=False)
+        two = gr.Number(value=2, visible=False)
+        three = gr.Number(value=3, visible=False)
+        four = gr.Number(value=4, visible=False)
+        color1_button.click(fn=demo.update_color, inputs=[color1_update, zero], outputs=[design, image1_gr, color1_update])
+        color2_button.click(fn=demo.update_color, inputs=[color2_update, one], outputs=[design, image2_gr, color2_update])
+        color3_button.click(fn=demo.update_color, inputs=[color3_update, two], outputs=[design, image3_gr, color3_update])
+        color4_button.click(fn=demo.update_color, inputs=[color4_update, three], outputs=[design, image4_gr, color4_update])
+        color5_button.click(fn=demo.update_color, inputs=[color5_update, four], outputs=[design, image5_gr, color5_update])
+
+        project_demo.launch()

color_palette/bash_scripts/training.sh

@@ -0,0 +1,24 @@
+FROM_FILE=0
+TO_FILE=11
+
+file=$FROM_FILE
+device_idx=2
+
+free_mem=$(nvidia-smi --query-gpu=memory.free --format=csv -i $device_idx | grep -Eo [0-9]+)
+
+while [ $file -le $TO_FILE ]
+do
+    # if [ $free_mem -lt 13000 ]; then
+    #     while [ $free_mem -lt 13000 ]; do
+    #         sleep 10
+    #         free_mem=$(nvidia-smi --query-gpu=memory.free --format=csv -i $device_idx | grep -Eo [0-9]+)
+    #     done
+    # fi
+
+    echo "Running experiment for conf$file.yaml"
+    #nohup python train.py --config_file config/conf$file.yaml > "config/out_$file.txt" &
+    python evaluate.py --config_file config/conf$file.yaml
+    file=$(($file+1))
+    #sleep 10
+    
+done

color_palette/cnn_dataset.py

@@ -0,0 +1,104 @@
+from torch.utils.data import Dataset, DataLoader
+from sklearn.model_selection import train_test_split
+import skimage.color as scicolor
+from utils import *
+
+class PreviewDataset(Dataset):
+    def __init__(self, root="../destijl_dataset/rgba_dataset/", 
+                 transform=None, test=False, color_space="RGB", 
+                 input_color_space="RGB", 
+                 is_classification=False, 
+                 normalize_cielab=True,
+                 normalize_rgb=True):
+
+        self.test = test
+        self.sample_filenames = os.listdir(root+"00_preview_cropped")
+        self.transform = transform
+        self.img_dir = root
+        self.color_space = color_space
+        self.is_classification = is_classification
+        self.input_color_space = input_color_space
+        self.normalize_cielab = normalize_cielab
+        self.normalize_rgb = normalize_rgb
+
+        self.train_filenames, self.test_filenames = train_test_split(self.sample_filenames,
+                                                                     test_size=0.2, 
+                                                                     random_state=42) 
+    def __len__(self):
+        if self.test:
+            return len(self.test_filenames)
+        else:
+            return len(self.train_filenames)
+        
+    def __getitem__(self, idx):
+
+        path_idx = "{:04d}".format(idx)
+        img_path = os.path.join(self.img_dir, "00_preview_cropped/" + self.sample_filenames[idx])
+
+        image = np.array(Image.open(img_path))
+        # Convert image to lab if the input space is CIELab.
+        # Image is a numpy array always. Convert to tensor at the end.
+        if self.input_color_space == "CIELab":
+            image = scicolor.rgb2lab(image)
+            image = torch.from_numpy(image)
+            # if self.normalize_cielab:
+            #     image = torch.from_numpy(image)
+            #     image = normalize_CIELab(image)
+        else:
+            image = torch.from_numpy(image)
+
+        
+        # Apply kmeans on RGB image always.
+        bg_path = os.path.join("../destijl_dataset/01_background/" + self.sample_filenames[idx])
+        # Most dominant color in RGB.
+        color = self.kmeans_for_bg(bg_path)[0]
+
+        # If output is in CIELab space but input is in RGB, convert target to CIELab also.
+        if self.color_space == "CIELab":
+            target_color = torch.squeeze(torch.tensor(RGB2CIELab(color.astype(np.int32))))
+            # if self.normalize_cielab:
+            #     target_color = normalize_CIELab(target_color)
+        # Input and output is in RGB space or input and output is in CIELab space.
+        # If Input is in CIELab and output is in RGB, than this is also valid since dataset is in RGB.
+        else:
+            target_color = torch.squeeze(torch.tensor(color))
+
+        if self.is_classification:
+            target_color = [torch.zeros(256), torch.zeros(256), torch.zeros(256)]
+            target_color[0][color[0]] = 1
+            target_color[1][color[1]] = 1
+            target_color[2][color[2]] = 1
+
+        if self.transform:
+            # Reshape the image if not in (C, H, W) form.
+            if image.shape[0] != 3:
+                image = image.reshape(-1, image.shape[0], image.shape[1]).type("torch.FloatTensor")
+            # Apply the transformation
+            image = self.transform(image)
+        
+        if self.normalize_rgb:
+            image /= 255
+
+        if self.color_space == "RGB" and self.normalize_rgb:
+            target_color /= 255
+
+        if self.normalize_cielab:
+            # we will only use lightness
+            target_color /= 100
+
+        return image, target_color
+    
+    def kmeans_for_bg(self, bg_path):
+        image = cv2.imread(bg_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        n_colors = 1
+
+        # Apply KMeans to the text area
+        pixels = np.float32(image.reshape(-1, 3))
+        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, .1)
+        flags = cv2.KMEANS_RANDOM_CENTERS
+
+        _, labels, palette = cv2.kmeans(pixels, n_colors, None, criteria, 10, flags)
+        palette = np.asarray(palette, dtype=np.int64) # RGB
+
+        return palette

color_palette/colorCNN.py

@@ -0,0 +1,238 @@
+"""
+
+    TODO: 
+    *   Make the white backgrounds transparent.
+    *   Locate the images in the bounding boxes in preview images in white background. 
+        Cut the images and paste to the location on the decoration layaer.
+    *   The pasting order: Decoration will have white bg. Paste image. Paste text.
+    *   When the white bg images are done, feed them to CNN.
+    *   The output will be the missing color which is the background color. 
+    *   Use CIELab distances to train
+
+"""
+
+import cv2
+import numpy as np
+from utils import *
+from dataset_processing import ProcessedDeStijl
+from PIL import Image
+
+class DestijlProcessorCNN():
+
+    def __init__(self, data_path):
+        self.path_dict = {
+            'preview': data_path + '/00_preview/',
+            'image': data_path + '/02_image/',
+            'decoration': data_path + '/03_decoration/',
+            'text': data_path + '/04_text/',
+        }
+
+        self.rgba_path_dict = {
+            'preview': data_path + '/rgba_dataset/00_preview/',
+            'image': data_path + '/rgba_dataset/02_image/',
+            'decoration': data_path + '/rgba_dataset/03_decoration/',
+            'text': data_path + '/rgba_dataset/04_text/',
+            'temporary': data_path + '/rgba_dataset/05_temporary/',
+            'cropped_preview': data_path + '/rgba_dataset/00_preview_cropped/',
+        }
+
+        self.xml_path_dict = {
+            'preview': data_path + '/xmls/00_preview/',
+            'image': data_path + '/xmls/02_image/',
+            'decoration': data_path + '/xmls/03_decoration/',
+            'text': data_path + '/xmls/04_text/',
+        }
+
+        self.processed_dataset = ProcessedDeStijl("../destijl_dataset")
+
+    def whitebg_to_transparent(self, img_path, layer):
+
+        """
+            WORKING
+        """
+        image_bgr = cv2.imread(img_path)
+        image_num = img_path[-8:]
+        # get the image dimensions (height, width and channels)
+        h, w, c = image_bgr.shape
+        # append Alpha channel -- required for BGRA (Blue, Green, Red, Alpha)
+        image_bgra = np.concatenate([image_bgr, np.full((h, w, 1), 255, dtype=np.uint8)], axis=-1)
+        # create a mask where white pixels ([255, 255, 255]) are True
+        white = np.all(image_bgr == [255, 255, 255], axis=-1)
+        # change the values of Alpha to 0 for all the white pixels
+        image_bgra[white, -1] = 0
+        # save the image
+        cv2.imwrite(self.rgba_path_dict[layer]+image_num, image_bgra)
+
+    def locate_images_in_image_layer(self, idx):
+
+        method = cv2.TM_SQDIFF_NORMED
+        path_idx = "{:04d}".format(idx)
+        preview_img = cv2.imread(self.path_dict["preview"]+path_idx+".png")
+        preview_bboxes = VOC2bbox(self.xml_path_dict["image"]+path_idx+".xml")[1]
+        image_img = cv2.imread(self.path_dict["image"]+path_idx+".png")
+        
+        boxes = []
+        design_boxes = []
+        for box in preview_bboxes:
+            xmin = box[0][0]
+            xmax = box[1][0]
+            ymin = box[0][1]
+            ymax = box[2][1]
+            cropped_img = preview_img[ymin:ymax, xmin:xmax]
+            
+            if(cropped_img.shape[0] > image_img.shape[0]):
+                diff_x = abs(cropped_img.shape[0] - image_img.shape[0])
+                image_img = cv2.copyMakeBorder(image_img, diff_x//2+5, diff_x//2+5, 0, 0, cv2.BORDER_CONSTANT, value=[255, 255, 255])
+
+            if(cropped_img.shape[1] > image_img.shape[1]):
+                diff_y = abs(cropped_img.shape[1] - image_img.shape[1])
+                image_img = cv2.copyMakeBorder(image_img, 0, 0, diff_y//2+5, diff_y//2+5, cv2.BORDER_CONSTANT, value=[255, 255, 255])
+
+            result = cv2.matchTemplate(cropped_img, image_img, method)
+            mn,_,mnLoc,_ = cv2.minMaxLoc(result)
+            MPx,MPy = mnLoc
+            trows,tcols = cropped_img.shape[:2]
+            boxes.append([MPx, MPx+tcols, MPy, MPy+trows])
+            design_boxes.append([xmin, xmax, ymin, ymax])
+
+        self.check_boxes(design_boxes, idx)
+        return boxes, design_boxes
+    
+    def check_boxes(self, bboxes, idx):
+        path_idx = "{:04d}".format(idx)
+        im = cv2.imread("../destijl_dataset/02_image/" + path_idx + ".png")
+        for box in bboxes:
+            # [[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]]
+            xmin = box[0]
+            xmax = box[1]
+            ymin = box[2]
+            ymax = box[3]
+            cv2.rectangle(im,(xmin, ymin),(xmax, ymax),(255,0,0),2)
+        cv2.imwrite("check_boxes.jpg", im)
+
+    def map_image_coordinates(self, text_coordinate, design_text_coordinate, design_img_coordinates, design_size, text_size):
+        prev_x, prev_y = design_size
+        text_x, text_y = text_size
+
+        design_x, design_y = design_text_coordinate[0]
+        text_x, text_y = text_coordinate[0]
+
+        diff_x = text_x - design_x
+        diff_y = text_y - design_y
+        
+        new_coordinates = []
+        for coordinate in design_img_coordinates:
+            for i in range(len(coordinate)):
+                if i < 2:
+                    coordinate[i] = int(coordinate[i] + diff_x)
+                else:
+                    coordinate[i] = int(coordinate[i] + diff_y)
+
+                if coordinate[i] < 0:
+                    coordinate[i] *= -1
+            new_coordinates.append(coordinate)
+        
+        return new_coordinates
+    
+    def convert_to_min_max_coordinate(self, box):
+        xmin, ymin = np.min(box, axis=0)
+        xmax, ymax = np.max(box, axis=0)
+        return [int(xmin), int(xmax), int(ymin), int(ymax)]
+
+    def paste_onto_decoration_layer(self, idx):
+        path_idx = "{:04d}".format(idx)
+        preview_path = self.path_dict["preview"] + path_idx + ".png"
+        img_path = self.path_dict["image"] + path_idx + ".png"
+        decoration_path = self.path_dict["decoration"] + path_idx + ".png"
+        text_path = self.path_dict["text"] + path_idx + ".png"
+        white_bg_text_path = self.rgba_path_dict["text"] + path_idx + ".png"
+        white_bg_img_path = self.rgba_path_dict["image"] + path_idx + ".png"
+
+        img = cv2.imread(white_bg_img_path)
+        decoration_img = cv2.imread(decoration_path)
+        prev = cv2.imread(preview_path)
+
+        design_size = (prev.shape[0], prev.shape[1])
+        text_size = (decoration_img.shape[0], decoration_img.shape[1])
+
+        image_boxes, design_image_boxes = self.locate_images_in_image_layer(idx)
+
+        text_bboxes, white_bg_text_boxes, texts = self.processed_dataset.extract_text_bbox(text_path, preview_path)
+        text_bboxes_from_design, composed_text_palettes = self.processed_dataset.extract_text_directly(preview_path, texts)
+
+        if not text_bboxes_from_design or not white_bg_text_boxes:
+            pass
+        else:
+            design_text_coordinate = text_bboxes_from_design[0]
+            text_coordinate = white_bg_text_boxes[0]
+            new_image_boxes = self.map_image_coordinates(text_coordinate, design_text_coordinate, design_image_boxes, design_size, text_size)
+
+            white_bg = np.zeros( [decoration_img.shape[0], decoration_img.shape[1], 3] ,dtype=np.uint8)
+            white_bg.fill(255)
+
+            cv2.imwrite('bg.jpg', white_bg)
+
+            white_bg = Image.open('bg.jpg')
+
+            decoration_overlay = Image.open(self.rgba_path_dict["decoration"] + path_idx + ".png")
+            text_overlay = Image.open(white_bg_text_path)
+            white_bg.paste(decoration_overlay, mask=decoration_overlay)
+            
+            for j, box in enumerate(new_image_boxes):
+                xmin1, xmax1, ymin1, ymax1 = box # box place on decoration
+                xmin2, xmax2, ymin2, ymax2 = image_boxes[j] # box place on image
+                cropped_img = img[ymin2:ymax2, xmin2:xmax2]
+
+                cv2.imwrite(self.rgba_path_dict["temporary"] + path_idx + ".png", cropped_img)
+                self.whitebg_to_transparent(self.rgba_path_dict["temporary"] + path_idx + ".png", "temporary")
+                cropped_img = Image.open(self.rgba_path_dict["temporary"] + path_idx + ".png")
+
+                offset = (xmin1, ymin1)
+                white_bg.paste(cropped_img, offset, mask=cropped_img)
+
+            white_bg.paste(text_overlay, mask=text_overlay)
+            white_bg.save(self.rgba_path_dict["preview"] + path_idx + ".png")
+
+    def pipeline(self):
+        for idx in range(550, 706):
+            print("Sample: ", idx)
+            path_idx = "{:04d}".format(idx)
+
+            img_path = self.path_dict["image"]+path_idx+".png"
+            text_path = self.path_dict["text"]+path_idx+".png"
+            decoration_path = self.path_dict["decoration"]+path_idx+".png"
+
+            self.whitebg_to_transparent(img_path, "image")
+            self.whitebg_to_transparent(text_path, "text")
+            self.whitebg_to_transparent(decoration_path, "decoration")
+
+            self.paste_onto_decoration_layer(idx)
+
+    def resize_images(self):
+        for idx in range(84, 705):
+            path_idx = "{:04d}".format(idx)
+
+            if os.path.exists(self.rgba_path_dict["preview"]+path_idx+".png"):
+                print(idx)
+
+                method = cv2.TM_SQDIFF_NORMED
+                big_image = cv2.imread(self.rgba_path_dict["preview"]+path_idx+".png")
+                small_image = cv2.imread(self.path_dict["preview"]+path_idx+".png")
+                if(small_image.shape[0] > big_image.shape[0]):
+                    diff_x = abs(big_image.shape[0] - small_image.shape[0])
+                    big_image = cv2.copyMakeBorder(big_image, diff_x//2+5, diff_x//2+5, 0, 0, cv2.BORDER_CONSTANT, value=[255, 255, 255])
+
+                if(small_image.shape[1] > big_image.shape[1]):
+                    diff_y = abs(big_image.shape[1] - small_image.shape[1])
+                    big_image = cv2.copyMakeBorder(big_image, 0, 0, diff_y//2+5, diff_y//2+5, cv2.BORDER_CONSTANT, value=[255, 255, 255])
+
+                result = cv2.matchTemplate(small_image, big_image, method)
+                mn,_,mnLoc,_ = cv2.minMaxLoc(result)
+                MPx,MPy = mnLoc
+                trows,tcols = small_image.shape[:2]
+                box = [MPx, MPx+tcols, MPy, MPy+trows]
+                new_img = big_image[box[2]:box[3], box[0]:box[1]]
+                cv2.imwrite(self.rgba_path_dict["cropped_preview"]+path_idx+".png", new_img)
+
+processor = DestijlProcessorCNN("../destijl_dataset")
+processor.resize_images()

color_palette/config.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass
+
+@dataclass
+class DataConfig:
+    
+    dataset = "shape_dataset_circle_image/"
+    model_name = "ColorAttentionCircleImageLayer_random_mask"
+    data_type = "processed_rgb_toy_dataset_circle_image_color_and_layer"
+
+    # dataset = "../shape_dataset_lightness_circle/"
+    # model_name = "ColorAttentionLightnessCircle_random_mask_class"
+    # data_type = "processed_rgb_toy_dataset_lightness_circle_color_and_layer"
+
+    # dataset = "../destijl_dataset/"
+    # model_name = "ColorGNN_random_mask_new"
+    # data_type = "processed_rgb_color_and_layer"
+
+    feature_size = 68
+    loss_function = "CrossEntropy"
+
+    device = "cpu"
+    lr = 0.005
+    batch_size = 1
+    weight_decay = 0
+    num_epoch = 300
+
+    node_to_mask = -1
+    is_classification = True
+    layers_to_consider = ["background", "text", "image"]

color_palette/config/conf.yaml

@@ -0,0 +1,15 @@
+model_name: "ColorAttentionColor_Layer"
+data_type: "processed_rgb_toy_dataset_color_and_layer"
+feature_size: 4
+
+threshold_for_neighbours: 1
+
+loss_function: "MSE"
+
+device: "cuda:2"
+lr: 0.01
+batch_size: 16
+weight_decay: 0
+num_epoch: 300
+
+node_to_mask: -1

color_palette/config/confCNN.yaml

@@ -0,0 +1,16 @@
+batch_size: 16
+color_space: CIELab
+device: cuda:1
+input_color_space: RGB
+input_size: 512
+is_classification: false
+loss_function: MSE
+out_features: 1
+lr: 0.05
+map_outputs: true
+model_name: ColorCNN_lightness_RGB_normalized
+num_epoch: 150
+step_size: 20
+weight_decay: 0
+normalize_rgb: True
+normalize_cielab: True

color_palette/config/grid_search_conf_generator.py

@@ -0,0 +1,24 @@
+import yaml
+
+lr_list = [0.01, 0.005, 0.001]
+weight_decay = [0, 0.01, 0.001, 0.0001]
+
+count = 0
+for i, lr in enumerate(lr_list):
+    for j, wd in enumerate(weight_decay):
+        with open("conf"+str(count)+".yaml", "w") as file:
+            data = {
+                "model_name": "ColorGNNEmbedding_lr"+str(lr)+"_wd"+str(wd),
+                "data_type": "processed_rgb",
+                "feature_size": 1005,
+
+                "device": "cuda:2",
+                "lr": lr,
+                "batch_size": 1,
+                "weight_decay": wd,
+                "num_epoch": 150,
+
+                "dataset_root": "../destijl_dataset",
+            }
+            documents = yaml.dump(data, file)
+        count+=1

color_palette/cube_num_one_hot_LR/test_gt.npy

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

color_palette/cube_num_one_hot_LR/test_preds.npy

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

color_palette/cube_num_one_hot_LR/test_preds_graph.npy

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

color_palette/cube_num_one_hot_LR/test_rgb_colors.npy

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

color_palette/cube_num_one_hot_LR_sequential/new_palettes.npy

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

color_palette/cube_num_one_hot_LR_sequential/new_palettes_purple.npy

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

color_palette/cube_num_one_hot_LR_sequential/original_palettes.npy

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

color_palette/cube_num_one_hot_LR_sequential/original_palettes_purple.npy

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

color_palette/cube_num_one_hot_LR_sequential/test_gt.npy

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

color_palette/cube_num_one_hot_LR_sequential/test_preds.npy

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

color_palette/dataset.py

@@ -0,0 +1,215 @@
+import torch
+import torch_geometric
+from torch_geometric.data import Dataset, Data
+from sklearn.model_selection import train_test_split
+import os
+import random
+import numpy as np
+import yaml
+from color_palette.utils import *
+from skimage.color import rgb2lab, lab2rgb
+from color_palette.config import *
+import math 
+
+# with open("config/conf.yaml", 'r') as f:
+#     config = yaml.load(f, Loader=yaml.FullLoader)
+
+config = DataConfig()
+
+data_type = config.data_type
+mask_type = -1
+dataset_root = config.dataset
+
+print(f"Torch version: {torch.__version__}")
+print(f"Cuda available: {torch.cuda.is_available()}")
+print(f"Torch geometric version: {torch_geometric.__version__}")
+
+class GraphDestijlDataset(Dataset):
+    def __init__(self, root, test=False, transform=None, pre_transform=None, cube_mapping=False, square_label=False):
+        """
+        root = Where the dataset should be stored. This folder is split
+        into raw_dir (downloaded dataset) and processed_dir (processed data). 
+        """
+        self.test = test
+        self.sample_filenames = os.listdir(root + data_type +'/')
+        self.processed_data_dir = root + data_type + '/'
+        self.square_label = square_label
+        self.map_to_cube = cube_mapping
+        self.cube_size = 4
+
+        # If you want to use less data than the whole dataset, you can specify the range here.
+        # Than it loads only samples up to that sample.
+        self.sample_filenames = ["data_{:04d}.pt".format(idx) for idx in range(0, 1000)]
+
+        # Train test filenames.
+        self.train_filenames, self.test_filenames = train_test_split(self.sample_filenames, 
+                                                            test_size=0.2, 
+                                                            random_state=42)
+        super().__init__(root, transform, pre_transform)
+        
+    @property
+    def raw_file_names(self):
+        return "empty"
+
+    @property
+    def processed_file_names(self):
+        """ If these files are found in raw_dir, processing is skipped"""
+
+        if self.test:
+            return self.test_filenames
+        else:
+            return self.train_filenames
+
+    def download(self):
+        pass
+
+    def process(self):
+        pass
+
+    def return_all_same_colors(self, input_data):
+        indices_list = [[],[],[],[],[]]
+        unique_colors = torch.from_numpy(np.unique(input_data.x[:, 4:], axis=0))
+        all_colors = input_data.x[:, 4:]
+        for idx, color in enumerate(unique_colors):
+            for node_num, element in enumerate(all_colors):
+                if torch.equal(color, element):
+                    indices_list[idx].append(node_num)
+
+
+        return indices_list, unique_colors
+
+
+    def test_train_mask(self, data):
+
+        '''
+            Input: graph data
+
+            Mask the color of one node. The ground truth color is the last 3 dimension of the feature vector.
+            Data is saved as RGB. 
+            If you want you can convert unnormalized RGB ground truth color to Lab.
+            (Conversion is done using COLORMATH)
+            Put mask on a random node's color information by setting that color to [0, 0, 0].
+
+            Return: new_data with masked RGB colors, color_to_hide in lab, node_to_mask scalar
+        '''
+
+        # Take number of nodes
+        n_nodes = len(data.x)
+        if mask_type == -1:
+            # Chose the color to mask randomly
+            node_to_mask = random.randint(0, n_nodes-1)
+            #node_to_mask = n_nodes-1
+        else:
+            # Mask the red each time
+            node_to_mask = n_nodes-2
+        # If you chosed a folder that has processed_rgb name in it, then all the colors are stores in (0, 255) RGB.
+        feature_vector = data.x
+        # This is our target.
+        
+        also_normal_values = []
+
+        for color in enumerate(feature_vector[:, -3:].clone()):
+            also_normal_values.append(color[1])
+
+        color_to_hide = feature_vector[node_to_mask, -3:].clone()
+
+        if self.map_to_cube:
+            # print("RGB")
+            # print(color_to_hide*255)
+            color_to_hide = self.cube_mapping(color_to_hide*255)
+            # print("CONVERSION")
+            # print(color_to_hide)
+
+        # Conversion to cielab. I do not use it anymore.
+        #color_to_hide = torch.tensor(RGB2CIELab(color_to_hide.numpy().astype(np.int32)))
+        #color_to_hide = torch.tensor(rgb2lab(color_to_hide.numpy().astype(np.int32)))
+        # Set node to mask in feature vector to zero.
+        feature_vector[node_to_mask, -3:] =  torch.Tensor([0, 0, 0]) #torch.Tensor([0.9, 0.1, 0.1])
+
+        # Seperate square and circle labels.
+
+        if self.square_label:
+
+            add_label = torch.zeros([5,1])
+            add_label[4, 0] = 1
+            labels = feature_vector[:, :3]
+            labels[4, 2] = 0
+            labels = torch.cat((labels, add_label), dim=1)
+            feature_vector = torch.cat((labels, feature_vector[:, -3:]), dim=1)
+
+        # Assing the new feature vector to the graph.
+        
+        new_data = data.clone()
+        new_data.x = feature_vector
+
+
+        # This code below is used if we want to apply a threshold while adding edges.
+        # It just removes the edges with a higher distance than the threshold.
+        
+        # new_edge_weight = []
+        # new_edge_index = []
+        # for k, edge in enumerate(new_data.edge_weight):
+        #     if edge.item() < threshold:
+        #         new_edge_weight.append(edge.item())
+        #         new_edge_index.append([new_data.edge_index[0][k], new_data.edge_index[1][k]])
+
+        # new_data.edge_index = torch.Tensor(new_edge_index).T
+        # new_data.edge_weight = torch.Tensor(new_edge_weight)
+        #print("New calculations")
+        #print(new_data.edge_index, new_data.edge_weight)
+
+        if self.map_to_cube:
+            cube_num_list = []
+            cube_colors = torch.zeros((feature_vector.shape[0], int(math.pow(self.cube_size, 3))))
+            for j, color in enumerate(feature_vector[:, -3:]):
+                if j != node_to_mask:
+                    cube_num = self.cube_mapping(color*255)
+                    cube_colors[j][cube_num] = 1
+                    cube_num_list.append(cube_num)
+                    
+            new_data.x = torch.cat((feature_vector[:, :-3], cube_colors), dim=1)
+            new_data.y = cube_num_list
+
+        return new_data, color_to_hide, node_to_mask, also_normal_values
+
+    def len(self):
+        if self.test:
+            return len(self.test_filenames)
+        else:
+            return len(self.train_filenames)
+
+    def get(self, idx):
+        if self.test:
+            data = torch.load(self.processed_data_dir+self.test_filenames[idx])
+            new_data, target_color, node_to_mask, also_normal_values = self.test_train_mask(data)
+        else:
+            data = torch.load(self.processed_data_dir+self.train_filenames[idx])
+            new_data, target_color, node_to_mask, also_normal_values = self.test_train_mask(data)
+        return new_data, target_color, node_to_mask, also_normal_values
+    
+    def cube_mapping(self, color):
+        intervals = np.arange(0, 256, 256//4)
+        cube_coordinates = []
+        for channel in color:
+            i = 0
+            for j, value in enumerate(intervals):
+                if value < channel:
+                    i = j
+            cube_coordinates.append(i)
+        
+        cube_num = cube_coordinates[0]*1 + cube_coordinates[1]*self.cube_size + cube_coordinates[2]*self.cube_size*self.cube_size
+        return cube_num
+    
+    def cube2rgb(self, cube_num):
+        """
+            Return the start of the ranges
+        """
+        cube_num = int(cube_num)
+        intervals = np.arange(0, 256, 256//4)
+        coor2 = cube_num // 16
+        coor1 = (cube_num - coor2*self.cube_size*self.cube_size) // 4
+        coor0 = cube_num - coor2*self.cube_size*self.cube_size - coor1*self.cube_size
+        return [intervals[coor0], intervals[coor1], intervals[coor2]]
+
+if __name__ == '__main__':
+    dataset_obj = GraphDestijlDataset(root=dataset_root)

color_palette/dataset_processing.py

@@ -0,0 +1,505 @@
+import torch
+import torch.nn as nn
+from torch_geometric.data import Dataset
+
+import os
+from utils import *
+
+from paddleocr import PaddleOCR, draw_ocr
+
+from PIL import Image, ImageFont
+import numpy as np
+import cv2
+from sklearn.cluster import KMeans
+from sklearn.metrics import silhouette_score
+from matplotlib.colors import hsv_to_rgb, rgb_to_hsv
+
+from torchvision.models import resnet50, ResNet50_Weights
+from model.CNN import Autoencoder
+
+from model.graph import DesignGraph
+from config import *
+
+class ProcessedDeStijl(Dataset):
+    def __init__(self, data_path):
+        self.path = data_path
+        self.path_dict = {
+            'preview': data_path + '/00_preview/',
+            'background': data_path + '/01_background/',
+            'image': data_path + '/02_image/',
+            'decoration': data_path + '/03_decoration/',
+            'text': data_path + '/04_text/',
+            'theme': data_path + '/05_theme/'
+        }
+
+        self.data_path = data_path
+        self.dataset_size = len(next(os.walk(self.path_dict['preview']))[2])
+        self.ocr = PaddleOCR(use_angle_cls=True, lang='en', use_gpu=False)
+
+        self.criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, .1)
+        self.flags = cv2.KMEANS_RANDOM_CENTERS
+
+        #self.layers = ['background', 'text', "decoration"] # Take this from config file later
+        self.layers = ['background', 'text', "decoration"]
+        
+        self.pretrained_model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
+        # self.pretrained_model = Autoencoder()
+        # self.pretrained_model.load_state_dict(torch.load("../CNN_models/CNNAutoencoder/weights/best.pth")["state_dict"])
+
+    def len(self):
+        return self.dataset_size
+
+    def get(self, idx):
+        '''
+            Return a graph object based on the information
+        '''
+        pass
+            
+    ######### RUNTIME EXTRACTION ###########
+
+
+    ######### PROCESSING AND ANNOTATING THE DATASET ###########
+
+    def extract_text_bbox(self, img_path, preview_image_path):
+        '''
+            Input: path to the text image
+            Extract text using paddleOCR.
+            Crop text from bounding box.
+            Extract colors using Kmeans inside the bbox.
+            Return the dominant color and the position.
+            
+            DONE: Try to combine very close lines as paragraph bbox. 
+            If the the distance between two bbox is smaller than the bbox height and color is the same,
+            we can group them as paragraphs.
+
+            Return: text color palettes, dominant colors for each text and position list (as bboxes).
+        '''
+        # Parameters for KMeans.
+        n_colors = 3
+
+        result = self.ocr.ocr(img_path, cls=True)[0]
+
+        image = Image.open(img_path).convert('RGB')
+        boxes = [line[0] for line in result]
+        texts = [line[1][0] for line in result]
+        image = cv2.imread(img_path)
+        preview_image = cv2.imread(preview_image_path)
+
+        palettes = []
+        dominants = []
+        new_bboxes = []
+
+        # Run KMeans for each text object
+        for bbox in boxes:
+            # Crop the text area
+            x, y = int(bbox[0][0]), int(bbox[0][1])
+            z, t = int(bbox[2][0]), int(bbox[2][1])
+            cropped_image = image[y:t, x:z]
+
+            # Do template matching to find the places at the actual image because not every image has the same size.
+            method = cv2.TM_SQDIFF_NORMED
+            result = cv2.matchTemplate(cropped_image, preview_image, method)
+            mn,_,mnLoc,_ = cv2.minMaxLoc(result)
+            MPx,MPy = mnLoc
+            trows,tcols = cropped_image.shape[:2]
+            # --> left top, right top, right bottom, left bottom
+            bbox = [[MPx,MPy], [MPx+tcols, MPy], [MPx+tcols, MPy+trows], [MPx, MPy+trows]]
+            new_bboxes.append(bbox)
+
+        return new_bboxes, boxes, texts
+
+    def compose_paragraphs(self, text_bboxes, text_palettes):
+
+        '''
+            Compose text data into paragraphs.
+            Return: Grouped indices of detected text elements.
+        '''
+
+        num_text_boxes = len(text_bboxes)
+        if num_text_boxes == 0:
+            return False
+        composed_text_idxs = [[0]]
+        for i in range(num_text_boxes-1):
+            palette1 = text_palettes[i]
+            palette2 = text_palettes[i+1]
+            if np.array_equal(palette1, palette2):
+                bbox1 = text_bboxes[i]
+                bbox2 = text_bboxes[i+1]
+                height1 = bbox1[0][1] - bbox1[3][1]
+                height2 = bbox2[0][1] - bbox2[3][1]
+                if abs(bbox1[0][1]-bbox2[0][1]) <= abs(height1)+30:
+                    if i != 0 and i not in composed_text_idxs[-1]:
+                        composed_text_idxs.append([i])
+                    composed_text_idxs[-1].append(i+1)
+                else:
+                    if i != 0 and i not in composed_text_idxs[-1]:
+                        composed_text_idxs.append([i])
+                    if i == num_text_boxes-2:
+                        composed_text_idxs.append([i+1])
+            else:
+                if i != 0 and i not in composed_text_idxs[-1]:
+                    composed_text_idxs.append([i])
+                if i == (num_text_boxes-2):
+                    composed_text_idxs.append([i+1])
+
+        return composed_text_idxs
+
+    def merge_bounding_boxes(self, composed_text_idxs, bboxes):
+        '''
+            openCV --> x: left-to-right, y: top--to-bottom
+            bbox coordinates --> [[256.0, 1105.0], [1027.0, 1105.0], [1027.0, 1142.0], [256.0, 1142.0]]
+                             --> left top, right top, right bottom, left bottom
+
+            TODO: Also return color palettes for each merged box.
+        '''
+        
+        biggest_borders = []
+        if len(bboxes) == 0:
+            return biggest_borders
+        for idxs in composed_text_idxs:
+            smallest_x = smallest_y = 10000
+            biggest_y = biggest_x = 0
+            if len(idxs) > 1:
+                for idx in idxs:
+                    bbox = bboxes[idx]
+                    bbox_smallest_x, bbox_smallest_y = np.min(bbox, axis=0)
+                    bbox_biggest_x, bbox_biggest_y = np.max(bbox, axis=0)
+
+                    if smallest_x > bbox_smallest_x:
+                        smallest_x = bbox_smallest_x
+                    if smallest_y > bbox_smallest_y:
+                        smallest_y = bbox_smallest_y
+                    if biggest_x < bbox_biggest_x:
+                        biggest_x = bbox_biggest_x
+                    if biggest_y < bbox_biggest_y:
+                        biggest_y =  bbox_biggest_y
+
+                biggest_border = [[smallest_x, smallest_y], [biggest_x, smallest_y], [biggest_x, biggest_y], [smallest_x, biggest_y]]
+                biggest_borders.append(biggest_border)
+            else:
+                biggest_borders.append(bboxes[idxs[0]])
+        return biggest_borders
+
+    def mini_kmeans(self, biggest_border, n_colors, image):
+        # for text
+        x, y = int(biggest_border[0][0]), int(biggest_border[0][1])
+        z, t = int(biggest_border[2][0]), int(biggest_border[2][1])
+        cropped_image = image[y:t, x:z]
+        pixels = np.float32(cropped_image.reshape(-1, 3))
+        _, labels, palette = cv2.kmeans(pixels, n_colors, None, self.criteria, 10, self.flags)
+        palette = np.asarray(palette, dtype=np.int64)
+
+        _, counts = np.unique(labels, return_counts=True)
+        color = palette[np.argmin(counts)]
+        return color
+
+
+    def extract_text_directly(self, img_path, white_bg_texts):
+        n_colors = 2
+
+        result = self.ocr.ocr(img_path, cls=True)[0]
+
+        image = Image.open(img_path).convert('RGB')
+        boxes = [line[0] for line in result]
+        texts = [line[1][0].replace(" ", "").lower() for line in result]
+        white_bg_texts = [elem.replace(" ", "").lower() for elem in white_bg_texts]
+        image = cv2.imread(img_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        same_idxs = []
+        new_boxes = []
+        
+        composed_text_palettes = []
+        for j, elem in enumerate(white_bg_texts):
+            for i, text in enumerate(texts):
+                if similar(elem, text) > 0.85:
+                    new_boxes.append(boxes[i])
+                    biggest_border = boxes[i]
+                    composed_text_palettes.append(self.mini_kmeans(biggest_border, n_colors, image))
+                elif i+1 != len(texts):
+                    if similar(elem, text + texts[i+1]) > 0.85:
+                        # merge boxes
+                        bboxes = [boxes[i], boxes[i+1]]
+                        smallest_x = 1000
+                        smallest_x = smallest_y = 10000
+                        biggest_y = biggest_x = 0
+                        for idx in [0, 1]:
+                            bbox = bboxes[idx]
+                            bbox_smallest_x, bbox_smallest_y = np.min(bbox, axis=0)
+                            bbox_biggest_x, bbox_biggest_y = np.max(bbox, axis=0)
+
+                            if smallest_x > bbox_smallest_x:
+                                smallest_x = bbox_smallest_x
+                            if smallest_y > bbox_smallest_y:
+                                smallest_y = bbox_smallest_y
+                            if biggest_x < bbox_biggest_x:
+                                biggest_x = bbox_biggest_x
+                            if biggest_y < bbox_biggest_y:
+                                biggest_y =  bbox_biggest_y
+
+                        biggest_border = [[smallest_x, smallest_y], [biggest_x, smallest_y], [biggest_x, biggest_y], [smallest_x, biggest_y]]
+                        new_boxes.append(biggest_border)
+                        composed_text_palettes.append(self.mini_kmeans(biggest_border, n_colors, image))
+
+        return new_boxes, composed_text_palettes
+
+    def extract_decor_elements(self, decoration_path, preview_path):
+        # Determine the number of dominant colors
+        num_colors = 6
+        
+        # Load the image
+        image = cv2.imread(decoration_path)
+        
+        # Convert the image to the RGB color space
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        image2 = image.copy()
+        
+        # Reshape the image to a 2D array of pixels
+        pixels = image.reshape(-1, 3)
+        
+        # Apply K-means clustering with the determined number of colors
+        kmeans = KMeans(n_clusters=num_colors)
+        kmeans.fit(pixels)
+        
+        # Get the RGB values of the dominant colors
+        colors = kmeans.cluster_centers_.astype(int)
+        
+        # Convert the colors to the HSV color space
+        hsv_colors = [] 
+
+        for i, color in enumerate(colors):
+            x, y, z = color
+            if not (252 < x < 256 and 252 < y < 256 and 252 < z < 256):
+                x, y, z = rgb_to_hsv([x/255, y/255, z/255])
+                hsv_colors.append([x*180, y*255, z*255])
+        # Convert the image to the HSV color space
+        hsv_image = cv2.cvtColor(image2, cv2.COLOR_RGB2HSV)
+        
+        # Create masks for each dominant color
+        masks = []
+        hsv_colors = np.asarray(hsv_colors, dtype=np.int32)
+        
+        colors = []
+        for i in range(len(hsv_colors)):
+            
+            h, s, v = hsv_colors[i, :]
+            lower_color = hsv_colors[i, :] - np.array([10, 50, 50])
+            upper_color = hsv_colors[i, :] + np.array([10, 255, 255])
+            mask = cv2.inRange(hsv_image, lower_color, upper_color)
+            colors.append([h,s,v])
+            masks.append(mask)
+        
+        # Find contours in each mask
+        contours = []
+        for mask in masks:
+            contours_color, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+            contours.append(contours_color)
+        
+        # Draw bounding boxes around the shapes
+        image_with_boxes = image.copy()
+        bboxes = []
+        for i, contour_color in enumerate(contours):
+            for contour in contour_color:
+                x, y, w, h = cv2.boundingRect(contour)
+                # left top, right top, right bottom, left bottom
+                bboxes.append([[x,y], [x+w, y], [x+w, y+h], [x,y+h]])
+
+        new_bboxes = delete_too_small_bboxes(np.asarray(bboxes))
+        return colors, new_bboxes
+
+    def map_decoration_coordinates(self, design_text_coordinate, text_coordinate, decoration_coordinates, prev_size, text_size):
+        # --> [[256.0, 1105.0], [1027.0, 1105.0], [1027.0, 1142.0], [256.0, 1142.0]]
+        # --> left top, right top, right bottom, left bottom
+
+        prev_x, prev_y = prev_size
+        text_x, text_y = text_size
+
+        design_x, design_y = design_text_coordinate[0]
+        text_x, text_y = text_coordinate[0]
+
+        diff_x = text_x - design_x
+        diff_y = text_y - design_y
+        
+        new_coordinates = []
+        for coordinate in decoration_coordinates:
+            new_coor = []
+            for elem in coordinate:
+                new_coor.append([elem[0]-diff_x, elem[1]-diff_y])
+            new_coordinates.append(new_coor)
+
+        return new_coordinates
+
+    def extract_image(self, preview_path, image_path):
+        '''
+            Use Template Matching the put a bounding box around the main image. Use it as the position.
+            Extract colors using KMeans.
+            Return: image color palettes and position list (as bboxes).
+        '''
+        
+        preview_image = cv2.imread(preview_path)
+        preview_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        cropped_image_path = trim_image(image_path, "02_image")
+        image = cv2.imread(cropped_image_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        if image.shape[0] > preview_image.shape[0]:
+            diff_x = image.shape[0] - preview_image.shape[0]
+            image = image[(diff_x//2+1):image.shape[0]-(diff_x//2+1), :]
+        
+        if image.shape[1] > preview_image.shape[1]:
+            diff_y = image.shape[1] - preview_image.shape[1]
+            image = image[:, (diff_y//2+1):image.shape[1]-(diff_y//2+1)]
+
+        method = cv2.TM_SQDIFF_NORMED
+        result = cv2.matchTemplate(image, preview_image, method)
+        mn,_,mnLoc,_ = cv2.minMaxLoc(result)
+        MPx,MPy = mnLoc
+        trows,tcols = image.shape[:2]
+        bbox = [[MPx,MPy], [MPx+tcols,MPy+trows]]
+        cropped_image = preview_image[MPx:MPx+tcols, MPy:MPy+trows]
+
+        pixels = np.float32(image.reshape(-1, 3))
+
+        n_colors = 6
+
+        _, labels, palette = cv2.kmeans(pixels, n_colors, None, self.criteria, 10, self.flags)
+        palette = np.asarray(palette, dtype=np.int64)
+
+        return [bbox], palette
+
+    def annotate_dataset(self):
+
+        for idx in range(388, self.dataset_size):
+            print("CURRENTLY AT: ", idx)
+            path_idx = "{:04d}".format(idx)
+            preview = self.path_dict['preview'] + path_idx + '.png'
+            decoration = self.path_dict['decoration'] + path_idx + '.png'
+            image = self.path_dict['image'] + path_idx + '.png'
+            text = self.path_dict['text'] + path_idx + '.png'
+
+            text_bboxes, white_bg_text_boxes, texts = self.extract_text_bbox(text, preview)
+            text_bboxes_from_design, composed_text_palettes = self.extract_text_directly(preview, texts)
+            composed_text_idxs = self.compose_paragraphs(text_bboxes_from_design, composed_text_palettes)
+            merged_bboxes = []
+            if composed_text_idxs != False:
+                merged_bboxes = self.merge_bounding_boxes(composed_text_idxs, text_bboxes_from_design)
+            image_bboxes, image_palette = self.extract_image(preview, image)
+
+            #decoration_hsv_xpalettes, decoration_bboxes = self.extract_decor_elements(decoration, preview)
+            # image_prev = cv2.imread(preview)
+            # image_text = cv2.imread(text)
+            #mapped_decoration_bboxes = self.map_decoration_coordinates(text_bboxes_from_design[0], white_bg_text_boxes[0], decoration_bboxes, (image_prev.shape[0], image_prev.shape[1]), (image_text.shape[0], image_text.shape[1]))
+
+            #create_xml("../destijl_dataset/xmls/03_decoration", path_idx+".xml", mapped_decoration_bboxes)
+            if len(merged_bboxes) == 0:
+                create_xml(self.data_path+"/xmls/04_text", path_idx+".xml", [[[0,0],[0,0],[0,0],[0,0]]])
+            else:
+                create_xml(self.data_path+"/xmls/04_text", path_idx+".xml", merged_bboxes)
+            create_xml(self.data_path+"/xmls/02_image", path_idx+".xml",  image_bboxes)
+
+    def process_dataset(self, idx):
+        '''
+            Process each node. Construct graph features and save the features as pt files.
+            This code should be used after we have an annotated dataset.
+        '''
+        path_idx = "{:04d}".format(idx)
+
+        img_path_dict = {
+            'preview': self.data_path + '/00_preview/' + path_idx + '.png',
+            'background': self.data_path + '/01_background/' + path_idx + '.png',
+            'image': self.data_path + '/02_image/' + path_idx + '.png',
+            'text': self.data_path + '/04_text/' + path_idx + '.png',
+            'decoration': self.data_path + '/03_decoration/' + path_idx + '.png',
+        }
+
+        annotation_path_dict = {
+            'preview': self.data_path + '/xmls' +'/00_preview/' + path_idx + '.xml',
+            'image': self.data_path + '/xmls' + '/02_image/' + path_idx + '.xml',
+            'text': self.data_path + '/xmls' + '/04_text/' + path_idx + '.xml',
+            'decoration': self.data_path + '/xmls' + '/03_decoration/' + path_idx + '.xml',
+        }
+       
+        all_bboxes = {
+            'image':[], 
+            'background':[], 
+            "decoration":[],
+            'text':[]
+        }
+        all_images = {
+            'image':[], 
+            'background':[], 
+            "decoration":[],
+            'text':[]
+        }
+
+        # For each layer:
+        #   * save all bounding boxes to all_bboxes
+        #   * save all paths of images in which we extract the objects from --> to all_images
+        #   * We generally extract all images from the preview image so that path is preview image path
+        #   * We save this information to use in DesignGraph. It extracts colors from bounding boxes
+        #   using the image we want to extract them from.
+
+        for i, layer in enumerate(self.layers):
+            # Check what is the layer, save information accordingly to dictionaries
+            if layer == "background":
+                # load image as CV image
+                self.preview_img = cv2.imread(img_path_dict[layer])
+                # save the layer image path
+                img = img_path_dict[layer]
+                all_images[layer] = img
+                # since it is background, just add a trivial bbox. This is not used
+                all_bboxes[layer] = [[[0, 0], [self.preview_img.shape[0], 0], [self.preview_img.shape[0], self.preview_img.shape[1]], [0, self.preview_img.shape[1]]]]
+            else:
+                if layer == 'text':
+                    # get the preview path since we extract text directly from preview image
+                    img_path = img_path_dict['preview']
+                    # get annotations from xml
+                    filename, bboxes = VOC2bbox(annotation_path_dict[layer])
+                    # assign bounding boxes and image path to extract colors later
+                    all_bboxes[layer] = bboxes
+                    all_images[layer] = img_path
+                elif layer == 'image' or layer == "decoration":
+                    # same logic is applied as text
+                    img_path = img_path_dict['preview']
+                    self.img_img = cv2.imread(img_path)
+                    filename, bboxes = VOC2bbox(annotation_path_dict[layer])
+                    """
+                    This comment below is for checking whether the annotation boxes work for
+                    each bounding box. It saves the annotated image.
+                    """
+                    # for k, box in enumerate(bboxes):
+                        # im = cv2.imread("../destijl_dataset/00_preview/"+path_idx+".png")
+                        # # [[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]]
+                        # xmin = box[0][0]
+                        # xmax = box[1][0]
+                        # ymin = box[0][1]
+                        # ymax = box[2][1]
+                        # print(xmin, xmax, ymin, ymax)
+                        # cv2.rectangle(im,(xmin, ymin),(xmax, ymax),(255,0,0),2)
+                        # cv2.imwrite("check_bboxes"+str(k)+".jpg", im)
+                    all_bboxes[layer] = bboxes
+                    all_images[layer] = img_path
+            
+
+        # Design graph constructs the graph object and saves it as a pt file.
+        design_graph = DesignGraph(self.pretrained_model, all_images, all_bboxes, self.layers, img_path_dict['preview'], idx)
+        return design_graph.get_all_colors_in_design()
+
+    def trial(self):
+        # Save the samples in the range (0, n)
+        for i in range(0, 3100): 
+            print("Sample: ", i)
+            all_colors = self.process_dataset(i)
+            for nested_list in all_colors:
+                color = nested_list[0]
+                # Fix if has unnecessary extra dimension
+                if len(color.shape) == 2:
+                    color = color[0]
+                color = color.tolist()
+
+if __name__ == "__main__":
+    config = DataConfig()
+    dataset_root = config.dataset
+    dataset = ProcessedDeStijl(data_path=dataset_root)
+    dataset.trial()
+
+

color_palette/deneme.py

@@ -0,0 +1,107 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from dataset import *
+from torch_geometric.loader import DataLoader
+from config import DataConfig
+
+config = DataConfig()
+model_name = config.model_name
+
+test_dataset = GraphDestijlDataset(root=dataset_root, test=True, cube_mapping=True)
+test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
+num_of_plots = len(test_loader)
+
+def my_palplot(pal, size=1, ax=None):
+    """Plot the values in a color palette as a horizontal array.
+    Parameters
+    ----------
+    pal : sequence of matplotlib colors
+        colors, i.e. as returned by seaborn.color_palette()
+    size :
+        scaling factor for size of plot
+    ax :
+        an existing axes to use
+    """
+
+    import numpy as np
+    import matplotlib as mpl
+    import matplotlib.pyplot as plt
+    import matplotlib.ticker as ticker
+
+    n = len(pal)
+    if pal[0][0] > 1:
+        pal = np.array(pal) / 255
+    ax.imshow(np.arange(n).reshape(1, n),
+              cmap=mpl.colors.ListedColormap(list(pal)),
+              interpolation="nearest", aspect="auto")
+    ax.set_xticks(np.arange(n) - .5)
+    ax.set_yticks([-.5, .5])
+    # Ensure nice border between colors
+    ax.set_xticklabels(["" for _ in range(n)])
+    # The proper way to set no ticks
+    ax.yaxis.set_major_locator(ticker.NullLocator())
+
+targets = np.load("targets.npy", allow_pickle=True)
+preds = np.load("preds.npy", allow_pickle=True)
+top_k_preds = np.squeeze(np.load("top_k_preds.npy", allow_pickle=True))
+node_to_masks = np.load("node_to_mask.npy", allow_pickle=True)
+
+rows = num_of_plots//3 + 1
+cols = 3
+fig, ax_array = plt.subplots(rows, cols, figsize=(60, 60), dpi=80, squeeze=False)
+
+column_titles = ["                                                                     Prediction | Target " for i in range(cols)]
+for ax, col in zip(ax_array[0], column_titles):
+    ax.set_title(col, fontdict={'fontsize': 30, 'fontweight': 'medium'})
+
+palettes = []
+for i in range(targets.shape[0]):
+    target = targets[i]
+    pred = preds[i]
+    top_k_pred = top_k_preds[i]
+
+    print(target.shape, pred.shape, top_k_pred.shape)
+
+    palette = np.concatenate((top_k_pred[1:], np.atleast_1d(pred), np.atleast_1d(target)))
+    palettes.append(palette)
+    ax = plt.subplot(rows, cols, i+1)
+
+    rgb_palette = []
+    for color in palette:
+        rgb_value = test_dataset.cube2rgb(color)
+        rgb_palette.append(rgb_value)
+
+    my_palplot(rgb_palette, ax=ax)
+
+    if i == num_of_plots - 1:
+        print("saviing")
+        plt.savefig("../models/"+model_name+"/top_k.jpg")
+
+
+unique_values = np.unique(node_to_masks)
+palettes = np.array(palettes)
+each_node_pred = []
+each_node_target = []
+for value in unique_values:
+    indices = np.where(node_to_masks == value)[0]
+    each_node_pred.append(palettes[indices])
+    each_node_target.append(targets[indices])
+
+targets_repeated = np.repeat(np.expand_dims(targets, axis=1), top_k_preds.shape[1], axis=1)
+total = np.sum(palettes[:, :-1] == targets_repeated)
+print("Total accuracy: ", total/600)
+
+for j, value in enumerate(unique_values):
+
+    targets_repeated = np.repeat(np.expand_dims(each_node_target[j], axis=1), top_k_preds.shape[1], axis=1)
+    total = np.sum(each_node_pred[j][:, :-1] == targets_repeated)
+    print("Accuracy of node " + str(value) + ": ", total)
+
+plt.close()
+plt.hist(preds, bins=64)
+plt.savefig("hist.jpg")
+
+plt.close()
+plt.hist(targets, bins=64)
+plt.savefig("hist_gt.jpg")
+

color_palette/evaluate.py

@@ -0,0 +1,188 @@
+import torch
+import torch.nn as nn
+from torch_geometric.loader import DataLoader
+
+from dataset import GraphDestijlDataset
+import yaml
+import argparse
+
+import seaborn as sns
+import matplotlib.pyplot as plt
+from utils import *
+from config import *
+
+######################## Set Parameters ########################
+
+os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+
+# you can specify the config file you want to provide
+# parser = argparse.ArgumentParser()
+# parser.add_argument("--config_file", type=str, default="config/conf.yaml", help="Path to the config file.")
+# args = parser.parse_args()
+# config_file = args.config_file
+
+# with open(config_file, 'r') as f:
+#     config = yaml.load(f, Loader=yaml.FullLoader)
+
+config = DataConfig()
+
+data_type = config.data_type
+model_name = config.model_name
+device = config.device
+feature_size = config.feature_size
+loss_function = config.loss_function
+dataset_root = config.dataset
+our_node_to_mask = config.node_to_mask
+
+model_weight_path = "../models/" + model_name + "/weights/best.pth"
+
+######################## Model ########################
+
+# Prepare dataset
+# Set test=True for testing on the test set. Otherwise it tests on the train set.
+test_dataset = GraphDestijlDataset(root=dataset_root, test=True)
+test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
+num_of_plots = len(test_loader)
+
+# Model is selected according to name you provide
+model = model_switch(model_name, feature_size).to(device)
+model.load_state_dict(torch.load(model_weight_path)["state_dict"])
+criterion = nn.MSELoss()
+######################## Helper Functions ########################
+
+# Skip black flag was used to eliminate examples that includes black colors.
+# Always set to False for default experiments
+
+skip_black_flag=False
+def test(data, target_color, node_to_mask):
+
+    model.eval()
+    out = model(data.x, data.edge_index.long(), data.edge_attr)
+
+    if skip_black_flag:
+        for color in data.y:
+            a, b, c = color
+            if 0 <= a <= 30 and 0 <= b <= 30 and 0 <= c <= 30:
+                return None, None
+            else:
+                loss = colormath_CIE2000(out[node_to_mask, :][0], target_color[0])
+    else:
+        # Loss is only in MSE now.
+        loss = criterion(out[node_to_mask, :][0], target_color[0]/255)
+    return loss, out
+
+def my_palplot(pal, size=1, ax=None):
+    """Plot the values in a color palette as a horizontal array.
+    Parameters
+    ----------
+    pal : sequence of matplotlib colors
+        colors, i.e. as returned by seaborn.color_palette()
+    size :
+        scaling factor for size of plot
+    ax :
+        an existing axes to use
+    """
+
+    import numpy as np
+    import matplotlib as mpl
+    import matplotlib.pyplot as plt
+    import matplotlib.ticker as ticker
+
+    n = len(pal)
+    if ax is None:
+        f, ax = plt.subplots(1, 1, figsize=(n * size, size))
+    ax.imshow(np.arange(n).reshape(1, n),
+              cmap=mpl.colors.ListedColormap(list(pal)),
+              interpolation="nearest", aspect="auto")
+    ax.set_xticks(np.arange(n) - .5)
+    ax.set_yticks([-.5, .5])
+    # Ensure nice border between colors
+    ax.set_xticklabels(["" for _ in range(n)])
+    # The proper way to set no ticks
+    ax.yaxis.set_major_locator(ticker.NullLocator())
+
+# Config for plot
+rows = num_of_plots//3 + 1
+cols = 3
+fig, ax_array = plt.subplots(rows, cols, figsize=(60, 60), dpi=80, squeeze=False)
+
+column_titles = ["                                                                     Prediction | Target " for i in range(cols)]
+for ax, col in zip(ax_array[0], column_titles):
+    ax.set_title(col, fontdict={'fontsize': 30, 'fontweight': 'medium'})
+
+fig.suptitle(model_name+" Test Palettes", fontsize=100)
+
+# Code for evaluation loop
+plot_count = 0
+val_losses = []
+palettes = []
+preds = []
+targets = []
+count = 0
+for i, (input_data, target_color, node_to_mask) in enumerate(test_loader):
+    loss, out = test(input_data.to(device), target_color.to(device), node_to_mask)
+    if loss != None:
+        val_losses.append(loss.item())
+
+        # Get predicton and other colors in the palette
+        ax = plt.subplot(rows, cols, plot_count+1)
+
+        # Get prediction for a masked node
+        prediction = out[node_to_mask, :]
+        print("which node: ", node_to_mask, "count: ", count)
+        preds.append(prediction.detach().cpu()[0])
+        targets.append(target_color.detach().cpu()[0])
+        # Concat unmasked colors with prediction and ground truth
+        other_colors = input_data.y.clone()
+        other_colors = torch.cat([other_colors[0:node_to_mask, :], other_colors[node_to_mask+1:, :]])
+        #print(other_colors[0:node_to_mask, :].shape, other_colors[node_to_mask+1:, :].shape, node_to_mask)
+        other_colors = other_colors.type(torch.float32).detach().cpu().numpy()/255
+        # Normalize since they are in (0, 255) range.
+        # other_colors /= 255
+
+        if loss_function == "CIELab":
+            palette = np.clip(np.concatenate([other_colors, CIELab2RGB(prediction), CIELab2RGB(target_color[0])]), a_min=0, a_max=1)
+        else:
+            # Concat palettes. All of them are between (0, 1)
+            palette = np.clip(np.concatenate([other_colors, prediction.detach().cpu().numpy(), target_color.detach().cpu().numpy()]), a_min=0, a_max=1)
+        # I commented out codes related to calculating results in CIELab
+
+        # if "embedding" in model_name.lower():
+        #     other_colors = other_colors.type(torch.float32).detach().cpu().numpy()
+        #     other_colors /= 255
+        #     palette = np.clip(np.concatenate([other_colors, CIELab2RGB(prediction), CIELab2RGB(target_color[0])]), a_min=0, a_max=1)
+        # else:
+        #     current_palette = torch.cat([other_colors, prediction, target_color.to(device)]).type(torch.float32).detach().cpu().numpy()
+        #     palette = CIELab2RGB(current_palette)
+
+        # Save all the palettes to use it for distribution histograms.
+        palettes.append(prediction.detach().tolist()[0])
+        my_palplot(palette, ax=ax)
+    else:
+        print("none")
+
+    plot_count+=1
+    print(plot_count)
+    if i == num_of_plots-1:
+        path = "../models/"+model_name
+        if not os.path.exists(path):
+            os.mkdir(path)
+
+        if our_node_to_mask == -1:
+            print("hello")
+            plt.savefig(path+"/palettes.jpg")
+        else:
+            print("why red")
+            plt.savefig(path+"/palettes_only_blue.jpg")
+        plt.close()
+
+# This is for checking prediction distribution
+# It is saved as a histogram.
+#check_distributions(palettes)
+
+criterion = nn.MSELoss()
+stacked_pred = torch.stack(preds)
+stacked_target = torch.stack(targets)
+random_results = np.random.random(size=(stacked_pred.shape[0], stacked_pred.shape[1]))
+print(criterion(stacked_pred, stacked_target))
+print(criterion(torch.Tensor(random_results), stacked_target))

color_palette/evaluate_CNN.py

@@ -0,0 +1,180 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+import yaml
+import argparse
+from torchvision import transforms
+
+import seaborn as sns
+import matplotlib.pyplot as plt
+from utils import *
+from model.CNN import *
+from cnn_dataset import *
+
+import pandas as pd
+
+######################## Set Parameters ########################
+
+os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+
+device = "cuda:1"
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model_name", type=str, default="ColorCNN", help="Give the name of the model.")
+args = parser.parse_args()
+model_name = args.model_name
+
+config_file = "../CNN_models/"+model_name+"/conf.yaml"
+
+with open(config_file, 'r') as f:
+    config = yaml.load(f, Loader=yaml.FullLoader)
+
+## Basic Training Parameters ##
+model_name = config["model_name"]
+device = config["device"]
+
+## Neural Network Parameters ##
+loss_function = config["loss_function"]
+out_features = config["out_features"]
+color_space = config["color_space"]
+input_color_space = config["input_color_space"]
+is_classification = config["is_classification"]
+input_size = config["input_size"]
+normalize_rgb = config["normalize_rgb"]
+normalize_cielab = config["normalize_cielab"]
+
+model_weight_path = "../CNN_models/" + model_name + "/weights/best.pth"
+
+if out_features == 1:
+    out_type = "Lightness"
+else:
+    out_type = "Color"
+
+print("Evaluating for the model: ", model_name, "\n",
+      "Loss function: ", loss_function, "\n",
+      "Output Color Space: ", color_space, "\n",
+      "Color or Lightness?: ", out_type, "\n",
+      "Device: ", device, "\n")
+######################## Model ########################
+
+transform = transforms.Compose([
+    transforms.Resize((input_size, input_size)),
+    #transforms.Normalize((255,), (255,))
+    ])
+
+test_dataset = PreviewDataset(transform=transform, 
+                               color_space=color_space, 
+                               input_color_space=input_color_space,
+                               normalize_rgb=normalize_rgb,
+                               normalize_cielab=normalize_cielab,
+                               test=True)
+
+test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
+num_of_plots = len(test_loader)
+
+model = model_switch_CNN(model_name, out_features).to(device)
+model.load_state_dict(torch.load(model_weight_path)["state_dict"])
+
+if loss_function == "MSE":
+    criterion = nn.MSELoss()
+elif loss_function == "Cross-Entropy":
+    criterion = nn.CrossEntropyLoss()
+elif loss_function == "MAE":
+    criterion = nn.L1Loss()
+
+######################## Helper Functions ########################
+def test(data, color):
+    model.eval()
+    out = model(data)
+
+    if out_features == 1:
+        if loss_function != "CIELab":
+            loss = criterion(out[0][0], color[0][0])
+        else:
+            loss = colormath_CIE2000(out[0][0], color[0][0])
+    else:
+        if loss_function != "CIELab":
+            loss = criterion(out, color)
+        else:
+            loss = colormath_CIE2000(out, color)
+    return loss, out
+
+
+def my_palplot(pal, size=1, ax=None):
+    """Plot the values in a color palette as a horizontal array.
+    Parameters
+    ----------
+    pal : sequence of matplotlib colors
+        colors, i.e. as returned by seaborn.color_palette()
+    size :
+        scaling factor for size of plot
+    ax :
+        an existing axes to use
+    """
+
+    import numpy as np
+    import matplotlib as mpl
+    import matplotlib.pyplot as plt
+    import matplotlib.ticker as ticker
+
+    n = len(pal)
+    if ax is None:
+        f, ax = plt.subplots(1, 1, figsize=(n * size, size))
+    ax.imshow(np.arange(n).reshape(1, n),
+              cmap=mpl.colors.ListedColormap(list(pal)),
+              interpolation="nearest", aspect="auto")
+    ax.set_xticks(np.arange(n) - .5)
+    ax.set_yticks([-.5, .5])
+    # Ensure nice border between colors
+    ax.set_xticklabels(["" for _ in range(n)])
+    # The proper way to set no ticks
+    ax.yaxis.set_major_locator(ticker.NullLocator())
+
+rows = num_of_plots//3 + 1
+cols = 3
+
+fig, ax_array = plt.subplots(rows, cols, figsize=(60, 60), dpi=80, squeeze=False)
+column_titles = ["Prediction                   Target" for i in range(cols)]
+for ax, col in zip(ax_array[0], column_titles):
+    ax.set_title(col, fontdict={'fontsize': 45, 'fontweight': 'medium'})
+
+fig.suptitle(model_name+" Test Palettes", fontsize=100)
+
+plot_count = 0
+val_losses = []
+
+outputs = []
+target_colors = []
+
+for i, (input_data, target_color) in enumerate(test_loader):
+    loss, out = test(input_data.to(device), target_color.to(device))
+    val_losses.append(loss.item())
+    # Get predicton and other colors in the palette
+    ax = plt.subplot(rows, cols, plot_count+1)
+
+    if color_space == "CIELab":
+        out = out.detach().cpu().numpy()
+        out = np.append(out, [[30.0, 30.0]], axis=1)
+        target_color = np.array([[target_color.detach().cpu().numpy()[0][0], 30.0, 30.0]])
+        palette = np.clip(np.concatenate([CIELab2RGB(out), CIELab2RGB(target_color)]), a_min=0, a_max=1)
+    else:
+        palette = np.clip(np.concatenate([out.detach().cpu().numpy(), target_color/255]), a_min=0, a_max=1)
+    outputs.append(out)
+    target_colors.append(target_color)
+
+    my_palplot(palette, ax=ax)
+
+    plot_count+=1
+
+    if i == num_of_plots-1:
+        path = "../CNN_models/"+model_name
+        if not os.path.exists(path):
+            os.mkdir(path)
+        plt.savefig(path+"/palettes.jpg")
+        plt.close()
+
+        cielab_dict = {'Output': outputs, 'Targets': target_colors}
+        df = pd.DataFrame(data=cielab_dict)
+
+#df.to_csv("trainset_predictions.csv")

color_palette/evaluate_classification.py

@@ -0,0 +1,217 @@
+import torch
+import torch.nn as nn
+from torch_geometric.loader import DataLoader
+
+from dataset import GraphDestijlDataset
+import yaml
+import argparse
+
+import seaborn as sns
+import matplotlib.pyplot as plt
+from utils import *
+from config import *
+
+from model.GNN import *
+
+######################## Set Parameters ########################
+
+os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+
+# you can specify the config file you want to provide
+# parser = argparse.ArgumentParser()
+# parser.add_argument("--config_file", type=str, default="config/conf.yaml", help="Path to the config file.")
+# args = parser.parse_args()
+# config_file = args.config_file
+
+# with open(config_file, 'r') as f:
+#     config = yaml.load(f, Loader=yaml.FullLoader)
+
+config = DataConfig()
+
+data_type = config.data_type
+model_name = config.model_name
+device = config.device
+feature_size = config.feature_size
+loss_function = config.loss_function
+dataset_root = config.dataset
+our_node_to_mask = config.node_to_mask
+
+model_weight_path = "../models/" + model_name + "/weights/best.pth"
+
+######################## Model ########################
+
+# Prepare dataset
+# Set test=True for testing on the test set. Otherwise it tests on the train set.
+test_dataset = GraphDestijlDataset(root=dataset_root, test=True, cube_mapping=True)
+test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
+num_of_plots = len(test_loader)
+
+# Model is selected according to name you provide
+model = ColorAttentionClassification(feature_size).to(device)
+model.load_state_dict(torch.load(model_weight_path)["state_dict"])
+criterion = nn.CrossEntropyLoss()
+######################## Helper Functions ########################
+
+# Skip black flag was used to eliminate examples that includes black colors.
+# Always set to False for default experiments
+
+skip_black_flag=False
+def test(data, target_color, node_to_mask):
+
+    model.eval()
+    out = model(data.x, data.edge_index.long(), data.edge_attr)
+    loss = criterion(out[node_to_mask, :], target_color)
+    return loss, out
+
+def my_palplot(pal, pred, tar, size=1, ax=None):
+    """Plot the values in a color palette as a horizontal array.
+    Parameters
+    ----------
+    pal : sequence of matplotlib colors
+        colors, i.e. as returned by seaborn.color_palette()
+    size :
+        scaling factor for size of plot
+    ax :
+        an existing axes to use
+    """
+
+    import numpy as np
+    import matplotlib as mpl
+    import matplotlib.pyplot as plt
+    import matplotlib.ticker as ticker
+
+    n = len(pal)
+    if pal[0][0] > 1:
+        pal = np.array(pal) / 255
+    if ax is None:
+        f, ax = plt.subplots(1, 1, figsize=(n * size, size))
+    ax.imshow(np.arange(n).reshape(1, n),
+              cmap=mpl.colors.ListedColormap(list(pal)),
+              interpolation="nearest", aspect="auto")
+    ax.set_xticks(np.arange(n) - .5)
+    ax.set_yticks([-.5, .5])
+    rgb_pred = test_dataset.cube2rgb(pred)
+    rgb_pred_str = " ".join(str(x) for x in rgb_pred)
+    rgb_tar = test_dataset.cube2rgb(tar)
+    rgb_tar_str = " ".join(str(x) for x in rgb_tar)
+    ax.set_ylabel("Pred: " + rgb_pred_str + " Tar: " + rgb_tar_str, rotation=0, labelpad=100)
+    # Ensure nice border between colors
+    ax.set_xticklabels(["" for _ in range(n)])
+    # The proper way to set no ticks
+    ax.yaxis.set_major_locator(ticker.NullLocator())
+
+# Config for plot
+rows = num_of_plots//3 + 1
+cols = 3
+fig, ax_array = plt.subplots(rows, cols, figsize=(60, 60), dpi=80, squeeze=False)
+
+
+column_titles = ["                                                                     Prediction | Target " for i in range(cols)]
+for ax, col in zip(ax_array[0], column_titles):
+    ax.set_title(col, fontdict={'fontsize': 30, 'fontweight': 'medium'})
+
+fig.suptitle(model_name+" Test Palettes", fontsize=100)
+
+# Code for evaluation loop
+plot_count = 0
+val_losses = []
+palettes = []
+preds = []
+targets = []
+count = 0
+top_k_all_preds = []
+node_names = []
+
+for i, (input_data, target_color, node_to_mask) in enumerate(test_loader):
+    loss, out = test(input_data.to(device), target_color.to(device), node_to_mask)
+    if loss != None:
+        val_losses.append(loss.item())
+
+        # Get predicton and other colors in the palette
+        ax = plt.subplot(rows, cols, plot_count+1)
+
+        # Get prediction for a masked node
+        
+        prediction = out[node_to_mask, :]
+        #preds.append(prediction.detach().cpu()[0])
+        #targets.append(target_color.detach().cpu()[0])
+        # Concat unmasked colors with prediction and ground truth
+        other_colors = torch.tensor(input_data.y)
+        # FIX THIS
+        other_colors = torch.cat((other_colors[0][0:node_to_mask], other_colors[0][node_to_mask+1:]))
+        #print(other_colors[0:node_to_mask, :].shape, other_colors[node_to_mask+1:, :].shape, node_to_mask)
+        other_colors = other_colors.type(torch.float32).detach().cpu().numpy()
+        # Normalize since they are in (0, 255) range.
+        # other_colors /= 255
+
+            # Concat palettes. All of them are between (0, 1)
+        a, top_k_preds = F.softmax(torch.tensor(prediction)).topk(5)
+        top_k_all_preds.append(top_k_preds.detach().numpy())
+        node_names.append(node_to_mask.detach().numpy())
+        prediction = torch.argmax(F.softmax(torch.tensor(prediction)), dim=1).numpy()
+        print("Pred: ", prediction, " Target: ", target_color)
+        preds.append(prediction.item())
+        print("pred cube to color: ", test_dataset.cube2rgb(prediction.item()))
+        print("target cube to color: ", test_dataset.cube2rgb(target_color.item()))
+        targets.append(target_color.item())
+        # print(other_colors, prediction, target_color)
+        # print(other_colors[0])
+        # print(np.atleast_1d(prediction), np.atleast_1d(target_color.detach().cpu().numpy()))
+        palette = np.concatenate((other_colors, np.atleast_1d(prediction.item()), np.atleast_1d(target_color.item())))
+        pred_palette = np.concatenate((top_k_preds[0], np.atleast_1d(target_color.item())))
+        # I commented out codes related to calculating results in CIELab
+        # if "embedding" in model_name.lower():
+        #     other_colors = other_colors.type(torch.float32).detach().cpu().numpy()
+        #     other_colors /= 255
+        #     palette = np.clip(np.concatenate([other_colors, CIELab2RGB(prediction), CIELab2RGB(target_color[0])]), a_min=0, a_max=1)
+        # else:
+        #     current_palette = torch.cat([other_colors, prediction, target_color.to(device)]).type(torch.float32).detach().cpu().numpy()
+        #     palette = CIELab2RGB(current_palette)
+
+        # Save all the palettes to use it for distribution histograms.
+       
+        all_colors = []
+        for e, num in enumerate(palette):
+            color = test_dataset.cube2rgb(num)
+            all_colors.append(color)
+
+        #palettes.append(prediction.detach().tolist()[0])
+        my_palplot(all_colors, prediction.item(), target_color.item(), ax=ax)
+    else:
+        print("none")
+
+    plot_count+=1
+    print(plot_count)
+    if i == num_of_plots-1:
+        path = "../models/"+model_name
+        if not os.path.exists(path):
+            os.mkdir(path)
+
+        if our_node_to_mask == -1:
+            print("hello")
+            plt.savefig(path+"/palettes.jpg")
+
+        else:
+            print("why red")
+            plt.savefig(path+"/palettes_only_red.jpg")
+
+N = len(test_dataset)
+print("Evaluation Accuracy: ", np.sum(np.array(preds) == np.array(targets))/N)
+random_results = (np.random.random(size=(N,))*64).astype(np.uint16)
+print("Evaluation Accuracy Random: ", np.sum(random_results == np.array(targets))/N)
+
+np.save("preds.npy", np.array(preds))
+np.save("targets.npy", np.array(targets))
+np.save("node_to_mask.npy", np.array(node_names))
+np.save("top_k_preds.npy", np.array(top_k_all_preds))
+
+# This is for checking prediction distribution
+# It is saved as a histogram.
+#check_distributions(palettes)
+
+# criterion = nn.MSELoss()
+# stacked_pred = torch.stack(preds)
+# stacked_target = torch.stack(targets)
+# random_results = np.random.random(size=(stacked_pred.shape[0], stacked_pred.shape[1]))
+# print(criterion(stacked_pred, stacked_target))
+# print(criterion(torch.Tensor(random_results), stacked_target))

color_palette/evaluate_recommend.py

@@ -0,0 +1,72 @@
+import numpy as np
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import matplotlib.ticker as ticker
+from regressor.config import config_to_use
+
+def my_palplot(pal, size=1, ax=None):
+
+    n = len(pal)
+    if ax is None:
+        f, ax = plt.subplots(1, 1, figsize=(n * size, size))
+    ax.imshow(np.arange(n).reshape(1, n),
+              cmap=mpl.colors.ListedColormap(list(pal)),
+              interpolation="nearest", aspect="auto")
+    ax.set_xticks(np.arange(n) - .5)
+    ax.set_yticks([-.5, .5])
+    # Ensure nice border between colors
+    ax.set_xticklabels(["" for _ in range(n)])
+    # The proper way to set no ticks
+    ax.yaxis.set_major_locator(ticker.NullLocator())
+
+# test_preds = np.load('all_one_hot_LR/test_preds_graph.npy')
+palettes = np.load(config_to_use.save_folder+'/new_palettes_purple.npy')
+original_palettes = np.load(config_to_use.save_folder+'/original_palettes_purple.npy')
+
+print("testing out stuff")
+
+# test_preds = np.expand_dims(test_preds, axis=1)
+
+# all_colors = np.concatenate((palettes, test_preds), axis=1)
+
+colors = np.clip(palettes, a_min=0, a_max=1)
+colors_org = np.clip(original_palettes, a_min=0, a_max=1)
+
+rows = 50
+cols = 2
+fig, ax_array = plt.subplots(rows, cols, figsize=(60, 60), dpi=80, squeeze=False)
+
+column_titles = ["Updated palettes" for i in range(cols)]
+for ax, col in zip(ax_array[0], column_titles):
+    ax.set_title(col, fontdict={'fontsize': 50, 'fontweight': 'medium'})
+
+fig.suptitle("Test Palettes", fontsize=100)
+plot_count = 0
+for i in range(len(colors)):
+    ax = plt.subplot(rows, cols, plot_count+1)
+    my_palplot(colors[i], ax=ax)
+    plot_count += 1
+    if plot_count == 100:
+        break
+
+plt.savefig(config_to_use.save_folder+'/recommended_purple.jpg')
+plt.clf()
+
+rows = 50
+cols = 2
+fig, ax_array = plt.subplots(rows, cols, figsize=(60, 60), dpi=80, squeeze=False)
+
+column_titles = ["Original Palettes" for i in range(cols)]
+for ax, col in zip(ax_array[0], column_titles):
+    ax.set_title(col, fontdict={'fontsize': 50, 'fontweight': 'medium'})
+
+fig.suptitle("Test Palettes", fontsize=100)
+plot_count = 0
+for i in range(len(colors)):
+    ax = plt.subplot(rows, cols, plot_count+1)
+    my_palplot(colors_org[i], ax=ax)
+    plot_count += 1
+    if plot_count == 100:
+        break
+
+plt.savefig(config_to_use.save_folder+'/original_purple.jpg')

color_palette/model/CNN.py

@@ -0,0 +1,209 @@
+import torch
+import torch.nn as nn
+from torchvision.models import resnet50, ResNet50_Weights, resnet18, ResNet18_Weights
+
+class Autoencoder(torch.nn.Module):
+    def __init__(self, num_channels=3, c_hid=16, latent_dim=256):
+        super().__init__()
+         
+        # Building an linear encoder with Linear
+        # layer followed by Relu activation function
+        # 784 ==> 9
+        self.encoder = torch.nn.Sequential(
+            nn.Conv2d(num_channels, c_hid, kernel_size=3, padding=1, stride=2),  # 256x256 => 128x128
+            nn.ReLU(),
+            nn.Conv2d(c_hid, c_hid, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(c_hid, 2 * c_hid, kernel_size=3, padding=1, stride=2),  # 128x128 => 64x64
+            nn.ReLU(),
+            nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1, stride=2),  # 64x64 => 32x32
+            nn.ReLU(),
+            nn.Conv2d(2 * c_hid, c_hid, kernel_size=3, padding=1, stride=2),  # 32x32 => 16x16
+           )
+         
+        # Building an linear decoder with Linear
+        # layer followed by Relu activation function
+        # The Sigmoid activation function
+        # outputs the value between 0 and 1
+        # 9 ==> 784
+        self.linear = nn.Sequential(
+            nn.Flatten(),  # Image grid to single feature vector
+            nn.Linear(16 * 16 * c_hid, num_channels*latent_dim),
+        )
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(
+                num_channels, 2 * c_hid, kernel_size=3, output_padding=1, padding=1, stride=2
+            ),  # 4x4 => 8x8
+            nn.ReLU(),
+            nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(2 * c_hid, c_hid, kernel_size=3, output_padding=1, padding=1, stride=2),  # 8x8 => 16x16
+            nn.ReLU(),
+            nn.Conv2d(c_hid, c_hid, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(
+                c_hid, c_hid, kernel_size=3, output_padding=1, padding=1, stride=2
+            ),  # 16x16 => 32x32
+            nn.ReLU(),
+            nn.ConvTranspose2d(
+                c_hid, num_channels, kernel_size=3, output_padding=1, padding=1, stride=2
+            ),  # 32x32 => 64x64
+            nn.Tanh(),  # The input images is scaled between -1 and 1, hence the output has to be bounded as well
+        )
+ 
+    def forward(self, x):
+        encoded = self.encoder(x)
+        linear = self.linear(encoded)
+        x = linear.reshape(linear.shape[0], -1, 16, 16)
+        decoded = self.decoder(x)
+        return decoded
+
+class FinetuneResNet18_classify(nn.Module):
+    def __init__(self, freeze_resnet=True):
+        super().__init__()
+
+        "Classify the color"
+
+        self.pretrained_model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
+
+        for param in self.pretrained_model.parameters():
+            param.requires_grad_ = False
+            
+        self.pretrained_model.fc = nn.Linear(in_features=512, out_features=1024)
+
+        self.color_head = nn.Sequential(
+            nn.Linear(in_features=1024, out_features=768),
+        )
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x):
+        x = self.pretrained_model(x)
+        x = self.color_head(x)
+        r = self.softmax(x[:, :256])
+        g = self.softmax(x[:, 256:512])
+        b = self.softmax(x[:, 512:])
+        return r, g, b
+    
+class ResNet18(nn.Module):
+    def __init__(self, freeze_resnet=True, map_outputs="RGB"):
+        super().__init__()
+
+        """
+            Just map to interval
+        """
+        self.pretrained_model = resnet18(weights=None)
+        self.map_outputs = map_outputs
+
+        for param in self.pretrained_model.parameters():
+            param.requires_grad_ = True
+            
+        self.pretrained_model.fc = nn.Linear(in_features=512, out_features=256)
+
+        self.color_head = nn.Sequential(
+            nn.Linear(in_features=256, out_features=128),
+            nn.Linear(in_features=128, out_features=64),
+            nn.Linear(in_features=64, out_features=3),
+        )
+
+        if self.map_outputs == "CIELab":
+            self.l_activation = nn.Sigmoid()
+            self.a_activation = nn.Tanh()
+            self.b_activation = nn.Tanh()
+
+        elif self.map_outputs == "RGB":
+            self.activation = nn.Sigmoid()
+
+    def forward(self, x):
+        x = self.pretrained_model(x)
+        x = self.color_head(x)
+        if self.map_outputs == "CIElab":
+            x[:, 0] = self.l_activation(x[:, 0]) * 100
+            x[:, 1] = self.a_activation(x[:, 1]) * 127
+            x[:, 2] = self.b_activation(x[:, 2]) * 127
+
+        elif self.map_outputs == "RGB":
+            x = x
+        return x
+    
+
+class ColorCNN(nn.Module):
+    def __init__(self, num_channels=3, c_hid=16, out_feature=3, reverse_normalize_output=True):
+        super().__init__()
+        self.encoder = torch.nn.Sequential(
+            nn.Conv2d(num_channels, c_hid, kernel_size=3, padding=1, stride=2),  # 512x512 -> 256x256
+            nn.BatchNorm2d(num_features=c_hid),
+            nn.ReLU(),
+            nn.Conv2d(c_hid, c_hid, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(c_hid, 2 * c_hid, kernel_size=3, padding=1, stride=2), # 256x256 -> 128x128
+            nn.BatchNorm2d(num_features=2*c_hid),
+            nn.ReLU(),
+            nn.Conv2d(2 * c_hid, 4 * c_hid, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(4 * c_hid, 2 * c_hid, kernel_size=3, padding=1, stride=2),  # 128x128 => 64x64
+            nn.BatchNorm2d(num_features=2*c_hid),
+            nn.ReLU(),
+            nn.Conv2d(2 * c_hid, c_hid, kernel_size=3, padding=1, stride=2),  # 64x64 => 32x32
+            nn.BatchNorm2d(num_features=c_hid),
+            nn.ReLU(),
+            nn.Conv2d(c_hid, 8, kernel_size=3, padding=1, stride=2),  # 32x32 => 16x16
+            nn.BatchNorm2d(num_features=8),
+            nn.ReLU(),
+           )
+        
+        self.color_head = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(in_features=16*16*8, out_features=out_feature)
+        )
+
+        self.reverse_normalize_output = reverse_normalize_output
+        #self.addition = Addition()
+        #self.multiplication =  torch.Tensor([100, 255, 255]).to("cuda:1")
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.color_head(x)
+        return x
+    
+class Addition(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        input += torch.Tensor([0, 127, 127]).to("cuda:1")
+        return input
+    
+class ColorCNNBigger(nn.Module):
+    def __init__(self, num_channels=3, c_hid=16):
+        super().__init__()
+        self.encoder = torch.nn.Sequential(
+            nn.Conv2d(num_channels, c_hid, kernel_size=3, padding=1, stride=2),  # 512x512 -> 256x256
+            nn.BatchNorm2d(num_features=c_hid),
+            nn.ReLU(),
+            nn.Conv2d(c_hid, 2 * c_hid, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(2 * c_hid, 4 * c_hid, kernel_size=3, padding=1, stride=2), # 256x256 -> 128x128
+            nn.BatchNorm2d(num_features=4*c_hid),
+            nn.ReLU(),
+            nn.Conv2d(4 * c_hid, 2 * c_hid, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(2 * c_hid, c_hid, kernel_size=3, padding=1, stride=2),  # 128x128 => 64x64
+            nn.BatchNorm2d(num_features=c_hid),
+            nn.ReLU(),
+            nn.Conv2d(c_hid, c_hid//2, kernel_size=3, padding=1, stride=2),  # 128x128 => 64x64
+            nn.BatchNorm2d(num_features=c_hid//2),
+           )
+        
+        self.color_head = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(in_features=32*32*8, out_features=16*16*4),
+            nn.Linear(in_features=16*16*4, out_features=3)
+        )
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.color_head(x)
+        return x
+