# from dataset.rico import Rico
from convertHTML.publaynet import PubLayNet
from convertHTML.rico import Rico25Dataset
from convertHTML.magazine import MagezineDataset

def get_dataset(name, datapath, split, transform=None, max_bbox_bins=32):
    if name == 'rico25':
        return Rico25Dataset(datapath, split, transform, max_bbox_bins=max_bbox_bins)

    elif name == 'publaynet':
        return PubLayNet(datapath, split, transform, max_bbox_bins=max_bbox_bins)

    elif name == 'magazine':
        return MagezineDataset(datapath, split, transform, max_bbox_bins=max_bbox_bins)

    raise NotImplementedError(name)
convertHTML/__init__.py
ProjectNUWA/LayoutNUWA
import torch
import seaborn as sns
from pathlib import Path
from torch.utils.data import Dataset
from torch_geometric.data import InMemoryDataset
import torch_geometric  
from torch_geometric.data import DataLoader  
  
class BaseDataset(InMemoryDataset):
    labels = []
    _label2index = None
    _index2label = None
    _colors = None

    def __init__(self, path, split, transform):
        assert split in ['train', 'val', 'test']
        self.path = path
        super().__init__(self.path, transform)
        
        idx = self.processed_file_names.index('{}.pt'.format(split))
        self.data, self.slices = torch.load(self.processed_paths[idx])

    @property
    def label2index(self):
        if self._label2index is None:
            self._label2index = dict()
            for idx, label in enumerate(self.labels):
                self._label2index[label] = idx
        return self._label2index

    @property
    def index2label(self):
        if self._index2label is None:
            self._index2label = dict()
            for idx, label in enumerate(self.labels):
                self._index2label[idx] = label
        return self._index2label

    @property
    def colors(self):
        if self._colors is None:
            n_colors = self.num_classes
            colors = sns.color_palette('husl', n_colors=n_colors)
            self._colors = [tuple(map(lambda x: int(x * 255), c))
                            for c in colors]
        return self._colors

    @property
    def raw_file_names(self):
        raw_dir = Path(self.path)
        if not raw_dir.exists():
            return []
        return [p.name for p in raw_dir.iterdir()]

    @property
    def processed_file_names(self):
        return ['train.pt', 'val.pt', 'test.pt']

    def process(self):
        raise NotImplementedError
    
    

convertHTML/base.py
ProjectNUWA/LayoutNUWA
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizer
import json
import random
import torch


class RawFileDataset(Dataset):  
    def __init__(self, args, file, tokenizer: PreTrainedTokenizer):  
        self.tokenizer = tokenizer  
        self.max_seq_len = args.max_seq_length
        with open(file, "r") as f:
            self.content = [json.loads(line) for line in f]
            
  
    def __len__(self):
        return len(self.content)
    
    def __getitem__(self, index):
        data = self.content[index]
        bbox_seq_modeling = data["bbox_seq_modeling"][0]
        cate_seq_modeling = data["cate_seq_modeling"][0]
        bbox_inputs = self.tokenizer(
            bbox_seq_modeling, 
            padding="max_length", 
            truncation=True, 
            max_length=self.max_seq_len,
            return_tensors="pt",
        )
        bbox_input_ids = bbox_inputs.input_ids[0]
        bbox_attention_mask = bbox_inputs.attention_mask[0]
        bbox_labels = torch.where(bbox_input_ids != self.tokenizer.pad_token_id, bbox_input_ids, -100)
        
        cate_inputs = self.tokenizer(
            cate_seq_modeling, 
            padding="max_length", 
            truncation=True, 
            max_length=self.max_seq_len,
            return_tensors="pt",
        )
        cate_input_ids = cate_inputs.input_ids[0]
        cate_attention_mask = cate_inputs.attention_mask[0]
        cate_labels = torch.where(cate_input_ids != self.tokenizer.pad_token_id, cate_input_ids, -100)
        
        if random.random() < 0.5:
            return {
                "input_ids": bbox_input_ids,
                "attention_mask": bbox_attention_mask,
                "labels": bbox_labels,
            }
        else:
            return {
                "input_ids": cate_input_ids,
                "attention_mask": cate_attention_mask,
                "labels": cate_labels,
            }
    
    
    
    
convertHTML/basedataset.py
ProjectNUWA/LayoutNUWA
import sys
import os
current_directory = os.path.dirname(os.path.abspath(__file__))
parent_directory = os.path.dirname(current_directory)
sys.path.append(parent_directory)  
import torch
import random
import torchvision.transforms as T
import os
import json
import copy
import argparse
from convertHTML.utils import LexicographicSort
from torch_geometric.loader import DataLoader
from torch_geometric.utils import to_dense_batch
from transformers import AutoTokenizer
from convertHTML import get_dataset
from helper.global_var import *
from collections import OrderedDict
from typing import List, Dict  
from tqdm import *
from helper.metrics import *

##################
### Global Config
##################
IGNORE_INDEX = -100
DEFAULT_PAD_TOKEN = "[PAD]"
DEFAULT_EOS_TOKEN = "</s>"
DEFAULT_BOS_TOKEN = "</s>"
DEFAULT_UNK_TOKEN = "</s>"
SPAN_MASK_TOKEN = "<FILL_{i}>"
SEP_TOKEN = "<sep>"
PLACE_HOLDER = "<MASK>"


def round_nested_list(nested_list, decimals):  
    result = []  
    for item in nested_list:  
        if isinstance(item, list):   
            result.append(round_nested_list(item, decimals))  
        else:  
            result.append(round(item, decimals))  
    return result 


