utils/utils.py

import PIL.Image as Image
import io
from collections import Counter

from .mamasnowflake import *


# 根据两个点坐标截取人物画像图片
def get_one_person(xyxy,
                   img):
    crop = img[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(-1)]
    result = output_to_binary(crop)
    return result


# 保存图像
def save_img_as_png(img,
                    file_path,
                    RGB=0):
    if RGB:
        Image.fromarray(img).save(file_path, format='PNG', quality=100, subsampling=0)
    else:
        Image.fromarray(img[..., ::-1]).save(file_path, format='PNG', quality=100, subsampling=0)


# 获取人物名称
def get_one_name():
    # 获取雪花算法对象
    snowflake = Snowflake(datacenter_id=1, worker_id=3)
    # 调用雪花算法获取文件名
    randID = snowflake.generate()
    randID = str(randID)
    file_name = f'{randID}.png'
    return file_name


# 从字符串获取其中的第一个数字
def extract_first_number(s):
    for char in s:
        if char.isdigit():
            return int(char)
    return -1  # 如果字符串中没有数字，返回-1


# 检查字符串中是否包含关键词列表中的关键词
def contains_any_keyword(string, keywords):
    for keyword in keywords:
        if keyword in string:
            return True
    return False


# 将PNG转换为二进制
def png_to_binary(file_path='../data/temp/temp.png'):
    with open(file_path, 'rb') as file:
        binary_data = file.read()
    return binary_data


# 将模型输出转换为二进制
def output_to_binary(img_array):
    # 创建PIL图像对象
    img = Image.fromarray(img_array)

    # 将图像保存为PNG格式到BytesIO对象
    binary_data = io.BytesIO()
    img.save(binary_data, format='PNG')

    # 获取二进制数据
    binary_data = binary_data.getvalue()

    return binary_data


# 计算预测值与目标值之间的function_loss、total_loss以及avg_loss
def calculate_loss(predictions, targets):
    # 计算每一对数值之间的损失的绝对值
    function_loss = [abs(int(pred) - int(target)) for pred, target in zip(predictions, targets)]

    # 计算总的损失
    total_loss = sum(function_loss)

    # 计算平均损失
    avg_loss = total_loss / len(predictions)

    result = {'function_loss': function_loss,
              'total_loss': total_loss,
              'avg_loss': avg_loss,
              'total_person_num': len(predictions)
              }
    return result


# 统计词汇次数
def count_words_in_strings(strings):
    word_counts = Counter()
    word_counts.update(strings)
    return dict(word_counts)