Table of contents

Extract necessary data from recognition results

foreword

When creative inspiration comes, we may first record the inspiration on the draft, and then realize it later. For example, one day, I suddenly came up with the inspiration for game creation, thinking that I could design some simple components and layout first, so I drew a few frames on the draft.

Note: L stands for Label component, B stands for Button component, and S stands for Sprite component.

A few days passed and I didn't want to turn on my computer even if the inspiration was good at the time, so the draft was still just a draft. I thought, if there is an AI that can recognize the draft I drew, and then automatically generate the corresponding components and layout, it would be great.

So, I decided to train an AI, to be precise, an image object detection AI model, which I will use to recognize and locate the boxes on my sketch.

train AI

First of all, the first step is to make a data set. It took me half a day to draw a lot of large and small boxes, and the number of drafts is 150. Some are drawn on post-it notes, some are drawn on paper towels, and some are drawn on white paper.

There are already pictures for training, and it took half a day to label them. I have to say that if you want to train AI, you must first train your perseverance.

The training set has been prepared, and the next step is to throw it to yolov5 for training. The training results are as follows.

After the verification, the recognition effect is still good, and now the AI can recognize the three SLB boxes I drew with a high probability.

Note: Due to limited time, I will implement three components first: Sprite, Label and Button. If more components are added, more data will be needed.

Extract necessary data from recognition results

AI can generate a txt file after recognition, which saves the data of each rectangle, as shown below.

1 0.682994 0.452703 0.127743 0.0981199
2 0.683777 0.452115 0.123041 0.0992949
2 0.562696 0.65188 0.12069 0.103995
2 0.5721 0.343713 0.125392 0.0963572
0 0.290752 0.341657 0.111285 0.0887192
1 0.302116 0.543772 0.0979624 0.0851939
1 0.320533 0.73913 0.111285 0.0951821

The numbers 0, 1, and 2 in the first column represent the Button box, Sprite box and Label box respectively; the second column is the x coordinate of the center of the rectangle (after normalization, the same below); the second column is the rectangle The y-coordinate of the center; the third column is the rectangle width; the fourth column is the rectangle height. Through these data, we can get the real coordinates and sizes of each rectangular box on the picture. The code is written as follows:

def get_rect_list_from_detect_result(pic_path, detect_result_path):
    """从预测结果中得出各个矩形的真实大小和位置数据"""
    with open(detect_result_path, "r", encoding="utf-8") as f:
        result = f.readlines()

    pic_width, pic_height = Image.open(pic_path).size
    class_dict = {"0": "B", "1": "S", "2": "L"}
    rect_list = []

    for line in result:
        _class, x_center_norm, y_center_norm, width_norm, height_norm = line.strip().split(" ")
        rect_width = round(float(width_norm)*pic_width, 2)
        rect_height = round(float(height_norm)*pic_height, 2)
        rect_x = round(float(x_center_norm)*pic_width, 2) - rect_width/2
        rect_y = round(float(y_center_norm)*pic_height, 2) - rect_height/2
        rect_list.append({
            "type": class_dict[_class],
            "width": rect_width,
            "height": rect_height,
            "x": rect_x,
            "y": rect_y
        })
    
    return rect_list

In order to generate components on the Canvas of the Cocos scene, we need to determine the size and orientation of the canvas. We can use 960 and 640 for the initial size, but what about the orientation? Is it horizontal or vertical? The solution I adopt here is: Find the leftmost and rightmost rectangles on the draft, and then get the horizontal distance, then find the top and bottom rectangles and get the vertical distance. Compare the horizontal distance and the vertical distance, if the former is larger, the canvas is horizontal, otherwise it is vertical. The code is written as follows:

def decide_canvas_size(rect_list):
    """通过各个矩形的最小和最大x/y值来确定画布尺寸"""
    
    # 获取所有矩形中的最小x值
    min_x = min(rect_list, key=lambda rect: rect["x"])["x"]
    max_x = max(rect_list, key=lambda rect: rect["x"])["x"]
    min_y = min(rect_list, key=lambda rect: rect["y"])["y"]
    max_y = max(rect_list, key=lambda rect: rect["y"])["y"]

    # 根据x和y的距离差判断是要横屏还是竖屏
    distance_x = max_x - min_x
    distance_y = max_y - min_y

    if distance_x >= distance_y:
        canvas_width = 960
        canvas_height = 640
    else:
        canvas_width = 640
        canvas_height = 960

    width_prop = distance_x / canvas_width if canvas_width > distance_x else canvas_width / distance_x
    height_prop = distance_y / canvas_height if canvas_height > distance_y else canvas_height / distance_y

    return canvas_width, canvas_height, width_prop, height_prop

