后端新增《蚕茧识别V2》模块

2025-11-10 18:08:50 +08:00
parent 9527cc2f1c
commit 625d185f69
15 changed files with 559 additions and 809 deletions
@@ -5,6 +5,7 @@ import config.minIO as minIO
 import db.postgres as pg
 from agent.licenseImageAgent import get_license_response
 from config.minIO import minio_client
+from config.yolo import YOLOSingleton
 from llm.ticketLLM import *
 from llm.ticketLLMv2 import get_ticket_response_v2

@@ -106,3 +107,86 @@ def process_license_image(
    )

    return json_data
+
+
+def process_silkworm_cocoon_image(
+    img_bytes=None,
+    file_name: str = None,
+    project_name: str = None,
+    user_id: UUID = None,
+):
+    # 上传到 OSS，使用 UUID 做对象名
+    if img_bytes is None:
+        img_bytes = []
+    pre_object_name = str(uuid.uuid4())
+    after_object_name = str(uuid.uuid4())
+    file_bytes = BytesIO(img_bytes)
+
+    bucket_name = "image-sca"
+    if not minio_client.bucket_exists(bucket_name):
+        minio_client.make_bucket(bucket_name)
+    minIO.push_file(
+        bucket_name, "raw/" + pre_object_name, file_bytes, img_bytes, "image/jpeg"
+    )
+
+    # YOLO检测
+    img_bytes_out, results_json = YOLOSingleton.detect(img_bytes)
+
+    # results_json = {
+    #     "total_objects": "",
+    #     "max_confidence": "",
+    #     "min_confidence": "",
+    #     "avg_confidence": "",
+    #     "class_counts": "",
+    #     "speed_ms": {
+    #         "preprocess": "",
+    #         "inference": "",
+    #         "postprocess": "",
+    #     },
+    # }
+    speed_json = results_json.get("speed_ms")
+
+    file_bytes_out = BytesIO(img_bytes_out)
+    minIO.push_file(
+        bucket_name,
+        "ai/" + after_object_name,
+        file_bytes_out,
+        img_bytes_out,
+        "image/jpeg",
+    )
+
+    # 获取图片分辨率和大小
+    img = Image.open(BytesIO(img_bytes))
+    resolution = f"{img.width}x{img.height}"
+    size_kb = round(len(img_bytes) / 1024, 2)
+
+    # 插入数据库
+    pg.insert_sca_image(
+        file_name=file_name,
+        resolution=resolution,
+        size=size_kb,
+        cocoon_count=results_json.get("total_objects"),
+        max_confidence=results_json.get("max_confidence"),
+        min_confidence=results_json.get("min_confidence"),
+        average_confidence=results_json.get("avg_confidence"),
+        other_info=results_json.get("class_counts"),
+        preprocess_time_ms=speed_json.get("preprocess"),
+        inference_time_ms=speed_json.get("inference"),
+        postprocess_time_ms=speed_json.get("postprocess"),
+        name=project_name if project_name else pre_object_name[:8],
+        image_pre=pre_object_name,
+        image_after=after_object_name,
+        created_by=user_id,
+    )
+    return {
+        "resolution": resolution,
+        "size": size_kb,
+        "cocoon_count": results_json.get("total_objects"),
+        "max_confidence": results_json.get("max_confidence"),
+        "min_confidence": results_json.get("min_confidence"),
+        "average_confidence": results_json.get("avg_confidence"),
+        "preprocess_time_ms": speed_json.get("preprocess"),
+        "inference_time_ms": speed_json.get("inference"),
+        "postprocess_time_ms": speed_json.get("postprocess"),
+        "details": results_json.get("class_counts"),
+    }