class CustomDataLoader(DataLoader):  
    def __init__(
        self, 
        args,
        tokenizer, 
        bbox_quantization, 
        dataset, 
        batch_size,
        shuffle=False, 
        split='train',
        **kwargs
    ):  
        super(CustomDataLoader, self).__init__(dataset, batch_size, shuffle, **kwargs)  
        self.split = split
        
        self.html_template = TEMPLATE_FORMAT.get("html_format")
        self.bbox_template = TEMPLATE_FORMAT.get("bbox_format")
        
        if args.infilling:
            self.cond_cate_to_size_pos = INFILLING_INSTRUCTION.get("cond_cate_to_size_pos")
            self.cond_cate_size_to_pos = INFILLING_INSTRUCTION.get("cond_cate_size_to_pos")
            self.cond_random_mask = INFILLING_INSTRUCTION.get("cond_random_mask")
        else:
            self.cond_cate_to_size_pos = INSTRUCTION.get("cond_cate_to_size_pos")
            self.cond_cate_size_to_pos = INSTRUCTION.get("cond_cate_size_to_pos")
            self.cond_random_mask = INSTRUCTION.get("cond_random_mask")
        if args.add_task_instruction:
            task_instruction = TASK_INSTRUCTION[args.dataset_name]
            self.cond_cate_to_size_pos = task_instruction + self.cond_cate_to_size_pos
            self.cond_cate_size_to_pos = task_instruction + self.cond_cate_size_to_pos
            self.cond_random_mask = task_instruction + self.cond_random_mask
        
        self.category_map = DATASET_META[dataset.dataset_name]
        self.glue_template_train_eval = SEP_SEQ[0]
        self.glue_template_test = SEP_SEQ[1]
        self.glue_template_codegen_train = SEP_SEQ[2]
        self.glue_template_codegen_test = SEP_SEQ[3]
        
        self.tokenizer = tokenizer
        self.N_category = dataset.N_category
        self.bbox_quantization = bbox_quantization  # quanlization approachs
        self.consistency_num = args.consistency_num
        self.infilling = args.infilling
        
        
        
    def filter_invalid_num(self, lst, mask):
        new_lst = []
        for i in range(len(lst)):
            new_lst.append(lst[i][:mask[i].sum().item()])
        return new_lst
    
    
    def build_input_with_ele_dict(self, ele_dict: Dict, type=None):
        answer_notepad = []
        ele_dict = copy.deepcopy(ele_dict)
        if type == "html_content":
            ele_dict = ele_dict
        elif type == "cate_masked_html":
            answer_notepad = ele_dict["c"]
            ele_dict["c"] = PLACE_HOLDER
        elif type == "size_pos_mask_html":
            c = ele_dict["c"]
            answer_notepad = [ele_dict[k] for k in ele_dict if k != "c"]
            ele_dict = dict([(k, PLACE_HOLDER) for k in ele_dict.keys()])
            ele_dict["c"] = c
        elif type == "pos_mask_html":
            answer_notepad = [ele_dict["x"], ele_dict["y"]]
            ele_dict["x"] = PLACE_HOLDER
            ele_dict["y"] = PLACE_HOLDER
        elif type == "random_mask_html":
            random_mask_num = random.choice([3, 4]) # mask up to 80% places (categoty is not masked)
            selected_mask_element = random.sample(['x', 'y', 'w', 'h'], random_mask_num)
            answer_notepad = []
            for key in selected_mask_element:
                answer_notepad.append(ele_dict[key])
                ele_dict[key] = PLACE_HOLDER
            
        return self.bbox_template.format(**ele_dict), answer_notepad
    
    
    def replace_order_mask(self, lst: List[str], ans_lst: List):
        '''
        replace the mask token and build corresponding results
        '''
        new_lst, new_ans = [], {}
        cnt = 1
        for line, ans in zip(lst, ans_lst):
            mask_count = line.count(PLACE_HOLDER)
            for i in range(mask_count):
                mask_token = SPAN_MASK_TOKEN.format(i=cnt)
                line = line.replace(PLACE_HOLDER, mask_token, 1)
                new_ans[mask_token] = ans[i]
                cnt += 1
            new_lst.append(line)
        return new_lst, new_ans
    
    
    def convert_num_to_html(self, coord_lst=None, category_lst=None, self_consistency=False, consistency_num=10):
        batched_html_lst = []  # target
        batched_cond_cate, batched_cond_bbox = [], []  # condition
        cond_cate_to_size_pos, cond_cate_size_to_pos = [], [] 
        random_mask = []
    
        if coord_lst is not None and category_lst is not None: # create the training data   
            for coords, categories in zip(coord_lst, category_lst):
                # store all the input code
                html_content, cate_masked_html = [], []
                size_pos_mask_html, pos_mask_html = [], []
                random_mask_html = []
                
                # store all the ans
                cate_masked_html_ans, random_masked_html_ans = [], []
                size_pos_mask_html_ans, pos_mask_html_ans = [], []
                
                all_category = OrderedDict([(i, 0) for i in range(self.N_category)])
                for coord, category in zip(coords, categories):
                    w, h = coord[2], coord[3]
                    x, y = int(coord[0] - w / 2), int(coord[1] - h / 2)
                    real_category = self.category_map[category]
                    all_category[category] += 1
                    ele_dict = {"c": real_category, "x": x, "y":y, "w":w, "h":h}
                    tmp1, _ = self.build_input_with_ele_dict(ele_dict, "html_content")
                    html_content.append(tmp1)
                    
                    tmp2, ans2 = self.build_input_with_ele_dict(ele_dict, "cate_masked_html")
                    cate_masked_html.append(tmp2)
                    cate_masked_html_ans.append(ans2)
                    
                    tmp3, ans3 = self.build_input_with_ele_dict(ele_dict, "size_pos_mask_html")
                    size_pos_mask_html.append(tmp3)
                    size_pos_mask_html_ans.append(ans3)
                    
                    tmp4, ans4 = self.build_input_with_ele_dict(ele_dict, "pos_mask_html")
                    pos_mask_html.append(tmp4)
                    pos_mask_html_ans.append(ans4)
                    
                    tmp5, ans5 = self.build_input_with_ele_dict(ele_dict, "random_mask_html")
                    random_mask_html.append(tmp5)
                    random_masked_html_ans.append(ans5)
               
                ### post process the mask token id
                cate_masked_html, cate_masked_ans = self.replace_order_mask(cate_masked_html, cate_masked_html_ans)
                size_pos_mask_html, size_pos_mask_ans = self.replace_order_mask(size_pos_mask_html, size_pos_mask_html_ans)
                pos_mask_html, pos_mask_ans = self.replace_order_mask(pos_mask_html, pos_mask_html_ans)
                random_mask_html, random_mask_ans = self.replace_order_mask(random_mask_html, random_masked_html_ans)
                
                verbal_all_categories = []
                for i in range(self.N_category):
                    if all_category[i] != 0:
                        verbal_category = self.category_map[i]
                        verbal_number = VERBALIZED_NUM[all_category[i]]
                        verbal_all_categories.append("{} {},".format(verbal_number, verbal_category))
                all_verbal_all_cates = " ".join(verbal_all_categories).rstrip(",")
                
                if self_consistency == True:  # random shuffle the condition, but stay with the target
                    shuffle_lst = [i for i in range(len(html_content))]
                    min_shuffle_num = min(len(shuffle_lst), consistency_num)
                    
                    def shuffle_list(input_list):  
                        random.shuffle(input_list)  
                        return input_list  
                    
                    shuffled_results = []  
                    for i in range(min_shuffle_num):  
                        shuffled_results.append(shuffle_list(shuffle_lst.copy()))
                    
                    for random_order in shuffled_results:
                        # new_html_content = [html_content[i] for i in random_order]
                        new_cate_masked_html = [cate_masked_html[i] for i in random_order]
                        new_size_pos_mask_html = [size_pos_mask_html[i] for i in random_order]
                        new_pos_mask_html = [pos_mask_html[i] for i in random_order]
                        new_random_mask_html = [random_mask_html[i] for i in random_order]

                        batched_cond_cate.append(all_verbal_all_cates)   
                        batched_html_lst.append("\n".join(html_content))  # save target
                        batched_cond_bbox.append('\n'.join(new_cate_masked_html))
                        cond_cate_to_size_pos.append("\n".join(new_size_pos_mask_html))
                        cond_cate_size_to_pos.append("\n".join(new_pos_mask_html))
                        random_mask.append("\n".join(new_random_mask_html))  
                        
                else:
                    # process all conditions
                    batched_cond_cate.append(all_verbal_all_cates)   
                    batched_cond_bbox.append('\n'.join(cate_masked_html))
                    batched_html_lst.append("\n".join(html_content))
                    cond_cate_to_size_pos.append("\n".join(size_pos_mask_html))
                    cond_cate_size_to_pos.append("\n".join(pos_mask_html))
                    random_mask.append("\n".join(random_mask_html))
                
        else:
            raise ValueError("Can not inplement to testing data")

        return {
            "batched_html_lst": batched_html_lst,
            "batched_cond_cate": batched_cond_cate,
            "batched_cond_bbox": batched_cond_bbox,
            "cond_cate_to_size_pos": cond_cate_to_size_pos,
            "cond_cate_size_to_pos": cond_cate_size_to_pos,
            "random_mask": random_mask,
            "codegen_ans": {
                "cate_masked_ans": cate_masked_ans,
                "size_pos_mask_ans": size_pos_mask_ans,
                "pos_mask_ans": pos_mask_ans,
                "random_mask_ans": random_mask_ans,
            },
        }
    
    
    def build_random_mask(self, lst):
        new_lst = lst.copy()
        num = random.sample([3, 4], 1)[0]  # mask up to 80% position
        pos = random.sample([0,1,2,3], num)
        for i in pos:
            new_lst[i] = PLACE_HOLDER
        return new_lst
    
    
    def generate_new_order(self, lst):
        shuffle_order = [i for i in range(len(lst))]
        random.shuffle(shuffle_order)
        return shuffle_order
        
                
    def custom_function(self, data, id_, self_consistency=True, consistency_num=10):  
        label, mask = to_dense_batch(data.y, data.batch)   # (B, S)
        bbox_real, _ = to_dense_batch(data.x, data.batch)  # (B, S, 4)
        W, H = data.attr["width"], data.attr["height"]
        
        size_ = torch.cat((W.unsqueeze(-1), H.unsqueeze(-1), W.unsqueeze(-1), H.unsqueeze(-1)), dim=-1)
        size_ = size_.unsqueeze(1)
        real_idx = size_ * bbox_real # [cx, cy, w, h]
        if self.bbox_quantization == "code":
            label = label.to(torch.int).tolist()
            label_lst = self.filter_invalid_num(label, mask)
            real_idx = real_idx.to(torch.float).tolist()
            
            real_idx = round_nested_list(real_idx, 1)
            bbox_lst = self.filter_invalid_num(real_idx, mask)
            preposed_res = self.convert_num_to_html(
                bbox_lst, label_lst, self_consistency=self_consistency, consistency_num=consistency_num
            )
            batched_html_lst = preposed_res.get("batched_html_lst")
            batched_cond_cate = preposed_res.get("batched_cond_cate")
            batched_cond_bbox = preposed_res.get("batched_cond_bbox")   
            
            cond_cate_to_size_pos = preposed_res.get("cond_cate_to_size_pos")
            cond_cate_to_size_pos_res_dict = preposed_res["codegen_ans"].get("size_pos_mask_ans")
            
            cond_cate_size_to_pos = preposed_res.get("cond_cate_size_to_pos")
            cond_cate_size_to_pos_res_dict = preposed_res["codegen_ans"].get("pos_mask_ans")
            
            random_mask = preposed_res.get("random_mask")
            random_mask_res_dict = preposed_res["codegen_ans"].get("random_mask_ans")
        
        elif self.bbox_quantization == "numerical":
            label = label.to(torch.int).tolist()
            label_lst = self.filter_invalid_num(label, mask)[0]
            real_idx = real_idx.to(torch.float).tolist()
            real_idx = round_nested_list(real_idx, 1)
            bbox_lst = self.filter_invalid_num(real_idx, mask)[0]
            label_lst = [str(i) for i in label_lst]
            bbox_lst = [[str(i) for i in item] for item in bbox_lst]
            
            if self.split == "train" or self.split == "val":
                consistency_num = min(len(label_lst), consistency_num)
            else:
                consistency_num = 1

            final_res = []
            for j in range(consistency_num):
                if j != 0:
                    shuffle_order = self.generate_new_order(label_lst)
                    new_bbox_lst = [bbox_lst[i] for i in shuffle_order]
                    new_label_lst = [label_lst[i] for i in shuffle_order]
                else:
                    new_bbox_lst = bbox_lst
                    new_label_lst = label_lst
            
                golden_layouts = [[new_label_lst[i]] + bbox_lst[i] for i in range(len(new_label_lst))]
                golden_layouts = " ".join([" ".join(item) for item in golden_layouts])
                
                # buld size and position mask
                cond_cate_to_size_pos_masks = [[PLACE_HOLDER] * 4 for i in range(len(new_label_lst))]
                cond_cate_to_size_pos_modeling = [[new_label_lst[i]] + cond_cate_to_size_pos_masks[i] for i in range(len(new_label_lst))]
                cond_cate_to_size_pos_modeling = [" ".join(item) for item in cond_cate_to_size_pos_modeling]
                cond_cate_to_size_pos_modeling = " ".join(cond_cate_to_size_pos_modeling)
                
                # build position mask
                cond_cate_to_size_pos_masks = [[PLACE_HOLDER] * 2 + new_bbox_lst[i][2:] for i in range(len(new_label_lst))]
                cond_cate_to_size_pos_masks_modeling = [[new_label_lst[i]] + cond_cate_to_size_pos_masks[i] for i in range(len(new_label_lst))]
                cond_cate_to_size_pos_masks_modeling = [" ".join(item) for item in cond_cate_to_size_pos_masks_modeling]
                cond_cate_to_size_pos_masks_modeling = " ".join(cond_cate_to_size_pos_masks_modeling)
                
                # build random mask
                cond_recover_masks = [self.build_random_mask(item) for item in new_bbox_lst]
                cond_recover_masks_modeling = [[new_label_lst[i]] + cond_recover_masks[i] for i in range(len(new_label_lst))]
                cond_recover_masks_modeling = [" ".join(item) for item in cond_recover_masks_modeling]
                cond_recover_masks_modeling = " ".join(cond_recover_masks_modeling)
                
                if self.split == "train" or self.split == "val":
                    cond_cate_to_size_pos_modeling = cond_cate_to_size_pos_modeling + " <MID> " + golden_layouts
                    cond_cate_to_size_pos_masks_modeling = cond_cate_to_size_pos_masks_modeling + " <MID> " + golden_layouts
                    cond_recover_masks_modeling = cond_recover_masks_modeling + " <MID> " + golden_layouts
                    final_res.append({
                        "cond_cate_to_size_pos_seq_modeling": cond_cate_to_size_pos_modeling,
                        "cond_cate_size_to_pos_seq_modeling": cond_cate_to_size_pos_masks_modeling,
                        "cond_recover_mask_seq_modeling": cond_recover_masks_modeling,
                    })
                else:
                    return {
                        "cond_cate_to_size_pos_seq_modeling": cond_cate_to_size_pos_modeling + " <MID>",
                        "cond_cate_size_to_pos_seq_modeling": cond_cate_to_size_pos_masks_modeling + " <MID>",
                        "cond_recover_mask_seq_modeling": cond_recover_masks_modeling + " <MID>",
                    }

            return final_res
            
            
        if self_consistency:  # resize W and H
            W = W.repeat(len(batched_html_lst))
            H = H.repeat(len(batched_html_lst))
        
        # construct the html input 
        batched_cond_bbox = [
            self.html_template.format(W=W[i], H=H[i], content=batched_cond_bbox[i])
            for i in range(len(batched_cond_bbox))                 
        ]
        cond_cate_to_size_pos = [
            self.html_template.format(W=W[i], H=H[i], content=cond_cate_to_size_pos[i])
            for i in range(len(cond_cate_to_size_pos))  
        ]
        cond_cate_size_to_pos = [
            self.html_template.format(W=W[i], H=H[i], content=cond_cate_size_to_pos[i])
            for i in range(len(cond_cate_size_to_pos))  
        ]
        cond_recover_mask = [
            self.html_template.format(W=W[i], H=H[i], content=random_mask[i])
            for i in range(len(random_mask))  
        ]
        
        # add task instructions
        cond_recover_mask = [
            self.cond_random_mask.format(bbox_html=bbox)
            for bbox in cond_recover_mask
        ]
        cond_cate_to_size_pos = [
            self.cond_cate_to_size_pos.format(bbox_html=bbox)
            for bbox in cond_cate_to_size_pos
        ]
        cond_cate_size_to_pos = [
            self.cond_cate_size_to_pos.format(bbox_html=bbox)
            for bbox in cond_cate_size_to_pos
        ]
        bbox_cond_seqs = [
            self.cond_bbox_prefix.format(categories=cate, bbox_html=bbox_html) 
            for cate, bbox_html in zip(batched_cond_cate, batched_cond_bbox)
        ]
        category_cond_seqs = [
            self.cond_cate_prefix.format(categories=batched_cond_cate[i], W=W[i], H=H[i]) 
            for i in range(len(batched_cond_cate))
        ]

        if self.infilling and self.split in ("train", "val"):  # do infilling task
            cond_cate_to_size_pos_golden = [f" {SEP_TOKEN} ".join(f"{key} {value}" for key, value in cond_cate_to_size_pos_res_dict.items())]
            cond_cate_size_to_pos_golden = [f" {SEP_TOKEN} ".join(f"{key} {value}" for key, value in cond_cate_size_to_pos_res_dict.items())]
            random_mask_res_dict_golden = [f" {SEP_TOKEN} ".join(f"{key} {value}" for key, value in random_mask_res_dict.items())]
        
        # build target seq
        if self.split == "train" or self.split == "val":
            if self.infilling:
                if self_consistency:
                    consistency_num = len(cond_cate_to_size_pos)
                    target_seqs = [
                        cond_cate_to_size_pos_golden * consistency_num, 
                        cond_cate_size_to_pos_golden * consistency_num, 
                        random_mask_res_dict_golden * consistency_num
                    ]
                else:
                    target_seqs = [cond_cate_to_size_pos_golden, cond_cate_size_to_pos_golden, random_mask_res_dict_golden]
                
                cond_cate_to_size_pos_seq_modeling = [
                    self.glue_template_codegen_train.format(instruct=sample, result=target)
                    for sample, target in zip(cond_cate_to_size_pos, target_seqs[0])
                ]
                
                cond_cate_size_to_pos_seq_modeling = [
                    self.glue_template_codegen_train.format(instruct=sample, result=target)
                    for sample, target in zip(cond_cate_size_to_pos, target_seqs[1])
                ]
                
                cond_recover_mask_seq_modeling = [
                    self.glue_template_codegen_train.format(instruct=sample, result=target)
                    for sample, target in zip(cond_recover_mask, target_seqs[2])
                ]

                return {
                    "cond_cate_to_size_pos_seq_modeling": cond_cate_to_size_pos_seq_modeling,
                    "cond_cate_size_to_pos_seq_modeling": cond_cate_size_to_pos_seq_modeling,
                    "cond_recover_mask_seq_modeling": cond_recover_mask_seq_modeling,
                }
                
            else:
                target_seqs = [
                    self.html_template.format(W=W[i], H=H[i], content=batched_html_lst[i])
                    for i in range(W.size(0))
                ]
            
                cond_recover_mask_seq_modeling = [
                    self.glue_template_codegen_train.format(instruct=sample, result=target)
                    for sample, target in zip(cond_recover_mask, target_seqs)
                ]
                
                cond_cate_to_size_pos_seq_modeling = [
                    self.glue_template_codegen_train.format(instruct=sample, result=target)
                    for sample, target in zip(cond_cate_to_size_pos, target_seqs)
                ]
                
                cond_cate_size_to_pos_seq_modeling = [
                    self.glue_template_codegen_train.format(instruct=sample, result=target)
                    for sample, target in zip(cond_cate_size_to_pos, target_seqs)
                ]
            
                return {
                    "cond_cate_to_size_pos_seq_modeling": cond_cate_to_size_pos_seq_modeling,
                    "cond_cate_size_to_pos_seq_modeling": cond_cate_size_to_pos_seq_modeling,
                    "cond_recover_mask_seq_modeling": cond_recover_mask_seq_modeling,
                }
            
        else:
            if self.infilling:
                cond_bbox_input_seqs = [
                    self.glue_template_codegen_test.format(instruct=instance)
                    for instance in bbox_cond_seqs
                ]
                
                continual_gen_input_seqs = [
                    self.glue_template_codegen_test.format(instruct=instance)
                    for instance in category_cond_seqs
                ]
                
                cond_cate_size_to_pos_input_seqs = [
                    self.glue_template_codegen_test.format(instruct=instance)
                    for instance in cond_cate_size_to_pos
                ]
                
                cond_cate_to_size_pos_input_seqs = [
                    self.glue_template_codegen_test.format(instruct=instance)
                    for instance in cond_cate_to_size_pos
                ]
                
                cond_recover_mask_input_seqs = [
                    self.glue_template_codegen_test.format(instruct=instance)
                    for instance in cond_recover_mask
                ]
            
            else:
                cond_bbox_input_seqs = [
                    self.glue_template_test.format(instruct=instance)
                    for instance in bbox_cond_seqs
                ]
                
                continual_gen_input_seqs = [
                    self.glue_template_test.format(instruct=instance)
                    for instance in category_cond_seqs
                ]
                
                cond_cate_size_to_pos_input_seqs = [
                    self.glue_template_test.format(instruct=instance)
                    for instance in cond_cate_size_to_pos
                ]
                
                cond_cate_to_size_pos_input_seqs = [
                    self.glue_template_test.format(instruct=instance)
                    for instance in cond_cate_to_size_pos
                ]
                
                cond_recover_mask_input_seqs = [
                    self.glue_template_test.format(instruct=instance)
                    for instance in cond_recover_mask
                ]
            
            
            labels = None
            if batched_html_lst is not None:
                
                labels = [
                    self.html_template.format(W=W[i], H=H[i], content=batched_html_lst[i])
                    for i in range(W.size(0))
                ]
            
            return {
                "cond_bbox_input_seqs": cond_bbox_input_seqs,
                "continual_gen_input_seqs": continual_gen_input_seqs,
                "cond_cate_size_to_pos_input_seqs": cond_cate_size_to_pos_input_seqs,
                "cond_cate_to_size_pos_input_seqs": cond_cate_to_size_pos_input_seqs,
                "cond_recover_mask_input_seqs": cond_recover_mask_input_seqs,
                "labels": labels,
                "raw_data": {
                    "category": bbox_lst[0],
                    "bbox": label_lst[0]
                },
                "id_": id_
            }
    
    def __iter__(self):
        for i, data in enumerate(super(CustomDataLoader, self).__iter__()):  
            if self.consistency_num > 1:
                self_consistency = True
            else:
                self_consistency = False
            yield self.custom_function(data, i, self_consistency=self_consistency)  

    
    @property
    def eos_token_id(self) -> int:
        return self.tokenizer.eos_token_id
    
    @property
    def bos_token_id(self) -> int:
        return self.tokenizer.bos_token_id
    
    @property
    def pad_token_id(self) -> int: 
        return self.tokenizer.pad_token_id
    
    @property
    def mask_token_id(self) -> int:
        return self.tokenizer.unk_token_id
    
    @property
    def unk_token_id(self) -> int:
        return self.tokenizer.unk_token_id
            


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate the Layout Code in HTML format")
    parser.add_argument("--dataset_name", type=str, default="rico25", help="dataset name")  
    parser.add_argument("--dataset_path", type=str, default="data/rico25-max25")
    parser.add_argument("--save_path", type=str, default="data/rico25-max25/html_format")
    parser.add_argument("--model_path_or_name", type=str, default=None, help="tokenizer model name")
    parser.add_argument("--bbox_quantization", type=str, default="code", choices=["code", "numerical"])
    parser.add_argument("--consistency_num", type=int, default=1, help="number of consistency num")
    parser.add_argument("--build_testing_set", action="store_true", help="whether to build the testing set")
    parser.add_argument("--infilling", action="store_true", help="whether to build the infilling data set")
    parser.add_argument("--add_task_instruction", action="store_true", help="whether to add the task instruction")
    args = parser.parse_args()
    
    if not os.path.exists(args.save_path):  
        os.makedirs(args.save_path)  
        print(f"Directory '{args.save_path}' created.")  
    
    if args.model_path_or_name is None:
        print("please specify the model name or path")
    
    tokenizer = AutoTokenizer.from_pretrained(args.model_path_or_name)
    transforms = [LexicographicSort()]
     
    if not args.build_testing_set:
        train_dataset = get_dataset(
            name=args.dataset_name,
            datapath=args.dataset_path,
            split='train',
            transform=T.Compose(transforms)
        )
        
        eval_dataset = get_dataset(
            name=args.dataset_name,
            datapath=args.dataset_path,
            split='val',
            transform=T.Compose(transforms)
        )
        
        train_dataloader = CustomDataLoader(
            args,
            tokenizer,
            bbox_quantization=args.bbox_quantization,
            dataset=train_dataset, 
            batch_size=1,
            split="train",
        )
        
        eval_dataloader = CustomDataLoader(
            args,
            tokenizer,
            bbox_quantization=args.bbox_quantization,
            dataset=eval_dataset, 
            batch_size=1,
            split="val",
        )
        
        all_train_data, all_eval_data = [], []
    
        train_file = os.path.join(args.save_path, "train_llama_numerical.jsonl")
        val_file = os.path.join(args.save_path, "val_llama_numerical.jsonl")
        
        print(f"begin to save train file >>> {args.save_path}")
        with tqdm(total=len(train_dataloader)) as pbar:
            for i, batch_inputs in enumerate(train_dataloader):
                if args.consistency_num > 1:
                    inner_batch = len(batch_inputs['cond_cate_to_size_pos_seq_modeling'])
                    new_batch_inputs = [{} for i in range(inner_batch)]
                    for k, v in batch_inputs.items():
                        for i, value in enumerate(v):
                            new_batch_inputs[i][k] = value    
                    batch_inputs = new_batch_inputs
                else:
                    batch_inputs = [batch_inputs]
                all_train_data.extend(batch_inputs)
                pbar.update(1)
        with open(train_file, "w") as f:
            for line in all_train_data:
                f.write(json.dumps(line) + "\n")
                
        print(f"training data saved done, begin to save eval dataset >>> {args.save_path}")
        with tqdm(total=len(eval_dataloader)) as pbar:
            for i, batch_inputs in enumerate(eval_dataloader):
                if args.consistency_num > 1:
                    inner_batch = len(batch_inputs['cond_cate_to_size_pos_seq_modeling'])
                    new_batch_inputs = [{} for i in range(inner_batch)]
                    for k, v in batch_inputs.items():
                        for i, value in enumerate(v):
                            new_batch_inputs[i][k] = value    
                    batch_inputs = new_batch_inputs
                else:
                    batch_inputs = [batch_inputs]
                all_eval_data.extend(batch_inputs)
                pbar.update(1)
        with open(val_file, "w") as f:
            for line in all_eval_data:
                f.write(json.dumps(line) + "\n")
    
    else:
        test_dataset = get_dataset(
            name=args.dataset_name,
            datapath=args.dataset_path,
            split='test',
            transform=T.Compose(transforms)
        )
        
        test_dataloader = CustomDataLoader(
            args,
            tokenizer,
            bbox_quantization=args.bbox_quantization,
            dataset=test_dataset, 
            batch_size=1,
            split="test",
        )
    
        all_test_data = []
        test_file = os.path.join(args.save_path, "test_numerical.jsonl")
        print("begin to save test file")
        with tqdm(total=len(test_dataloader)) as pbar:
            for i, batch_inputs in enumerate(test_dataloader):
                all_test_data.append(batch_inputs)
                pbar.update(1)
        with open(test_file, "w") as f:
            for line in all_test_data:
                f.write(json.dumps(line) + "\n")
    
    
    
convertHTML/build_code.py
ProjectNUWA/LayoutNUWA
import numpy as np
from einops import rearrange

EPS = 1e-12


class Percentile:

    def __init__(self, n_clusters: int = 32, v_min: float = 0.0, v_max: float = 1.0):
        self.n_clusters = n_clusters
        self.v_min = v_min
        self.v_max = v_max

    def fit(self, X: np.ndarray):
        assert X.ndim == 2  # (B, 1)
        X = X[:, 0].clip(self.v_min, self.v_max)

        X = np.sort(np.unique(X))
        thresholds = np.linspace(0.0, 1.0, self.n_clusters + 1)[:-1]
        thresholds = [X[int(t * len(X))] for t in thresholds]
        ids = (thresholds <= X.reshape(-1, 1)).sum(axis=1).astype(np.uint64) - 1

        self.cluster_centers_ = np.full(
            (self.n_clusters, 1), -1.0, dtype=np.float32
        )  
        for i in range(self.n_clusters):
            values = X[ids == i]
            if len(values) > 0:
                self.cluster_centers_[i, 0] = values.mean().astype(np.float32)
        return self

    def predict(self, X):
        if not hasattr(self, "cluster_centers_"):
            raise NotImplementedError

        assert X.ndim == 2
        X = X.clip(self.v_min, self.v_max)
        dist = np.fabs(self.cluster_centers_ - rearrange(X, "s 1 -> 1 s"))
        ids = np.argmin(dist, axis=0)

        return ids
convertHTML/cluster.py
ProjectNUWA/LayoutNUWA
import sys
import os
current_directory = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_directory)  
import copy
import json
import numpy as np
import torch
from pathlib import Path
from PIL import Image, ImageDraw
from torch_geometric.data import Data
from convertHTML.utils import sparse_to_dense
from convertHTML.utils import convert_xywh_to_ltrb

from convertHTML.base import BaseDataset

magazine_labels = [
    "text",
    "image",
    "headline",
    "text-over-image",
    "headline-over-image",
]

label2index = {
    "text": 0,
    "image": 1,
    "headline": 2,
    "text-over-image": 3,
    "headline-over-image": 4,
}



def append_child(element, elements):
    if "children" in element.keys():
        for child in element["children"]:
            elements.append(child)
            elements = append_child(child, elements)
    return elements


