# -*- coding: utf-8 -*-
# @Time : 2021/1/15 9:06
# @Author : smw
# @Site : jnsenter
# @File : remove_img_without_jsonlabel.py
# @Software: PyCharm

# 功能：1. 统一标注和图像，将没有标注的图像删除；
#      2. 将不同文件夹下的图像和标注文件汇总到一个文件夹下

import os
import shutil

def scandir(path, file_list):
    for item in os.scandir(path):
        if item.is_dir():
            scandir(item.path, file_list)
        elif item.is_file():
            file_list.append(item.path)

def image_remove(image_file_path, delete_num=0, residue_num=0):
    move_path_img = "E:\image_dataset\yaban_train_images\JPEGImages"
    move_path_anno = "E:\image_dataset\yaban_train_images\Annotations"
    os.makedirs(move_path_img, exist_ok=True)
    os.makedirs(move_path_anno, exist_ok=True)

    for root, dir, file_names in os.walk(image_file_path):
        if not file_names:  # 空
            for _dir in dir:
                next_dir = os.path.join(root, _dir)
                delete_num, residue_num = image_remove(next_dir, delete_num, residue_num)
        else:   # 非空
            for file_name in file_names:
                image_path = os.path.join(root, file_name)
                if image_path.endswith(".jpg"):
                    print("Process on {}".format(image_path))
                    json_path = image_path.replace(".jpg", ".json")
                    if not os.path.isfile(json_path):   # 如果不是一个文件则删除图像
                        os.remove(image_path)
                        delete_num += 1
                    else:    # 如果是一个文件 就把它们统一汇总在一起
                        img_dir = os.path.join(move_path_img, file_name)
                        anno_dir = os.path.join(move_path_anno, file_name.replace(".jpg", ".json"))
                        shutil.move(image_path, img_dir)
                        shutil.move(json_path, anno_dir)
                        residue_num += 1
                else:    # 非jpg文件跳过不处理
                    continue
    return delete_num, residue_num

if __name__ == '__main__':
    path = "E:\image_dataset\yaban"
    delete_num, residue_num = image_remove(path)
    print(delete_num, residue_num)

# -*- coding: utf-8 -*-
# @Time : 2021/1/15 16:06
# @Author : smw
# @Site : jnsenter
# @File : create_txt.py
# @Software: PyCharm

import json
import os 
import os.path as osp
import sys
import shutil
import math
import numpy as np

from functools import reduce
import logging
import logging.handlers
import logging.config

# 功能： 将数据集划分为训练集、验证集、测试集，并生成对应的txt于train_2021, val_2021, test_2021

classses = ['XX1', 'XX2']
sets = [('2022', 'train'), ('2022', 'val'), ("2022", "test")]  # float为划分比例 需保证比例总和为1
ratio = [0.8, 0.1, 0.1]
names_reflect = {"XX1": "YY1", "XX2": "YY2"}
# print(names_reflect.values())
# print("")

def initLog(logfile):
    logger = logging.getLogger()
    # 指定日志的最低输出级别，默认为WARN级别
    logger.setLevel(logging.INFO)
    fileHandeler = logging.handlers.TimedRotatingFileHandler(logfile, 'M', 1, 0, encoding="utf-8")
    # rotate_handler = ConcurrentRotatingFileHandler(logfile, "a", 1024 * 1024 * 100 * 6, backupCount=5, encoding="utf-8")
    # 指定logger输出格式
    formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
    fileHandeler.setFormatter(formatter)
    # 为logger添加的日志处理器
    logger.addHandler(fileHandeler)
    return logger

