ASJPush/Controller/DrawBoxesOnImageController.py

# -*- encoding: utf-8 -*-
"""
@File    : DrawBoxesOnImageController.py
@Time    : 2025/9/2 09:21
@Author  : stephen
@Email   : zhangdongming@asj6.wecom.work
@Software: PyCharm
"""
import base64
import io
import json
import os
import time

from PIL import Image, ImageDraw, ImageFont
from django.views import View

from AnsjerPush.config import BASE_DIR
from Object.ResponseObject import ResponseObject


class DrawBoxesOnImageView(View):
    def get(self, request, *args, **kwargs):
        request.encoding = 'utf-8'
        operation = kwargs.get('operation')
        return self.validation(request.GET, request, operation)

    def post(self, request, *args, **kwargs):
        request.encoding = 'utf-8'
        operation = kwargs.get('operation')
        return self.validation(request.POST, request, operation)

    def validation(self, request_dict, request, operation):
        response = ResponseObject()
        if operation is None:
            return response.json(444, 'error path')
        elif operation == 'drawBoxesFromRequest':
            return self.draw_boxes_from_request(request_dict, response)
        return response.json(414)

    def draw_boxes_from_request(self, request_dict, response: ResponseObject):
        """
        处理绘制边界框的API请求。
        预期请求体 (JSON):
        {
            "fileOne": "base64_string_1",
            "fileTwo": "base64_string_2",
            "fileThree": "base64_string_3", // 可选
            "coordinates": { ... } // 包含所有图片坐标的JSON对象
        }
        """
        # 1. 定义图片保存的相对路径和绝对路径
        #    请确保 settings.BASE_DIR 指向您的Django项目根目录
        relative_save_dir = os.path.join('static', 'ai')
        absolute_save_dir = os.path.join(BASE_DIR, relative_save_dir)

        # 确保目标目录存在
        os.makedirs(absolute_save_dir, exist_ok=True)

        # 2. 从请求中提取图片和坐标数据
        base64_images = {
            "file_0": request_dict.get("fileOne"),
            "file_1": request_dict.get("fileTwo"),
            "file_2": request_dict.get("fileThree"),
        }

        # 坐标可以直接是字典，也可以是JSON字符串，这里做兼容处理
        coordinates_data = request_dict.get("coordinates")
        if isinstance(coordinates_data, str):
            try:
                coordinates_data = json.loads(coordinates_data)
            except json.JSONDecodeError:
                return response.json(400, "coordinates字段不是一个有效的JSON字符串")

        if not coordinates_data or not isinstance(coordinates_data, dict):
            return response.json(400, "缺少或无效的 'coordinates' 数据")

        processed_files = []
        errors = []

        # 3. 循环处理每张图片
        for i in range(3):
            image_key = f"file_{i}"

            base64_str = base64_images.get(image_key)
            detections = coordinates_data.get(image_key)

            # 如果图片和对应的坐标都存在，则进行处理
            if base64_str and detections:
                # 生成一个唯一的文件名以避免冲突
                timestamp = int(time.time() * 1000)
                output_filename = f"result_{image_key}_{timestamp}.jpg"
                output_path = os.path.join(absolute_save_dir, output_filename)

                # 调用绘图函数
                success = self.draw_boxes_on_image(base64_str, detections, output_path)

                if success:
                    # 返回可供前端访问的静态文件URL
                    static_url = os.path.join(BASE_DIR, 'static', 'ai', output_filename).replace("\\", "/")
                    processed_files.append({"source": image_key, "url": static_url})
                else:
                    errors.append(f"处理 {image_key} 失败")

        # 4. 根据处理结果返回响应
        if not processed_files and not errors:
            return response.json(400, "未提供任何有效的图片和坐标数据进行处理")
        elif errors:
            return response.json(500, {"errors": errors, "success": processed_files})
        else:
            return response.json(0, {"processed_files": processed_files})

    def draw_boxes_on_image(self, base64_image_string: str, detections: list, output_path: str) -> bool:
        """
        在一个Base64编码的图片上根据提供的检测坐标绘制边界框，并保存到指定路径。

        Args:
            base64_image_string (str): 图片的Base64编码字符串。
            detections (list): 一个包含检测对象信息的字典列表。
                               每个字典应包含 'class' 和 'Left', 'Top', 'Width', 'Height' 键。
            output_path (str): 带有标注的图片的完整保存路径。

        Returns:
            bool: 如果成功保存图片则返回 True，否则返回 False。
        """
        if not base64_image_string or not detections:
            print("错误：未提供图片数据或检测坐标。")
            return False

        try:
            # 1. 解码Base64图片并加载
            image_bytes = base64.b64decode(base64_image_string)
            image = Image.open(io.BytesIO(image_bytes))
            draw = ImageDraw.Draw(image)
            img_width, img_height = image.size

            # 2. 尝试加载字体，如果失败则使用默认字体
            try:
                # 确保服务器上有这个字体文件，或者换成一个您确定存在的字体路径
                font = ImageFont.truetype("arial.ttf", size=20)
            except IOError:
                print("警告: 未找到arial.ttf字体，将使用默认字体。")
                font = ImageFont.load_default()

            colors = ["red", "blue", "green", "yellow", "purple", "orange", "cyan", "magenta"]

            # 3. 遍历所有检测框并绘制
            for idx, item in enumerate(detections):
                label = item.get("class", "unknown")

                left_ratio = float(item["Left"])
                top_ratio = float(item["Top"])
                width_ratio = float(item["Width"])
                height_ratio = float(item["Height"])

                x1 = int(left_ratio * img_width)
                y1 = int(top_ratio * img_height)
                x2 = int((left_ratio + width_ratio) * img_width)
                y2 = int((top_ratio + height_ratio) * img_height)

                color = colors[idx % len(colors)]
                draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
                draw.text((x1 + 4, y1 + 2), label, fill=color, font=font)

            # 4. 保存绘制好的图片
            image.save(output_path)
            print(f"检测结果图片已成功保存到: {output_path}")
            return True

        except Exception as e:
            print(f"在绘制和保存图片时发生错误: {e}")
            return False