class MagezineDataset(BaseDataset):
    def __init__(self, datapath, split='train', transform=None, max_bbox_bins=32):
        super().__init__(datapath, split, transform)
        self.N_category = 5
        self.max_box_bins = max_bbox_bins
        self.dataset_name = "magazine"

    def process(self):
        for split_publaynet in ['train', 'val']:
            data_list = []
            with open(os.path.join(self.path, f'{split_publaynet}.json'), "r") as f:
                content = json.load(f)
            for k, v in content.items():
                W = 225.0
                H = 300.0
                name = k
                
                def is_valid(element):
                    x1, y1, width, height = element['bbox']
                    x2, y2 = x1 + width, y1 + height
                    if x1 < 0 or y1 < 0 or W < x2 or H < y2:
                        return False

                    if x2 <= x1 or y2 <= y1:
                        return False

                    return True
                
                N = len(v)
                if N == 0 or 32 < N:
                    continue
                
                boxes, labels = [], []
                
                for ele_name, ele_val in v.items():
                    # bbox
                    for ele in ele_val:
                        if len(ele) > 4:
                            import pdb; pdb.set_trace()
                        x1, y1 = ele[0][0], ele[0][1]
                        width = abs(ele[2][0] - ele[0][0]) 
                        height = abs(ele[2][1] - ele[0][1])
                        xc = x1 + width / 2.
                        yc = y1 + height / 2.
                        b = [xc / W, yc / H, width / W, height / H]
                        boxes.append(b)
                        labels.append(label2index[ele_name])

                boxes = torch.tensor(boxes, dtype=torch.float)  # xc, yc, W, H
                labels = torch.tensor(labels, dtype=torch.long)

                data = Data(x=boxes, y=labels)
                data.attr = {
                    'name': name,
                    'width': W,
                    'height': H,
                    'filtered': None,
                    'has_canvas_element': False,
                    "NoiseAdded": False,
                }
                data_list.append(data)
            
            s_t_v = int(len(data_list) * .95)
            print(s_t_v)
            if split_publaynet == 'train':
                train_list = data_list[:s_t_v]
            else:
                val_list = data_list[s_t_v:]

        # shuffle train with seed
        generator = torch.Generator().manual_seed(0)
        indices = torch.randperm(len(train_list), generator=generator)
        train_list = [train_list[i] for i in indices]
        
        # train_list -> train 95% / val 5%
        # val_list -> test 100%
        s = int(len(train_list) * .95)
        torch.save(self.collate(train_list), self.processed_paths[0])
        torch.save(self.collate(train_list[s:]), self.processed_paths[1])
        torch.save(self.collate(val_list), self.processed_paths[2])


    def get_original_resource(self, batch) -> Image:
        assert not self.raw_dir.startswith("gs://")
        bbox, _, _, _ = sparse_to_dense(batch)

        img_bg, img_original, cropped_patches = [], [], []
        names = batch.attr["name"]
        if isinstance(names, str):
            names = [names]

        for i, name in enumerate(names):
            name = Path(name).name.replace(".json", ".jpg")
            img = Image.open(Path(self.raw_dir) / "combined" / name)
            img_original.append(copy.deepcopy(img))

            W, H = img.size
            ltrb = convert_xywh_to_ltrb(bbox[i].T.numpy())
            left, right = (ltrb[0] * W).astype(np.uint32), (ltrb[2] * W).astype(
                np.uint32
            )
            top, bottom = (ltrb[1] * H).astype(np.uint32), (ltrb[3] * H).astype(
                np.uint32
            )
            draw = ImageDraw.Draw(img)
            patches = []
            for (l, r, t, b) in zip(left, right, top, bottom):
                patches.append(img.crop((l, t, r, b)))
                draw.rectangle([(l, t), (r, b)], fill=(255, 255, 255))
            img_bg.append(img)
            cropped_patches.append(patches)


        return {
            "img_bg": img_bg,
            "img_original": img_original,
            "cropped_patches": cropped_patches,
        }

convertHTML/magazine.py
ProjectNUWA/LayoutNUWA
import sys
sys.path.append("/workspace/zecheng/layout-generation")
import torch
from pathlib import Path
from pycocotools.coco import COCO
from torch_geometric.data import Data
from convertHTML.base import BaseDataset
import os

class PubLayNet(BaseDataset):
    labels = [
        'text',
        'title',
        'list',
        'table',
        'figure',
    ]

    def __init__(self, datapath, split='train', transform=None, max_bbox_bins=32):
        super().__init__(datapath, split, transform)
        self.N_category = self.num_classes
        self.max_box_bins = max_bbox_bins
        self.dataset_name = "publaynet"

    def process(self):
        for split_publaynet in ['train', 'val']:
            data_list = []
            coco = COCO(os.path.join(self.path, f'{split_publaynet}.json'))
            for img_id in sorted(coco.getImgIds()):
                ann_img = coco.loadImgs(img_id)
                W = float(ann_img[0]['width'])
                H = float(ann_img[0]['height'])
                name = ann_img[0]['file_name']
                if H < W:
                    continue

                def is_valid(element):
                    x1, y1, width, height = element['bbox']
                    x2, y2 = x1 + width, y1 + height
                    if x1 < 0 or y1 < 0 or W < x2 or H < y2:
                        return False

                    if x2 <= x1 or y2 <= y1:
                        return False

                    return True

                elements = coco.loadAnns(coco.getAnnIds(imgIds=[img_id]))
                _elements = list(filter(is_valid, elements))
                filtered = len(elements) != len(_elements)
                elements = _elements

                N = len(elements)
                if N == 0 or self.max_bbox_bins < N:  # max remain one bbox pos for eos and bos
                    continue

                boxes = []
                labels = []

                for element in elements:
                    # bbox
                    x1, y1, width, height = element['bbox']
                    xc = x1 + width / 2.
                    yc = y1 + height / 2.
                    b = [xc / W, yc / H, width / W, height / H]
                    boxes.append(b)

                    # label
                    l = coco.cats[element['category_id']]['name']
                    labels.append(self.label2index[l])

                boxes = torch.tensor(boxes, dtype=torch.float)  # xc, yc, W, H
                labels = torch.tensor(labels, dtype=torch.long)

                data = Data(x=boxes, y=labels)
                data.attr = {
                    'name': name,
                    'width': W,
                    'height': H,
                    'filtered': filtered,
                    'has_canvas_element': False,
                    "NoiseAdded": False,
                }
                data_list.append(data)

            if split_publaynet == 'train':
                train_list = data_list
            else:
                val_list = data_list

        # shuffle train with seed
        generator = torch.Generator().manual_seed(0)
        indices = torch.randperm(len(train_list), generator=generator)
        train_list = [train_list[i] for i in indices]
        
        # train_list -> train 95% / val 5%
        # val_list -> test 100%
        s = int(len(train_list) * .95)
        torch.save(self.collate(train_list[:s]), self.processed_paths[0])
        torch.save(self.collate(train_list[s:]), self.processed_paths[1])
        torch.save(self.collate(val_list), self.processed_paths[2])
        

if __name__ == "__main__":
    train_dataset = PubLayNet(
        datapath="/zecheng/dataset/layout/publaynet",
        split='val',
        transform=None,
        N_category=5,
        max_bbox_bins=32
    )
    print(train_dataset)
convertHTML/publaynet.py
ProjectNUWA/LayoutNUWA
import sys
sys.path.append("/workspace/zecheng/layout-generation/")

import copy
import json
import os
from pathlib import Path
from zipfile import ZipFile
import numpy as np
import torch
from fsspec.core import url_to_fs
from PIL import Image, ImageDraw
from torch_geometric.data import Data
from tqdm import tqdm
from convertHTML.utils import sparse_to_dense
from convertHTML.utils import convert_xywh_to_ltrb

from .base import BaseDataset

_rico5_labels = [
    "Text",
    "Text Button",
    "Toolbar",
    "Image",
    "Icon",
]

_rico13_labels = [
    "Toolbar",
    "Image",
    "Text",
    "Icon",
    "Text Button",
    "Input",
    "List Item",
    "Advertisement",
    "Pager Indicator",
    "Web View",
    "Background Image",
    "Drawer",
    "Modal",
]

_rico25_labels = [
    "Text",
    "Image",
    "Icon",
    "Text Button",
    "List Item",
    "Input",
    "Background Image",
    "Card",
    "Web View",
    "Radio Button",
    "Drawer",
    "Checkbox",
    "Advertisement",
    "Modal",
    "Pager Indicator",
    "Slider",
    "On/Off Switch",
    "Button Bar",
    "Toolbar",
    "Number Stepper",
    "Multi-Tab",
    "Date Picker",
    "Map View",
    "Video",
    "Bottom Navigation",
]


def append_child(element, elements):
    if "children" in element.keys():
        for child in element["children"]:
            elements.append(child)
            elements = append_child(child, elements)
    return elements


class _RicoDataset(BaseDataset):
    def __init__(self, datapath, split='train', transform=None, max_bbox_bins=32):
        super().__init__(datapath, split, transform)
        self.N_category = self.num_classes
        self.max_box_bins = max_bbox_bins
        self.dataset_name = "rico25"

    def process(self):
        data_list = []
        raw_file = os.path.join(
            self.raw_dir, "rico_dataset_v0.1_semantic_annotations.zip"
        )
        fs, _ = url_to_fs(self.raw_dir)
        with fs.open(raw_file, "rb") as f, ZipFile(f) as z:
            names = sorted([n for n in z.namelist() if n.endswith(".json")])
            for name in tqdm(names):
                ann = json.loads(z.open(name).read())

                B = ann["bounds"]
                W, H = float(B[2]), float(B[3])
                if B[0] != 0 or B[1] != 0 or H < W:
                    continue

                def is_valid(element):
                    if element["componentLabel"] not in set(self.labels):
                        print(element["componentLabel"])
                        return False

                    x1, y1, x2, y2 = element["bounds"]
                    if x1 < 0 or y1 < 0 or W < x2 or H < y2:
                        return False

                    if x2 <= x1 or y2 <= y1:
                        return False

                    return True

                elements = append_child(ann, [])
                _elements = list(filter(is_valid, elements))
                filtered = len(elements) != len(_elements)
                elements = _elements
                N = len(elements)
                if N == 0 or self.max_seq_length < N:
                    continue

                boxes = []
                labels = []

                for element in elements:
                    # bbox
                    x1, y1, x2, y2 = element["bounds"]
                    xc = (x1 + x2) / 2.0
                    yc = (y1 + y2) / 2.0
                    width = x2 - x1
                    height = y2 - y1
                    b = [xc / W, yc / H, width / W, height / H]
                    boxes.append(b)

                    # label
                    l = element["componentLabel"]
                    labels.append(self.label2index[l])

                boxes = torch.tensor(boxes, dtype=torch.float)
                labels = torch.tensor(labels, dtype=torch.long)

                data = Data(x=boxes, y=labels)
                data.attr = {
                    "name": name,
                    "width": W,
                    "height": H,
                    "filtered": filtered,
                    "has_canvas_element": False,
                    "NoiseAdded": False,
                }
                data_list.append(data)

        # shuffle with seed
        generator = torch.Generator().manual_seed(0)
        indices = torch.randperm(len(data_list), generator=generator)
        data_list = [data_list[i] for i in indices]

        # train 85% / val 5% / test 10% (validation data)
        N = len(data_list)
        s = [int(N * 0.85), int(N * 0.90)]

        with fs.open(self.processed_paths[0], "wb") as file_obj:
            torch.save(self.collate(data_list[: s[0]]), file_obj)
        with fs.open(self.processed_paths[1], "wb") as file_obj:
            torch.save(self.collate(data_list[s[0] : s[1]]), file_obj)
        with fs.open(self.processed_paths[2], "wb") as file_obj:
            torch.save(self.collate(data_list[s[1] :]), file_obj)


    def get_original_resource(self, batch) -> Image:
        assert not self.raw_dir.startswith("gs://")
        bbox, _, _, _ = sparse_to_dense(batch)

        img_bg, img_original, cropped_patches = [], [], []
        names = batch.attr["name"]
        if isinstance(names, str):
            names = [names]

        for i, name in enumerate(names):
            name = Path(name).name.replace(".json", ".jpg")
            img = Image.open(Path(self.raw_dir) / "combined" / name)
            img_original.append(copy.deepcopy(img))

            W, H = img.size
            ltrb = convert_xywh_to_ltrb(bbox[i].T.numpy())
            left, right = (ltrb[0] * W).astype(np.uint32), (ltrb[2] * W).astype(
                np.uint32
            )
            top, bottom = (ltrb[1] * H).astype(np.uint32), (ltrb[3] * H).astype(
                np.uint32
            )
            draw = ImageDraw.Draw(img)
            patches = []
            for (l, r, t, b) in zip(left, right, top, bottom):
                patches.append(img.crop((l, t, r, b)))
                # draw.rectangle([(l, t), (r, b)], fill=(255, 0, 0))
                draw.rectangle([(l, t), (r, b)], fill=(255, 255, 255))
            img_bg.append(img)
            cropped_patches.append(patches)


        return {
            "img_bg": img_bg,
            "img_original": img_original,
            "cropped_patches": cropped_patches,
        }


class Rico5Dataset(_RicoDataset):
    name = "rico5"
    labels = _rico5_labels

    def __init__(self, dir: str, split: str, max_seq_length: int, transform=None):
        super().__init__(dir, split, max_seq_length, transform)


# Constrained Graphic Layout Generation via Latent Optimization (ACMMM2021)
class Rico13Dataset(_RicoDataset):
    name = "rico13"
    labels = _rico13_labels

    def __init__(self, dir: str, split: str, max_seq_length: int, transform=None):
        super().__init__(dir, split, max_seq_length, transform)


class Rico25Dataset(_RicoDataset):
    name = "rico25"
    labels = _rico25_labels

    def __init__(self, datapath, split='train', transform=None, max_bbox_bins=32):
        super().__init__(datapath, split=split, transform=transform, max_bbox_bins=max_bbox_bins)
convertHTML/rico.py
ProjectNUWA/LayoutNUWA
import torch
import random
from enum import IntEnum
from itertools import product, combinations
from torch_geometric.utils import to_dense_batch
from typing import Tuple, List
from torch import FloatTensor, LongTensor, BoolTensor, BoolTensor
import numpy as np


def convert_xywh_to_ltrb(bbox):
    xc, yc, w, h = bbox
    x1 = xc - w / 2
    y1 = yc - h / 2
    x2 = xc + w / 2
    y2 = yc + h / 2
    return [x1, y1, x2, y2]


class RelSize(IntEnum):
    UNKNOWN = 0
    SMALLER = 1
    EQUAL = 2
    LARGER = 3


class RelLoc(IntEnum):
    UNKNOWN = 4
    LEFT = 5
    TOP = 6
    RIGHT = 7
    BOTTOM = 8
    CENTER = 9

REL_SIZE_ALPHA = 0.1


def detect_size_relation(b1, b2):
    a1, a2 = b1[2] * b1[3], b2[2] * b2[3]
    a1_sm = (1 - REL_SIZE_ALPHA) * a1
    a1_lg = (1 + REL_SIZE_ALPHA) * a1

    if a2 <= a1_sm:
        return RelSize.SMALLER

    if a1_sm < a2 and a2 < a1_lg:
        return RelSize.EQUAL

    if a1_lg <= a2:
        return RelSize.LARGER

    raise RuntimeError(b1, b2)


def detect_loc_relation(b1, b2, canvas=False):
    if canvas:
        yc = b2[1]
        y_sm, y_lg = 1. / 3, 2. / 3

        if yc <= y_sm:
            return RelLoc.TOP

        if y_sm < yc and yc < y_lg:
            return RelLoc.CENTER

        if y_lg <= yc:
            return RelLoc.BOTTOM

    else:
        l1, t1, r1, b1 = convert_xywh_to_ltrb(b1)
        l2, t2, r2, b2 = convert_xywh_to_ltrb(b2)

        if b2 <= t1:
            return RelLoc.TOP

        if b1 <= t2:
            return RelLoc.BOTTOM

        if t1 < b2 and t2 < b1:
            if r2 <= l1:
                return RelLoc.LEFT

            if r1 <= l2:
                return RelLoc.RIGHT

            if l1 < r2 and l2 < r1:
                return RelLoc.CENTER

    raise RuntimeError(b1, b2, canvas)


def get_rel_text(rel, canvas=False):
    if type(rel) == RelSize:
        index = rel - RelSize.UNKNOWN - 1
        if canvas:
            return [
                'within canvas',
                'spread over canvas',
                'out of canvas',
            ][index]

        else:
            return [
                'larger than',
                'equal to',
                'smaller than',
            ][index]

    else:
        index = rel - RelLoc.UNKNOWN - 1
        if canvas:
            return [
                '', 'at top',
                '', 'at bottom',
                'at middle',
            ][index]

        else:
            return [
                'right to', 'below',
                'left to', 'above',
                'around',
            ][index]


class LexicographicSort():
    def __call__(self, data):
        assert not data.attr['has_canvas_element']
        l, t, _, _ = convert_xywh_to_ltrb(data.x.t())
        _zip = zip(*sorted(enumerate(zip(t, l)), key=lambda c: c[1:]))
        idx = list(list(_zip)[0])
        data.x_orig, data.y_orig = data.x, data.y
        data.x, data.y = data.x[idx], data.y[idx]
        return data

class AddNoiseToBBox:
    def __init__(self, std: float = 0.05):
        self.std = float(std)

    def __call__(self, data):
        noise = torch.normal(0, self.std, size=data.x.size(), device=data.x.device)
        data.x_orig = data.x.clone()
        # print(f"add a noise >>>>>> {noise}\n original data is {data.x}")
        data.x = data.x + noise
        # prevent data.x smaller than 0 or greater than one
        data.x = data.x.clamp(0, 1)
        data.attr = data.attr.copy()
        data.attr["NoiseAdded"][0] = True
        return data
    
class LexicographicOrder:
    def __call__(self, data):
        assert not data.attr["has_canvas_element"]
        x, y, _, _ = convert_xywh_to_ltrb(data.x.t())
        _zip = zip(*sorted(enumerate(zip(y, x)), key=lambda c: c[1:]))
        idx = list(list(_zip)[0])
        data.x_orig, data.y_orig = data.x, data.y
        data.x, data.y = data.x[idx], data.y[idx]
        return data    
    

class HorizontalFlip():
    def __call__(self, data):
        data.x = data.x.clone()
        data.x[:, 0] = 1 - data.x[:, 0]
        return data


class AddCanvasElement():
    def __init__(self):
        self.x = torch.tensor([[.5, .5, 1., 1.]], dtype=torch.float)
        self.y = torch.tensor([0], dtype=torch.long)

    def __call__(self, data):
        if not data.attr['has_canvas_element']:
            data.x = torch.cat([self.x, data.x], dim=0)
            data.y = torch.cat([self.y, data.y + 1], dim=0)
            data.attr = data.attr.copy()
            data.attr['has_canvas_element'] = True
        return data


class AddRelation():
    def __init__(self, seed=None, ratio=0.1):
        self.ratio = ratio
        self.generator = random.Random()
        if seed is not None:
            self.generator.seed(seed)

    def __call__(self, data):
        N = data.x.size(0)
        has_canvas = data.attr['has_canvas_element']

        rel_all = list(product(range(2), combinations(range(N), 2)))
        size = int(len(rel_all) * self.ratio)
        rel_sample = set(self.generator.sample(rel_all, size))

        edge_index, edge_attr = [], []
        rel_unk = 1 << RelSize.UNKNOWN | 1 << RelLoc.UNKNOWN
        for i, j in combinations(range(N), 2):
            bi, bj = data.x[i], data.x[j]
            canvas = data.y[i] == 0 and has_canvas

            if (0, (i, j)) in rel_sample:
                rel_size = 1 << detect_size_relation(bi, bj)
            else:
                rel_size = 1 << RelSize.UNKNOWN

            if (1, (i, j)) in rel_sample:
                rel_loc = 1 << detect_loc_relation(bi, bj, canvas)
            else:
                rel_loc = 1 << RelLoc.UNKNOWN

            rel = rel_size | rel_loc
            if rel != rel_unk:
                edge_index.append((i, j))
                edge_attr.append(rel)

        data.edge_index = torch.as_tensor(edge_index).long()
        data.edge_index = data.edge_index.t().contiguous()
        data.edge_attr = torch.as_tensor(edge_attr).long()

        return data


