sonar/sonar_process/count_fish_jh5_v1_bushu.py

# -*- coding:utf-8 -*-
import math
import os
import tarfile
import time
import cv2
import numpy as np
from tqdm import tqdm
from ultralytics import YOLO
from hard_disk_storage import HardDiskStorage
import shutil
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
model = YOLO('20240820jh5best.pt')  # load a custom model

def image24to30(image):
    height, width, channels = image.shape
    new_b_range = height * 30 / 24
    canjian = round(new_b_range - height)
    new_width = width + canjian
    new_height = height + canjian
    black_image = np.zeros((new_height, new_width, 3), dtype=np.uint8)
    black_image[:height, :width] = image
    resized_image = cv2.resize(black_image, (width, height), cv2.INTER_LINEAR)
    return resized_image


def image30to24(image):
    height, width, channels = image.shape
    new_b_range = height * 24 / 30
    canjian = round(height - new_b_range)
    image_caijian = image[:-canjian, :-canjian]
    resized_image = cv2.resize(image_caijian, (width, height), cv2.INTER_LINEAR)

    # cv2.imshow('1',resized_image)
    # cv2.waitKey(0)

    return resized_image


def make_mask(mask, image):
    background = np.zeros(image.shape, dtype=np.uint8) + 0
    image[mask] = background[mask]
    return image


def onload_mask(mask_path, IMAGE_SHAPE):
    mask_dict = {}
    all_flie_name = os.listdir(mask_path)
    for file_name in all_flie_name:
        if file_name.endswith(".bin") and file_name.startswith("mask"):
            key = file_name[file_name.index("_") + 1: -4]
            path = os.path.join(mask_path, file_name)
            mask = np.fromfile(path, dtype=np.bool_)
            mask.shape = IMAGE_SHAPE
            mask_dict[key] = mask
    return mask_dict


def list_split(info_list):
    s_list, b_list = [], []
    for image_info in info_list:

        if "s" in image_info.name:
            s_list.append(image_info)
        elif "b" in image_info.name:
            b_list.append(image_info)
        else:
            print("error")
    return s_list, b_list


def rotate(image, angle, center=None, scale=1.0):
    (h, w) = image.shape[:2]
    if center is None:
        center = (w // 2, h // 2)

    M = cv2.getRotationMatrix2D(center, angle, scale)
    rotated = cv2.warpAffine(image, M, (w, h))
    return rotated


def get_contours_aear(contour):
    """
    计算轮廓大小
    """

    x, y, w, h = cv2.boundingRect(contour)
    contour = contour - [x, y]
    p = np.zeros(shape=(h, w))
    # print(contour)
    cv2.drawContours(p, [contour], -1, 255, -1)
    # cv2.imshow('p',p)
    # cv2.waitKey(0)
    count_pixels = np.count_nonzero(p)
    # count_pixels = cv2.contourArea(contour=contour, oriented=False)
    return count_pixels


def get_fish_count(image, image_name, fish_size=130, YUZHI=0):
    """
    计算10m的鱼群数据
    """

    # 寻找轮廓
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    H, S, V = cv2.split(hsv)
    # 去噪
    # dst_V = cv2.fastNlMeansDenoising(V, None, 10, 7, 21)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    dst_V = cv2.morphologyEx(V, cv2.MORPH_OPEN, kernel, iterations=2)
    ret, dst_V = cv2.threshold(dst_V, YUZHI, 255, cv2.THRESH_BINARY)
    contours, hierarchy = cv2.findContours(
        dst_V, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )
    fish_count = 0

    for c in contours:
        data = np.array(get_contours_aear(c))
        if data > 20 and data < fish_size:
            fish_count = fish_count + 1
        elif data > fish_size:
            fish_count = fish_count + int(data / fish_size)

    return fish_count


def Calculate_density(fish_count_list, r=10, angel_big=100, angel_small=20):
    """
    :param fish_count_list:  30张鱼数量列表
    :param r:           范围半径
    :param angel_big:   声呐上下角度
    :param angel_small: 声呐左右角度
    :return:            10m鱼的密度

    ps：
    目前图像上下角度为 100° ， 左右角度为 20°
    体积计算过程：
    （1） 首先将声呐的上下范围当做180°（看作一个半圆），半圆旋转360°为球体，声呐左右角度为 20° ，此时上下角度180度，左右角度 20°的体积为球体积的 20/360。
    （2） 已知 上下角度180度，20， 那么上下角度100°的体积为    （上下角度180度，左右角度20度的体积） * 100/180
    """
    print("最大：", max(fish_count_list))
    print("最小：", min(fish_count_list))

    mean_fish_count = sum(fish_count_list) // len(fish_count_list)

    ball_volumes = math.pi * 10 * 10 * (135 / 360)

    density = mean_fish_count / ball_volumes

    print(
        "[ 10米 ] 鱼数量 ： {} ， 水面积 ： {}立方米 ， 单位体积鱼的数量（鱼密度）： {}".format(
            mean_fish_count, ball_volumes, density
        )
    )

    return density, mean_fish_count


def distance(x1, y1, x2, y2):
    return math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)


