From c219e554e2b67aa1407cb0657abb5a217f3dd75c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=A1=9C=E8=8F=AF?= <huan0016@gmail.com>
Date: Wed, 21 Jun 2023 18:37:27 +0800
Subject: [PATCH] =?UTF-8?q?ws=E6=9B=B4=E6=96=B0=E8=A7=A6=E5=8F=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 models/Model.go   | 17 +++++++++++------
 routers/images.go | 11 ++++-------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/models/Model.go b/models/Model.go
index b84dcf1..4b02d78 100644
--- a/models/Model.go
+++ b/models/Model.go
@@ -66,7 +66,7 @@ func init() {
 	//}()
 }
 
-func (model *Model) Inference(image_list []Image, callback func()) {
+func (model *Model) Inference(image_list []Image, callback func(Image)) {
 	log.Println(image_list)
 
 	// 模型未部署到推理機
@@ -188,17 +188,22 @@ func (model *Model) Inference(image_list []Image, callback func()) {
 			if err := goreq.Post(url).SetJsonBody(data).Do().BindJSON(&rest); err != nil {
 				log.Println("API 查询失败:", err)
 			}
-			log.Println("API 查询成功:", rest)
-			for _, img := range rest.Images {
-				log.Println("保存图片:", img)
-				// 将base64编码的图片保存到本地webp
+			for index, img := range rest.Images {
 				var filename = fmt.Sprintf("%x", md5.Sum([]byte(img)))
+				log.Println("保存图片:", filename)
 				if err := SaveBase64Image(img, "data/images/"+filename+".webp"); err != nil {
 					log.Println(err)
 				}
+				image_list[index].Name = filename
+				image_list[index].Path = "data/images/" + filename + ".webp"
+				image_list[index].Hash = filename
+				image_list[index].Type = "image/webp"
+				image_list[index].Width = 512
+				image_list[index].Height = 512
+				image_list[index].Format = "webp"
+				callback(image_list[index])
 			}
 		}
-		callback()
 		return
 	}
 	log.Println("模型未部署到推理機, 取消推理模型")
diff --git a/routers/images.go b/routers/images.go
index 2ae1cf4..5dd61bf 100644
--- a/routers/images.go
+++ b/routers/images.go
@@ -154,11 +154,6 @@ func ImagesPost(w http.ResponseWriter, r *http.Request) {
 				return
 			}
 
-			// 推理图像
-			go model.Inference([]models.Image{}, func() {
-				log.Println("推理完成")
-			})
-
 			// 直接创建一组图片
 			var image_list []models.Image
 			var task string = uuid.New().String()
@@ -177,8 +172,10 @@ func ImagesPost(w http.ResponseWriter, r *http.Request) {
 				image_list = append(image_list, image)
 			}
 
-			go model.Inference(image_list, func() {
-				images_websocket_manager.NotifyTaskChange(task, image_list)
+			// 推理图像
+			go model.Inference([]models.Image{}, func(img models.Image) {
+				log.Println("推理完成")
+				images_websocket_manager.NotifyTaskChange(task, img)
 			})
 
 			// 存储图片信息到数据库