import os
import xml
import xml.etree.ElementTree as ET
from lxml.etree import Element,SubElement,tostring
import pprint
import numpy as np
from xml.dom.minidom import parseString
import cv2
from PIL import Image
import tqdm
root='/VOC2012'
save_root='Downloads/data/VOC20122'
if not os.path.exists(save_root):
    os.mkdir(save_root)
id='2' #2或者7，为了给后面生成的文件做标记  2007：7 2012:2
split='train'  #用于标记是train还是val
root_txt='VOC2012/ImageSets/Main/val.txt'
c_root=['Annotations', 'JPEGImages', 'ImageSets']
# for i in c_root:
#     if not os.path.exists(root+'/'+i):
#         # print('i  ',i)
#         os.mkdir(root+'/'+i)



old_image_root=root+'/'+c_root[1]
save_image_root=save_root+'/'+c_root[1]
if not os.path.exists(save_image_root):
    os.mkdir(save_image_root)
old_xml_root=root+'/'+c_root[0]
save_xml_root=save_root+'/'+c_root[0]
if not os.path.exists(save_xml_root):
    os.mkdir(save_xml_root)
if not os.path.exists(save_root+'/'+c_root[2]):
    os.mkdir(save_root+'/'+c_root[2])
old_txt_root=root+'/'+c_root[2]+'/'+'Main'
save_txt_root=save_root+'/'+c_root[2]+'/'+'Main'
if not os.path.exists(save_txt_root):
    os.mkdir(save_txt_root)
class_name_train=[]
class_name_val=[]
class_name_test=[]

#### 生成txt文件的文件名
for class_txt in os.listdir(old_txt_root):
    new_txt=save_txt_root+'/'+class_txt
    if not os.path.exists(save_txt_root):
        os.mkdir(save_txt_root)
    # if not os.path.exists(new_txt):
    #     os.mknod(new_txt)

for f in os.listdir(save_txt_root):
    if split in new_txt:
        if f=='test.tx' or f=='train.txt' or f=='trainval.txt' or f=='val.txt':
            continue
        if split == 'train' and 'train' in f:
            if f=='train_val.txt' or f=='train_test.txt':
                continue
            class_name_train.append(f)
        elif split == 'val' and 'val' in f:
            class_name_val.append(f)
        elif split == 'test' and 'test' in f :
            class_name_test.append(f)

#切图规则
#1、面积<thera:不切
#2、
def read_xml(filename):
    """Parse a PASCAL VOC xml file."""
    tree = ET.parse(filename)
    objects = []
    for obj in tree.findall("object"):
        obj_struct = {}
        obj_struct["name"] = obj.find("name").text
        obj_struct["pose"] = obj.find("pose").text
        obj_struct["truncated"] = int(obj.find("truncated").text)
        obj_struct["difficult"] = int(obj.find("difficult").text)
        bbox = obj.find("bndbox")
        obj_struct["bbox"] = [
            int(bbox.find("xmin").text),
            int(bbox.find("ymin").text),
            int(bbox.find("xmax").text),
            int(bbox.find("ymax").text),
        ]
        objects.append(obj_struct)
    return objects
def gen_obj_struct(name,pose,truncated,difficult,xmin,ymin,xmax,ymax):
    obj_struct = {}
    obj_struct["name"] = name
    obj_struct["pose"] = pose
    obj_struct["truncated"] = truncated
    obj_struct["difficult"] = difficult
    obj_struct["bbox"] = [
        xmin,
        ymin,
        xmax,
        ymax,
    ]
    return obj_struct