def Calculate_volume(range_shang, range_xia, rotation_angle, new_point):
    volume_list = []
    volume = 0
    for i in range(len(new_point)):
        rang = new_point[i][1]
        if range_shang <= rang * 24 / 630 <= range_xia:
            radius = (
                (new_point[i][0] * 24 / 630, new_point[i][0] * 24 / 630)
                if i == 0
                else ((new_point[i - 1][0] * 24 / 630, new_point[i][0] * 24 / 630))
            )
            height = (
                new_point[i][1] if i == 0 else new_point[i][1] - new_point[i - 1][1]
            )
            volume += (
                              np.pi
                              * (radius[0] ** 2 + radius[1] ** 2 + radius[0] * radius[1])
                              * (height * 24 / 630)
                      ) / 3

    if rotation_angle > 0 and range_xia == 6:
        jian_volume = (
                (
                        np.pi
                        * (new_point[0][0] * 24 / 630) ** 2
                        * (
                                math.tan(math.radians(rotation_angle))
                                * (new_point[0][0] * 24 / 630)
                        )
                )
                * 2
                / 3
        )
    else:
        jian_volume = 0
    volume_list.append(volume)
    volume_list.append(jian_volume)
    return volume_list


def Detection_Clothing(A, B):
    y_min_A = np.min(A[:, 1])
    y_max_A = np.max(A[:, 1])

    # 在数组B中选择y值在最大值和最小值中间的坐标
    mask = (B[:, 1] >= y_min_A) & (B[:, 1] <= y_max_A)
    C = B[mask]

    # 判断A中任意一个坐标与C中任意一个坐标的距离是否小于100个像素
    distances = np.sqrt(np.sum((A[:, None] - C) ** 2, axis=-1))
    result_distance = np.any(distances < (0.5 * 630 / 24))

    # 判断A中任意一个坐标的x值是否大于C中任意一个坐标的x值
    result_x = np.any(A[:, 0][:, None] > C[:, 0])

    result = result_distance or result_x  # 两个条件同时满足才返回True

    return result


def get_fish_ratio_all(image, image_name, mask_dict, CAGE_LENGTH, fish_size=0, min_size=0):
    # 初始化一个列表用于存储符合条件的轮廓
    temp_con = []
    # 初始化一个计数器，用于记录符合条件的鱼的数量
    tmp_total_s = 0

    # 将输入图像从BGR颜色空间转换为HSV颜色空间
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    # 将HSV图像分割成三个通道：H（色调）、S（饱和度）、V（明度）
    H, S, V = cv2.split(hsv)
    # 对V通道应用阈值操作，生成二值图像g
    # 阈值为60，超过阈值的像素设为255，否则设为0
    ret, g = cv2.threshold(V, 60, 255, cv2.THRESH_BINARY)

    # 在二值图像g中查找轮廓
    # 使用RETR_EXTERNAL模式只检测外部轮廓
    # 使用CHAIN_APPROX_SIMPLE模式压缩水平、垂直和对角线方向的轮廓点
    contours, hierarchy = cv2.findContours(
        g, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )

    # 遍历每一个检测到的轮廓
    for one_contours in contours:
        # 计算轮廓的面积
        # get_contours_aear是一个函数，用于计算轮廓的面积
        data = np.array(get_contours_aear(one_contours))
        # 如果轮廓面积在min_size和fish_size之间
        if data > min_size and data < fish_size:
            # 将该轮廓添加到符合条件的轮廓列表中
            temp_con.append(one_contours)
            # 计数器加1
            tmp_total_s = tmp_total_s + 1
        # 如果轮廓面积大于fish_size
        elif data > fish_size:
            # 将该轮廓添加到符合条件的轮廓列表中
            temp_con.append(one_contours)
            # 计数器加上面积除以fish_size的值，并向上取整
            tmp_total_s = tmp_total_s + round(data / fish_size)

    # 将符合条件的鱼的计数赋值给one_image_fish_count
    one_image_fish_count = tmp_total_s
    # 计算最终的鱼数量
    # 这里假设每张图片代表一个20度的区域
    # 360度表示一个完整的环形区域，所以需要乘以18（360 / 20）来计算整个区域内的鱼数
    # 再乘以2可能是某种修正系数
    finall_num = one_image_fish_count * (360 / 20) * 2
    # 返回最终计算的鱼的数量，并将其转换为整数
    return int(finall_num)