When AI recognizes, it takes the upper left corner of the picture as the origin, the right direction is the positive direction of the x-axis, and the downward direction is the positive direction of the y-axis.

But the coordinate origin of Cocos Creator is in the center of the canvas, so we must convert the obtained rectangle coordinates, otherwise the generated components will be arranged in the upper right corner of the canvas.

In addition, when drawing a box on the draft, we may want to align several rectangles, and the sizes of several rectangles should be the same. But it is not too accurate when drawing, so we should also adjust the rectangles with similar positions and sizes to make their x or y values or width and height equal. The code is written as follows:

Note: Regarding other adjustments, I will not go into details, readers can read the notes.

def adjust_rect_data(rect_list, canvas_width, canvas_height, width_prop, height_prop):
    """调整各个矩形的值"""
    # 找到最小x和最小y值（也就是找到最左上角的矩形的x和y值）
    # 然后其他将其他矩形的x和y坐标减去最小x和最小y，求出相对距离
    # 所有相对距离包括宽高全部乘以宽高比例
    # 同时将坐标转换为Cocos类型，以画布中心为原点
    min_x = min(rect_list, key=lambda rect: rect["x"])["x"]
    min_y = min(rect_list, key=lambda rect: rect["y"])["y"]

    for rect in rect_list:
        rect["x"] = (rect["x"] - min_x) * width_prop - canvas_width/2
        rect["y"] = canvas_height/2 - (rect["y"] - min_y) * height_prop
        rect["width"] *= width_prop
        rect["height"] *= height_prop

    # 算出下边和右边的空白距离，将所有矩形往下和往右平移空白距离/2个像素点
    max_x = max(rect_list, key=lambda rect: rect["x"])["x"]
    min_y = min(rect_list, key=lambda rect: rect["y"])["y"]
    right_distance = (canvas_width/2 - max_x) / 2
    bottom_distance = abs(-canvas_height/2 - min_y) / 2

    for rect in rect_list:
        rect["x"] += right_distance
        rect["y"] -= bottom_distance

    # 将x或y坐标距离不相差15像素的矩形对齐
    diff = 15
    for rect1 in rect_list:
        for rect2 in rect_list:
            if rect1 == rect2:
                continue
                
            if abs(rect1["x"] - rect2["x"]) <= diff:
                average_x = (rect1["x"] + rect2["x"]) / 2
                rect1["x"] = average_x
                rect2["x"] = average_x
            
            if abs(rect1["y"] - rect2["y"]) <= diff:
                average_y = (rect1["y"] + rect2["y"]) / 2
                rect1["x"] = average_y
                rect2["x"] = average_y

            if abs(rect1["width"] - rect2["width"]) <= diff:
                average_width = (rect1["width"] + rect2["width"]) / 2
                rect1["width"] = average_width
                rect2["width"] = average_width
            
            if abs(rect1["height"] - rect2["height"]) <= diff:
                average_height= (rect1["height"] + rect2["height"]) / 2
                rect1["height"] = average_height
                rect2["height"] = average_height

    # 四舍五入保留整数
    for rect in rect_list:
        rect["x"] = round(rect["x"])
        rect["y"] = round(rect["y"])
        rect["width"] = round(rect["width"])
        rect["height"] = round(rect["height"])
    
    return rect_list

After the data is adjusted, we can return the canvas data and each rectangle data together. I have developed a backend with the flask framework here, and the complete code is as follows:

import os
import uuid
from PIL import Image
from pathlib import Path
from flask import Flask, request


app = Flask(__name__)
app.config["UPLOAD_FOLDER"] = str(Path(__file__).parent / "upload")
app.config["SECRET_KEY"] = "SECRET_KEY"


