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,18 +1,26 @@
 ### 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文件夹
 
-> 采用代码持久化，方便后续更新代码及模型而不需要重新创建容器，以及统一配置配置文件
-> 
-> 后续更新，如果只更新了核心文件，则仅需git pull后重新启动容器
-> 
-> 如果更新了requirements.txt，则需要重新构建镜像
+  ```bash
+  tar -xvf core.tar
+  ```
+
+4. 使用指令运行docker镜像
 
   ```bash
   sudo docker run -d \
@@ -20,11 +28,29 @@ sudo docker run -d \
       --gpus all \
       -p [local_port]:80 \
       -v $(pwd)/uploads:/code/uploads \
-  -v $(pwd)/app:/code/app \
+      -v $(pwd)/core:/code/app/core \
       -v $(pwd)/.env:/code/.env \
       wall
   ```
 
-> TIPS：由于部分文件采用GIT LFS 管理，请先安装GIT LFS
-> 
-> 在clone或者pull时，建议先clone代码文件，然后停掉，再用git lfs pull可查看大文件下载进度
+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
+	```
+
+	

app/core/yolo_detect/detect.py

@@ -98,7 +98,6 @@ 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'))
@@ -127,13 +126,8 @@ class YOLODetect(YOLO):
                     if keep:
                         color = self.colors[int(c)]
                         mask[top:bottom, left:right] = color
-                        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))]
-                            )
-                        )
+                        coords.append((self.classes.get(predicted_class),
+                                       [(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 = "predict"
+    mode = "dir_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"], score=mask["score"], coords=mask["coords"])
+                MaskInfo(name=mask["name"], coords=mask["coords"])
                 for mask in coords_data
             ]
 

app/schemas/analyze_result.py

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

app/services/worker.py

@@ -42,8 +42,7 @@ class Worker:
                     print(f"处理任务 {task_id}, 处理图片 {input_img_path}...")
 
                     img_res, coords_res = self.detection.detect(input_img_path)
-                    coords_res = [{"name": name, "score": score, "coords": coords} for name, score, coords in coords_res]
-                    print(coords_res)
+                    coords_res = [{"name": name, "coords": coords} for name, coords in 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)