AIGC草稿转真实 | 征稿活动V6

引言：本文为参赛文章，作者「last forever」。他将介绍通过 AI 识别草稿内容，最后得到了一个可以生成真实组件和布局的代码。以下为开发实录。

前言

当创作灵感来的时候，我们可能会先把灵感记录在草稿上，之后再去实现它。比方说有一天，我突然来了游戏创作灵感，想着那可以先把一些简单的组件和布局设计出来，于是就在草稿上画了几个框。L 代表 Label 组件，B 代表 Button 组件，S 代表 Sprite 组件。

(资料图片)

几天过去了，就算当时的灵感再好，我也不想打开电脑，所以草稿还是只是草稿。我想着，如果有一个 AI 能够识别我画的草稿，然后自动生成对应组件以及布局的话该有多好啊。

于是，我决定训练一个 AI，准确来说是一个图像目标检测 AI 模型，我将用它来识别并定位我草稿上的各个方框。

训练AI

首先第一步，制作数据集，花了我半天时间画了很多大大小小的方框，草稿数量为 150 张。有在便利贴上画的，有在纸巾上画的，还有在白纸上画的。

用于训练的图片已经有了，接下来又花了半天时间去进行标注，不得不说想要训练 AI，首先就要训练自己的毅力。

训练集已经做好，接下来就是丢给 yolov5 训练，训练结果如下。

验证之后，识别效果还是可以的，现在 AI 可以大概率将我画的 SLB 三个方框识别出来。

注：由于时间有限，我就先实现三种组件：Sprite，Label 以及 Button，如果加入更多组件的话就需要更多数据。

从识别结果中提取必要数据

AI 在识别之后可以生成一个 txt 文件，其中保存着各个矩形的数据，如下所示。

10.6829940.4527030.1277430.098119920.6837770.4521150.1230410.099294920.5626960.651880.120690.10399520.57210.3437130.1253920.096357200.2907520.3416570.1112850.088719210.3021160.5437720.09796240.085193910.3205330.739130.1112850.0951821

第一列中的数字 0，1，2 分别代表 Button 方框，Sprite 方框和 Label 方框；第二列是矩形中心的 x 坐标（归一化后的，下同）；第二列是矩形中心的 y 坐标；第三列是矩形宽度；第四列是矩形高度。通过这些数据我们能够得出各个矩形方框在图片上的真实坐标以及大小。代码编写如下：

defget_rect_list_from_detect_result(pic_path,detect_result_path):\"\"\"从预测结果中得出各个矩形的真实大小和位置数据\"\"\"withopen(detect_result_path,\"r\",encoding=\"utf-8\")asf:result=f.readlines()pic_width,pic_height=Image.open(pic_path).sizeclass_dict={\"0\":\"B\",\"1\":\"S\",\"2\":\"L\"}rect_list=[]forlineinresult:_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/2rect_y=round(float(y_center_norm)*pic_height,2)-rect_height/2rect_list.append({\"type\":class_dict[_class],\"width\":rect_width,\"height\":rect_height,\"x\":rect_x,\"y\":rect_y})returnrect_list

为了在 Cocos 场景的 Canvas 画布上生成组件，我们需要先确定画布大小和方向。初始大小我们可以用 960 和 640 这两个值，但是方向呢？是采用横向还是竖向？我这里采用的解决方案是：找出草稿上最左边和最右边的矩形，然后得出横向距离，接着找出最上边和最下边的矩形并得出纵向距离。比较横向距离和纵向距离，假如前者大，则画布为横向，反之则为竖向。代码编写如下：

defdecide_canvas_size(rect_list):\"\"\"通过各个矩形的最小和最大x/y值来确定画布尺寸\"\"\"#获取所有矩形中的最小x值min_x=min(rect_list,key=lambdarect:rect[\"x\"])[\"x\"]max_x=max(rect_list,key=lambdarect:rect[\"x\"])[\"x\"]min_y=min(rect_list,key=lambdarect:rect[\"y\"])[\"y\"]max_y=max(rect_list,key=lambdarect:rect[\"y\"])[\"y\"]#根据x和y的距离差判断是要横屏还是竖屏distance_x=max_x-min_xdistance_y=max_y-min_yifdistance_x>=distance_y:canvas_width=960canvas_height=640else:canvas_width=640canvas_height=960width_prop=distance_x/canvas_widthifcanvas_width>distance_xelsecanvas_width/distance_xheight_prop=distance_y/canvas_heightifcanvas_height>distance_yelsecanvas_height/distance_yreturncanvas_width,canvas_height,width_prop,height_prop

AI 在识别时，是以图片左上角为原点的，往右为 x 轴正方向，往下为 y 轴正方向。

但是 Cocos Creator 的坐标原点在画布中心，所以我们肯定要将得到的矩形坐标进行转换，否则生成的组件将会都排布在画布右上角。

