rendering.py

import os
import time
import fitz
import json
import random
import argparse
import datetime
import multiprocessing
from utils.process import *

random.seed(datetime.datetime.now().timestamp())


def render_layout(layouts):
    """
    Render layouts to images and save in the yolo format.

    Args:
        layouts (list): List of generated layouts information.
    """
    doc = fitz.open()
    w, h = sample_hw(
        width_range=[1200, 2000],
        ratio_range=[0.7,1.5],
        max_height=3000,
    )
    page = doc.new_page(width=w, height=h)

    annotation_json = {"bbox": [], "labels": [], "width":w, "height":h}
    for bbox, category, relpath in zip(layouts["boxes"], layouts["categories"], layouts["relpaths"]):
        bbox[0], bbox[2] = bbox[0]*w, bbox[2]*w
        bbox[1], bbox[3] = bbox[1]*h, bbox[3]*h
        rect = fitz.Rect([bbox[i] for i in range(4)])

        abs_filepath = os.path.join('./element_pool',relpath)
        start_str = abs_filepath.rsplit('/',1)[1].split('.')[0]
        sampled_path = random.choice(INSTANCE2PATHLIST[start_str])

        page.insert_image(rect, filename=sampled_path, keep_proportion=False)

        annotation_json["bbox"].append(bbox)
        annotation_json["labels"].append(category)  
                  
    _id = str(time.time()).replace(".", "_")
    pix = page.get_pixmap()
    pix.save(os.path.join(IMAGE_DIR, f"{_id}.jpg"))
    anno_txt = open(os.path.join(ANNO_DIR, f"{_id}.txt"), "w")
    for bbox, category_id in zip(annotation_json["bbox"], annotation_json["labels"]):
        w, h = annotation_json["width"], annotation_json["height"]
        x0, y0, x1, y1 = bbox
        x0, y0 = x0/w, y0/h
        x1, y1 = x1/w, y1/h
        anno_txt.write(f"{category_id} {x0} {y0} {x1} {y0} {x1} {y1} {x0} {y1}\n")
    anno_txt.close()


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--save_dir', default="./generated_dataset", type=str, help='planned root dir for generated dataset')
    parser.add_argument('--n_jobs', default=None, required=True, type=int, help='number of processes to use in multiprocessing')
    parser.add_argument('--json_path', default=None, required=True, type=str, help='json path for layouts generated by the Mesh-candidate Bestfit alogorithm')
    parser.add_argument('--map_dict_path', default=None, required=True, type=str, help='json path for element to pathlist map dict')
    args = parser.parse_args()
    
    # Setting save path
    IMAGE_DIR = os.path.join(args.save_dir, "images")
    ANNO_DIR = os.path.join(args.save_dir, "labels")
    os.makedirs(IMAGE_DIR, exist_ok=True)
    os.makedirs(ANNO_DIR, exist_ok=True)
    
    # Load layout data
    layout_json = json.load(open(args.json_path))
    INSTANCE2PATHLIST = json.load(open(args.map_dict_path))
    
    # Using multiprocessing to accelerate rendering        
    n_jobs = args.n_jobs
    with multiprocessing.Pool(n_jobs) as p:
        result = p.starmap(
            render_layout, [(layout,) for layout in layout_json]
        )
    p.close()
    p.join()