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

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

Dockerfile

@@ -19,11 +19,11 @@ COPY requirements.txt /code/requirements.txt
 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

@@ -8,11 +8,11 @@ docker build -t wall .
 
 #### 启动方式
 
-> 采用模型文件持久化，方便后续更新模型而不需要重新创建容器，以及统一配置配置文件
+> 采用代码持久化，方便后续更新代码及模型而不需要重新创建容器，以及统一配置配置文件
 > 
 > 后续更新，如果只更新了核心文件，则仅需git pull后重新启动容器
 > 
-> 如果更新了代码文件，则需要重新构建镜像
+> 如果更新了requirements.txt，则需要重新构建镜像
 
 ```bash
 sudo docker run -d \
@@ -20,10 +20,11 @@ sudo docker run -d \
   --gpus all \
   -p [local_port]:80 \
   -v $(pwd)/uploads:/code/uploads \
-  -v $(pwd)/app/core:/code/app/core \
+  -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)