另外在草稿上画方框时，我们可能是想着让几个矩形对齐，还有几个矩形的大小应该是一样的。但是在画的时候并不会太精准，所以我们还应该将位置和大小相近的矩形进行调整，让它们的 x 或 y 值或宽高相等。代码编写如下：

注：关于其他方面的调整，我就不再赘述了，读者可以看下注释。

defadjust_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=lambdarect:rect[\"x\"])[\"x\"]min_y=min(rect_list,key=lambdarect:rect[\"y\"])[\"y\"]forrectinrect_list:rect[\"x\"]=(rect[\"x\"]-min_x)*width_prop-canvas_width/2rect[\"y\"]=canvas_height/2-(rect[\"y\"]-min_y)*height_proprect[\"width\"]*=width_proprect[\"height\"]*=height_prop#算出下边和右边的空白距离，将所有矩形往下和往右平移空白距离/2个像素点max_x=max(rect_list,key=lambdarect:rect[\"x\"])[\"x\"]min_y=min(rect_list,key=lambdarect:rect[\"y\"])[\"y\"]right_distance=(canvas_width/2-max_x)/2bottom_distance=abs(-canvas_height/2-min_y)/2forrectinrect_list:rect[\"x\"]+=right_distancerect[\"y\"]-=bottom_distance#将x或y坐标距离不相差15像素的矩形对齐diff=15forrect1inrect_list:forrect2inrect_list:ifrect1==rect2:continueifabs(rect1[\"x\"]-rect2[\"x\"])<=diff:average_x=(rect1[\"x\"]+rect2[\"x\"])/2rect1[\"x\"]=average_xrect2[\"x\"]=average_xifabs(rect1[\"y\"]-rect2[\"y\"])<=diff:average_y=(rect1[\"y\"]+rect2[\"y\"])/2rect1[\"x\"]=average_yrect2[\"x\"]=average_yifabs(rect1[\"width\"]-rect2[\"width\"])<=diff:average_width=(rect1[\"width\"]+rect2[\"width\"])/2rect1[\"width\"]=average_widthrect2[\"width\"]=average_widthifabs(rect1[\"height\"]-rect2[\"height\"])<=diff:average_height=(rect1[\"height\"]+rect2[\"height\"])/2rect1[\"height\"]=average_heightrect2[\"height\"]=average_height#四舍五入保留整数forrectinrect_list:rect[\"x\"]=round(rect[\"x\"])rect[\"y\"]=round(rect[\"y\"])rect[\"width\"]=round(rect[\"width\"])rect[\"height\"]=round(rect[\"height\"])returnrect_list

数据调整完毕后，我们就可以把画布数据和各个矩形数据一同返回。我这里用 flask 框架开发了一个后端，完整代码如下：

importosimportuuidfromPILimportImagefrompathlibimportPathfromflaskimportFlask,requestapp=Flask(__name__)app.config[\"UPLOAD_FOLDER\"]=str(Path(__file__).parent/\"upload\")app.config[\"SECRET_KEY\"]=\"SECRET_KEY\"@app.route(\"/d2r\",methods=[\"POST\"])defdraft_to_reality():\"\"\"将草稿转成真实布局所需要的数据格式\"\"\"#从前端获取图片file=request.files.get(\"file\")ifnotfile.filename.endswith(\".png\")andnotfile.filename.endswith(\".jpg\")andnotfile.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)ifnotis_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}defdetect(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\"ifPath(detect_result_path).exists():returnTrue,detect_result_pathelse:returnFalse,Nonedefget_rect_list_from_detect_result(pic_path,detect_result_path):\"\"\"从预测结果中得出各个矩形的真实大小和位置数据\"\"\"withopen(detect_result_path,\"r\",encoding=\"utf-8\")asf:result=f.readlines()pic_width,pic_height=Image.open(pic_path).sizeclass_dict={\"0\":\"B\",\"1\":\"S\",\"2\":\"L\"}rect_list=[]forlineinresult:_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/2rect_y=round(float(y_center_norm)*pic_height,2)-rect_height/2rect_list.append({\"type\":class_dict[_class],\"width\":rect_width,\"height\":rect_height,\"x\":rect_x,\"y\":rect_y})returnrect_listdefdecide_canvas_size(rect_list):\"\"\"通过各个矩形的最小和最大x/y值来确定画布尺寸\"\"\"#获取所有矩形中的最小x值min_x=min(rect_list,key=lambdarect:rect[\"x\"])[\"x\"]max_x=max(rect_list,key=lambdarect:rect[\"x\"])[\"x\"]min_y=min(rect_list,key=lambdarect:rect[\"y\"])[\"y\"]max_y=max(rect_list,key=lambdarect:rect[\"y\"])[\"y\"]#根据x和y的距离差判断是要横屏还是竖屏distance_x=max_x-min_xdistance_y=max_y-min_yifdistance_x>=distance_y:canvas_width=960canvas_height=640else:canvas_width=640canvas_height=960width_prop=distance_x/canvas_widthifcanvas_width>distance_xelsecanvas_width/distance_xheight_prop=distance_y/canvas_heightifcanvas_height>distance_yelsecanvas_height/distance_yreturncanvas_width,canvas_height,width_prop,height_propdefadjust_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=lambdarect:rect[\"x\"])[\"x\"]min_y=min(rect_list,key=lambdarect:rect[\"y\"])[\"y\"]forrectinrect_list:rect[\"x\"]=(rect[\"x\"]-min_x)*width_prop-canvas_width/2rect[\"y\"]=canvas_height/2-(rect[\"y\"]-min_y)*height_proprect[\"width\"]*=width_proprect[\"height\"]*=height_prop#算出下边和右边的空白距离，将所有矩形往下和往右平移空白距离/2个像素点max_x=max(rect_list,key=lambdarect:rect[\"x\"])[\"x\"]min_y=min(rect_list,key=lambdarect:rect[\"y\"])[\"y\"]right_distance=(canvas_width/2-max_x)/2bottom_distance=abs(-canvas_height/2-min_y)/2forrectinrect_list:rect[\"x\"]+=right_distancerect[\"y\"]-=bottom_distance#将x或y坐标距离不相差15像素的矩形对齐diff=15forrect1inrect_list:forrect2inrect_list:ifrect1==rect2:continueifabs(rect1[\"x\"]-rect2[\"x\"])<=diff:average_x=(rect1[\"x\"]+rect2[\"x\"])/2rect1[\"x\"]=average_xrect2[\"x\"]=average_xifabs(rect1[\"y\"]-rect2[\"y\"])<=diff:average_y=(rect1[\"y\"]+rect2[\"y\"])/2rect1[\"x\"]=average_yrect2[\"x\"]=average_yifabs(rect1[\"width\"]-rect2[\"width\"])<=diff:average_width=(rect1[\"width\"]+rect2[\"width\"])/2rect1[\"width\"]=average_widthrect2[\"width\"]=average_widthifabs(rect1[\"height\"]-rect2[\"height\"])<=diff:average_height=(rect1[\"height\"]+rect2[\"height\"])/2rect1[\"height\"]=average_heightrect2[\"height\"]=average_height#四舍五入保留整数forrectinrect_list:rect[\"x\"]=round(rect[\"x\"])rect[\"y\"]=round(rect[\"y\"])rect[\"width\"]=round(rect[\"width\"])rect[\"height\"]=round(rect[\"height\"])returnrect_listdefmake_final_data(rect_list,canvas_width,canvas_height):return{\"canvas\":{\"width\":canvas_width,\"height\":canvas_height},\"rects\":rect_list}if__name__==\"__main__\":app.run()

发送图片并生成最终代码

后端代码已经开发完毕，前端功能的话我们就用 Cocos 插件来实现，插件界面如下。

选择一张草稿图片，然后点击生成。

此时后端就会识别上传的图片，然后返回画布和矩形数据。

{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}]}

