我生成了一个（200 x 200 x 3）图像的数据集，其中每个图像包含一个40 X 40的不同颜色的框。使用张流创建一个模型，可以预测这个40 x 40框的坐标。enter image description here
代码我用来生成这些图像：

from PIL import Image, ImageDraw
from random import randrange

colors = ["#ffd615", "#f9ff21", "#00d1ff", 
"#0e153a", "#fc5c9c", "#ac3f21",
"#40514e", "#492540", "#ff8a5c",
"#000000", "#a6fff2", "#f0f696",
"#d72323", "#dee1ec", "#fcb1b1"]

def genrate_image(color):
    img = Image.new(mode="RGB", size=(200, 200), color=color)
    return img

def save_image(img, imgname):
    img.save(imgname)

def draw_rect(image, color, x, y):
    draw = ImageDraw.Draw(image)
    coords = ((x, y), (x+40, y), (x+40, y+40), (x, y+40))
    draw.polygon(coords, fill=color)
    #return image, str(coords)
    return image, coords[0][0], coords[2][0], coords[0][1], coords[2][1]

FILE_NAME = "train_annotations.txt"

for i in range(0, 100):
    img = genrate_image(colors[randrange(0, len(colors))])
    img, x0, x1, y0, y1 = draw_rect(img, colors[randrange(0, len(colors))], randrange(200 - 50), randrange(200 - 50))
    save_image(img, "dataset/train_images/img"+str(i)+".png")
    with open(FILE_NAME, "a+") as f:
        f.write(f"{x0} {x1} {y0} {y1}\n")
        f.close()

有谁能帮我一个忙，建议我怎样才能建立一个模型，可以预测一个新的图像的坐标。