def sparse_to_dense(
    batch,
    device: torch.device = torch.device("cpu"),
    remove_canvas: bool = False,
) -> Tuple[FloatTensor, LongTensor, BoolTensor, BoolTensor]:
    batch = batch.to(device)
    bbox, _ = to_dense_batch(batch.x, batch.batch)
    label, mask = to_dense_batch(batch.y, batch.batch)

    if remove_canvas:
        bbox = bbox[:, 1:].contiguous()
        label = label[:, 1:].contiguous() - 1  # cancel +1 effect in transform
        label = label.clamp(min=0)
        mask = mask[:, 1:].contiguous()

    padding_mask = ~mask
    return bbox, label, padding_mask, mask


def loader_to_list(
    loader: torch.utils.data.dataloader.DataLoader,
) -> List[Tuple[np.ndarray, np.ndarray]]:
    layouts = []
    for batch in loader:
        bbox, label, _, mask = sparse_to_dense(batch)
        for i in range(len(label)):
            valid = mask[i].numpy()
            layouts.append((bbox[i].numpy()[valid], label[i].numpy()[valid]))
    return layouts


def split_num_samples(N: int, batch_size: int) -> List[int]:
    quontinent = N // batch_size
    remainder = N % batch_size
    dataloader = quontinent * [batch_size]
    if remainder > 0:
        dataloader.append(remainder)
    return dataloader
convertHTML/utils.py
ProjectNUWA/LayoutNUWA
import sys
import os
current_directory = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_directory)  
import json
import re
import torch
import fsspec
import pickle
import numpy as np
import argparse
from helper.util import *
from helper.global_var import *
from helper.fid_model import *
from helper.metrics import *
from convertHTML import get_dataset
from convertHTML.utils import LexicographicSort
from collections import defaultdict
from typing import Dict
from helper.visualization import *
from tqdm import *
from torch_geometric.data import Data, Batch, DataLoader
from torch_geometric.loader import DataLoader


def remove_repeat(bbox, label):
    if bbox.size(0) == 0:
        return bbox, label 
    bbox_label = torch.cat((bbox, label.unsqueeze(1)), dim=1)  
    unique_bbox_label = []  
    for item in bbox_label:   
        same_bbox_label_exists = False  
        for unique_item in unique_bbox_label:  
            if torch.all(torch.eq(item, unique_item)):  
                same_bbox_label_exists = True  
                break  
        if not same_bbox_label_exists:  
            unique_bbox_label.append(item) 
     
    unique_bbox_label = torch.stack(unique_bbox_label) 
    unique_bbox = unique_bbox_label[:, :-1]  
    unique_label = unique_bbox_label[:, -1] 
    
    return unique_bbox, unique_label


def cluster_corrd(_clustering_models, bbox, keys):
    all_res = []
    device_ = bbox.device
    if bbox.size(0) == 0:
        return bbox
    elif bbox.dim() > 2:
        for item in bbox:
            itermediate_res = []
            for i, k in enumerate(keys):
                model = _clustering_models[k]
                ins_ = item[..., i][:, None].cpu().numpy().astype(np.float32)
                cluster_i = model.predict(ins_)
                new_bbox = model.cluster_centers_[cluster_i]
                itermediate_res.append(torch.from_numpy(new_bbox).to(device_))
            all_res.append(torch.concatenate(itermediate_res, dim=1))
        all_res = torch.stack(all_res, dim=0)
    else:
        for i, k in enumerate(keys):
            model = _clustering_models[k]
            bbox_ = bbox[..., i][:, None].cpu().numpy().astype(np.float32)
            cluster_i = model.predict(bbox_)
            new_bbox = model.cluster_centers_[cluster_i]
            all_res.append(torch.from_numpy(new_bbox).to(device_))
        all_res = torch.concatenate(all_res, dim=1)
    return all_res
        

def preprocess_batch(cluster_model, layouts, max_len: int, device: torch.device, k_means=True):
    layout = defaultdict(list)
    empty_ids = []  # 0: empty 1: full
    for sample in layouts:
        if not isinstance(sample["bbox"], torch.Tensor):
            bbox, label = torch.tensor(sample["bbox"]), torch.tensor(sample["categories"])
        else:
            bbox, label = sample["bbox"], sample["categories"]
        bbox, label = remove_repeat(bbox, label)
        if k_means == True:
            bbox =  cluster_corrd(cluster_model, bbox, ['x-64', 'y-64', 'w-64', 'h-64'])
        pad_len = max_len - label.size(0)

        if pad_len == max_len:
            empty_ids.append(0)
            pad_bbox = torch.tensor(np.full((max_len, 4), 0.0), dtype=torch.float)
            pad_label = torch.tensor(np.full((max_len,), 0), dtype=torch.long)
            mask = torch.tensor([False for _ in range(max_len)])
        else:
            empty_ids.append(1)  # not empty
            pad_bbox = torch.tensor(
                np.concatenate([bbox, np.full((pad_len, 4), 0.0)], axis=0),
                dtype=torch.float,
            )
            pad_label = torch.tensor(
                np.concatenate([label, np.full((pad_len,), 0)], axis=0),
                dtype=torch.long,
            )
            mask = torch.tensor(
                [True for _ in range(bbox.shape[0])] + [False for _ in range(pad_len)]
            )

        layout["bbox"].append(pad_bbox)
        layout["label"].append(pad_label)
        layout["mask"].append(mask)
        
    bbox = torch.stack(layout["bbox"], dim=0).to(device)
    label = torch.stack(layout["label"], dim=0).to(device)
    mask = torch.stack(layout["mask"], dim=0).to(device)
    
    padding_mask = ~mask.bool()  
    return bbox, label, padding_mask, mask, empty_ids     
        

def extract_WH(s):
    pattern = r'<svg width="([\d.]+)" height="([\d.]+)">'  
    match = re.search(pattern, s)  
    W, H = match.groups()  
    return  W, H  


def extract_xywh(s):
    bboxs, labels = [], []  
    pattern = r'<rect data-category="([^"]+)".*?x="(\d+(\.\d+)?)".*?y="(\d+(\.\d+)?)".*?width="(\d+(\.\d+)?)".*?height="(\d+(\.\d+)?)".*?/>'  
    matches = re.findall(pattern, s)  
    for match in matches:  
        data_category, x, _, y, _, width, _, height, _ = match  
        label = label_to_int.get(data_category)  
        # bbox = [x, y, x + width, y + height]  
        if label is None:
            continue
        bboxs.append([eval(x), eval(y), eval(width), eval(height)])
        labels.append(label)
    return bboxs, labels


def preprocess(sample: str):
    """
    key_lst can contain:
        - cond_bbox_input
        - continual_gen_input
        - cond_cate_to_size_pos
        - cond_cate_size_to_pos
        - cond_recover_mask_input
    """
    instance = sample.split("##Here is the result:")
    if len(instance) > 2:
        res = instance[-2]
    else:
        res = instance[1]
    bboxs, categories = extract_xywh(res)
    return bboxs, categories


def convert_to_array(generations: List[str]):
    final_res = []
    for sample in generations:
        bboxs, cates = sample.get("bbox"), sample.get("categories")
        bboxs, cates = bboxs.tolist(), cates.tolist()
        bboxs = np.array(bboxs, dtype=np.float32)
        cates = np.array(cates, dtype=np.int32)
        final_res.append((bboxs, cates))
    return final_res
        
        
def print_scores(scores: Dict):
    scores = {k: scores[k] for k in sorted(scores)}

    tex = ""
    for k, v in scores.items():
        # if k == "Alignment" or k == "Overlap" or "Violation" in k:
        #     v = [_v * 100 for _v in v]
        mean, std = np.mean(v), np.std(v)
        stdp = std * 100.0 / mean
        print(f"\t{k}: {mean:.4f} ({stdp:.4f}%)")
        tex += f"& {mean:.4f}\\std{{{stdp:.1f}}}\% "
    print(tex)

def average(scores):
    return sum(scores) / len(scores)

def main(args):
    
    ## load cluster model
    _clustering_models = dict()
    valid_keys = [f"{k}-64" for k in ['x', 'y', 'w', 'h']]
    with fsspec.open(args.cluster_model, "rb") as f:
        for key, model in pickle.load(f).items():
            if key not in valid_keys:
                continue
            # sort cluster center in 1d case
            var_name = key.split("-")[0]
            if len(var_name) == 1:
                cluster_centers = np.sort(
                    model.cluster_centers_, axis=0
                )  # (N, 1)
                model.cluster_centers_ = cluster_centers
            _clustering_models[key] = model   
    
    # load val data
    print(f">>> load validation data")
    transforms = [LexicographicSort()]
    val_dataset = get_dataset(
            name=args.dataset_name,
            datapath=args.dataset_path,
            split='val',
            transform=T.Compose(transforms)
        )
    val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, pin_memory=True, shuffle=False, num_workers=0)

    print(">>> load testing data")
    transforms = [LexicographicSort()]
    test_dataset = get_dataset(
            name=args.dataset_name,
            datapath=args.dataset_path,
            split='test',
            transform=T.Compose(transforms)
        )
    test_dataloader = DataLoader(test_dataset, batch_size=args.batch_size, pin_memory=True, shuffle=False, num_workers=0)
    layouts_main = loader_to_list(test_dataloader)

    # load golden data
    print(f">>> load golden file from {args.golden_file}")
    with open(args.golden_file, "r") as f:
        goldens = [json.loads(line) for line in f]
    
    # load fid model
    print(">>> load fid model")
    num_classes = len(label_to_int)
    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    fid_model = load_fidnet_v3_simple(args.fid_model_name_or_path, device, num_label=num_classes)
    
    print(">>> Extract features from testing file")
    feats_test = []
    for i, batch in enumerate(test_dataloader):
        bbox, label, padding_mask, mask = sparse_to_dense(batch, device)
        bbox = cluster_corrd(_clustering_models, bbox, valid_keys)
        with torch.set_grad_enabled(False):
            feat = fid_model.extract_features(bbox, label, padding_mask)
        feats_test.append(feat.cpu())
        
    print(">>> Extract features from eval file")
    feats_val = []
    for i, batch in enumerate(val_dataloader):
        bbox, label, padding_mask, mask = sparse_to_dense(batch, device)
        with torch.set_grad_enabled(False):
            feat = fid_model.extract_features(bbox, label, padding_mask)
        feats_val.append(feat.cpu())
    
    real_fid = compute_generative_model_scores(feats_test, feats_val)
    print("real FID score is")
    print(real_fid)
    
    if args.test_others:  # implement your own evaluation code here
        pass
    
    ### check and save all generations
    if os.path.exists(args.intermediate_saved_path):  
        print(f"{args.intermediate_saved_path} already exists, read directly!")
        all_generations = torch.load(args.intermediate_saved_path)
    else:  
        print(f"{args.intermediate_saved_path} not exist, process and save the results to this path!")
        with open(args.gen_res_path, "r") as f:
            generations = [json.loads(line) for line in f]    
        all_generations = []
        for i, sample in enumerate(generations):
            if "id" in sample:
                id_ = sample.pop("id")
            else:
                id_ = i
            
            # extract corresponding golden bbox and categories
            golden_bboxs, golden_labels = extract_xywh(goldens[id_].get('labels')[0])
            W, H = extract_WH(goldens[id_].get('labels')[0])
            # construct golden features
            golden_bboxs = [convert_real_to_xywh(bbox, W, H) for bbox in golden_bboxs]
            golden_bboxs = torch.tensor(golden_bboxs, dtype=torch.float)
            golden_labels = torch.tensor(golden_labels, dtype=torch.long)
            per_sample = {
                "id": id_, 
                "golden": {
                    "bbox": golden_bboxs,
                    "categories": golden_labels,
                },
            } 
            for k, v in sample.items():
                gen_bboxs, gen_categories = preprocess(v)
                gen_bboxs = [convert_real_to_xywh(bbox, W, H) for bbox in gen_bboxs]
                gen_bboxs = torch.tensor(gen_bboxs, dtype=torch.float)
                
                gen_categories = torch.tensor(gen_categories, dtype=torch.long)
                per_sample[k] = {"bbox": gen_bboxs, "categories": gen_categories}
            all_generations.append(per_sample)
        
        torch.save(all_generations, args.intermediate_saved_path)
    
    all_keys = list(all_generations[0].keys())
    block_lst = ["id", "golden"]
    print(f">>> All keys are {all_keys} | Begin to extract features from generations")
    
    scores_all = dict()
    for gk in all_keys:
        # create saver for each k
        if gk in block_lst:
            continue
        feats_gen = [] 
        batch_metrics = defaultdict(float)
        filter_ids = [] # filter the empty bbox 
        k_generations = [tmp[gk] for tmp in all_generations]  
        for i in range(0, len(k_generations), args.batch_size):
            i_end = min(i + args.batch_size, len(k_generations))
            batch = k_generations[i:i_end]
            max_len = max(len(g["categories"]) for g in batch)
            if max_len == 0:  # prevent not generations
                max_len == 1
            bbox, label, padding_mask, mask, empty_ids = preprocess_batch(_clustering_models, batch, max_len, device, k_means=True)
            filter_ids.extend(empty_ids)
            
            mask_empty_ids = torch.tensor(empty_ids, dtype=torch.bool, device=device)
            bbox = torch.masked_select(bbox, mask_empty_ids.unsqueeze(1).unsqueeze(2)).reshape(-1, bbox.size(1), bbox.size(2)).contiguous()
            label = torch.masked_select(label, mask_empty_ids.unsqueeze(1)).reshape(-1, label.size(1)).contiguous()
            padding_mask = torch.masked_select(padding_mask, mask_empty_ids.unsqueeze(1)).reshape(-1, padding_mask.size(1)).contiguous()
            mask = torch.masked_select(mask, mask_empty_ids.unsqueeze(1)).reshape(-1, mask.size(1)).contiguous()
            with torch.set_grad_enabled(False):
                feat = fid_model.extract_features(bbox, label, padding_mask)
            feats_gen.append(feat.cpu())    
                
            for k, v in compute_alignment(bbox, mask).items():
                batch_metrics[k] += v.sum().item()
            for k, v in compute_overlap(bbox, mask).items():
                batch_metrics[k] += v.sum().item()  

        print(">>> Extract features from testing file again (this time filter the failed bbox)")
        filter_feats_test = []
        end_flag = False
        for i, batch in enumerate(test_dataloader):
            bbox, label, padding_mask, mask = sparse_to_dense(batch, device)
            b_i = i * args.batch_size
            e_i = b_i + bbox.size(0)
            empty_ids = filter_ids[b_i: e_i]
            if len(empty_ids) < bbox.size(0):
                bbox = bbox[:len(empty_ids)]
                label = label[:len(empty_ids)]
                padding_mask = padding_mask[:len(empty_ids)]
                end_flag = True
            mask_empty_ids = torch.tensor(empty_ids, dtype=torch.bool, device=device)
            bbox = torch.masked_select(bbox, mask_empty_ids.unsqueeze(1).unsqueeze(2)).reshape(-1, bbox.size(1), bbox.size(2)).contiguous()
            label = torch.masked_select(label, mask_empty_ids.unsqueeze(1)).reshape(-1, label.size(1)).contiguous()
            padding_mask = torch.masked_select(padding_mask, mask_empty_ids.unsqueeze(1)).reshape(-1, padding_mask.size(1)).contiguous()
            with torch.set_grad_enabled(False):
                feat = fid_model.extract_features(bbox, label, padding_mask)
            filter_feats_test.append(feat.cpu())
            if end_flag:
                break
        
        # calculate IOU (max, average) and docsim score
        gen_data = convert_to_array(k_generations)
        
        scores = {}
        for k, v in batch_metrics.items():
            scores[k] = v / len(k_generations)
        
        scores.update(compute_average_iou(gen_data))
        scores.update(compute_generative_model_scores(filter_feats_test, feats_gen))
        scores["maximum_iou"] = compute_maximum_iou(layouts_main, gen_data)
        scores["DocSim"] = compute_docsim(layouts_main, gen_data)
        
        scores["Failed_rate"] = (len(filter_ids) - sum(filter_ids)) / len(filter_ids) * 100
        
        scores_all[gk] = defaultdict(list)
        for k, v in scores.items():
            scores_all[gk][k].append(v)
            
        
        print(f">>> Logging Scores for condition {gk}\n")
        print("-"*20 + "\n")
        print_scores(scores_all[gk])
     
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluation Generated Layout Code")
    parser.add_argument("--test_others", action='store_true', help="custom your evaluation own code here")  
    parser.add_argument("--file_dir", type=str, default="data/generated_results/rico")  
    parser.add_argument("--intermediate_saved_path", type=str, default=None) 
    parser.add_argument("--golden_file", type=str, default="data/generated_results/rico/golden.jsonl")  
    parser.add_argument("--fid_model_name_or_path", type=str, default="models/rico25-max25",)  
    parser.add_argument("--cluster_model", type=str, default="models/rico25-max25/rico25_max25_kmeans_train_clusters.pkl")
    parser.add_argument("--dataset_name", type=str, default="rico25")
    parser.add_argument("--dataset_path", type=str, default="data/rico25-max25")
    parser.add_argument("--gen_res_path", type=str, default=None, help="generated html code")  
    parser.add_argument("--batch_size", type=int, default=32)  
    parser.add_argument("--device", type=str, default="cuda:0")  

    args = parser.parse_args()
    
    int_to_lable = DATASET_META.get(args.dataset_name)
    label_to_int = dict([(v, k) for k, v in int_to_lable.items()])

    main(args)
evaluate.py
ProjectNUWA/LayoutNUWA
import sys
import os
current_directory = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_directory)  

import argparse
import logging
import pickle
import time
import torch
from sklearn.cluster import KMeans
from convertHTML.cluster import Percentile
from convertHTML import get_dataset

logger = logging.getLogger(__name__)

KEYS = ["x", "y", "w", "h"]

parser = argparse.ArgumentParser()
parser.add_argument("--dataset_name", type=str, default='publaynet')
parser.add_argument("--dataset_dir", type=str, default=None)
parser.add_argument("--max_seq_length", type=int, default=25)  
parser.add_argument("--algorithm", '-a', type=str, choices=["kmeans", "percentile"])
parser.add_argument("--result_dir", '-rd', type=str, default=None)
parser.add_argument("--random_state", type=int, default=0)
parser.add_argument(
    "--max_bbox_num",
    type=int,
    default=int(1e5),
    help="filter number of bboxes to avoid too much time consumption in kmeans",
)

args = parser.parse_args()

n_clusters_list = [2**i for i in range(1, 9)]

dataset = get_dataset(args.dataset_name, args.dataset_dir, split="train", transform=None)
bboxes = torch.cat([e.x for e in dataset], axis=0)