def convert(size, box):
    """坐标转换为中心点坐标+宽高"""
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[1])/2.0 - 1
    y = (box[2] + box[3])/2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def image_datasets_split(year, image_set, json_path, data, names_reflect, f_main, logger):
    """移动图像数据到相应的训练 测试 验证文件夹下"""
    move_path_image = "/home/smw/Project/yolov5/COCOdevkit/COCO{}/images/{}".format(year, image_set+year)
    os.makedirs(move_path_image, exist_ok=True)
    save_label_path = "/home/smw/Project/yolov5/COCOdevkit/COCO{}/labels/{}".format(year, image_set+year)
    os.makedirs(save_label_path, exist_ok=True)
    no_object_num = 0
    main_info_list = []   # 汇总的txt
    for path in data:  # 某中数据集的循环
        image_name = os.path.split(path)[1]
        if not image_name.endswith(".jpg"):
            logger.warning("{} is not endwith the .jpg".format(path))
            print("{} is not endwith the .jpg".format(path))
            continue
        
        # label txt
        save_txt_dir = os.path.join(save_label_path, image_name.replace(".jpg", ".txt"))
        # 1. 读取json 信息
        json_name = image_name.replace(".jpg", ".json")
        json_dir = os.path.join(json_path, json_name)
        # 再次验证json文件是否存在
        if not os.path.isfile(json_dir):
            logger.info("{} is not exist".format(json_dir))
            print("{} is not exist".format(json_dir))
            continue
        
        # 2. 过滤标签
        label_names = list(names_reflect.keys())
        with open(json_dir, "r", encoding="utf-8") as f:
            json_info = json.load(f)
        shapes_info = json_info.get("shapes")
        
        info_num = 0
        info_list = []  # 单个的label txt
        main_info_list_pick = [image_name]
        h = json_info.get("imageHeight")
        w = json_info.get("imageWidth")
        img_size = (w, h)

        for label_info in shapes_info:
            if label_info.get("shape_type") != 'rectangle':   # 过滤掉标错的 类别名称相同的内容
                continue

            label_name = label_info.get("label")
            if not label_name in label_names:    # 标注类型符合后 过滤标签名字
                continue
            # 准备写txt的内容
            info_num += 1  # 符合数量增
            label_reflect_name = names_reflect[label_name]

            label_class_id = label_names.index(label_name)   # info 1
            points = label_info.get("points")
            
            x, y, _w, _h = convert(img_size, [points[0][0], points[1][0],points[0][1], points[1][1]])   # 转xywh
            info_list.append([label_class_id, x, y, _w, _h])   # 记录信息
            main_info_list_pick.append([str(math.floor(points[0][0])), str(math.floor(points[0][1])), str(math.floor(points[1][0])), str(math.floor(points[1][1])), str(label_class_id)])

        
        if info_num > 0:  # 存在目标对象， 开始写txt，移动图像， 否则不移动图像
            # ToDo 移动图像，写txt
            # 如果不存在，则移动
            if not os.path.isfile(os.path.join(move_path_image, image_name)):
                # 移动图像
                shutil.copy(path, os.path.join(move_path_image, image_name))
            # shutil.move(path, os.path.join(move_path_image, image_name))
            # 写txt
            if not os.path.isfile(save_txt_dir):
                with open(save_txt_dir, "w") as f_txt:
                    for write_info in info_list:
                        write_info_str = map(str, write_info)
                        str_info = " ".join(write_info_str)
                        f_txt.write(str_info)
                        f_txt.write("\r\n")
        else:  # 不存在目标     
            # ToDo 不移动目标也不写txt 加日志输出
            no_object_num += 1
            logger.info("{} file does not have object!".format(image_name))
            print("{} file does not have object!".format(image_name))
        if len(main_info_list_pick) != 1:
            main_info_list.append(main_info_list_pick)  # 汇总每张图的
    # 所有图像汇总好后，写train2021.txt
    main_txt_write(main_info_list, f_main)
    return no_object_num

def main_txt_write(info_list, f):
    def add_dou(x, y):
        return x+","+y
    def add_space(x, y):
        return x+" "+y

    for info in info_list:
        image_name = info[0]
        bbox_list = []
        for info_bbox in info[1:]:
            info_single = reduce(add_dou, info_bbox)
            bbox_list.append(info_single)
        info_str_all = image_name + " " + reduce(add_space, bbox_list) 
        f.write(info_str_all)
        f.write("\r\n")
    print("")
    

