123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479 |
- # -*- 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)
|