models = {}
# name = Path(args.dataset_yaml).stem
weight_path = f"{args.result_dir}/{args.dataset_name}_max{args.max_seq_length}_{args.algorithm}_train_clusters.pkl"

if bboxes.size(0) > args.max_bbox_num and args.algorithm == "kmeans":
    text = f"{bboxes.size(0)} -> {args.max_bbox_num}"
    logger.warning(
        f"Subsampling bboxes because there are too many for kmeans: ({text})"
    )
    generator = torch.Generator().manual_seed(args.random_state)
    indices = torch.randperm(bboxes.size(0), generator=generator)
    bboxes = bboxes[indices[: args.max_bbox_num]]

for n_clusters in n_clusters_list:
    start_time = time.time()
    if args.algorithm == "kmeans":
        kwargs = {"n_clusters": n_clusters, "random_state": args.random_state}
        # one variable
        for i, key in enumerate(KEYS):
            key = f"{key}-{n_clusters}"
            models[key] = KMeans(**kwargs).fit(bboxes[..., i : i + 1].numpy())
    elif args.algorithm == "percentile":
        kwargs = {"n_clusters": n_clusters}
        for i, key in enumerate(KEYS):
            key = f"{key}-{n_clusters}"
            models[key] = Percentile(**kwargs).fit(bboxes[..., i : i + 1].numpy())
    print(
        f"{args.dataset_name} ({args.algorithm} {n_clusters} clusters): {time.time() - start_time}s"
    )

with open(weight_path, "wb") as f:
    pickle.dump(models, f, protocol=pickle.HIGHEST_PROTOCOL)
helper/clustering_corrdinates.py
ProjectNUWA/LayoutNUWA
import os
import fsspec
import torch
import torch.nn as nn


class TransformerWithToken(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward, num_layers):
        super().__init__()

        self.token = nn.Parameter(torch.randn(1, 1, d_model))
        token_mask = torch.zeros(1, 1, dtype=torch.bool)
        self.register_buffer("token_mask", token_mask)

        self.core = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=d_model,
                nhead=nhead,
                dim_feedforward=dim_feedforward,
            ),
            num_layers=num_layers,
        )

    def forward(self, x, src_key_padding_mask):
        # x: [N, B, E]
        # padding_mask: [B, N]
        #   `False` for valid values
        #   `True` for padded values

        B = x.size(1)

        token = self.token.expand(-1, B, -1)
        x = torch.cat([token, x], dim=0)

        token_mask = self.token_mask.expand(B, -1)
        padding_mask = torch.cat([token_mask, src_key_padding_mask], dim=1)

        x = self.core(x, src_key_padding_mask=padding_mask)

        return x


class FIDNet(nn.Module):
    def __init__(self, num_label):
        super().__init__()

        self.emb_label = nn.Embedding(num_label, 32)
        self.fc_bbox = nn.Linear(4, 32)
        self.transformer = TransformerWithToken(
            d_model=64, nhead=4, dim_feedforward=32, num_layers=4
        )
        self.fc_out = nn.Linear(64, 1)

    def extract_features(self, bbox, label, padding_mask):
        l = self.emb_label(label)
        b = self.fc_bbox(bbox)
        x = torch.cat([l, b], dim=-1).permute(1, 0, 2)
        x = self.transformer(x, padding_mask)
        return x[0]

    def forward(self, bbox, label, padding_mask):
        x = self.extract_features(bbox, label, padding_mask)
        x = self.fc_out(x)
        return x.squeeze(-1)


class FIDNetV2(nn.Module):
    def __init__(self, num_label, max_bbox=50):
        super().__init__()

        self.emb_label = nn.Embedding(num_label, 128)
        self.fc_bbox = nn.Linear(4, 128)
        self.encoder = TransformerWithToken(
            d_model=256, nhead=4, dim_feedforward=128, num_layers=8
        )

        self.fc_out = nn.Sequential(
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Linear(64, 1),
        )

        self.token = nn.Parameter(torch.rand(max_bbox, 1, 256))
        te = nn.TransformerEncoderLayer(d_model=256, dim_feedforward=128, nhead=4)
        self.decoder = nn.TransformerEncoder(te, num_layers=8)
        self.fc_out_cls = nn.Linear(256, num_label)
        self.fc_out_bbox = nn.Linear(256, 4)

    def extract_features(self, bbox, label, padding_mask):
        l = self.emb_label(label)
        b = self.fc_bbox(bbox)
        x = torch.cat([l, b], dim=-1).permute(1, 0, 2)
        x = self.encoder(x, padding_mask)
        return x[0]

    def forward(self, bbox, label, padding_mask):
        B, N, _ = bbox.size()
        x = self.extract_features(bbox, label, padding_mask)

        logit = self.fc_out(x).squeeze(-1)

        t = self.token[:N].expand(-1, B, -1)
        x = torch.cat([x.unsqueeze(0), t], dim=0)

        token_mask = self.encoder.token_mask.expand(B, -1)
       
        _padding_mask = torch.cat([token_mask, padding_mask], dim=1)

        x = self.decoder(x, src_key_padding_mask=_padding_mask)
        # x = x[1:].permute(1, 0, 2)[~padding_mask]
        x = x[1:].permute(1, 0, 2)

        logit_cls = self.fc_out_cls(x)
        bbox = torch.sigmoid(self.fc_out_bbox(x))

        return logit, logit_cls, bbox


class FIDNetV3(nn.Module):
    def __init__(self, num_label, d_model=256, nhead=4, num_layers=4, max_bbox=50):
        super().__init__()

        # encoder
        self.emb_label = nn.Embedding(num_label, d_model)
        self.fc_bbox = nn.Linear(4, d_model)
        self.enc_fc_in = nn.Linear(d_model * 2, d_model)

        self.enc_transformer = TransformerWithToken(
            d_model=d_model,
            dim_feedforward=d_model // 2,
            nhead=nhead,
            num_layers=num_layers,
        )

        self.fc_out_disc = nn.Linear(d_model, 1)

        # decoder
        self.pos_token = nn.Parameter(torch.rand(max_bbox, 1, d_model))
        self.dec_fc_in = nn.Linear(d_model * 2, d_model)

        te = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model // 2
        )
        self.dec_transformer = nn.TransformerEncoder(te, num_layers=num_layers)

        self.fc_out_cls = nn.Linear(d_model, num_label)
        self.fc_out_bbox = nn.Linear(d_model, 4)

    def extract_features(self, bbox, label, padding_mask):
        b = self.fc_bbox(bbox)
        l = self.emb_label(label)
        x = self.enc_fc_in(torch.cat([b, l], dim=-1))
        x = torch.relu(x).permute(1, 0, 2)
        x = self.enc_transformer(x, padding_mask)
        return x[0]

    def forward(self, bbox, label, padding_mask):
        B, N, _ = bbox.size()
        x = self.extract_features(bbox, label, padding_mask)

        logit_disc = self.fc_out_disc(x).squeeze(-1)

        x = x.unsqueeze(0).expand(N, -1, -1)
        t = self.pos_token[:N].expand(-1, B, -1)
        x = torch.cat([x, t], dim=-1)
        x = torch.relu(self.dec_fc_in(x))

        x = self.dec_transformer(x, src_key_padding_mask=padding_mask)
        # x = x.permute(1, 0, 2)[~padding_mask]
        x = x.permute(1, 0, 2)

        # logit_cls: [B, N, L]    bbox_pred: [B, N, 4]
        logit_cls = self.fc_out_cls(x)
        bbox_pred = torch.sigmoid(self.fc_out_bbox(x))

        return logit_disc, logit_cls, bbox_pred
    

class LayoutNet(nn.Module):
    def __init__(self, num_label):
        super().__init__()

        d_model = 256
        nhead = 4
        num_layers = 4
        max_bbox = 50

        # encoder
        self.emb_label = nn.Embedding(num_label, d_model)
        self.fc_bbox = nn.Linear(4, d_model)
        self.enc_fc_in = nn.Linear(d_model * 2, d_model)

        self.enc_transformer = TransformerWithToken(d_model=d_model,
                                                    dim_feedforward=d_model // 2,
                                                    nhead=nhead, num_layers=num_layers)

        self.fc_out_disc = nn.Linear(d_model, 1)

        # decoder
        self.pos_token = nn.Parameter(torch.rand(max_bbox, 1, d_model))
        self.dec_fc_in = nn.Linear(d_model * 2, d_model)

        te = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead,
                                        dim_feedforward=d_model // 2)
        self.dec_transformer = nn.TransformerEncoder(te, num_layers=num_layers)

        self.fc_out_cls = nn.Linear(d_model, num_label)
        self.fc_out_bbox = nn.Linear(d_model, 4)

    def extract_features(self, bbox, label, padding_mask):
        b = self.fc_bbox(bbox)
        l = self.emb_label(label)
        x = self.enc_fc_in(torch.cat([b, l], dim=-1))
        x = torch.relu(x).permute(1, 0, 2)
        x = self.enc_transformer(x, padding_mask)
        return x[0]

    def forward(self, bbox, label, padding_mask):
        B, N, _ = bbox.size()
        x = self.extract_features(bbox, label, padding_mask)

        logit_disc = self.fc_out_disc(x).squeeze(-1)

        x = x.unsqueeze(0).expand(N, -1, -1)
        t = self.pos_token[:N].expand(-1, B, -1)
        x = torch.cat([x, t], dim=-1)
        x = torch.relu(self.dec_fc_in(x))

        x = self.dec_transformer(x, src_key_padding_mask=padding_mask)
        x = x.permute(1, 0, 2)[~padding_mask]

        # logit_cls: [M, L]    bbox_pred: [M, 4]
        logit_cls = self.fc_out_cls(x)
        bbox_pred = torch.sigmoid(self.fc_out_bbox(x))

        return logit_disc, logit_cls, bbox_pred
    

def load_fidnetsimple(model_name_or_path: str, device: torch.device, num_label=5, max_bbox=25) -> nn.Module:
    fid_model = LayoutNet(num_label).to(device)
    with fsspec.open(model_name_or_path, "rb") as file_obj:
        x = torch.load(file_obj, map_location=device)
    fid_model.load_state_dict(x)
    fid_model.eval()
    return fid_model


def load_fidnet_v3(dataset, weight_dir: str, device: torch.device) -> nn.Module:
    prefix = f"{dataset.name}-max{dataset.max_seq_length}"
    ckpt_path = os.path.join(weight_dir, prefix, "model_best.pth.tar")
    fid_model = FIDNetV3(
        num_label=dataset.num_classes, max_bbox=dataset.max_seq_length
    ).to(device)
    with fsspec.open(ckpt_path, "rb") as file_obj:
        x = torch.load(file_obj, map_location=device)
    fid_model.load_state_dict(x["state_dict"])
    fid_model.eval()
    return fid_model

def load_fidnet_v3_simple(model_name_or_path: str, device: torch.device, num_label, max_bbox=25) -> nn.Module:
    ckpt_path = os.path.join(model_name_or_path, "model_best.pth.tar")
    fid_model = FIDNetV3(num_label=num_label, max_bbox=max_bbox).to(device)
    with fsspec.open(ckpt_path, "rb") as file_obj:
        x = torch.load(file_obj, map_location=device)
    fid_model.load_state_dict(x["state_dict"])
    fid_model.eval()
    return fid_model

helper/fid_model.py
ProjectNUWA/LayoutNUWA
# this is global vars

KEY_MULT_DICT = {
    "x-y-w-h": {"x":0, "y": 1, "w": 2, "h": 3},
    "xywh": {},
}


MULTI_CHOICE = {
    "kmeans": "x-y-w-h",
    "linear": "xywh",
    "int2str": "",
    "float2str": "",
}

SPECIAL_TOKENS = ["<MASK>", "<PAD>", "<EOS>", "<BOS>", "<SEP>"]

IGNORE_INDEX = -100

TEMPLATE_FORMAT = {
    "html_format": " <body> <svg width=\"{W}\" height=\"{H}\"> {content} </svg> </body>",
    "bbox_format": "<rect data-category=\"{c}\", x=\"{x}\", y=\"{y}\", width=\"{w}\", height=\"{h}\"/>",
}

TASK_INSTRUCTION = {
    "rico25": "I want to generate layout in the mobile app design format. ",
    "publaynet": "I want to generate layout in the document design format. ",
    "magazine": "I want to generate layout in the magazine design format. ",
}

INSTRUCTION = {
    "cond_cate_to_size_pos": "please generate the layout html according to the categories I provide (in html format):\n###bbox html: {bbox_html}",
    "cond_cate_size_to_pos": "please generate the layout html according to the categories and size I provide (in html format):\n###bbox html: {bbox_html}",
    "cond_random_mask": "please recover the layout html according to the bbox, categories and size I provide (in html format):\n###bbox html: {bbox_html}"
}

INFILLING_INSTRUCTION = {
    "cond_cate_to_size_pos": "please fulfilling the layout html according to the categories I provide (in html format):\n###bbox html: {bbox_html}",
    "cond_cate_size_to_pos": "please fulfilling the layout html according to the categories and size I provide (in html format):\n###bbox html: {bbox_html}",
    "cond_random_mask": "please recover the layout html according to the bbox, categories and size I provide (in html format):\n###bbox html: {bbox_html}"
}

SEP_SEQ = [
    "{instruct}\n\n##Here is the result:\n\n```{result}```",
    "{instruct}\n\n##Here is the result:",
    "{instruct} <MID> {result}",
    "{instruct} <MID>",
]

DATASET_META = {
    "magazine": {
        0: 'text',
        1: 'image',
        2: 'headline',
        3: 'text-over-image',
        4: 'headline-over-image',
    },
    "publaynet": {
        0: 'text',
        1: 'title',
        2: 'list',
        3: 'table',
        4: 'figure',
    },
    "rico25": {
        0: "Text",
        1: "Image",
        2: "Icon",
        3: "Text Button",
        4: "List Item",
        5: "Input",
        6: "Background Image",
        7: "Card",
        8: "Web View",
        9: "Radio Button",
        10: "Drawer",
        11: "Checkbox",
        12: "Advertisement",
        13: "Modal",
        14: "Pager Indicator",
        15: "Slider",
        16: "On/Off Switch",
        17: "Button Bar",
        18: "Toolbar",
        19: "Number Stepper",
        20: "Multi-Tab",
        21: "Date Picker",
        22: "Map View",
        23: "Video",
        24: "Bottom Navigation",
    },
}

# verborsed number
VERBALIZED_NUM = {  
    0: 'zero', 1: 'one', 2: 'two', 3: 'three', 4: 'four', 5: 'five', 6: 'six', 7: 'seven', 8: 'eight', 9: 'nine',  
    10: 'ten', 11: 'eleven', 12: 'twelve', 13: 'thirteen', 14: 'fourteen', 15: 'fifteen', 16: 'sixteen',  
    17: 'seventeen', 18: 'eighteen', 19: 'nineteen', 20: 'twenty',  
    21: 'twenty-one', 22: 'twenty-two', 23: 'twenty-three', 24: 'twenty-four', 25: 'twenty-five',  
    26: 'twenty-six', 27: 'twenty-seven', 28: 'twenty-eight', 29: 'twenty-nine', 30: 'thirty',  
    31: 'thirty-one', 32: 'thirty-two'  
}  
  
 

helper/global_var.py
ProjectNUWA/LayoutNUWA
import sys
import os
current_directory = os.path.dirname(os.path.abspath(__file__))
parent_directory = os.path.dirname(current_directory)
sys.path.append(parent_directory)  

import multiprocessing
from functools import partial
from itertools import chain
from typing import Dict, List, Optional, Tuple, Union
from helper.fid_model import *
import numpy as np
import torch
from einops import rearrange, reduce, repeat
from prdc import compute_prdc
from pytorch_fid.fid_score import calculate_frechet_distance
from scipy.optimize import linear_sum_assignment
from scipy.stats import wasserstein_distance
from torch import BoolTensor, FloatTensor
from convertHTML.utils import convert_xywh_to_ltrb


Feats = Union[FloatTensor, List[FloatTensor]]
Layout = Tuple[np.ndarray, np.ndarray]

# set True to disable parallel computing by multiprocessing (typically for debug)
# DISABLED = False
DISABLED = True


def __to_numpy_array(feats: Feats) -> np.ndarray:
    if isinstance(feats, list):
        # flatten list of batch-processed features
        if isinstance(feats[0], FloatTensor):
            feats = [x.detach().cpu().numpy() for x in feats]
    else:
        feats = feats.detach().cpu().numpy()
    return np.concatenate(feats)


def compute_generative_model_scores(
    feats_real: Feats,
    feats_fake: Feats,
) -> Dict[str, float]:
    """
    Compute precision, recall, density, coverage, and FID.
    """
    feats_real = __to_numpy_array(feats_real)
    feats_fake = __to_numpy_array(feats_fake)
    
    mu_real = np.mean(feats_real, axis=0)
    sigma_real = np.cov(feats_real, rowvar=False)
    mu_fake = np.mean(feats_fake, axis=0)
    sigma_fake = np.cov(feats_fake, rowvar=False)
    
    results = compute_prdc(
        real_features=feats_real, fake_features=feats_fake, nearest_k=5
    )

    results["fid"] = calculate_frechet_distance(
        mu_real, sigma_real, mu_fake, sigma_fake
    )
    return results


def compute_alignment(bbox: FloatTensor, mask: BoolTensor) -> Dict[str, FloatTensor]:
    """
    Computes some alignment metrics that are different to each other in previous works.
    Attribute-conditioned Layout GAN for Automatic Graphic Design (TVCG2020)
    https://arxiv.org/abs/2009.05284
    """
    S = bbox.size(1)

    bbox = bbox.permute(2, 0, 1)
    xl, yt, xr, yb = convert_xywh_to_ltrb(bbox)
    xc, yc = bbox[0], bbox[1]
    X = torch.stack([xl, xc, xr, yt, yc, yb], dim=1)
    X = X.unsqueeze(-1) - X.unsqueeze(-2)
    idx = torch.arange(X.size(2), device=X.device)
    X[:, :, idx, idx] = 1.0
    X = X.abs().permute(0, 2, 1, 3)
    X[~mask] = 1.0
    X = X.min(-1).values.min(-1).values
    X.masked_fill_(X.eq(1.0), 0.0)
    X = -torch.log(1 - X)

    # original
    # return X.sum(-1) / mask.float().sum(-1)

    score = reduce(X, "b s -> b", reduction="sum")
    score_normalized = score / reduce(mask, "b s -> b", reduction="sum")
    score_normalized[torch.isnan(score_normalized)] = 0.0

    Y = torch.stack([xl, xc, xr], dim=1)
    Y = rearrange(Y, "b x s -> b x 1 s") - rearrange(Y, "b x s -> b x s 1")

    batch_mask = rearrange(~mask, "b s -> b 1 s") | rearrange(~mask, "b s -> b s 1")
    idx = torch.arange(S, device=Y.device)
    batch_mask[:, idx, idx] = True
    batch_mask = repeat(batch_mask, "b s1 s2 -> b x s1 s2", x=3)
    Y[batch_mask] = 1.0

    # Y = rearrange(Y.abs(), "b x s1 s2 -> b s1 x s2")
    # Y = reduce(Y, "b x s1 s2 -> b x", "min")
    # Y = rearrange(Y.abs(), " -> b s1 x s2")
    Y = reduce(Y.abs(), "b x s1 s2 -> b s1", "min")
    Y[Y == 1.0] = 0.0
    score_Y = reduce(Y, "b s -> b", "sum")

    results = {
        "alignment-ACLayoutGAN": score,
        "alignment-LayoutGAN++": score_normalized,
        "alignment-NDN": score_Y,
    }
    return results