if __name__ == "__main__":
    pwd = os.getcwd()
    print(pwd)
    logfile = os.path.join(pwd, "run.log")
    logger = initLog(logfile)

    image_path = "/home/smw/Project/yolov5/COCOdevkit/COCO2021/images_ori"
    json_path = "/home/smw/Project/yolov5/COCOdevkit/COCO2021/labels_json"

    main_txt_path = "/home/smw/Project/yolov5/COCOdevkit/COCO{}".format(sets[0][0])
    image_dir = [os.path.join(image_path, image_name) for image_name in os.listdir(image_path)]
    json_dir = [os.path.join(json_path, json_name) for json_name in os.listdir(json_path)]
    image_info = np.array(image_dir)   # image_dir 和 json_dir 的序号直接拼接对不上  因此仅以图像为主！！
    # 判断数量一致
    image_num = len(image_dir)
    json_num = len(json_dir)
    if image_num != json_num:
        logger.error(" The num of images is not equal to json, please excute remove_img_without_jsonlabel.py first")
        print(" The num of images is not equal to json, please excute remove_img_without_jsonlabel.py first")
        sys.exit(0)   # 此处可以进行删除无json的图像的代码 代替退出程序！！
    else:
        logger.info("The num of images is equal to json")
        logger.info("The num of images is {}, the number of images is {}".format(image_num, json_num))
        print(" The num of images is not equal to json, please excute remove_img_without_jsonlabel.py first")
    # 数据集划分  只看数据 json文件没有再说吧
    image_num_index = list(range(image_num))
    np.random.shuffle(image_num_index)
    train_end_num = int(image_num * ratio[0])
    val_end_num = int(image_num * (ratio[0] + ratio[1]))     #  没有对是否有测试 验证集的兼容！！
    
    print(image_info[0])
    print(image_info.shape)
    print(type(image_num_index))
    train_info = image_info[image_num_index[:train_end_num]]
    val_info = image_info[image_num_index[train_end_num: val_end_num]]
    test_info = image_info[image_num_index[val_end_num:]]
    logger.info("train num: {}".format(train_info.shape[0]))
    logger.info("val num:{}".format(val_info.shape[0]))
    logger.info("test num{}".format(test_info.shape[0]))
    print("train num: {}".format(train_info.shape[0]))
    print("val num:{}".format(val_info.shape[0]))
    print("test num:{}".format(test_info.shape[0]))
    print("The number of total split data is {}".format(train_info.shape[0]+val_info.shape[0]+test_info.shape[0]))   # 目前来看划分后的数据和总数是相等的
    no_object_list = []
    for path_info, data in zip(sets, (train_info, val_info, test_info)):
        year = path_info[0]
        image_set = path_info[1]
        main_txt_dir = os.path.join(main_txt_path, image_set+year+".txt")
        f = open(main_txt_dir, "w")
        no_object_num = image_datasets_split(year, image_set, json_path, data, names_reflect, f, logger)
        no_object_list.append(no_object_num)
        f.close()
        print("{} finished!".format(image_set))
    print("No object Number:")
    print("train: {}".format(no_object_list[0]))
    print("val: {}".format(no_object_list[1]))
    print("test: {}".format(no_object_list[2]))

import numpy as np

