import sys
sys.path.extend(['..','../AIlib' ])

from AI import AI_process,AI_process_forest,get_postProcess_para
import cv2,os,time
from segutils.segmodel import SegModel
from models.experimental import attempt_load
from utils.torch_utils import select_device
from utilsK.queRiver import get_labelnames,get_label_arrays
import numpy as np
import torch
from utilsK.masterUtils import get_needed_objectsIndex
def river_demo():
    ##预先设置的参数
    device_='1'  ##选定模型，可选 cpu,'0','1'
    
    ##以下参数目前不可改
    #Detweights = "../AIlib/weights/yolov5/class5/best_5classes.pt" 
    #labelnames = "../AIlib/weights/yolov5/class5/labelnames.json"
    
    Detweights = "../AIlib/weights/yolov5/class8/bestcao.pt" ##升级后的检测模型
    labelnames = "../AIlib/weights/yolov5/class8/labelnames.json" ##对应类别表
    
    seg_nclass = 2 
    #Segweights = "../AIlib/weights/BiSeNet/checkpoint.pth"
    Segweights = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640.pth' ##升级的分割模型
    
    postFile= '../AIlib/conf/para.json'
    
    conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile)
   
    ####模型选择参数用如下：
    mode_paras=[
        {"id":"0","config":{"k1":"v1","k2":"v2"}},
        {"id":"1","config":{"k1":"v1","k2":"v2"}},
        {"id":"2","config":{"k1":"v1","k2":"v2"}},
        {"id":"3","config":{"k1":"v1","k2":"v2"}},
        {"id":"4","config":{"k1":"v1","k2":"v2"}},
        {"id":"5","config":{"k1":"v1","k2":"v2"}},
        {"id":"6","config":{"k1":"v1","k2":"v2"}},
        {"id":"7","config":{"k1":"v1","k2":"v2"}},
    ]

    allowedList,allowedList_string=get_needed_objectsIndex(mode_paras)
    #allowedList=[0,1,2,3]
    
    

    ##加载模型，准备好显示字符
    device = select_device(device_)
    names=get_labelnames(labelnames)
    
    label_arraylist = get_label_arrays(names,rainbows,outfontsize=40,fontpath="../AIlib/conf/platech.ttf")
    half = device.type != 'cpu'  # half precision only supported on CUDA
    model = attempt_load(Detweights, map_location=device)  # load FP32 model
    if half: model.half()
    segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device)
    
    
    
    ##图像测试
    #url='images/examples/20220624_响水河_12300_1621.jpg'
    impth = 'images/slope/'
    outpth = 'images/results/'
    folders = os.listdir(impth)
    for i in range(len(folders)):
        imgpath = os.path.join(impth, folders[i])
        im0s=[cv2.imread(imgpath)]
        time00 = time.time()
        p_result,timeOut = AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList)
        time11 = time.time()
        image_array = p_result[1]
        cv2.imwrite(  os.path.join( outpth,folders[i] ) ,image_array )
        print('----AIprocess:%s, %.1f, %s'%(folders[i], (time11 - time00) * 1000,timeOut))
        
def forest_demo():
    ##预先设置的参数
    device_='1'  ##选定模型，可选 cpu,'0','1'
    
    ##以下参数目前不可改
    Detweights = "../AIlib/weights/forest/best.pt"
    labelnames = "../AIlib/weights/forest/labelnames.json"
    postFile= '../AIlib/conf/para.json'
    
    conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile)
    
    ####模型选择参数用如下：
    mode_paras=[
        {
            "id":"0",
            "config":{
            "k1":"v1",
            "k2":"v2"
            }
        },
        {
            "id":"1",
            "config":{
            "k1":"v1",
            "k2":"v2"
            }
        }
    ]

    allowedList,allowedList_string=get_needed_objectsIndex(mode_paras)
    #allowedList=[0,1,2,3]
    
    print('####line108###')

    ##只加载检测模型，准备好显示字符
    device = select_device(device_)
    names=get_labelnames(labelnames)
    imageW=4915;###默认是1920，在森林巡检的高清图像中是4920
    outfontsize=int(imageW/1920*40);###
    label_arraylist = get_label_arrays(names,rainbows,outfontsize=outfontsize,fontpath="../AIlib/conf/platech.ttf")
    half = device.type != 'cpu'  # half precision only supported on CUDA
    model = attempt_load(Detweights, map_location=device)  # load FP32 model
    if half: model.half()
    segmodel = None 
    
    
    
    ##图像测试
    #url='images/examples/20220624_响水河_12300_1621.jpg'
    impth = 'images/forest/'
    outpth = 'images/results/'
    folders = os.listdir(impth)
    for i in range(len(folders)):
    #for i in range(2):
        imgpath = os.path.join(impth, folders[i])
        im0s=[cv2.imread(imgpath)]
        time00 = time.time()
        #使用不同的函数。每一个领域采用一个函数
        p_result,timeOut = AI_process_forest(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList)
        time11 = time.time()
        image_array = p_result[1]
        cv2.imwrite(  os.path.join( outpth,folders[i] ) ,image_array )
        
        print('----image:%s, process:%s ,save:%s, %s'%(folders[i],(time11-time00) * 1000, (time.time() - time11) * 1000,timeOut) )

    
    
    
        
if __name__=="__main__":
    river_demo()
    #forest_demo()