def compute_overlap(bbox: FloatTensor, mask: BoolTensor) -> Dict[str, FloatTensor]:
    """
    Based on
    (i) Attribute-conditioned Layout GAN for Automatic Graphic Design (TVCG2020)
    https://arxiv.org/abs/2009.05284
    (ii) LAYOUTGAN: GENERATING GRAPHIC LAYOUTS WITH WIREFRAME DISCRIMINATORS (ICLR2019)
    https://arxiv.org/abs/1901.06767
    "percentage of total overlapping area among any two bounding boxes inside the whole page."
    At least BLT authors seems to sum. (in the MSCOCO case, it surpasses 1.0)
    """
    B, S = mask.size()
    bbox = bbox.masked_fill(~mask.unsqueeze(-1), 0)
    bbox = bbox.permute(2, 0, 1)

    l1, t1, r1, b1 = convert_xywh_to_ltrb(bbox.unsqueeze(-1))
    l2, t2, r2, b2 = convert_xywh_to_ltrb(bbox.unsqueeze(-2))
    a1 = (r1 - l1) * (b1 - t1)

    # intersection
    l_max = torch.maximum(l1, l2)
    r_min = torch.minimum(r1, r2)
    t_max = torch.maximum(t1, t2)
    b_min = torch.minimum(b1, b2)
    cond = (l_max < r_min) & (t_max < b_min)
    ai = torch.where(cond, (r_min - l_max) * (b_min - t_max), torch.zeros_like(a1[0]))

    # diag_mask = torch.eye(a1.size(1), dtype=torch.bool, device=a1.device)
    # ai = ai.masked_fill(diag_mask, 0)
    batch_mask = rearrange(~mask, "b s -> b 1 s") | rearrange(~mask, "b s -> b s 1")
    idx = torch.arange(S, device=ai.device)
    batch_mask[:, idx, idx] = True
    ai = ai.masked_fill(batch_mask, 0)

    ar = torch.nan_to_num(ai / a1)  # (B, S, S)

    # original
    # return ar.sum(dim=(1, 2)) / mask.float().sum(-1)

    # fixed to avoid the case with single bbox
    score = reduce(ar, "b s1 s2 -> b", reduction="sum")
    score_normalized = score / reduce(mask, "b s -> b", reduction="sum")
    score_normalized[torch.isnan(score_normalized)] = 0.0

    ids = torch.arange(S)
    ii, jj = torch.meshgrid(ids, ids, indexing="ij")
    ai[repeat(ii >= jj, "s1 s2 -> b s1 s2", b=B)] = 0.0
    overlap = reduce(ai, "b s1 s2 -> b", reduction="sum")

    results = {
        "overlap-ACLayoutGAN": score,
        "overlap-LayoutGAN++": score_normalized,
        "overlap-LayoutGAN": overlap,
    }
    return results


def compute_iou(
    box_1: Union[np.ndarray, FloatTensor],
    box_2: Union[np.ndarray, FloatTensor],
    generalized: bool = False,
) -> Union[np.ndarray, FloatTensor]:
    # box_1: [N, 4]  box_2: [N, 4]

    if isinstance(box_1, np.ndarray):
        lib = np
    elif isinstance(box_1, FloatTensor):
        lib = torch
    else:
        raise NotImplementedError(type(box_1))

    l1, t1, r1, b1 = convert_xywh_to_ltrb(box_1.T)
    l2, t2, r2, b2 = convert_xywh_to_ltrb(box_2.T)
    a1, a2 = (r1 - l1) * (b1 - t1), (r2 - l2) * (b2 - t2)

    # intersection
    l_max = lib.maximum(l1, l2)
    r_min = lib.minimum(r1, r2)
    t_max = lib.maximum(t1, t2)
    b_min = lib.minimum(b1, b2)
    cond = (l_max < r_min) & (t_max < b_min)
    ai = lib.where(cond, (r_min - l_max) * (b_min - t_max), lib.zeros_like(a1[0]))

    au = a1 + a2 - ai
    iou = ai / au

    if not generalized:
        return iou

    # outer region
    l_min = lib.minimum(l1, l2)
    r_max = lib.maximum(r1, r2)
    t_min = lib.minimum(t1, t2)
    b_max = lib.maximum(b1, b2)
    ac = (r_max - l_min) * (b_max - t_min)

    giou = iou - (ac - au) / ac

    return giou


def compute_perceptual_iou(
    box_1: Union[np.ndarray, FloatTensor],
    box_2: Union[np.ndarray, FloatTensor],
) -> Union[np.ndarray, FloatTensor]:
    """
    Computes 'Perceptual' IoU [Kong+, BLT'22]
    """
    # box_1: [N, 4]  box_2: [N, 4]

    if isinstance(box_1, np.ndarray):
        lib = np
    elif isinstance(box_1, FloatTensor):
        lib = torch
    else:
        raise NotImplementedError(type(box_1))

    l1, t1, r1, b1 = convert_xywh_to_ltrb(box_1.T)
    l2, t2, r2, b2 = convert_xywh_to_ltrb(box_2.T)
    a1, a2 = (r1 - l1) * (b1 - t1), (r2 - l2) * (b2 - t2)

    # intersection
    l_max = lib.maximum(l1, l2)
    r_min = lib.minimum(r1, r2)
    t_max = lib.maximum(t1, t2)
    b_min = lib.minimum(b1, b2)
    cond = (l_max < r_min) & (t_max < b_min)
    ai = lib.where(cond, (r_min - l_max) * (b_min - t_max), lib.zeros_like(a1[0]))

    # numpy-only procedure in this part
    if isinstance(box_1, FloatTensor):
        unique_box_1 = np.unique(box_1.numpy(), axis=0)
    else:
        unique_box_1 = np.unique(box_1, axis=0)
    N = 32
    l1, t1, r1, b1 = [
        (x * N).round().astype(np.int32).clip(0, N)
        for x in convert_xywh_to_ltrb(unique_box_1.T)
    ]
    canvas = np.zeros((N, N))
    for (l, t, r, b) in zip(l1, t1, r1, b1):
        canvas[t:b, l:r] = 1
    global_area_union = canvas.sum() / (N**2)

    if global_area_union > 0.0:
        iou = ai / global_area_union
        return iou
    else:
        return lib.zeros((1,))


def __compute_maximum_iou_for_layout(layout_1: Layout, layout_2: Layout) -> float:
    score = 0.0
    (bi, li), (bj, lj) = layout_1, layout_2
    N = len(bi)
    for l in list(set(li.tolist())):
        _bi = bi[np.where(li == l)]
        _bj = bj[np.where(lj == l)]
        n = len(_bi)
        ii, jj = np.meshgrid(range(n), range(n))
        ii, jj = ii.flatten(), jj.flatten()
        iou = compute_iou(_bi[ii], _bj[jj]).reshape(n, n)
        # note: maximize is supported only when scipy >= 1.4
        ii, jj = linear_sum_assignment(iou, maximize=True)
        score += iou[ii, jj].sum().item()
    return score / N


def __compute_maximum_iou(layouts_1_and_2: Tuple[List[Layout]]) -> np.ndarray:
    layouts_1, layouts_2 = layouts_1_and_2
    N, M = len(layouts_1), len(layouts_2)
    ii, jj = np.meshgrid(range(N), range(M))
    ii, jj = ii.flatten(), jj.flatten()
    scores = np.asarray(
        [
            __compute_maximum_iou_for_layout(layouts_1[i], layouts_2[j])
            for i, j in zip(ii, jj)
        ]
    ).reshape(N, M)
    ii, jj = linear_sum_assignment(scores, maximize=True)
    return scores[ii, jj]


def __get_cond2layouts(layout_list: List[Layout]) -> Dict[str, List[Layout]]:
    out = dict()
    for bs, ls in layout_list:
        cond_key = str(sorted(ls.tolist()))
        if cond_key not in out.keys():
            out[cond_key] = [(bs, ls)]
        else:
            out[cond_key].append((bs, ls))
    return out


def compute_maximum_iou(
    layouts_1: List[Layout],
    layouts_2: List[Layout],
    disable_parallel: bool = DISABLED,
    n_jobs: Optional[int] = None,
):
    """
    Computes Maximum IoU [Kikuchi+, ACMMM'21]
    """
    c2bl_1 = __get_cond2layouts(layouts_1)
    keys_1 = set(c2bl_1.keys())
    c2bl_2 = __get_cond2layouts(layouts_2)
    keys_2 = set(c2bl_2.keys())
    keys = list(keys_1.intersection(keys_2))
    args = [(c2bl_1[key], c2bl_2[key]) for key in keys]
    # to check actual number of layouts for evaluation
    # ans = 0
    # for x in args:
    #     ans += len(x[0])
    if disable_parallel:
        scores = [__compute_maximum_iou(a) for a in args]
    else:
        with multiprocessing.Pool(n_jobs) as p:
            scores = p.map(__compute_maximum_iou, args)
    scores = np.asarray(list(chain.from_iterable(scores)))
    if len(scores) == 0:
        return 0.0
    else:
        return scores.mean().item()


def __compute_average_iou(layout: Layout, perceptual: bool = False) -> float:
    bbox, _ = layout
    N = bbox.shape[0]
    if N in [0, 1]:
        return 0.0  # no overlap in principle

    ii, jj = np.meshgrid(range(N), range(N))
    ii, jj = ii.flatten(), jj.flatten()
    is_non_diag = ii != jj  # IoU for diag is always 1.0
    ii, jj = ii[is_non_diag], jj[is_non_diag]

    if perceptual:
        iou = compute_perceptual_iou(bbox[ii], bbox[jj])
    else:
        iou = compute_iou(bbox[ii], bbox[jj])

    # pick all pairs of overlapped objects
    cond = iou > np.finfo(np.float32).eps  # to avoid very-small nonzero
    # return iou.mean().item()
    if len(iou[cond]) > 0:
        return iou[cond].mean().item()
    else:
        return 0.0


def compute_average_iou(
    layouts: List[Layout],
    disable_parallel: bool = DISABLED,
    n_jobs: Optional[int] = None,
) -> Dict[str, float]:
    """
    Compute IoU between overlapping objects for each layout.
    Note that the lower is better unlike popular IoU.

    Reference:
        Variational Transformer Networks for Layout Generation (CVPR2021)
        https://arxiv.org/abs/2104.02416
    Reference: (perceptual version)
        BLT: Bidirectional Layout Transformer for Controllable Layout Generation (ECCV2022)
        https://arxiv.org/abs/2112.05112
    """
    func1 = partial(__compute_average_iou, perceptual=True)
    func2 = partial(__compute_average_iou, perceptual=False)

    # single-thread process for debugging
    if disable_parallel:
        scores1 = [func1(l) for l in layouts]
        scores2 = [func2(l) for l in layouts]
    else:
        with multiprocessing.Pool(n_jobs) as p1:
            scores1 = p1.map(func1, layouts)
        with multiprocessing.Pool(n_jobs) as p2:
            scores2 = p2.map(func2, layouts)
    results = {
        "average_iou-BLT": np.array(scores1).mean().item(),
        "average_iou-VTN": np.array(scores2).mean().item(),
    }
    return results


def __compute_bbox_sim(
    bboxes_1: np.ndarray,
    category_1: np.int64,
    bboxes_2: np.ndarray,
    category_2: np.int64,
    C_S: float = 2.0,
    C: float = 0.5,
) -> float:
    # bboxes from diffrent categories never match
    if category_1 != category_2:
        return 0.0

    cx1, cy1, w1, h1 = bboxes_1
    cx2, cy2, w2, h2 = bboxes_2

    delta_c = np.sqrt(np.power(cx1 - cx2, 2) + np.power(cy1 - cy2, 2))
    delta_s = np.abs(w1 - w2) + np.abs(h1 - h2)
    area = np.minimum(w1 * h1, w2 * h2)
    alpha = np.power(np.clip(area, 0.0, None), C)

    weight = alpha * np.power(2.0, -1.0 * delta_c - C_S * delta_s)
    return weight


def __compute_docsim_between_two_layouts(
    layouts_1_layouts_2: Tuple[List[Layout]],
    max_diff_thresh: int = 3,
) -> float:
    layouts_1, layouts_2 = layouts_1_layouts_2
    bboxes_1, categories_1 = layouts_1
    bboxes_2, categories_2 = layouts_2

    N, M = len(bboxes_1), len(bboxes_2)
    if N >= M + max_diff_thresh or N <= M - max_diff_thresh:
        return 0.0

    ii, jj = np.meshgrid(range(N), range(M))
    ii, jj = ii.flatten(), jj.flatten()
    scores = np.asarray(
        [
            __compute_bbox_sim(
                bboxes_1[i], categories_1[i], bboxes_2[j], categories_2[j]
            )
            for i, j in zip(ii, jj)
        ]
    ).reshape(N, M)
    ii, jj = linear_sum_assignment(scores, maximize=True)

    if len(scores[ii, jj]) == 0:
        # sometimes, predicted bboxes are somehow filtered.
        return 0.0
    else:
        return scores[ii, jj].mean()


def compute_docsim(
    layouts_gt: List[Layout],
    layouts_generated: List[Layout],
    disable_parallel: bool = DISABLED,
    n_jobs: Optional[int] = None,
) -> float:
    """
    Compute layout-to-layout similarity and average over layout pairs.
    Note that this is different from layouts-to-layouts similarity.
    """
    args = list(zip(layouts_gt, layouts_generated))
    if disable_parallel:
        scores = []
        for arg in args:
            scores.append(__compute_docsim_between_two_layouts(arg))
    else:
        with multiprocessing.Pool(n_jobs) as p:
            scores = p.map(__compute_docsim_between_two_layouts, args)
    return np.array(scores).mean()


def _compute_wasserstein_distance_class(
    layouts_1: List[Layout],
    layouts_2: List[Layout],
    n_categories: int = 25,
) -> float:
    categories_1 = np.concatenate([l[1] for l in layouts_1])
    counts = np.array(
        [np.count_nonzero(categories_1 == i) for i in range(n_categories)]
    )
    prob_1 = counts / np.sum(counts)

    categories_2 = np.concatenate([l[1] for l in layouts_2])
    counts = np.array(
        [np.count_nonzero(categories_2 == i) for i in range(n_categories)]
    )
    prob_2 = counts / np.sum(counts)
    return np.absolute(prob_1 - prob_2).sum()


def _compute_wasserstein_distance_bbox(
    layouts_1: List[Layout],
    layouts_2: List[Layout],
) -> float:
    bboxes_1 = np.concatenate([l[0] for l in layouts_1]).T
    bboxes_2 = np.concatenate([l[0] for l in layouts_2]).T

    # simple 1-dimensional wasserstein for (cx, cy, w, h) independently
    N = 4
    ans = 0.0
    for i in range(N):
        ans += wasserstein_distance(bboxes_1[i], bboxes_2[i])
    ans /= N

    return ans


def compute_wasserstein_distance(
    layouts_1: List[Layout],
    layouts_2: List[Layout],
    n_classes: int = 25,
) -> Dict[str, float]:
    w_class = _compute_wasserstein_distance_class(layouts_1, layouts_2, n_classes)
    w_bbox = _compute_wasserstein_distance_bbox(layouts_1, layouts_2)
    return {
        "wdist_class": w_class,
        "wdist_bbox": w_bbox,
    }

helper/metrics.py
ProjectNUWA/LayoutNUWA
import sys
import os
current_directory = os.path.dirname(os.path.abspath(__file__))
parent_directory = os.path.dirname(current_directory)
sys.path.append(parent_directory) 

import argparse
from pycocotools.coco import COCO
from collections import defaultdict 
from tqdm import trange


def cal_box_bins(coco_instance):
      
    annotations = coco_instance.getAnnIds()
    image_layout_count = defaultdict(int)  
  
    for i in trange(len(annotations)):  
        ann_id = annotations[i]
        ann = coco_instance.loadAnns(ann_id)[0]  
        image_layout_count[ann["image_id"]] += 1  
  
    max_layouts_image_id = max(image_layout_count, key=image_layout_count.get)  
    max_layouts_count = image_layout_count[max_layouts_image_id] 
    
    return max_layouts_image_id, max_layouts_count 
    


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="A script to statistic the layout dataset")  
  
    parser.add_argument("--input_file", "-i", type=str, required=True, help="Path to the input file")  
    parser.add_argument("--output_file", "-o", type=str, default=None, help="Path to the output file")  
    parser.add_argument("--num_lines", "-n", type=int, default=10, help="Number of lines to extract from the input file (default: 10)")  
  
    args = parser.parse_args()  
    coco_instance = COCO(args.input_file)

    id_, max_bbox_bins = cal_box_bins(coco_instance)
    print(f">>> max bbox bins {max_bbox_bins} | id is: {id_}")
    
    
    
    


helper/stastic.py
ProjectNUWA/LayoutNUWA
import sys
import os
current_directory = os.path.dirname(os.path.abspath(__file__))
parent_directory = os.path.dirname(current_directory)
sys.path.append(parent_directory) 

import argparse
import shutil
from pathlib import Path
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as T
from torch_geometric.loader import DataLoader
from torch_geometric.utils import to_dense_batch
from convertHTML.utils import AddNoiseToBBox, LexicographicOrder
from helper.fid_model import FIDNetV3
from helper.visualization import save_image
from convertHTML import get_dataset

def save_checkpoint(state, is_best, out_dir):
    out_path = Path(out_dir) / "checkpoint.pth.tar"
    torch.save(state, out_path)

    if is_best:
        best_path = Path(out_dir) / "model_best.pth.tar"
        shutil.copyfile(out_path, best_path)


