"""
author:Hatcher
e-mail:hcwang@semi.ac.cn
功能：扫描文件夹下面的所有图像，并以文件夹为单位统计鸟类数量
文件夹格式：
    images
        data
            images.jpg
输出：每个文件夹中保存一个txt文件，每个种类鸟类的数量，所有鸟类种类的数量

"""
import os
import torch
import cv2
from System_I import Detection_Y5,Fine_grained_T
from common import global_para as gp
import json
from tqdm import tqdm
from creat_xml_fordetection import make_xml_scan
root_path = "/data5/Project/scanData/YNDC/scanData/"
save_root = "/data5/Project/scanData/YNDC/result/"
save_json = "/data5/Project/scanData/YNDC/JSON/"
os.environ['CUDA_VISIBLE_DEVICES'] = '7'
devices = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
Detection = Detection_Y5(devices,root_path)
Fine = Fine_grained_T(devices)
if not os.path.exists(os.path.join(save_json, "ForDataBase")):
    os.makedirs(os.path.join(save_json, "ForDataBase"))
if not os.path.exists(os.path.join(save_json, "ForTest")):
    os.makedirs(os.path.join(save_json, "ForTest"))
for file_paths in tqdm(os.listdir(root_path)):
    box_all = []
    label_all = []
    record_time = file_paths.replace("_","-",2)
    record_time = record_time.replace("_"," ",1)
    record_time = record_time.replace("_",":")
    file_path = os.path.join(root_path,file_paths)
    labels_fineall = []
    resultDict = {}
    allimageresult =[0 for i in range(len(gp.BHG_chineseName))]
    if not os.path.exists(os.path.join(save_root,file_paths)):
        os.makedirs(os.path.join(save_root,file_paths))
    for imagepath in os.listdir(file_path):
        image = cv2.imread(os.path.join(file_path,imagepath))
        img_show,boxes_fine, labels_fine = Detection.inference_scan(image=image,fine_grained=Fine)
        for m in range(len(labels_fine)):
            box_all.append(boxes_fine[m])
            label_all.append(gp.BHG_chineseName[int(labels_fine[m])])
        cv2.imwrite(os.path.join(save_root,file_paths,imagepath),img_show)
        oneimageresultcdict = {}
        oneimageresult =[0 for i in range(len(gp.BHG_chineseName))]
        for i in labels_fine:
            oneimageresult[int(i)]+=1
            allimageresult[int(i)]+=1
        for i in range(len(oneimageresult)):
            if not (int(oneimageresult[i]) == 0):
                oneimageresultcdict[gp.BHG_chineseName[i]]=int(oneimageresult[i])
        resultDict[imagepath] = oneimageresultcdict
    allimageresultcdict = {}
    for i in range(len(allimageresult)):
            if not (int(allimageresult[i]) == 0):
                allimageresultcdict[gp.BHG_chineseName[i]]=int(allimageresult[i])
    resultDict[file_paths] = allimageresultcdict
    json_str = json.dumps(resultDict, indent=4)
    with open(os.path.join(save_json,"ForTest",file_paths + "_Test.json"), 'w',encoding='utf-8') as json_file:
        #json_file.write(json_str)
        json.dump(resultDict,json_file,indent=4,ensure_ascii=False)
    json_conversion = make_xml_scan(record_time=record_time, serial_number="iDS-2DF8C840IXS-AFW/SP20220507CCCHJ90488961",boxes=box_all,labels=label_all, savename=os.path.join(save_json,"ForDataBase" ,file_paths + "_DataBase.json"))