@app.route("/d2r", methods=["POST"])
def draft_to_reality():
    """将草稿转成真实布局所需要的数据格式"""

    # 从前端获取图片
    file = request.files.get("file")
    if not file.filename.endswith(".png") and not file.filename.endswith(".jpg") and not file.filename.endswith("jpeg"):
        return {
            "code": "1",
            "message": "图片格式错误"
        }
    
    # 保存图片
    pic_path = Path(app.config["UPLOAD_FOLDER"]) / f"{uuid.uuid4()}.jpg"
    file.save(pic_path)

    # 目标识别
    is_ok, detect_result_path = detect(pic_path)
    if not is_ok:
        return {
             "code": "2",
             "message": "图片识别失败"
        }
    
    # 制作数据
    rect_list = get_rect_list_from_detect_result(pic_path, detect_result_path)
    canvas_width, canvas_height, width_prop, height_prop = decide_canvas_size(rect_list)
    rect_list = adjust_rect_data(rect_list, canvas_width, canvas_height, width_prop, height_prop)
    final_data = make_final_data(rect_list, canvas_width, canvas_height)
    
    return {
        "code": "0",
        "message": final_data
    }


def detect(pic_path):
    os.system(f"python ./yolov5/detect.py --weights ./yolov5/best.pt --source {pic_path} --save-txt --exist-ok")
    
    # 如果识别成功，则会生成一个txt文件
    detect_result_path = f"./yolov5/runs/detect/exp/labels/{Path(pic_path).name.split('.')[0]}.txt"
    if Path(detect_result_path).exists():
        return True, detect_result_path
    else:
        return False, None
    

def get_rect_list_from_detect_result(pic_path, detect_result_path):
    """从预测结果中得出各个矩形的真实大小和位置数据"""
    with open(detect_result_path, "r", encoding="utf-8") as f:
        result = f.readlines()

    pic_width, pic_height = Image.open(pic_path).size
    class_dict = {"0": "B", "1": "S", "2": "L"}
    rect_list = []

    for line in result:
        _class, x_center_norm, y_center_norm, width_norm, height_norm = line.strip().split(" ")
        rect_width = round(float(width_norm)*pic_width, 2)
        rect_height = round(float(height_norm)*pic_height, 2)
        rect_x = round(float(x_center_norm)*pic_width, 2) - rect_width/2
        rect_y = round(float(y_center_norm)*pic_height, 2) - rect_height/2
        rect_list.append({
            "type": class_dict[_class],
            "width": rect_width,
            "height": rect_height,
            "x": rect_x,
            "y": rect_y
        })
    
    return rect_list


def decide_canvas_size(rect_list):
    """通过各个矩形的最小和最大x/y值来确定画布尺寸"""
    
    # 获取所有矩形中的最小x值
    min_x = min(rect_list, key=lambda rect: rect["x"])["x"]
    max_x = max(rect_list, key=lambda rect: rect["x"])["x"]
    min_y = min(rect_list, key=lambda rect: rect["y"])["y"]
    max_y = max(rect_list, key=lambda rect: rect["y"])["y"]

    # 根据x和y的距离差判断是要横屏还是竖屏
    distance_x = max_x - min_x
    distance_y = max_y - min_y

    if distance_x >= distance_y:
        canvas_width = 960
        canvas_height = 640
    else:
        canvas_width = 640
        canvas_height = 960

    width_prop = distance_x / canvas_width if canvas_width > distance_x else canvas_width / distance_x
    height_prop = distance_y / canvas_height if canvas_height > distance_y else canvas_height / distance_y

    return canvas_width, canvas_height, width_prop, height_prop


