"""
# File       : mask2small.py
# Time       ：2024-04-23 16:40
# Author     ：FEANGYANG
# version    ：python 3.7
# Contact    ：1071082183@qq.com
# Description：
"""
"""
# File       : mask2iou.py
# Time       ：2024-04-08 13:35
# Author     ：FEANGYANG
# version    ：python 3.7
# Contact    ：1071082183@qq.com
# Description：
"""
import os
import cv2
import numpy as np


def cut_small_images(image_path, mask_path, pic_target='./result_class/'):
    if not os.path.exists(pic_target):  # 判断是否存在文件夹如果不存在则创建为文件夹
        os.makedirs(pic_target)
    image = cv2.imread(image_path)
    mask = cv2.imread(mask_path)
    name = os.path.split(image_path)[-1]

    (width, length, depth) = image.shape
    w, h, _ = image.shape
    # _w = w % pixes
    # _h = h % pixes

    cut_width = cut_length = 128

    # 预处理生成0矩阵
    pic = np.zeros((cut_width, cut_length, depth))
    # 计算可以划分的横纵的个数
    num_width = int(width / cut_width)
    num_length = int(length / cut_length)
    # for循环迭代生成
    lines_list = []
    print(image_path)
    for i in range(0, num_width):
        for j in range(0, num_length):
            im = image[i * cut_width: (i + 1) * cut_width, j * cut_length: (j + 1) * cut_length, :]
            ma = mask[i * cut_width: (i + 1) * cut_width, j * cut_length: (j + 1) * cut_length, :]
            agv = np.sum(ma)/(ma.shape[0]*ma.shape[1]*ma.shape[2])
            if agv < 25:
                fire_name = 'nosmoke' #'./data/images/fire (1747).png'
                nosmoke_dir = os.path.join(pic_target, fire_name)
                if not os.path.exists(nosmoke_dir):  # 判断是否存在文件夹如果不存在则创建为文件夹
                    os.makedirs(nosmoke_dir)
                result_path = os.path.join(nosmoke_dir, fire_name + '_' + os.path.splitext(name)[0] + '_{}_{}_image.jpg'.format(i + 1, j + 1))
            else:
                fire_name = 'smoke'
                smoke_dir = os.path.join(pic_target, fire_name)
                if not os.path.exists(smoke_dir):  # 判断是否存在文件夹如果不存在则创建为文件夹
                    os.makedirs(smoke_dir)
                result_path = os.path.join(smoke_dir, fire_name + '_' + os.path.splitext(name)[0] + '_{}_{}_image.jpg'.format(i + 1, j + 1))

            # result_path = pic_target + fire_name + '_'+os.path.splitext(name)[0] +'_{}_{}_mask.jpg'.format(i + 1, j + 1)
            # cv2.imwrite(result_path, ma)



            cv2.imwrite(result_path, im)

if __name__ == '__main__':
    masks_dir_path = './data/mask'
    images_dir_path = './data/images'
    # masks_dir_path = './data/test/mask'
    # images_dir_path = './data/fire'
    images_list_path = os.listdir(images_dir_path)
    for image_path in images_list_path:
        cut_small_images(os.path.join(images_dir_path, image_path), os.path.join(masks_dir_path, os.path.splitext(image_path)[0]+'.jpg'))