feat(model): 添加检测结果置信度分数支持、修改Dockerfile以支持代码持久化部署

- 在MaskInfo模型中添加score字段用于存储检测置信度
- 修改YOLO检测逻辑以提取和传递预测分数
- 更新坐标数据结构以包含置信度信息
- 调整数据处理流程以正确传输分数数据
- 修改Dockerfile以支持代码持久化部署
- 更新README文档说明代码持久化配置方式

Dockerfile

@@ -15,15 +15,15 @@ WORKDIR /code
 # 复制并安装 Python 依赖
 COPY requirements.txt /code/requirements.txt
 
-# 安装 Python 依赖（加速源）
+# 安装 Python 依赖
 RUN pip install --no-cache-dir --upgrade -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple \
     && pip install --no-cache-dir --upgrade torch torchvision --index-url https://download.pytorch.org/whl/cu130
 
-# 复制应用代码
-COPY ./app /code/app
+# 复制应用代码（正式上线后取消注释）
+# COPY ./app /code/app
 
-# 删除无用的文件，避免占用磁盘空间
-RUN rm -rf /code/app/core
+# 删除核心文件，减小体积（正式上线后取消注释）
+# RUN rm -rf /code/app/core
 
 # 暴露端口并启动应用
 CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "80"]

README.md

@@ -1,56 +1,30 @@
 ### Wall Docker 镜像使用教程
 
----
+#### 构建方式
 
-> 采用模型文件持久化，方便后续更新模型而不需要重新创建容器，以及统一配置配置文件
+```bash
+docker build -t wall .
+```
 
-1. 导入docker images `docker load -i wall.tar`
+#### 启动方式
 
-2. 进入持久化目录，新建.env，添加相应内容
-
-  ```env
-  UPLOAD_DIR=uploads
-  MOCK=false
-  MODEL=segformer #segformer, yolo，目前打包模型只有segformer
-  ```
-
-3. 解压算法模型目录到core文件夹
-
-  ```bash
-  tar -xvf core.tar
-  ```
-
-4. 使用指令运行docker镜像
-
-  ```bash
-  sudo docker run -d \
-      --name [docker_container_name] \
-      --gpus all \
-      -p [local_port]:80 \
-      -v $(pwd)/uploads:/code/uploads \
-      -v $(pwd)/core:/code/app/core \
-      -v $(pwd)/.env:/code/.env \
-      wall
-  ```
-
-5. 如果后续需要更新模型，只需要覆盖掉core内的文件，更改.env配置文件后，即可继续运行
-
----
-
-> 如果不想要模型文件持久化，则不需要解压算法文件了
-
-1. 导入docker images `docker load -i wall.tar`
-
-2. 使用指令运行docker镜像
-
-	```bash
-	sudo docker run -d \
-	    --name [docker_container_name] \
-	    --gpus all \
-	    -p [local_port]:80 \
-	    -v $(pwd)/core:/code/app/core \
-	    -v $(pwd)/.env:/code/.env \
-	    wall
-	```
+> 采用代码持久化，方便后续更新代码及模型而不需要重新创建容器，以及统一配置配置文件
+> 
+> 后续更新，如果只更新了核心文件，则仅需git pull后重新启动容器
+> 
+> 如果更新了requirements.txt，则需要重新构建镜像
 
+```bash
+sudo docker run -d \
+  --name [docker_container_name] \
+  --gpus all \
+  -p [local_port]:80 \
+  -v $(pwd)/uploads:/code/uploads \
+  -v $(pwd)/app:/code/app \
+  -v $(pwd)/.env:/code/.env \
+  wall
+```
 
+> TIPS：由于部分文件采用GIT LFS 管理，请先安装GIT LFS
+> 
+> 在clone或者pull时，建议先clone代码文件，然后停掉，再用git lfs pull可查看大文件下载进度

app/core/yolo_detect/detect.py

@@ -98,6 +98,7 @@ class YOLODetect(YOLO):
                 predicted_class = self.class_names[int(c)]
                 if predicted_class != "wall":
                     box = top_boxes[i]
+                    score = top_conf[i]
                     top, left, bottom, right = box
                     top = max(0, np.floor(top).astype('int32'))
                     left = max(0, np.floor(left).astype('int32'))
@@ -126,8 +127,13 @@ class YOLODetect(YOLO):
                     if keep:
                         color = self.colors[int(c)]
                         mask[top:bottom, left:right] = color
-                        coords.append((self.classes.get(predicted_class),
-                                       [(int(left), int(top)), (int(right), int(top)), (int(right), int(bottom)), (int(left), int(bottom))]))
+                        coords.append(
+                            (
+                                self.classes.get(predicted_class),
+                                float(score),
+                                [(int(left), int(top)), (int(right), int(top)), (int(right), int(bottom)), (int(left), int(bottom))]
+                            )
+                        )
             mask = cv2.cvtColor(mask, cv2.COLOR_RGB2BGR)
             # print("coords:", coords)
 

app/core/yolo_detect/predict.py

@@ -23,7 +23,7 @@ if __name__ == "__main__":
     #   'heatmap'           表示进行预测结果的热力图可视化，详情查看下方注释。
     #   'export_onnx'       表示将模型导出为onnx，需要pytorch1.7.1以上。
     #----------------------------------------------------------------------------------------------------------#
-    mode = "dir_predict"
+    mode = "predict"
     #-------------------------------------------------------------------------#
     #   crop                指定了是否在单张图片预测后对目标进行截取
     #   count               指定了是否进行目标的计数

app/routes/analyze_result.py

@@ -52,7 +52,7 @@ async def get_task_result(task_id: str, response: Response):
 
             # 构建mask信息
             masks = [
-                MaskInfo(name=mask["name"], coords=mask["coords"])
+                MaskInfo(name=mask["name"], score=mask["score"], coords=mask["coords"])
                 for mask in coords_data
             ]
 

app/schemas/analyze_result.py

@@ -11,6 +11,7 @@ class ImageInfo(BaseModel):
 
 class MaskInfo(BaseModel):
     name: str
+    score: float
     coords: List[List[int]]
 
 

app/services/worker.py

@@ -42,7 +42,8 @@ class Worker:
                     print(f"处理任务 {task_id}, 处理图片 {input_img_path}...")
 
                     img_res, coords_res = self.detection.detect(input_img_path)
-                    coords_res = [{"name": name, "coords": coords} for name, coords in coords_res]
+                    coords_res = [{"name": name, "score": score, "coords": coords} for name, score, coords in coords_res]
+                    print(coords_res)
                     coords_json = json.dumps(coords_res, ensure_ascii=False)
                     out_img_path = os.path.join(str(output_dir), f"{idx}.jpg")
                     cv2.imwrite(out_img_path, img_res)