def adjust_rect_data(rect_list, canvas_width, canvas_height, width_prop, height_prop):
    """调整各个矩形的值"""
    # 找到最小x和最小y值（也就是找到最左上角的矩形的x和y值）
    # 然后其他将其他矩形的x和y坐标减去最小x和最小y，求出相对距离
    # 所有相对距离包括宽高全部乘以宽高比例
    # 同时将坐标转换为Cocos类型，以画布中心为原点
    min_x = min(rect_list, key=lambda rect: rect["x"])["x"]
    min_y = min(rect_list, key=lambda rect: rect["y"])["y"]

    for rect in rect_list:
        rect["x"] = (rect["x"] - min_x) * width_prop - canvas_width/2
        rect["y"] = canvas_height/2 - (rect["y"] - min_y) * height_prop
        rect["width"] *= width_prop
        rect["height"] *= height_prop

    # 算出下边和右边的空白距离，将所有矩形往下和往右平移空白距离/2个像素点
    max_x = max(rect_list, key=lambda rect: rect["x"])["x"]
    min_y = min(rect_list, key=lambda rect: rect["y"])["y"]
    right_distance = (canvas_width/2 - max_x) / 2
    bottom_distance = abs(-canvas_height/2 - min_y) / 2

    for rect in rect_list:
        rect["x"] += right_distance
        rect["y"] -= bottom_distance

    # 将x或y坐标距离不相差15像素的矩形对齐
    diff = 15
    for rect1 in rect_list:
        for rect2 in rect_list:
            if rect1 == rect2:
                continue
                
            if abs(rect1["x"] - rect2["x"]) <= diff:
                average_x = (rect1["x"] + rect2["x"]) / 2
                rect1["x"] = average_x
                rect2["x"] = average_x
            
            if abs(rect1["y"] - rect2["y"]) <= diff:
                average_y = (rect1["y"] + rect2["y"]) / 2
                rect1["x"] = average_y
                rect2["x"] = average_y

            if abs(rect1["width"] - rect2["width"]) <= diff:
                average_width = (rect1["width"] + rect2["width"]) / 2
                rect1["width"] = average_width
                rect2["width"] = average_width
            
            if abs(rect1["height"] - rect2["height"]) <= diff:
                average_height= (rect1["height"] + rect2["height"]) / 2
                rect1["height"] = average_height
                rect2["height"] = average_height

    # 四舍五入保留整数
    for rect in rect_list:
        rect["x"] = round(rect["x"])
        rect["y"] = round(rect["y"])
        rect["width"] = round(rect["width"])
        rect["height"] = round(rect["height"])
    
    return rect_list


def make_final_data(rect_list, canvas_width, canvas_height):
    return {
        "canvas": {
            "width": canvas_width,
            "height": canvas_height
        },
        "rects": rect_list
    }



if __name__ == "__main__":
    app.run()

Send pictures and generate final code

The back-end code has been developed. For the front-end function, we use the Cocos plug-in to realize it. The plug-in interface is as follows.

Choose a draft image and click Generate.

At this point, the backend will recognize the uploaded image and return the canvas and rectangle data.

{ 
  canvas: { height: 960, width: 640 }, 
  rects: [ 
    { height: 93, type: 'S', width: 122, x: 215, y: 128 }, 
    { height: 208, type: 'B', width: 241, x: 193, y: -165 }, 
    { height: 119, type: 'S', width: 148, x: -171, y: -56 }, 
    { height: 119, type: 'L', width: 148, x: -215, y: 165 } 
  ] 
}

Finally, the plug-in will generate a d2r.ts file in assets according to the data. We mount the ts file on the canvas and run it to see the effect.

Note: The resources folder is automatically generated when the plugin is started, which contains the initial images for the Sprite and Button components.

The running effect is as follows:

Summary and Improvement

We identified the draft content through AI, then extracted and adjusted the necessary data, and finally got a code that can generate real components and layouts.

Points to be improved:

1. The amount of training data used for AI can be larger, which can improve the recognition accuracy and generate more types of components.

2. The recognition results of individual drafts are not very satisfactory. Some components are separated on the draft, but they overlap after generation, and the size ratio of the box needs to be more precise. Therefore, when adjusting the extracted data, the adjustment algorithm needs to be improved.

3. At present, the role of AI is to identify, which cannot well reflect the charm of AIGC. You can add natural language processing (NLP) technology to obtain user intentions, and then use crawlers to crawl corresponding image resources to generate better-looking and more complete content.

Download

Back-end code download address, the library that needs to be installed has been written in the yolov5/requirements.txt file. Because it contains AI models and training pictures, the file will be relatively large:

Link: https://pan.baidu.com/s/1Z-q2mc2jsX5h_fWD_QjaOA
Extraction code: cr9t

Front-end plug-in download address, Cocos Creator version >=3.6.0:

Link: https://pan.baidu.com/s/141gpeSjGunKMf9SlqY0H7Q
Extraction code: 9e6t

"Cocos Creator Game Combat" AIGC turns draft content into real content

foreword

train AI

Extract necessary data from recognition results

Send pictures and generate final code

Summary and Improvement

Download

Guess you like