Boen_Shi
b054769b40
- 新增yolo_detect模块，包含backbone、nets、utils等组件 - 在模型配置中添加yolo_detect选项，支持新的检测模型 - 移除SAM3预处理相关代码和配置项 - 更新Dockerfile删除core目录下所有文件以减少镜像体积 - 修改worker服务移除图像标签预处理逻辑，直接进行模型检测
65 lines
2.2 KiB
Python
65 lines
2.2 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.services.model import TaskStatus, TaskStore
|
|
|
|
|
|
class Worker:
|
|
def __init__(self):
|
|
self.detection = Model().getModel()
|
|
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)
|
|
|
|
for idx, input_img_path in enumerate(task.images):
|
|
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_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)
|
|
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()
|