def split_image(old_image,ob,objects,):
    #输入
    split_ans=[]
    img_1 = Image.open(old_image)
    old_width,old_height=img_1.size[0],img_1.size[1]
    obj_xmin,obj_ymin,obj_xmax,obj_ymax=ob['bbox']
    cent_x=abs(obj_xmax-obj_xmin)//2+obj_xmin
    cent_y=abs(obj_ymax-obj_ymin)//2+obj_ymin
    xmin_x=[0,cent_x,0,cent_x]
    ymin_y=[0,0,cent_y,cent_y]
    xmax_x=[cent_x,old_width,cent_x,old_width]
    ymax_y=[cent_y,cent_y,old_height,old_height]
    for i in range(4):
        sub_image={}
        cur_xmin=xmin_x[i]
        cur_ymin=ymin_y[i]
        cur_xmax=xmax_x[i]
        cur_ymax=ymax_y[i]

        cur_objects=[]

        for obb in objects:
            other_xmin,other_ymin,other_xmax,other_ymax=obb['bbox']
            U_xmin=max(cur_xmin,other_xmin)
            U_ymin=max(cur_ymin,other_ymin)
            U_xmax=min(cur_xmax,other_xmax)
            U_ymax=min(cur_ymax,other_ymax)
            if U_xmin>U_xmax:
                continue
            if U_ymin>U_ymax:
                continue
            #####要使用相对坐标

            re_xmin=U_xmin-cur_xmin
            re_xmax=U_xmax-cur_xmin
            re_ymin=U_ymin-cur_ymin
            re_ymax=U_ymax-cur_ymin
            if (abs(re_xmax-re_xmin)*abs(re_ymax-re_ymin)<16*16):
                continue
            cur_objects.append(gen_obj_struct(obb['name'],obb['pose'],obb['truncated'],obb['difficult'],re_xmin,re_ymin,re_xmax,re_ymax))
        sub_image['xmin']=cur_xmin
        sub_image['ymin']=cur_ymin
        sub_image['xmax']=cur_xmax
        sub_image['ymax']=cur_ymax
        sub_image['width']=abs(cur_xmax-cur_xmin)
        sub_image['height']=abs(cur_ymax-cur_ymin)
        sub_image['ans_object']=cur_objects
        split_ans.append(sub_image)

    return split_ans
def f_save_xml(width,height,image,split_object_ans,save_xml):
    # 创建文件的头信息
    node_root = Element('annotation')
    node_folder = SubElement(node_root, 'folder')
    node_folder.text = 'VOC20077'
    node_filename = SubElement(node_root, 'filename')
    node_filename.text = image
    #####################################################################
    # 图片的长宽
    node_size = SubElement(node_root, 'size')
    node_width = SubElement(node_size, 'width')
    node_width.text = str(width)
    node_height = SubElement(node_size, 'height')
    node_height.text = str(height)
    node_depth = SubElement(node_size, 'depth')
    node_depth.text = '3'
    #####################################################################
    # 物体
    for sub_obj in split_object_ans:
        # print(sub_obj)
        node_object = SubElement(node_root, 'object')
        node_name = SubElement(node_object, 'name')
        node_name.text = sub_obj['name']
        node_pose = SubElement(node_object, 'pose')
        node_pose.text = 'Left'
        node_truncates = SubElement(node_object, 'truncated')
        node_truncates.text = '1'
        node_difficult = SubElement(node_object, 'difficult')
        node_difficult.text = '0'
        node_bndbox = SubElement(node_object, 'bndbox')
        node_xmin = SubElement(node_bndbox, 'xmin')
        node_xmin.text = str(sub_obj['bbox'][0])
        node_ymin = SubElement(node_bndbox, 'ymin')
        node_ymin.text = str(sub_obj['bbox'][1])
        node_xmax = SubElement(node_bndbox, 'xmax')
        node_xmax.text = str(sub_obj['bbox'][2])
        node_ymax = SubElement(node_bndbox, 'ymax')
        node_ymax.text = str(sub_obj['bbox'][3])
    xml = tostring(node_root, pretty_print=True)
    dom = parseString(xml)
    with open(save_xml, 'w', encoding='UTF-8') as fh:
        dom.writexml(fh, indent='', addindent='\t', encoding='UTF-8')
# print(trainval)
def f_save_image(old_image,image,xmin,ymin,xmax,ymax):
    # python裁剪图片并保存
    # 读取图片
    img_1 = Image.open(old_image)
    # 设置裁剪的位置
    crop_box = (xmin, ymin, xmax,ymax)  # xmin,ymin,xmax,ymax
    # 裁剪图片
    img_2 = img_1.crop(crop_box)
    crop_save_path=save_image_root+'/'+image
    img_2.save(crop_save_path)


def f_save_txt(save_txt_root,ans_object,image):
    image=image.split('.')[0]
    # print('image ',image)
    cl_txt_name=[]
    # print(split)
    # print('class_name_train  ',class_name_train)
    if split=='train':
        cl_txt_name=class_name_train
    elif split=='val':
        cl_txt_name=class_name_val
    have_cl=[]
    for sub_obj in ans_object:
        have_cl.append(sub_obj['name'])
    # print('have_cl  ',have_cl)
    # print('cl_txt_name ',cl_txt_name,len(cl_txt_name))
    for cl_path in cl_txt_name:
        txt_path = save_txt_root + '/' + cl_path
        # print('txt_path  ',txt_path)
        with open(txt_path,'a') as f:
            cur_txt_cls=cl_path.split('_')[0]
            # print(cur_txt_cls)
            strr=''
            if cur_txt_cls in have_cl:
                strr=image+'  '+str(1)+'\n'
            else:
                strr=image+'  '+str(0)+'\n'
            # print(strr)
            f.write(strr)
            f.close()