def main():
    parser = argparse.ArgumentParser()
    # parser.add_argument("dataset_yaml", type=str)
    parser.add_argument("--out_dir", type=str, default="models/datasetname")
    parser.add_argument("--batch_size", type=int, default=64, help="input batch size")
    parser.add_argument(
        "--iteration",
        type=int,
        default=int(2e5),
        help="number of iterations to train for",
    )
    parser.add_argument(
        "--lr", type=float, default=3e-4, help="learning rate, default=3e-4"
    )
    parser.add_argument("--seed", type=int, default=23, help="manual seed")
    args = parser.parse_args()
    print(args)

    prefix = "FIDNetV3"
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    transform = T.Compose(
        [
            T.RandomApply([AddNoiseToBBox()], 0.5),
            LexicographicOrder(),
        ]
    )
    
    train_dataset = get_dataset('publaynet', '/zecheng/dataset/layout/publaynet', split="train", transform=transform)
    val_dataset = get_dataset('publaynet', '/zecheng/dataset/layout/publaynet', split="test", transform=transform)
    categories = train_dataset.labels

    kwargs = {
        "batch_size": args.batch_size,
        "num_workers": os.cpu_count(),  # os.cpu_count()
        "pin_memory": True,
    }

    train_dataloader = DataLoader(train_dataset, shuffle=True, **kwargs)
    val_dataloader = DataLoader(val_dataset, shuffle=False, **kwargs)

    device = torch.device("cuda:7" if torch.cuda.is_available() else "cpu")
    
    # import pdb; pdb.set_trace()
    
    model = FIDNetV3(num_label=len(categories), max_bbox=25).to(device)   # publaynet: 25

    # setup optimizer
    optimizer = optim.Adam(model.parameters(), lr=args.lr)

    criterion_bce = nn.BCEWithLogitsLoss(reduction="none")
    criterion_label = nn.CrossEntropyLoss(reduction="none")
    criterion_bbox = nn.MSELoss(reduction="none")

    def proc_batch(batch):
        batch = batch.to(device)
        bbox, _ = to_dense_batch(batch.x, batch.batch)
        label, mask = to_dense_batch(batch.y, batch.batch)
        padding_mask = ~mask

        is_real = batch.attr["NoiseAdded"].float()
        return bbox, label, padding_mask, mask, is_real

    iteration = 0
    best_loss = 1e8
    max_epoch = args.iteration * args.batch_size / len(train_dataset)
    max_epoch = torch.ceil(torch.tensor(max_epoch)).int().item()
    for epoch in range(max_epoch):
        model.train()
        train_loss = {
            "Loss_BCE": 0,
            "Loss_Label": 0,
            "Loss_BBox": 0,
        }
        
        # for i, batch in enumerate(train_dataloader):
        #     bbox, label, padding_mask, mask, is_real = proc_batch(batch)
        #     import pdb; pdb.set_trace()
        #     model.zero_grad()

        #     logit, logit_cls, bbox_pred = model(bbox, label, padding_mask)

        #     loss_bce = criterion_bce(logit, is_real)
        #     loss_label = criterion_label(logit_cls[mask], label[mask])
        #     loss_bbox = criterion_bbox(bbox_pred[mask], bbox[mask]).sum(-1)
        #     loss = loss_bce.mean() + loss_label.mean() + 10 * loss_bbox.mean()
        #     loss.backward()

        #     optimizer.step()

        #     loss_bce_mean = loss_bce.mean().item()
        #     train_loss["Loss_BCE"] += loss_bce.sum().item()
        #     loss_label_mean = loss_label.mean().item()
        #     train_loss["Loss_Label"] += loss_label.sum().item()
        #     loss_bbox_mean = loss_bbox.mean().item()
        #     train_loss["Loss_BBox"] += loss_bbox.sum().item()
            
        #     if i % 50 == 0:
        #         log_prefix = f"[{epoch}/{max_epoch}][{i}/{len(train_dataset) // args.batch_size}]"
        #         log = f"Loss: {loss.item():E}\tBCE: {loss_bce_mean:E}\tLabel: {loss_label_mean:E}\tBBox: {loss_bbox_mean:E}"
        #         print(f"{log_prefix}\t{log}")

        #     iteration += 1

        # for key in train_loss.keys():
        #     train_loss[key] /= len(train_dataset)
        
        model.eval()
        with torch.no_grad():
            val_loss = {
                "Loss_BCE": 0,
                "Loss_Label": 0,
                "Loss_BBox": 0,
            }

            for i, batch in enumerate(val_dataloader):
                bbox, label, padding_mask, mask, is_real = proc_batch(batch)

                # logit, logit_cls, bbox_pred = model(bbox, label, padding_mask)

                # loss_bce = criterion_bce(logit, is_real)
                # loss_label = criterion_label(logit_cls[mask], label[mask])
                # loss_bbox = criterion_bbox(bbox_pred[mask], bbox[mask]).sum(-1)

                # val_loss["Loss_BCE"] += loss_bce.sum().item()
                # val_loss["Loss_Label"] += loss_label.sum().item()
                # val_loss["Loss_BBox"] += loss_bbox.sum().item()

                if i == 0 and epoch % 10 == 0:
                    import pdb; pdb.set_trace()
                    save_image(
                        bbox,
                        label,
                        mask,
                        val_dataset.colors,
                        out_dir / f"samples_{epoch}.png",
                    )
                    cls_pred = logit_cls.argmax(dim=-1)
                    save_image(
                        bbox_pred,
                        cls_pred,
                        mask,
                        val_dataset.colors,
                        out_dir / f"recon_samples_{epoch}.png",
                    )

            for key in val_loss.keys():
                val_loss[key] /= len(val_dataset)

        tag_scalar_dict = {
            "train": sum(train_loss.values()),
            "val": sum(val_loss.values()),
        }
        
        for key in train_loss.keys():
            tag_scalar_dict = {"train": train_loss[key], "val": val_loss[key]}
            print(tag_scalar_dict)

        # do checkpointing
        val_loss = sum(val_loss.values())
        is_best = val_loss < best_loss
        best_loss = min(val_loss, best_loss)

        save_checkpoint(
            {
                "epoch": epoch + 1,
                "state_dict": model.state_dict(),
                "best_loss": best_loss,
                "optimizer": optimizer.state_dict(),
            },
            is_best,
            out_dir,
        )


if __name__ == "__main__":
    main()
helper/train_fid.py
ProjectNUWA/LayoutNUWA
import json
import random
import shutil
import numpy as np
from pathlib import Path
from datetime import datetime
from PIL import Image, ImageDraw

import torch
import torchvision.utils as vutils
import torchvision.transforms as T
from torch_geometric.utils import to_dense_batch

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    print("Random Seed:", seed)


def init_experiment(args, prefix):
    if args.seed is None:
        args.seed = random.randint(0, 10000)

    set_seed(args.seed)

    if not args.name:
        args.name = datetime.now().strftime('%Y%m%d%H%M%S%f')

    out_dir = Path('output') / args.dataset / prefix / args.name
    out_dir.mkdir(parents=True, exist_ok=True)

    json_path = out_dir / 'args.json'
    with json_path.open('w') as f:
        json.dump(vars(args), f, indent=2)

    return out_dir


def save_checkpoint(state, is_best, out_dir):
    out_path = Path(out_dir) / 'checkpoint.pth.tar'
    torch.save(state, out_path)

    if is_best:
        best_path = Path(out_dir) / 'model_best.pth.tar'
        shutil.copyfile(out_path, best_path)


def loader_to_list(loader: torch.utils.data.dataloader.DataLoader):
    layouts = []
    for batch in loader:
        bbox, label, _, mask = sparse_to_dense(batch)
        for i in range(len(label)):
            valid = mask[i].numpy()
            layouts.append((bbox[i].numpy()[valid], label[i].numpy()[valid]))
    return layouts

def sparse_to_dense(
    batch,
    device: torch.device = torch.device("cpu"),
    remove_canvas: bool = False,
):
    batch = batch.to(device)
    
    bbox, _ = to_dense_batch(batch.x, batch.batch)
    label, mask = to_dense_batch(batch.y, batch.batch)

    if remove_canvas:
        bbox = bbox[:, 1:].contiguous()
        label = label[:, 1:].contiguous() - 1  # cancel +1 effect in transform
        label = label.clamp(min=0)
        mask = mask[:, 1:].contiguous()

    padding_mask = ~mask
    return bbox, label, padding_mask, mask


def convert_xywh_to_ltrb(bbox):
    xc, yc, w, h = bbox
    x1 = xc - w / 2
    y1 = yc - h / 2
    x2 = xc + w / 2
    y2 = yc + h / 2
    return [x1, y1, x2, y2]

def convert_real_to_xywh(bbox, W, H):
    bbox = [float(i) for i in bbox]
    W = eval(W)
    H = eval(H)
    x1, y1, width, height = bbox
    
    xc = x1 + width / 2.
    yc = y1 + height / 2.
    b = [xc / W, yc / H, width / W, height / H]
    return b


def convert_layout_to_image(boxes, labels, colors, canvas_size):
    H, W = canvas_size
    img = Image.new('RGB', (int(W), int(H)), color=(255, 255, 255))
    draw = ImageDraw.Draw(img, 'RGBA')

    # draw from larger boxes
    area = [b[2] * b[3] for b in boxes]
    indices = sorted(range(len(area)),
                     key=lambda i: area[i],
                     reverse=True)

    for i in indices:
        bbox, color = boxes[i], colors[labels[i]]
        c_fill = color + (100,)
        x1, y1, x2, y2 = convert_xywh_to_ltrb(bbox)
        x1, x2 = x1 * (W - 1), x2 * (W - 1)
        y1, y2 = y1 * (H - 1), y2 * (H - 1)
        draw.rectangle([x1, y1, x2, y2],
                       outline=color,
                       fill=c_fill)
    return img


def save_image(batch_boxes, batch_labels, batch_mask,
               dataset_colors, out_path, canvas_size=(60, 40),
               nrow=None):
    # batch_boxes: [B, N, 4]
    # batch_labels: [B, N]
    # batch_mask: [B, N]

    imgs = []
    B = batch_boxes.size(0)
    to_tensor = T.ToTensor()
    for i in range(B):
        mask_i = batch_mask[i]
        boxes = batch_boxes[i][mask_i]
        labels = batch_labels[i][mask_i]
        img = convert_layout_to_image(boxes, labels,
                                      dataset_colors,
                                      canvas_size)
        imgs.append(to_tensor(img))
    image = torch.stack(imgs)

    if nrow is None:
        nrow = int(np.ceil(np.sqrt(B)))

    vutils.save_image(image, out_path, normalize=False, nrow=nrow)
helper/util.py
ProjectNUWA/LayoutNUWA
import pathlib
from collections import Counter
from typing import Dict, List, Optional, Tuple, Union

import numpy as np
import torch
import torchvision.transforms as T
import torchvision.utils as vutils
from einops import rearrange
from numpy.typing import NDArray
from PIL import Image, ImageDraw, ImageFont
from torch import BoolTensor, FloatTensor, LongTensor


def convert_xywh_to_ltrb(bbox: Union[np.ndarray, FloatTensor]):
    xc, yc, w, h = bbox
    # print(xc, yc, w, h)
    x1 = xc - w / 2
    y1 = yc - h / 2
    x2 = xc + w / 2
    y2 = yc + h / 2
    return [x1, y1, x2, y2]

