【yolov5】接收他人传输的图片（base64），通过flask服务 返回图片物种的坐标，类别，和裁剪后的图片片段（base64）

代码如下，没有优化（只实现）

test2_trans.py

import torch
import cv2
from PIL import Image
from io import BytesIO
import os


pt_path = r'E:\Code\Python\yolov5py38\yolov5\best.pt'
# Model
model = torch.hub.load(r"E:\Code\Python\YoLov5\yolov5",'custom', path=pt_path, source='local')  # or yolov5n - yolov5x6, custom



# 遍历目录内所有图片
def list_files_in_folder(folder_path):
    files = os.listdir(folder_path)
    for file in files:
        file_path = os.path.join(folder_path, file)
        if os.path.isfile(file_path):
            print("File:", rf'{folder_path}\{file}')
            yield rf'{folder_path}\{file}'
        # elif os.path.isdir(file_path):
        #     print("Subfolder:", file)


    



# Image  获取图片物体的坐标，可信度，类别， 图像的三维数组
def get_info(img_path=r"E:\Code\Python\yolov5py38\dataset\dog_and_cat\images\val\119.jpg"):
    point_ls, cls_ls, b64_ls = [], [], []  # 一个图片的所有类别坐标组，类组， im的一维数组的base64组
    import base64
    # # img_path = r"E:\Code\Python\yolov5py38\dataset\dog_and_cat\images\val\119.jpg"  # or file, Path, PIL, OpenCV, numpy, list
    # folder_to_list = r"E:\Code\Python\yolov5py38\dataset\dog_and_cat\images\val"
    # for img_path in list_files_in_folder(folder_to_list):
    with open(img_path, 'rb') as f:
        img_data = f.read()
    img_obj = Image.open(BytesIO(img_data))

    # print(type(img_obj))

    # Inference
    results = model(img_obj)

    # Results
    # results.print()

    s = results.crop(save=False)  # or .show(), .save(), .crop(), .pandas(), etc.
    # print(s)
    for item in s:

        # need
        point = item['box']
        point_ls.append(point)

        cls = item['cls']
        cls_ls.append(cls)


        # print(item)
        im = item['im']
        # a = cv2.imencode(".jpg", im)  # 三维转一维
        success, enc_img = cv2.imencode(".jpg", im)
        # print(success, enc_img)  # enc_img 为一维数组
        # print(type(success), type(enc_img))
        img_data = enc_img.tobytes()  # 转字节码
        # print(f'aaaa{type(img_data)}')
        # print(f'bbbbb', enc_img)
        # b64_code = base64.b64encode(img_data)
        # print(b64_code)


        # ls.append(b64_code)
        b64_code1 = base64.b64encode(img_data).decode()
        b64_ls.append(b64_code1)

        # print(b64_code1)
    print(len(point_ls), len(b64_ls))
    return point_ls, cls_ls, b64_ls



# 测试
cc = get_info()
print(cc[0])
print(cc[1])

for point, clss,  b64 in zip(cc[0], cc[1], cc[2]):
    tensor_values = [tensor.item() for tensor in point]  # 使用 .item() 方法来提取Tensor的数值部分
    cls = clss.item() # 与上同理
    print(tensor_values, cls,  b64)

# with open('a.txt', 'a') as f:
#     for index, bs64 in enumerate(get_info()):
#         print(f"第{index}张图--{bs64}")
#         f.write(f"第{index}张图--{bs64}\n")

con_api.py [主文件]

import base64
from test2_trans import get_info

from flask import Flask, render_template, request, jsonify

app = Flask(__name__)


@app.route('/', methods=['GET', 'POST'])
def index():
    dic = None
    CLS, POINT, Bc64 = [], [], []

    if request.method == 'POST':
        try:
            json_data = request.get_json()  # 获取POST请求中的JSON数据

            if 'image_data' in json_data:
                # 假设用户发送的JSON中有一个名为 "image_data" 的字段
                image_data = json_data['image_data']

                output_path = "output_image.jpg"

                # 对用户发送的图片（base64）进行解码
                base64_to_image(image_data, output_path)


                res = get_info(output_path)   # 返回 point_ls, cls_ls, b64_ls

                for point, clss, b64 in zip(res[0], res[1], res[2]):
                    Bc64.append(b64)
                    tensor_values = [tensor.item() for tensor in point]  # 使用 .item() 方法来提取Tensor的数值部分
                    POINT.append(tensor_values) # 坐标
                    cls = clss.item()  # 与上同理
                    CLS.append(cls)  # 类别
                    print(tensor_values, cls, b64)
                    dic = {
                        "result": Bc64,
                        "point": POINT
                        # "b64": b64
                    }


                # # 在这里可以对图像数据进行处理
                # processed_data = process_image_data(image_data)

        except Exception as e:
            error_message = str(e)
            return jsonify({"error": error_message})

    if dic:
        dic['message'] = 'success'
        return jsonify(dic)
    else:
        return jsonify({"message": "No data processed."})


# base64转图片
def base64_to_image(base64_data, output_path):
    image_data = base64.b64decode(base64_data)
    with open(output_path, "wb") as image_file:
        image_file.write(image_data)
        print('图片保存已保存为', output_path)  # 每次替换

base64_encoded_image = "base64_encoded_data_here"  # 替换为实际的Base64编码
output_path = "output_image.jpg"  # 替换为输出图片的路径和文件名

# base64_to_image(base64_encoded_image, output_path)
# print("Image saved:", output_path)




def process_image_data(image_data):
    # 在这里进行图像数据处理的实际操作
    # 这个函数是一个示例，根据实际需求来实现图像数据的处理逻辑
    processed_image = image_data.upper()  # 示例：将图像数据转换为大写

    return processed_image


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5555, debug=True)