最后插件会根据该数据在 assets 中生成一个 d2r.ts 文件，我们将该 ts 文件挂载到画布上然后运行就可以看到效果了。

注：resources 文件夹是插件在开启时自动生成的，其中包含着用于 Sprite 和 Button 组件的初始图片。

运行效果如下：

运行视频：

总结与提高

我们通过 AI 将草稿内容识别了出来，然后提取并调整了必要数据，最后得到了一个可以生成真实组件和布局的代码。

有待提高的点：

1. 用于 AI 的训练数据量可以再大一些，这样可以提高识别精度，也可以生成更多种类的组件。

2. 个别草稿识别出来的结果不是很满意，有些组件在草稿上是分开的，但生成后却是重合的，还有方框的大小比例还需要更精确。所以在对提取的数据做调整时，调整方面的算法还有待改进。

3. 目前 AI 的作用就在识别，并不能很好的体现 AIGC 的魅力。可以加入自然语言处理(NLP)技术，获取用户意向，然后再用爬虫爬取相应的图片资源，生成一个更好看更完善的内容。

资源下载

后端代码下载地址，需要安装的库已经写在 yolov5/requirements.txt 文件中。因为包含 AI 模型以及训练图片，所以文件会比较大：

下载链接：提取码【cr9t】

https://pan.baidu.com/s/1Z-q2mc2jsX5h_fWD_QjaOA

前端插件下载地址，Cocos Creator 版本要>=3.6.0：提取码【9e6t】

https://pan.baidu.com/s/141gpeSjGunKMf9SlqY0H7Q

点击文末【阅读原文】前往原文查看。

往期精彩