# 定义函数，第一个参数是缩放比例，第二个参数是需要显示的图片组成的元组或者列表
def ManyImgs(scale, imgarray):
    rows = len(imgarray)         # 元组或者列表的长度
    cols = len(imgarray[0])      # 如果imgarray是列表，返回列表里第一幅图像的通道数，如果是元组，返回元组里包含的第一个列表的长度
    # print("rows=", rows, "cols=", cols)

    # 判断imgarray[0]的类型是否是list
    # 是list，表明imgarray是一个元组，需要垂直显示
    rowsAvailable = isinstance(imgarray[0], list)

    # 第一张图片的宽高
    width = imgarray[0][0].shape[1]
    height = imgarray[0][0].shape[0]
    # print("width=", width, "height=", height)


    # 如果传入的是一个元组
    if rowsAvailable:
        for x in range(0, rows):
            for y in range(0, cols):
                # 遍历元组，如果是第一幅图像，不做变换
                if imgarray[x][y].shape[:2] == imgarray[0][0].shape[:2]:
                    imgarray[x][y] = cv2.resize(imgarray[x][y], (0, 0), None, scale, scale)
                # 将其他矩阵变换为与第一幅图像相同大小，缩放比例为scale
                else:
                    imgarray[x][y] = cv2.resize(imgarray[x][y], (imgarray[0][0].shape[1], imgarray[0][0].shape[0]), None, scale, scale)
                # 如果图像是灰度图，将其转换成彩色显示
                if  len(imgarray[x][y].shape) == 2:
                    imgarray[x][y] = cv2.cvtColor(imgarray[x][y], cv2.COLOR_GRAY2BGR)

        # 创建一个空白画布，与第一张图片大小相同
        imgBlank = np.zeros((height, width, 3), np.uint8)
        hor = [imgBlank] * rows   # 与第一张图片大小相同，与元组包含列表数相同的水平空白图像
        for x in range(0, rows):
            # 将元组里第x个列表水平排列
            hor[x] = np.hstack(imgarray[x])
        ver = np.vstack(hor)   # 将不同列表垂直拼接
    # 如果传入的是一个列表
    else:
        # 变换操作，与前面相同
        for x in range(0, rows):
            if imgarray[x].shape[:2] == imgarray[0].shape[:2]:
                imgarray[x] = cv2.resize(imgarray[x], (0, 0), None, scale, scale)
            else:
                imgarray[x] = cv2.resize(imgarray[x], (imgarray[0].shape[1], imgarray[0].shape[0]), None, scale, scale)
            if len(imgarray[x].shape) == 2:
                imgarray[x] = cv2.cvtColor(imgarray[x], cv2.COLOR_GRAY2BGR)
        # 将列表水平排列
        # hor = np.hstack(imgarray)
        hor = np.vstack(imgarray)
        ver = hor
    return ver

