Boen_Shi
6a2e046884
- 新增 analyze、analyze_result、analyze_status 和 health 路由 - 实现图像上传和任务提交功能 - 添加任务状态查询和结果获取接口 - 集成 segformer 和 yolo 模型进行图像检测 - 实现 SAM3 预处理功能用于图像预处理判断 - 添加模型选择配置支持 segformer 和 yolo - 实现任务队列管理和异步处理机制 - 添加 Dockerfile 用于容器化部署 - 配置环境变量和 gitignore 规则 - 创建数据模型定义 API 响应结构
81 lines
2.9 KiB
Python
81 lines
2.9 KiB
Python
import json
|
||
import os
|
||
import cv2
|
||
import threading
|
||
from datetime import datetime
|
||
from queue import Queue
|
||
from typing import Dict
|
||
|
||
from app.core.model import Model
|
||
from app.core.preprocess import Preprocess
|
||
from app.services.model import TaskStatus, TaskStore
|
||
|
||
|
||
class Worker:
|
||
def __init__(self):
|
||
self.detection = Model().getModel()
|
||
self.preprocess = Preprocess().getPreprocess()
|
||
self.task_queue = Queue()
|
||
self.task_store: Dict[str, TaskStore] = {}
|
||
|
||
threading.Thread(target=self.worker, daemon=True).start()
|
||
|
||
def worker(self):
|
||
from app.main import UPLOAD_DIR
|
||
while True:
|
||
task_id = self.task_queue.get()
|
||
if task_id is None:
|
||
break
|
||
|
||
task = self.task_store.get(task_id)
|
||
if not task:
|
||
continue
|
||
|
||
try:
|
||
task.status = TaskStatus.PROCESSING.name
|
||
task.progress = 0
|
||
print(f"开始处理任务 {task_id}...")
|
||
|
||
# 创建输出目录
|
||
output_dir = os.path.join(UPLOAD_DIR, task_id, "outputs")
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 获取图像的标签列表
|
||
image_labels = self.preprocess.preprocess(task.images) # 返回一个0和1的列表，0代表跳过，1代表进行检测
|
||
|
||
for idx, (input_img_path, label) in enumerate(zip(task.images, image_labels)):
|
||
print(f"处理任务 {task_id}, 处理图片 {input_img_path}...")
|
||
|
||
if label == 0:
|
||
# 如果标签是0，跳过模型检测，输出路径和坐标为空
|
||
task.result.append(
|
||
{"input_img_path": input_img_path, "output_img_path": "", "coords": "[]"}
|
||
)
|
||
else:
|
||
# 进行模型检测
|
||
img_res, coords_res = self.detection.detect(input_img_path)
|
||
|
||
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(output_dir, f"{idx}.jpg")
|
||
cv2.imwrite(out_img_path, img_res)
|
||
|
||
task.result.append(
|
||
{"input_img_path": input_img_path, "output_img_path": out_img_path, "coords": coords_json}
|
||
)
|
||
|
||
task.progress = int((idx + 1) / len(task.images) * 100)
|
||
|
||
task.status = TaskStatus.COMPLETED.name
|
||
task.completedAt = datetime.now()
|
||
task.message = "处理完成"
|
||
|
||
print(f"任务 {task_id} 处理完成")
|
||
|
||
except Exception as e:
|
||
task.status = TaskStatus.FAILED.name
|
||
task.message = str(e)
|
||
finally:
|
||
self.task_queue.task_done()
|