def move_gz(gz_path):

    target_folder = os.path.basename(gz_path)
    target_base_path = '/data/transfer_data'
    target_path = os.path.join(target_base_path, target_folder)

    # 检查目标文件夹是否存在，如果不存在则创建
    if not os.path.exists(target_path):
        os.makedirs(target_path)
    # 移动文件到目标文件夹
    try:
        shutil.move(image_tar, target_path)
        print(f"File moved from {image_tar} to {target_path}")
    except Exception as e:
        print(f"Error: {e}")

if __name__ == "__main__":

    harddisk_db = HardDiskStorage()
    while True:
        flag_30to24 = 0
        flag_24to30 = 0

        # UUID = "jh5_0605_s16b38"
        # # cage_num = UUID.split("yue")[-1]
        # sql = "SELECT uuid,cage_length,fish_weight FROM jh5_measurement_data WHERE UUID = %s ORDER BY id DESC LIMIT 1;"
        # var = UUID
        # res = harddisk_db.execute_sql(sql, var)
        #
        # if len(res) == 0:
        #     print("暂无数据")
        #     time.sleep(30)
        #     continue
        #
        # res = res[0]
        # UUID = res["uuid"]

        sql = "SELECT uuid,cage_length,fish_weight FROM measurement_data ORDER BY id DESC LIMIT 1;"
        res = harddisk_db.get_all(sql)
        print(res)
        if len(res) == 0:
            print("暂无数据")
            time.sleep(30)
            continue

        res = res[0]
        UUID = res["uuid"]
        CAGE_LENGTH = res["cage_length"]
        FISH_WEIGHT = res["fish_weight"]

        # =============================== |  1. 参数配置   | ================================
        fishsze = 60
        min_size = 10

        print("*" * 50)
        print("uuid:", UUID)
        print("cage_length:", CAGE_LENGTH)
        print("weight:", FISH_WEIGHT)
        print("fish_size:", fishsze)
        if UUID == None:
            print("uuid none")
            time.sleep(30)
            continue

        sql = "SELECT id,package_name, gz_path FROM package_data WHERE UUID = %s AND is_processed = 0 ORDER BY id ASC LIMIT 1;"
        var = UUID
        res = harddisk_db.execute_sql(sql, var)
        if len(res) == 0:
            print("暂无数据")
            time.sleep(30)
            continue

        res = res[0]
        ID = res["id"]
        package_name = res["package_name"]
        print("package_name:", package_name)
        print("*" * 50)

        gz_path = res["gz_path"]
        image_tar = os.path.join(gz_path, package_name)

        if not os.path.exists(image_tar):
            print('不存在gz文件')
            error_text = "不存在gz文件"
            is_abnormal = 1
            fish_num = None
            sql = "UPDATE package_data SET is_processed = 1 , fish_count = %s,is_abnormal = %s,error_text=%s WHERE id = %s; "
            var = (fish_num, is_abnormal, error_text, ID)
            res = harddisk_db.execute_sql(sql, var)
            
            move_gz(gz_path)
            
            continue

        data_start = time.time()
        # =============================== |  3. 初始化参数   | ================================
        # 默认参数
        IMAGE_SHAPE = (631, 1078, 3)
        mask_dict = onload_mask("mask/", IMAGE_SHAPE)

        all_fish_count_list = []
        total_num_list = []
        Volume_Process = [[], [], [], [], [], []]
        try:
            tar = tarfile.open(image_tar)
            info_list = tar.getmembers()
            s_list, b_list = list_split(info_list)
        except:
            error_text = "打开压缩包gz失败"
            is_abnormal = 1
            fish_num = None
            sql = "UPDATE package_data SET is_processed = 1 , fish_count = %s,is_abnormal = %s,error_text=%s WHERE id = %s; "
            var = (fish_num, is_abnormal, error_text, ID)
            res = harddisk_db.execute_sql(sql, var)
            
            move_gz(gz_path)
            
            continue

        if len(b_list) < 60:
            error_text = f"len(s_list): {len(s_list)}, len(b_list): {len(b_list)}"
            is_abnormal = 1
            fish_num = None
            sql = "UPDATE package_data SET is_processed = 1 , fish_count = %s,is_abnormal = %s,error_text=%s WHERE id = %s; "
            var = (fish_num, is_abnormal, error_text, ID)
            res = harddisk_db.execute_sql(sql, var)
            
            move_gz(gz_path)
            continue
        # =============================== |  4. 开始计算   | ================================
        image_path = os.path.split(image_tar)[0]
        fish_count_list = []
        error_num = 0
        for image_info_24m in tqdm(b_list):
            (image_name2, ext) = os.path.splitext(image_info_24m.name)

            f = tar.extractfile(image_info_24m)
            img_buffer = f.read()
            nparr = np.frombuffer(img_buffer, np.uint8)
            try:
                image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            except:
                continue
            if flag_30to24 == 1:
                image = image30to24(image)
            elif flag_24to30 == 1:
                image = image24to30(image)

            # 预测鱼群轮廓
            results1 = model(image, conf=0.90)
            if results1[0].masks is None:
                error_num += 1
                print(image_name2,'yolo error')
                continue
            mask = results1[0].masks.xy
            result_tuple = (mask[0].reshape(len(mask[0]), 1, 2).astype(int),)
            # 创建一个空白图像，大小与原图像相同
            fish_mask = np.zeros_like(image)
            # 将掩码绘制到空白图像上
            for contour in mask:
                cv2.drawContours(fish_mask, [contour.reshape(len(contour), 1, 2).astype(int)], -1, (255, 255, 255), thickness=cv2.FILLED)

            # 使用掩码提取鱼的部分
            fish_image = cv2.bitwise_and(image, fish_mask)
            finall_num = get_fish_ratio_all(fish_image, image_name2, mask_dict, CAGE_LENGTH, fish_size=fishsze, min_size=min_size)
            total_num_list.append(finall_num)

        mean_value = np.median(total_num_list)
        std_dev = np.std(total_num_list)
        threshold_max = mean_value + 2 * std_dev
        threshold_min = mean_value - 2 * std_dev
        if threshold_min < 0:
            threshold_min = 0

        total_num_list = [
            x for x in total_num_list if x > threshold_min and x < threshold_max
        ]


        if len(total_num_list) < 20:
            print(image_name2, "有效数据不足")
            error_text = f"len(total_num_list) < 20:{len(total_num_list)}"
            is_abnormal = 1
            fish_num = None
            sql = "UPDATE package_data SET is_processed = 1 , fish_count = %s,is_abnormal = %s,error_text=%s WHERE id = %s; "
            var = (fish_num, is_abnormal, error_text, ID)
            res = harddisk_db.execute_sql(sql, var)
            move_gz(gz_path)
            continue

        print("这一组情况：")
        print(total_num_list)
        print("max:{}".format(max(total_num_list)))
        print("min:{}".format(min(total_num_list)))
        print("mean:{}".format(sum(total_num_list) // len(total_num_list)))
        date_off = time.time()
        print("花费时间", date_off - data_start)
        print("异常图片数量", error_num)

        fish_num = int(sum(total_num_list) // len(total_num_list))

        is_abnormal = 0

        sql = "UPDATE package_data SET is_processed = 1 , fish_count = %s,is_abnormal = %s WHERE id = %s; "
        var = (fish_num, is_abnormal, ID)
        res = harddisk_db.execute_sql(sql, var)
        
        move_gz(gz_path)