def convert_layout_to_image(
    boxes: FloatTensor,
    labels: LongTensor,
    colors: List[Tuple[int]],
    canvas_size: Optional[Tuple[int]] = (60, 40),
    resources: Optional[Dict] = None,
    names: Optional[List[str]] = None,
    **kwargs,
):
    H, W = canvas_size
    if names:
        font = ImageFont.truetype("LiberationSerif-Regular", W // 10)

    if resources:
        img = resources["img_bg"].resize((W, H))
    else:
        img = Image.new("RGB", (int(W), int(H)), color=(255, 255, 255))
    draw = ImageDraw.Draw(img, "RGBA")

    # draw from larger boxes
    a = [b[2] * b[3] for b in boxes]
    indices = sorted(range(len(a)), key=lambda i: a[i], reverse=True)

    for i in indices:
        bbox, label = boxes[i], labels[i]
        if isinstance(label, LongTensor):
            label = label.item()

        c_fill = colors[label] + (100,)
        x1, y1, x2, y2 = convert_xywh_to_ltrb(bbox)
        x1, x2 = x1 * (W - 1), x2 * (W - 1)
        y1, y2 = y1 * (H - 1), y2 * (H - 1)

        if resources:
            patch = resources["cropped_patches"][i]
            # round coordinates for exact size match for rendering images
            x1, x2 = int(x1), int(x2)
            y1, y2 = int(y1), int(y2)
            w, h = x2 - x1, y2 - y1
            patch = patch.resize((w, h))
            img.paste(patch, (x1, y1))
        else:
            # print([x1, y1, x2, y2])  # TODO: debug
            draw.rectangle([x1, y1, x2, y2], outline=colors[label], fill=c_fill)
            if names:
                # draw.text((x1, y1), names[label], colors[label], font=font)
                draw.text((x1, y1), names[label], "black", font=font)
    return img


def save_image(
    batch_boxes: FloatTensor,
    batch_labels: LongTensor,
    batch_mask: BoolTensor,
    colors: List[Tuple[int]],
    out_path: Optional[Union[pathlib.PosixPath, str]] = None,
    canvas_size: Optional[Tuple[int]] = (60, 40),
    nrow: Optional[int] = None,
    batch_resources: Optional[Dict] = None,
    use_grid: bool = False,
    draw_label: bool = False,
    **kwargs,
):
    # batch_boxes: [B, N, 4]
    # batch_labels: [B, N]
    # batch_mask: [B, N]
    if isinstance(out_path, pathlib.PosixPath):
        out_path = str(out_path)
    if not draw_label and "names" in kwargs:
        kwargs.pop("names")

    imgs = []
    B = batch_boxes.size(0)
    to_tensor = T.ToTensor()
    for i in range(B):
        mask_i = batch_mask[i]
        boxes = batch_boxes[i][mask_i]
        labels = batch_labels[i][mask_i]
        if batch_resources:
            resources = {k: v[i] for (k, v) in batch_resources.items()}
            img = convert_layout_to_image(
                boxes, labels, colors, canvas_size, resources, **kwargs
            )
        else:
            img = convert_layout_to_image(boxes, labels, colors, canvas_size, **kwargs)
        imgs.append(to_tensor(img))
    image = torch.stack(imgs)

    if nrow is None:
        nrow = int(np.ceil(np.sqrt(B)))

    if out_path:
        vutils.save_image(image, out_path, normalize=False, nrow=nrow)
        print(f">>> save layout images to {out_path}")
    else:
        if use_grid:
            return torch_to_numpy_image(
                vutils.make_grid(image, normalize=False, nrow=nrow)
            )
        else:
            return image


def save_label(
    labels: Union[LongTensor, list],
    names: List[str],
    colors: List[Tuple[int]],
    out_path: Optional[Union[pathlib.PosixPath, str]] = None,
    **kwargs
    # canvas_size: Optional[Tuple[int]] = (60, 40),
):
    space, pad = 12, 12
    x_offset, y_offset = 500, 100

    img = Image.new("RGBA", (1000, 1000))
    fnt = ImageFont.truetype("LiberationSerif-Regular", 40)
    fnt_sm = ImageFont.truetype("LiberationSerif-Regular", 32)
    d = ImageDraw.Draw(img)

    if isinstance(labels, LongTensor):
        labels = labels.tolist()

    cnt = Counter(labels)
    for l in range(len(colors)):
        if l not in cnt.keys():
            continue

        text = names[l]
        use_multiline = False

        if cnt[l] > 1:
            add_width = d.textsize(f" × {cnt[l]}", font=fnt)[0] + pad
        else:
            add_width = 0

        width = d.textsize(text, font=fnt)[0]
        bbox = d.textbbox(
            (x_offset - (width + add_width) / 2, y_offset), text, font=fnt
        )
        bbox = (bbox[0] - pad, bbox[1] - pad, bbox[2] + pad, bbox[3] + pad)
        d.rectangle(bbox, fill=colors[l])

        _x_offset = x_offset - (width + add_width) / 2
        if cnt[l] > 1:
            d.text(
                (_x_offset + width + pad, y_offset),
                f" × {cnt[l]}",
                font=fnt,
                fill=(0, 0, 0),
            )

        d.text((_x_offset, y_offset), text, font=fnt, fill=(255, 255, 255))

        y_offset = y_offset + bbox[3] - bbox[1] + space

    # crop
    bbox = img.getbbox()
    bbox = (bbox[0] - pad, bbox[1] - pad, bbox[2] + pad, bbox[3] + pad)
    img = img.crop(bbox)

    # add white background
    out = Image.new("RGB", img.size, color=(255, 255, 255))
    out.paste(img, mask=img)
    # pil_size = canvas_size[::-1]  # (H, W) -> (W, H)
    # out = out.resize(pil_size)
    if out_path:
        out.save(out_path)
    else:
        return np.array(out)


def save_label_with_size(
    labels: LongTensor,
    boxes: FloatTensor,
    names: List[str],
    colors: List[Tuple[int]],
    out_path: Optional[Union[pathlib.PosixPath, str]] = None,
    # canvas_size: Optional[Tuple[int]] = (60, 40),
    **kwargs,
):
    space, pad = 12, 12
    x_offset, y_offset = 500, 100
    B = 32

    img = Image.new("RGBA", (1000, 1000))
    fnt = ImageFont.truetype("LiberationSerif-Regular", 40)
    fnt_sm = ImageFont.truetype("LiberationSerif-Regular", 32)
    d = ImageDraw.Draw(img)

    for i, l in enumerate(labels):
        w, h = [int(x) for x in (boxes[i].clip(1 / B, 1.0) * B).long()[2:]]
        text = f"{names[l]} ({w},{h})"
        add_width = 0

        width = d.textsize(text, font=fnt)[0]
        bbox = d.textbbox(
            (x_offset - (width + add_width) / 2, y_offset), text, font=fnt
        )
        bbox = (bbox[0] - pad, bbox[1] - pad, bbox[2] + pad, bbox[3] + pad)
        d.rectangle(bbox, fill=colors[l])

        _x_offset = x_offset - (width + add_width) / 2
        d.text((_x_offset, y_offset), text, font=fnt, fill=(255, 255, 255))
        y_offset = y_offset + bbox[3] - bbox[1] + space

    # crop
    bbox = img.getbbox()
    bbox = (bbox[0] - pad, bbox[1] - pad, bbox[2] + pad, bbox[3] + pad)
    img = img.crop(bbox)

    # add white background
    out = Image.new("RGB", img.size, color=(255, 255, 255))
    out.paste(img, mask=img)
    # pil_size = canvas_size[::-1]  # (H, W) -> (W, H)
    # out = out.resize(pil_size)
    if out_path:
        out.save(out_path)
    else:
        return np.array(out)


def torch_to_numpy_image(input_th: FloatTensor) -> np.ndarray:
    """
    Args
        input_th: (C, H, W), in [0.0, 1/0], torch image
    Returns
        output_npy: (H, W, C), in {0, 1, ..., 255}, numpy image
    """
    x = (input_th * 255.0).clamp(0, 255)
    x = rearrange(x, "c h w -> h w c")
    output_npy = x.numpy().astype(np.uint8)
    return output_npy


def save_relation(
    label_with_canvas: LongTensor,
    edge_attr: LongTensor,
    names: List[str],
    colors: List[Tuple[int]],
    out_path: Optional[Union[pathlib.PosixPath, str]] = None,
    **kwargs,
):
    from trainer.data.util import (  # lazy load to avoid circular import
        RelLoc,
        RelSize,
        get_rel_text,
    )

    pairs, triplets = [], []
    relations = list(RelSize) + list(RelLoc)
    for rel_value in relations:
        if rel_value in [RelSize.UNKNOWN, RelLoc.UNKNOWN]:
            continue
        cond = edge_attr & 1 << rel_value
        ii, jj = np.where(cond.numpy() > 0)
        for i, j in zip(ii, jj):
            li = label_with_canvas[i] - 1
            lj = label_with_canvas[j] - 1

            if i == 0:
                rel = get_rel_text(rel_value, canvas=True)
                pairs.append((lj, rel, None))
            else:
                rel = get_rel_text(rel_value, canvas=False)
                triplets.append((li, rel, lj))

    triplets = pairs + triplets

    space, pad = 6, 6
    img = Image.new("RGBA", (1000, 1000))
    fnt = ImageFont.truetype("LiberationSerif-Regular", 20)
    fnt_sm = ImageFont.truetype("LiberationSerif-Regular", 16)
    d = ImageDraw.Draw(img)

    def draw_text(x_offset, y_offset, text, color=None, first=False):
        if color is None:
            d.text((x_offset, y_offset), text, font=fnt, fill=(0, 0, 0))
            x_offset = x_offset + d.textsize(text, font=fnt)[0] + space
        else:
            x_offset = x_offset + pad

            use_multiline = False
            bbox = d.textbbox((x_offset, y_offset), text, font=fnt)
            if bbox[2] - bbox[0] > 120 and " " in text:
                use_multiline = True
                h_old = d.textsize(text, font=fnt)[1]
                text = text.replace(" ", "\n")
                h_new = d.multiline_textsize(text, font=fnt_sm)[1]
                h_diff = h_new - h_old
                if first:
                    y_offset = y_offset + h_diff / 2
                bbox = d.multiline_textbbox(
                    (x_offset, y_offset - h_diff / 2), text, font=fnt_sm
                )

            bbox = (bbox[0] - pad, bbox[1] - pad, bbox[2] + pad, bbox[3] + pad)
            d.rectangle(bbox, fill=color)

            if use_multiline:
                d.multiline_text(
                    (x_offset, y_offset - h_diff / 2),
                    text,
                    align="center",
                    font=fnt_sm,
                    fill=(255, 255, 255),
                )
                text_width = d.multiline_textsize(text, font=fnt_sm)[0]
            else:
                d.text((x_offset, y_offset), text, font=fnt, fill=(255, 255, 255))
                text_width = d.textsize(text, font=fnt)[0]

            x_offset = x_offset + text_width + space + pad
        return x_offset, y_offset

    for i, (l1, rel, l2) in enumerate(triplets):
        x_offset, y_offset = 20, 40 * (i + 1)
        x_offset, y_offset = draw_text(
            x_offset, y_offset, names[l1], colors[l1], first=True
        )
        x_offset, y_offset = draw_text(x_offset, y_offset, rel)
        if l2 is not None:
            draw_text(x_offset, y_offset, names[l2], colors[l2])

    # crop
    bbox = img.getbbox()
    if bbox is not None:
        bbox = (bbox[0] - pad, bbox[1] - pad, bbox[2] + pad, bbox[3] + pad)
        img = img.crop(bbox)

        # add white background
        out = Image.new("RGB", img.size, color=(255, 255, 255))
        out.paste(img, mask=img)

        if out_path:
            out.save(out_path)
        else:
            return np.array(out)


def save_gif(
    images: List[NDArray],
    out_path: str,
    **kwargs,
):
    assert images[0].ndim == 4
    to_pil = T.ToPILImage()
    for i in range(len(images[0])):
        tmp_images = [to_pil(image[i]) for image in images]
        tmp_images[0].save(
            # f"tmp/animation/{i}.gif",
            out_path.format(i),
            save_all=True,
            append_images=tmp_images[1:],
            optimize=False,
            duration=200,
            loop=0,
        )
helper/visualization.py
ProjectNUWA/LayoutNUWA
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizer
import json
import random
import torch
import transformers


class RawFileDataset(Dataset):  
    def __init__(self, args, file, tokenizer: PreTrainedTokenizer):  
        self.tokenizer = tokenizer  
        self.max_seq_len = args.model_max_length
        with open(file, "r") as f:
            self.content = [json.loads(line) for line in f]
        self.cond_type = ["cond_cate_size", "cond_category", "cond_mask"]
        self.sample_prob = [1, 1, 1]  # custom your sampling weight here
  
    def __len__(self):
        return len(self.content)
    
    def __getitem__(self, index):
        data = self.content[index]

        cond_cate_to_size_pos_seq_modeling = data["cond_cate_to_size_pos_seq_modeling"]
        cond_cate_size_to_pos_seq_modeling = data["cond_cate_size_to_pos_seq_modeling"]
        cond_recover_mask_seq_modeling = data["cond_recover_mask_seq_modeling"]
        
        ## category -> size + pos
        cond_cate_to_size_pos_inputs = self.tokenizer(
            cond_cate_to_size_pos_seq_modeling, 
            padding="max_length", 
            truncation=True, 
            max_length=self.max_seq_len,
            return_tensors="pt",
        )
        cond_cate_ids = cond_cate_to_size_pos_inputs.input_ids[0]
        cond_cate_attention_mask = cond_cate_to_size_pos_inputs.attention_mask[0]
        cond_cate_labels = torch.where(
            cond_cate_ids != self.tokenizer.pad_token_id, cond_cate_ids, -100
        )
        
        ## categoty + size -> pos
        cond_cate_size_inputs = self.tokenizer(
            cond_cate_size_to_pos_seq_modeling, 
            padding="max_length", 
            truncation=True, 
            max_length=self.max_seq_len,
            return_tensors="pt",
        )
        cond_cate_size_ids = cond_cate_size_inputs.input_ids[0]
        cond_cate_size_attention_mask = cond_cate_size_inputs.attention_mask[0]
        cond_cate_size_labels = torch.where(
            cond_cate_size_ids != self.tokenizer.pad_token_id, cond_cate_size_ids, -100
        )
        
        ## random mask
        random_mask_inputs = self.tokenizer(
            cond_recover_mask_seq_modeling, 
            padding="max_length", 
            truncation=True, 
            max_length=self.max_seq_len,
            return_tensors="pt",
        )
        random_mask_input_ids = random_mask_inputs.input_ids[0]
        random_mask_attention_mask = random_mask_inputs.attention_mask[0]
        random_mask_labels = torch.where(
            random_mask_input_ids != self.tokenizer.pad_token_id, random_mask_input_ids, -100
        )
        
        instances = {
            "cond_cate_size":{
                "input_ids": cond_cate_size_ids,
                "attention_mask": cond_cate_size_attention_mask,
                "labels":cond_cate_size_labels,
            },
            "cond_category": {
                "input_ids": cond_cate_ids,
                "attention_mask": cond_cate_attention_mask,
                "labels": cond_cate_labels,
            },
            "cond_mask": {
                "input_ids": random_mask_input_ids,
                "attention_mask": random_mask_attention_mask,
                "labels": random_mask_labels,
            }
        }
        
        ## random sampling
        selected_types = random.choices(self.cond_type, self.sample_prob, k=2)  # joint loss (t=2 here)
        instance_1 = instances[selected_types[0]]
        instance_2 = instances[selected_types[1]]

        return {
            "instance_1": instance_1,
            "instance_2": instance_2,
        }
trainer/src/custom_dataset.py
ProjectNUWA/LayoutNUWA
import sys
import os
import fire
import torch
import transformers
import json
from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
from tqdm import *

def main(
    file_path: str = None,
    base_model: str = "/zecheng/model_hub/Llama-2-7b-multinode/checkpoint-1800",
    device: int=None,
    output_dir: str=None,
    max_new_tokens: int=1024,
):
    assert base_model, (
        "Please specify a --base_model, e.g. --base_model='decapoda-research/llama-7b-hf'"
    )
    
    assert file_path, (
        "Please specify a --file_path, e.g. --file_path='/path/to/json_file'"
    )
    
    assert device is not None, (
        "Please specify a --device, e.g. --device='0'"
    )
    
    assert output_dir, (
        "Please specify a --output_dir, e.g. --output_dir='/path/to/output_dir'"
    )
    
    tokenizer = LlamaTokenizer.from_pretrained("/zecheng/model_hub/Llama-2-7b-hf")
    model = LlamaForCausalLM.from_pretrained(
            base_model,
            load_in_8bit=False,
            torch_dtype=torch.float16,
            device_map={"": device},
        )
    
    model.half()
    model.eval()
    
    # unwind broken decapoda-research config
    model.config.pad_token_id = tokenizer.pad_token_id = 0  # unk
    model.config.bos_token_id = 1
    model.config.eos_token_id = 2

    def evaluate(
        instruction,
        temperature=0.6,
        top_p=0.9,
        top_k=40,
        num_beams=4,
        do_sample=True,
        max_new_tokens=max_new_tokens,
        **kwargs,
    ):
        prompt = instruction
        inputs = tokenizer(prompt, return_tensors="pt")
        input_ids = inputs["input_ids"].to(device)
        generation_config = GenerationConfig(
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            num_beams=num_beams,
            **kwargs,
        )
        # import pdb; pdb.set_trace()
        with torch.no_grad():
            generation_output = model.generate(
                input_ids=input_ids,
                generation_config=generation_config,
                return_dict_in_generate=True,
                output_scores=True,
                max_new_tokens=max_new_tokens,
            )
        s = generation_output.sequences[0]
        output = tokenizer.decode(s)
        return output
    
    with open(file_path, "r") as f:
        content = [json.loads(line) for line in f]
    
    res = []
    with tqdm(total=len(content)) as pbar:
        for samples in content:
            cond_cate_to_size_pos = samples.get("cond_cate_to_size_pos_seq_modeling")[0]
            cond_cate_size_to_pos = samples.get("cond_cate_size_to_pos_seq_modeling")[0]
            cond_recover_mask_input = samples.get("cond_recover_mask_seq_modeling")[0]
            id_ = samples.get("id_")

          
            cond_cate_to_size_pos = evaluate(cond_cate_to_size_pos, max_new_tokens=max_new_tokens)
            cond_cate_size_to_pos = evaluate(cond_cate_size_to_pos, max_new_tokens=max_new_tokens)
            cond_recover_mask = evaluate(cond_recover_mask_input, max_new_tokens=max_new_tokens)
            
            res.append({
                "cond_cate_to_size_pos": cond_cate_to_size_pos,
                "cond_cate_size_to_pos": cond_cate_size_to_pos,
                "cond_recover_mask": cond_recover_mask,
                "id": id_
            })
            pbar.update(1)
        
    output_file = os.path.join(output_dir, file_path.split("/")[-1])
    
    with open(output_file, "w") as f:
        for line in res:
            f.write(json.dumps(line) + "\n")
            
    print("save done !!")
            
        
if __name__ == "__main__":
    fire.Fire(main)
trainer/src/generate.py
ProjectNUWA/LayoutNUWA
import sys
import os
current_directory = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_directory) 

from src.custom_dataset import RawFileDataset
import copy
import random
from dataclasses import dataclass, field
from typing import Optional, Dict, Sequence
import os

import torch
import torch.distributed
import transformers
from transformers import Trainer

IGNORE_INDEX = -100
DEFAULT_PAD_TOKEN = "[PAD]"
DEFAULT_EOS_TOKEN = "</s>"
DEFAULT_BOS_TOKEN = "</s>"
DEFAULT_UNK_TOKEN = "</s>"



@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")


@dataclass
class DataArguments:
    data_path: str = field(default=None, metadata={"help": "Path to the training data."})
    train_file: str = field(default=None, metadata={"help": "train file name"})
    val_file: str = field(default=None, metadata={"help": "val file name"})

@dataclass
class TrainingArguments(transformers.TrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    optim: str = field(default="adamw_torch")
    model_max_length: int = field(
        default=512,
        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
    )


def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
    """Collects the state dict and dump to disk."""
    state_dict = trainer.model.state_dict()
    if trainer.args.should_save:
        cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
        del state_dict
        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa


def smart_tokenizer_and_embedding_resize(
    special_tokens_dict: Dict,
    tokenizer: transformers.PreTrainedTokenizer,
    model: transformers.PreTrainedModel,
):
    """Resize tokenizer and embedding.

    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
    """
    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
    model.resize_token_embeddings(len(tokenizer))

    if num_new_tokens > 0:
        input_embeddings = model.get_input_embeddings().weight.data
        output_embeddings = model.get_output_embeddings().weight.data

        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)

        input_embeddings[-num_new_tokens:] = input_embeddings_avg
        output_embeddings[-num_new_tokens:] = output_embeddings_avg


def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
    """Tokenize a list of strings."""
    tokenized_list = [
        tokenizer(
            text,
            return_tensors="pt",
            padding="longest",
            max_length=tokenizer.model_max_length,
            truncation=True,
        )
        for text in strings
    ]
    input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
    input_ids_lens = labels_lens = [
        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
    ]
    return dict(
        input_ids=input_ids,
        labels=labels,
        input_ids_lens=input_ids_lens,
        labels_lens=labels_lens,
    )


def preprocess(
    sources: Sequence[str],
    targets: Sequence[str],
    tokenizer: transformers.PreTrainedTokenizer,
) -> Dict:
    """Preprocess the data by tokenizing."""
    examples = [s + t for s, t in zip(sources, targets)]
    examples_tokenized, sources_tokenized = [_tokenize_fn(strings, tokenizer) for strings in (examples, sources)]
    input_ids = examples_tokenized["input_ids"]
    labels = copy.deepcopy(input_ids)
    for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
        label[:source_len] = IGNORE_INDEX
    return dict(input_ids=input_ids, labels=labels)


@dataclass
class DataCollatorForSupervisedDataset(object):
    """Collate examples for supervised fine-tuning."""

    tokenizer: transformers.PreTrainedTokenizer

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        ### one can customize here, since we set the T for joint loss as 2
        
        batch_input_ids1, batch_input_ids2 = [], []
        batch_attention_mask1, batch_attention_mask2 = [], []
        batch_labels1, batch_labels2 = [], []

        for instance in instances:
            instance1, instance2 = instance["instance_1"], instance["instance_2"]
            batch_input_ids1.append(instance1["input_ids"])
            batch_input_ids2.append(instance2["input_ids"])
            batch_attention_mask1.append(instance1["attention_mask"])
            batch_attention_mask2.append(instance2["attention_mask"])
            batch_labels1.append(instance1["labels"])
            batch_labels2.append(instance2["labels"])
        
        batch_input_ids1 = torch.stack(batch_input_ids1, dim=0)
        batch_input_ids2 = torch.stack(batch_input_ids2, dim=0)
        batch_attention_mask1 = torch.stack(batch_attention_mask1, dim=0)
        batch_attention_mask2 = torch.stack(batch_attention_mask2, dim=0)
        batch_labels1 = torch.stack(batch_labels1, dim=0)
        batch_labels2 = torch.stack(batch_labels2, dim=0)
        
        return {
            "batch_input_ids1": batch_input_ids1,
            "batch_input_ids2": batch_input_ids2,
            "batch_attention_mask1": batch_attention_mask1,
            "batch_attention_mask2": batch_attention_mask2,
            "batch_labels1": batch_labels1,
            "batch_labels2": batch_labels2,
        }
        

class CustomTrainier(Trainer):
    def __init__(self, model, args, train_dataset, eval_dataset, tokenizer, **kwargs):
        super().__init__(
            model=model, 
            args=args, 
            train_dataset=train_dataset, 
            eval_dataset=eval_dataset, 
            tokenizer=tokenizer,
            **kwargs,
        )
        
    def compute_loss(self, model, inputs, return_outputs=False):
        input_ids1 = inputs.get("batch_input_ids1")
        input_ids2 = inputs.get("batch_input_ids2")
        batch_attention_mask1 = inputs.get("batch_attention_mask1")
        batch_attention_mask2 = inputs.get("batch_attention_mask2")
        batch_labels1 = inputs.get("batch_labels1")
        batch_labels2 = inputs.get("batch_labels2")
        
        outputs1 = model(
            input_ids=input_ids1,
            attention_mask=batch_attention_mask1,
            labels=batch_labels1,
        )
        outputs2 = model(
            input_ids=input_ids2,
            attention_mask=batch_attention_mask2,
            labels=batch_labels2,
        )
        
        outputs = (outputs1, outputs2)
        loss = outputs1.loss + outputs2.loss

        return (loss, outputs) if return_outputs else loss 

              
def train():
    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
        
    model = transformers.AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
    )
    model.config.pad_token_id = 0

    tokenizer = transformers.AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
        model_max_length=training_args.model_max_length,
        padding_side="right",
        use_fast=True,
    )
    if tokenizer.pad_token is None:
        smart_tokenizer_and_embedding_resize(
            special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
            tokenizer=tokenizer,
            model=model,
        )
    tokenizer.pad_token_id = tokenizer.unk_token_id
    if "llama" in model_args.model_name_or_path:
        tokenizer.add_special_tokens(
            {
                "eos_token": DEFAULT_EOS_TOKEN,
                "bos_token": DEFAULT_BOS_TOKEN,
                "unk_token": DEFAULT_UNK_TOKEN,
            }
        )
    
    train_file = os.path.join(data_args.data_path, data_args.train_file)
    val_file = os.path.join(data_args.data_path, data_args.val_file)
    
    train_dataset = RawFileDataset(training_args, train_file, tokenizer)
    val_dataset = RawFileDataset(training_args, val_file, tokenizer)


    if training_args.local_rank == 0:
        print(len(train_dataset))
        for index in random.sample(range(len(train_dataset)), 3):
            print(f"Sample {index} of the training set: {train_dataset[index]}.")
    
    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
    data_module = dict(train_dataset=train_dataset, eval_dataset=val_dataset, data_collator=data_collator)

    model.is_parallelizable = True
    model.model_parallel = True

    trainer = CustomTrainier(model=model, tokenizer=tokenizer, args=training_args, **data_module)
    model.config.use_cache = False

    trainer.train()
    trainer.save_state()
    safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)


if __name__ == "__main__":
    train()

trainer/src/train.py
ProjectNUWA/LayoutNUWA