with open(root_txt,'r') as f:
    num=0
    while True:
        name=f.readline()
        if not name:
            break
        name=name.replace('\n','')
        num=num+1
        # if num>10:  #先切分几个，小批量处理处理，方便调试
        #     break
        old_xml=old_xml_root+'/'+name+'.xml'
        old_image=old_image_root+'/'+name+'.jpg'
        objects=read_xml(old_xml)
        # print(object)
        #以每张图片中的每一个物体为中心，进行切分
        show=1   #对切分结果进行可视化展示
        sun_obj=1
        show_tuple=[]


        if show:
            print('old_image  ',old_image,os.path.exists(old_image))
            save_show_root='Downloads/data/test_crop_show'
            im=cv2.imread('Downloads/data/voc/VOCtrainval_11-May-2012/VOCdevkit/VOC2012/JPEGImages/2008_000008.jpg')
            print(im.size)
            if not os.path.exists(save_show_root):
                os.mkdir(save_show_root)
            image_whole=cv2.imread(old_image)
            # print(image_whole.size())
            for ob in objects:
                cv2.rectangle(image_whole,(ob['bbox'][0],ob['bbox'][1]),(ob['bbox'][2],ob['bbox'][3]),(0,255,0),2)
            if not os.path.exists(save_show_root+'/'+name):
                os.mkdir(save_show_root+'/'+name)
            cv2.imwrite(save_show_root+'/'+name+'/'+'whole.jpg', image_whole)
            show_tuple.append(image_whole)




        for whole_image_ob_idx,ob in enumerate(objects):  #图中的每个instance都会当一次中心处理点
            ##backbone下采样16，面积＜16*16的物体不进行切分
            if(whole_image_ob_idx%50==0):
                print('{}/{}'.format(whole_image_ob_idx,len(objects)))
            if int(ob['difficult'])==1:
                continue
            area=(abs(int(ob['bbox'][2])-int(ob['bbox'][0]))*(abs(int(ob['bbox'][3])-int(ob['bbox'][1]))))
            if (16*16)>area:
                continue
            ####开始切分，切分结果
            split_ans=split_image(old_image,ob,objects)
            for cur_ob_pathch,sub_split_ans in enumerate(split_ans): #以当前instance为中心处理点，会产生几个子图
                xmin, ymin, xmax, ymax, width, height, ans_object =sub_split_ans['xmin'],sub_split_ans['ymin'],sub_split_ans['xmax'],sub_split_ans['ymax'],sub_split_ans['width'],sub_split_ans['height'],sub_split_ans['ans_object']
                ########保存的地址和信息
                save_xml = save_xml_root + '/' + name + '_' + str(id) + '_' + str(whole_image_ob_idx)+'_'+str(cur_ob_pathch) + '.xml'
                image = name + '_' + str(id) + '_' + str(whole_image_ob_idx) +'_'+str(cur_ob_pathch)+ '.jpg'
                save_image = save_image_root + '/' + name + '_' + str(id) + '_' + str(whole_image_ob_idx) +'_'+str(cur_ob_pathch)+ '.jpg'
                #####################################################################
                f_save_xml(width, height, image, ans_object, save_xml)
                f_save_image(old_image, image, xmin, ymin, xmax, ymax)
                # f_save_txt(save_txt_root, ans_object, image)
                save_txt_con=name + '_' + str(id) + '_' + str(whole_image_ob_idx) +'_'+str(cur_ob_pathch)
                if split=='train':
                    with open(save_txt_root+'/'+'train.txt','a') as train_f:
                        train_f.write(save_txt_con+'\n')
                        train_f.close()
                    with open(save_txt_root+'/'+'trainval.txt','a') as train_val_f:
                        train_val_f.write(save_txt_con+'\n')
                        train_val_f.close()
                elif split=='val':
                    with open(save_txt_root+'/'+'val.txt','a')as val_f:
                        val_f.write(save_txt_con+'\n')
                        val_f.close()
                    with open(save_txt_root+'/'+'trainval.txt','a') as train_val_f:
                        train_val_f.write(save_txt_con+'\n')
                        train_val_f.close()
                elif split=='test':
                    with open(save_txt_root+'/'+'test.txt','a')as test_f:
                        test_f.write(save_txt_con+'\n')
                        test_f.close()
                if show:
                    # save_show_root =
                    r=save_image
                    print(r)
                    image_sub = cv2.imread(r)
                    for ob in ans_object:
                        cv2.rectangle(image_sub, (ob['bbox'][0], ob['bbox'][1]), (ob['bbox'][2], ob['bbox'][3]),
                                      (0, 255, 0), 2)
                    s=save_show_root + '/' + name+'/'+str(whole_image_ob_idx)
                    # print(s)
                    # print(s +'/'+str(cur_ob_pathch) + '_sub.jpg')
                    # print(image_sub.size)
                    if not os.path.exists(s ):
                        os.mkdir(s )
                    cv2.imwrite(s +'/'+str(cur_ob_pathch) + '_sub.jpg', image_sub)
                    show_tuple.append(image_sub)
        # 第二种方法，调用我们自己编写的一个函数实现任意大小、任意通道数的图像堆叠显示
        # 以原来的0.5倍并排显示两张图片
        # stackedimageh = ManyImgs(0.2, show_tuple)

        # 垂直显示
        # stackedimagev = ManyImgs(0.2, ([imgflower, imgwomen], [imgwomen, imgflower]))

        # 如果只有图片数量为奇数，并希望能够垂直显示，可以创建一个空白图像
        Blankimg = np.zeros((800, 800), np.uint8)  # 大小可以任意函数会将其强制转换
        stackedimageb = ManyImgs(0.5, show_tuple)

        # cv2.namedWindow("stackedimageh")
        # cv2.imshow("stackedimageh", stackedimageh)
        #
        # cv2.namedWindow("stackedimagev")
        # cv2.imshow("stackedimagev", stackedimagev)

        cv2.namedWindow("stackedimageb")
        cv2.imshow("stackedimageb", stackedimageb)

        cv2.waitKey(0)
        cv2.destroyAllWindows()



    f.close()