class YOLO_Kmeans:
    def __init__(self, cluster_number, filename):
    self.cluster_number = cluster_number
    self.filename = filename

    def iou(self, boxes, clusters):  # 1 box -> k clusters
        n = boxes.shape[0]
        k = self.cluster_number

        box_area = boxes[:, 0] * boxes[:, 1]    # 把要聚类的框的宽高相乘，作为了一个box_area
        box_area = box_area.repeat(k)     # 要算到k个类中心的距离，需要搞一个每个都有k个的矩阵
        box_area = np.reshape(box_area, (n, k))

        cluster_area = clusters[:, 0] * clusters[:, 1]
        cluster_area = np.tile(cluster_area, [1, n])
        cluster_area = np.reshape(cluster_area, (n, k))
        # 把box和cluster的宽都整理成n行k列的形式，并把两者做比较，最后还是一个n行k列的形式，这个
        # 过程其实在比较box和两个cluster的宽，并选出小的
        box_w_matrix = np.reshape(boxes[:, 0].repeat(k), (n, k))
        cluster_w_matrix = np.reshape(np.tile(clusters[:, 0], (1, n)), (n, k))
        min_w_matrix = np.minimum(cluster_w_matrix, box_w_matrix)
        # 把box和cluster的高都整理成n行k列的形式，并把两者做比较，最后还是一个n行k列的形式，这个
        # 过程其实在比较box和两个cluster的高，并选出小的
        box_h_matrix = np.reshape(boxes[:, 1].repeat(k), (n, k))
        cluster_h_matrix = np.reshape(np.tile(clusters[:, 1], (1, n)), (n, k))
        min_h_matrix = np.minimum(cluster_h_matrix, box_h_matrix)
        # 将筛选出来的小的宽高 相乘
        inter_area = np.multiply(min_w_matrix, min_h_matrix)

        result = inter_area / (box_area + cluster_area - inter_area)
        return result

    def avg_iou(self, boxes, clusters):
        accuracy = np.mean([np.max(self.iou(boxes, clusters), axis=1)])
        return accuracy

    def kmeans(self, boxes, k, dist=np.median):
        box_number = boxes.shape[0]
        distances = np.empty((box_number, k))
        last_nearest = np.zeros((box_number,))
        np.random.seed()
        clusters = boxes[np.random.choice(
            box_number, k, replace=False)]  # 随机选择k个类中心
        while True:

            distances = 1 - self.iou(boxes, clusters)

            current_nearest = np.argmin(distances, axis=1)
            if (last_nearest == current_nearest).all():
                break  # clusters won't change
            for cluster in range(k):
                # print(clusters[cluster])
                # print(boxes[current_nearest == cluster])
                # print(np.mean(boxes[current_nearest == cluster][:, 0]))
                clusters[cluster] = dist(boxes[current_nearest == cluster], axis=0)  # update clusters
                # 类中心的修改选取的是中位数 不是平均值
            last_nearest = current_nearest

        return clusters

	def result2txt(self, data):
        f = open("/home/smw/Project/yolov5/COCOdevkit/COCO2022/huikonggui2022.txt", 'w')
        row = np.shape(data)[0]
        for i in range(row):
            if i == 0:
                x_y = "%d,%d" % (data[i][0], data[i][1])
            else:
                x_y = ", %d,%d" % (data[i][0], data[i][1])
            f.write(x_y)
        f.close()

    def txt2boxes(self):
        f = open(self.filename, 'r')
        dataSet = []
        for line in f:
            infos = line.split(" ")
            length = len(infos)
            for i in range(1, length):
                width = int(infos[i].split(",")[2]) - \
                    int(infos[i].split(",")[0])
                height = int(infos[i].split(",")[3]) - \
                    int(infos[i].split(",")[1])
                dataSet.append([width, height])
        result = np.array(dataSet)
        f.close()
        return result

    def txt2clusters(self):
        all_boxes = self.txt2boxes()      # 将txt中数值信息转化为图像标记框的宽高，并返回
        result = self.kmeans(all_boxes, k=self.cluster_number)
        result = result[np.lexsort(result.T[0, None])]
        self.result2txt(result)
        print("K anchors:\n {}".format(result))
        print("Accuracy: {:.2f}%".format(
            self.avg_iou(all_boxes, result) * 100)

if  __name__ == "__main__":
    cluster_number = 9
    filename = "/home/smw/Project/yolov5/COCOdevkit/COCO2022/train2022.txt"
    kmeans = YOLO_Kmeans(cluster_number, filename)
    kmeans.txt2clusters()