#!/usr/bin/env python 
# -*- coding: utf-8 -*-
"""
@project: 
@File    : roofpv_gender
@Author  : qiqq
@create_time    : 2023/7/2 16:28
"""

#根据屋頂和光伏结果结果合成多目标图

from collections import namedtuple  #namedtuple创建一个和tuple类似的对象，而且对象拥有可访问的属性

import  os
from tqdm import tqdm
import matplotlib.pyplot as  plt
import  numpy as np
from PIL import Image
Image.MAX_IMAGE_PIXELS = None
Cls = namedtuple('cls', ['name', 'id', 'color'])
Clss = [
    Cls('c0', 0, (0, 0, 0)),
    Cls('c1', 1, (0, 255, 0)),
    Cls('c2', 1, (255, 0, 0)),
        ]



def get_putpalette(Clss, color_other=[0, 0, 0]):
    '''
    灰度图转8bit彩色图
    :param Clss:颜色映射表
    :param color_other:其余颜色设置
    :return:
    '''
    putpalette = []
    for cls in Clss:
        putpalette += list(cls.color)
    putpalette += color_other * (255 - len(Clss))
    return putpalette

def Convert_get_color8bit(gt):
    '''
    灰度图转8bit彩色图
    :param grays_path:  灰度图文件路径
    :param colors_path: 彩色图文件路径
    :return:
    '''
    bin_colormap = get_putpalette(Clss)
    gt = Image.fromarray(gt)
    gt = gt.convert("P")
    gt.putpalette(bin_colormap)
    return gt

bin_colormap = get_putpalette(Clss)
PALETTE = [(0, 0, 0), (255, 0, 0), (0, 255, 0),(0,0,255)]
palette = [0, 0, 0, 255, 0, 0, 0, 255, 0,0, 0, 255]

def roofpvgener(roofresult,pvresult,size):
    '''



    大概的思路：
    屋顶数据集r0 背景 r1屋顶
    光伏数据集p0 背景 p1光伏

    生成一张空白图，逐个像素的操作
    r0  p0 -->背景       0 黑色
     r1   p0  --->屋顶    1(应该是空白屋顶了)  #红色
    r0   p1 -->非屋顶光伏   2 #绿色

    r1   p1--- 屋顶光伏   3 #蓝色
    '''
    roof_nroof_pv = np.zeros(size)
    # 使用numpy的条件判断语句进行逐元素的判断和赋值
    roof_nroof_pv[(pvresult == 0) & (roofresult == 0)] = 0 #黑色
    roof_nroof_pv[(pvresult == 0) & (roofresult == 1)] = 1 #红色不带光伏的屋顶
    roof_nroof_pv[(pvresult == 1) & (roofresult == 0)] = 2
    roof_nroof_pv[(pvresult == 1) & (roofresult == 1)] = 3

    return roof_nroof_pv


if __name__ == '__main__':

    ##读取
    model_name="manet"
    pv_path="/"#光伏分割结果
    roof_path="/"#屋顶分割结果
    imagedir="/"##原图路径

    rnrpv_outdir="./"+model_name+"/"+model_name+"_output"
    rnrpv_fusindir="./"+model_name+"/"+model_name+"_fusion"


    useFusion=False


    pv_list= os.listdir(pv_path)
    roof_list= os.listdir(roof_path)
    assert len(pv_list)==len(roof_list)

    if not os.path.exists(model_name):
        os.makedirs(model_name)

    if not os.path.exists(rnrpv_outdir):
        os.makedirs(rnrpv_outdir)

    for i in tqdm(pv_list):
        imagename= i.split(".")[0]
        pvimage = Image.open(os.path.join(pv_path,i))
        roofimage=Image.open(os.path.join(roof_path,i))
        pv =np.array(pvimage)
        roof= np.array(roofimage)
        orininal_h, orininal_w=pv.shape

        result =roofpvgener(roof,pv,size=(orininal_h, orininal_w))

        result=np.uint8(result)
        result = Image.fromarray(result).convert('P')  # 原来的
        result.putpalette(palette)
        result.save(os.path.join(rnrpv_outdir, imagename + ".png"))

        if  useFusion:
            if not os.path.exists(rnrpv_fusindir):
                os.makedirs(rnrpv_fusindir)

            sourceimage =Image.open(os.path.join(imagedir,imagename+".png"))
            result = np.reshape(np.array(PALETTE, np.uint8)[np.reshape(result, [-1])],
                                 [orininal_h, orininal_w, -1])


            image = Image.fromarray(np.uint8(result))
            fusion = Image.blend(sourceimage, image, 0.3)
            fusion.save(os.path.join(rnrpv_fusindir,imagename+".png"))

    print("完成")