from genericpath import exists
import os
import re
from posixpath import join
from sys import path
import numpy as np
from numpy.lib.histograms import _ravel_and_check_weights
from models import spinal_net
import torch
from decoder import DecDecoder
from data.dataset import BaseDataset_ap, BaseDataset_lat, BaseDataset_csvl
import cv2
from cobb_evaluate import *
from data.pre2post import pre2post
from models import regression
from torch.autograd import Variable
import matplotlib.pyplot as plt
from modeling.deeplab import *
import data.custom_transforms as tr
from PIL import Image
from torchvision import transforms
import cv2 as cv

def load_model(model, resume):
    checkpoint = torch.load(resume, map_location=lambda storage, loc: storage)
    print('loaded weights from {}, epoch {}'.format(resume, checkpoint['epoch']))
    state_dict_ = checkpoint['state_dict']
    model.load_state_dict(state_dict_, strict=False)
    return model

def ap_report():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    heads = {'hm': 1,
            'reg': 2*1,
            'wh': 2*4,}

    model_ap = spinal_net.SpineNet(heads=heads,
                                    pretrained=True,
                                    down_ratio=4,
                                    final_kernel=1,
                                    head_conv=256)

    model_lat = spinal_net.SpineNet(heads=heads,
                                    pretrained=True,
                                    down_ratio=4,
                                    final_kernel=1,
                                    head_conv=256)

    # self.num_classes = args.num_classes
    decoder = DecDecoder(K=17, conf_thresh=0.2)
    # dataset = BaseDataset

    ap_save_path = './models/pretrain/ap_model_best.pth'
    model_ap = load_model(model_ap, ap_save_path)
    model_ap = model_ap.to(device)
    model_ap.eval()

    lat_save_path = './models/pretrain/lat_model_best.pth'
    model_lat = load_model(model_lat, lat_save_path)
    model_lat = model_lat.to(device)
    model_lat.eval()

    # dataset_module = dataset
    dsets = BaseDataset_ap(data_dir='./data/DATAEXPORT_TEST/',
                            phase='test',
                            input_h=1024,
                            input_w=512,
                            down_ratio=4
                            )

    data_loader = torch.utils.data.DataLoader(dsets,
                                            batch_size=1,
                                            shuffle=False,
                                            num_workers=1,
                                            pin_memory=True)

    for cnt, data_dict in enumerate(data_loader):
        images = data_dict['images'][0]
        patient_id = data_dict['patient_id'][0]
        images = images.to(device)
        print("###############################################")
        print('processing {}/{} patient ... {}'.format(cnt, len(data_loader), patient_id))

        index = ['ap', 'bending_l', 'bending_r', 'lat']
        angles = []
        pts_all = []
        ori_image_points_all = []
        for i in range(4):

            print('processing image {}'.format(index[i]))
            if i == 3:
                model = model_lat
            else:
                model = model_ap
            #####################处理 ap#################
            # print(images.shape)
            with torch.no_grad():
                output = model(images[i])
                hm = output['hm']
                wh = output['wh']
                reg = output['reg']

            torch.cuda.synchronize(device)
            pts2 = decoder.ctdet_decode(hm, wh, reg)   # 17, 11
            pts0 = pts2.copy()
            pts0[:,:10] *= 4

            # print('totol pts num is {}'.format(len(pts2)))

            ori_images = dsets.load_images(dsets.patient_ids.index(patient_id))
            ori_image = ori_images[i]
            k = 256
            ori_image = ori_image * k

            ori_image_regress = cv2.resize(ori_image, (512, 1024))
            ori_image_points = ori_image_regress.copy()

            h,w,c = ori_image.shape
            pts0 = np.asarray(pts0, np.float32)
            # pts0[:,0::2] = pts0[:,0::2]/args.input_w*w
            # pts0[:,1::2] = pts0[:,1::2]/args.input_h*h
            sort_ind = np.argsort(pts0[:,1])
            pts0 = pts0[sort_ind]

            # ori_image_regress, ori_image_points = draw_points.draw_landmarks_regress_test(pts0,
            #                                                                             ori_image_regress,
            #                                                                             ori_image_points)

            pts3 = []
            for p in pts0:
                point1 = [p[2], p[3]]
                point2 = [p[4], p[5]]
                point3 = [p[6], p[7]]
                point4 = [p[8], p[9]]
                pts3.extend([point1,point2,point3,point4])
            pts3 = np.array(pts3)
            pts_all.append(pts3)
            ori_image_points_all.append(ori_image_points)

        angles, ptss, image = cal_angles(pts_all[0], ori_image_points_all[0])

        if not os.path.exists(os.path.join('./results/', patient_id)):
            os.makedirs(os.path.join('./results/', patient_id))
        cv2.imwrite(os.path.join('./results/', patient_id, index[0]+'.jpg'), image)

        report = open(os.path.join('./results', patient_id, 'report.txt'), 'a')

        ptss2 = []
        for x in ptss:
            ptss2.append((x[0]+1,x[1]+1))
        print('三段脊椎角度和对应脊椎段：\n',angles, ptss2, file=report)

        ################处理bending, lat################
        scoliosis_type = []
        scoliosis_seg = ['PT', 'MT', 'TL']
        for i, angle in enumerate(angles):
            print("-----{}段诊断流程".format(scoliosis_seg[i]), file=report)
            if angle < 25:
                print("#####{}为非结构弯：脊椎{}-{}：{}".format(scoliosis_seg[i], ptss[i][0]+1, ptss[i][1]+1, round(angles[i],2)), file=report)
                scoliosis_type.append(False)
            else:
                angle1, image1 = cal_angles_with_pos(pts_all[1], ori_image_points_all[1], ptss[i])
                angle2, image2 = cal_angles_with_pos(pts_all[2], ori_image_points_all[2], ptss[i])

                angle_bending = angle1 if angle1 < angle2 else angle2
                image = image1 if angle1 < angle2 else image2
                i_lr = 1 if angle1 < angle2 else 2
                cv2.imwrite(os.path.join('./results/', patient_id, scoliosis_seg[i]+'_'+index[i_lr]+'.jpg'), image)

                if angle_bending > 25:
                    print("#####{}为结构弯：直立位{}-{}：{}，bending位：{}".format(scoliosis_seg[i],ptss[i][0]+1, ptss[i][1]+1, round(angles[i],2), round(angle_bending,2)), file=report)
                    scoliosis_type.append(True)
                else:
                    if i == 0:
                        pos = (1, 4)
                    else:
                        pos = (9, 13)
                    angle_lat, image = cal_angles_with_pos(pts_all[3], ori_image_points_all[3], pos)
                    cv2.imwrite(os.path.join('./results/', patient_id, scoliosis_seg[i]+'_'+index[3]+'.jpg'), image)
                    if angle_lat < 20:
                        print("#####{}为非结构弯：脊椎{}-{}：{}，bending位：{}，矢状面：{}-{}：{}".format(scoliosis_seg[i],ptss[i][0]+1, ptss[i][1]+1, round(angles[i],2), round(angle_bending,2), pos[0],pos[1],round(angle_lat,2)), file=report)
                        scoliosis_type.append(False)
                    else:
                        print("#####{}为结构弯：脊椎{}-{}：{}，bending位：{}，矢状面：{}-{}：{}".format(scoliosis_seg[i],ptss[i][0]+1, ptss[i][1]+1, round(angles[i],2), round(angle_bending,2), pos[0],pos[1],round(angle_lat,2)), file=report)
                        scoliosis_type.append(True)

def ap_report_v2():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    heads = {'hm': 1,
            'reg': 2*1,
            'wh': 2*4,}
    print(device)

    model_ap = spinal_net.SpineNet(heads=heads,
                                    pretrained=True,
                                    down_ratio=4,
                                    final_kernel=1,
                                    head_conv=256)

    model_lat = spinal_net.SpineNet(heads=heads,
                                    pretrained=True,
                                    down_ratio=4,
                                    final_kernel=1,
                                    head_conv=256)

    # self.num_classes = args.num_classes
    decoder = DecDecoder(K=17, conf_thresh=0.2)
    # dataset = BaseDataset

    ap_save_path = './models/pretrain/ap_model_best.pth'
    model_ap = load_model(model_ap, ap_save_path)
    model_ap = model_ap.to(device)
    model_ap.eval()

    lat_save_path = './models/pretrain/lat_model_best.pth'
    model_lat = load_model(model_lat, lat_save_path)
    model_lat = model_lat.to(device)
    model_lat.eval()

    # dataset_module = dataset
    dsets = BaseDataset_ap(data_dir='./data/DATAEXPORT_TEST/',
                            phase='test',
                            input_h=1024,
                            input_w=512,
                            down_ratio=4
                            )

    data_loader = torch.utils.data.DataLoader(dsets,
                                            batch_size=1,
                                            shuffle=False,
                                            num_workers=1,
                                            pin_memory=True)

    for cnt, data_dict in enumerate(data_loader):
        images = data_dict['images'][0]
        patient_id = data_dict['patient_id'][0]
        images = images.to(device)
        print("###############################################")
        print('processing {}/{} patient ... {}'.format(cnt, len(data_loader), patient_id))

        index = ['ap', 'bending_l', 'bending_r', 'lat']
        angles = []
        pts_all = []
        ori_image_points_all = []
        for i in range(4):

            print('processing image {}'.format(index[i]))
            if i == 3:
                model = model_lat
            else:
                model = model_ap
            #####################处理 ap#################
            # print(images.shape)
            with torch.no_grad():
                output = model(images[i])
                hm = output['hm']
                wh = output['wh']
                reg = output['reg']

            # torch.cuda.synchronize(device)
            pts2 = decoder.ctdet_decode(hm, wh, reg)   # 17, 11
            pts0 = pts2.copy()
            pts0[:,:10] *= 4

            # print('totol pts num is {}'.format(len(pts2)))

            ori_images = dsets.load_images(dsets.patient_ids.index(patient_id))
            ori_image = ori_images[i]
            k = 256
            ori_image = ori_image * k

            ori_image_regress = cv2.resize(ori_image, (512, 1024))
            ori_image_points = ori_image_regress.copy()

            h,w,c = ori_image.shape
            pts0 = np.asarray(pts0, np.float32)
            # pts0[:,0::2] = pts0[:,0::2]/args.input_w*w
            # pts0[:,1::2] = pts0[:,1::2]/args.input_h*h
            sort_ind = np.argsort(pts0[:,1])
            pts0 = pts0[sort_ind]

            # ori_image_regress, ori_image_points = draw_points.draw_landmarks_regress_test(pts0,
            #                                                                             ori_image_regress,
            #                                                                             ori_image_points)

            pts3 = []
            for p in pts0:
                point1 = [p[2], p[3]]
                point2 = [p[4], p[5]]
                point3 = [p[6], p[7]]
                point4 = [p[8], p[9]]
                pts3.extend([point1,point2,point3,point4])
            pts3 = np.array(pts3)
            pts_all.append(pts3)
            ori_image_points_all.append(ori_image_points)

        angles, ptss, image = cal_angles(pts_all[0], ori_image_points_all[0])

        if not os.path.exists(os.path.join('./results/', patient_id)):
            os.makedirs(os.path.join('./results/', patient_id))
        cv2.imwrite(os.path.join('./results/', patient_id, index[0]+'.jpg'), image)

        report = open(os.path.join('./results', patient_id, 'report.txt'), 'a')

        ptss2 = []
        for x in ptss:
            ptss2.append((x[0]+1,x[1]+1))
        print('三段脊椎角度和对应脊椎段：\n',angles, ptss2, file=report)

        t2t5angle_lat, t2t5image = cal_angles_with_pos(pts_all[3], ori_image_points_all[3].copy(), (1,4))
        cv2.imwrite(os.path.join('./results/', patient_id, 'lat_t2_t5.jpg'), t2t5image)
        t10l2angle_lat, t10l2image = cal_angles_with_pos(pts_all[3], ori_image_points_all[3].copy(), (9,13))
        cv2.imwrite(os.path.join('./results/', patient_id, 'lat_t10_l2.jpg'), t10l2image)
        t5t12angle_lat, t5t12image = cal_angles_with_pos(pts_all[3], ori_image_points_all[3].copy(), (4,11))
        cv2.imwrite(os.path.join('./results/', patient_id, 'lat_t5_t12.jpg'), t5t12image)
        ################处理bending, lat################
        scoliosis_type = []
        scoliosis_seg = ['PT', 'MT', 'TL']
        for i, angle in enumerate(angles):
            print("-----{}段诊断流程".format(scoliosis_seg[i]), file=report)
            bl_angle, bl_image = cal_angles_with_pos(pts_all[1], ori_image_points_all[1].copy(), ptss[i])
            br_angle, br_image = cal_angles_with_pos(pts_all[2], ori_image_points_all[2].copy(), ptss[i])
            cv2.imwrite(os.path.join('./results/', patient_id, scoliosis_seg[i]+'_bl.jpg'), bl_image)
            cv2.imwrite(os.path.join('./results/', patient_id, scoliosis_seg[i]+'_br.jpg'), br_image)
            angle_bending = bl_angle if bl_angle < br_angle else br_angle
            j1 = segType_Trans(ptss[i][0] + 1)
            j2 = segType_Trans(ptss[i][1] + 1)
            if i == 0:
                pos = ("T2", "T5")
                angle_lat = t2t5angle_lat
            else:
                pos = ("T10", "L2")
                angle_lat = t10l2angle_lat

            if angle < 25:
                print("#####{}为非结构弯：脊椎{}-{}：{}，bending位：{}，矢状面：{}-{}：{}".format(scoliosis_seg[i],j1, j2, round(angles[i],2), round(angle_bending,2), pos[0],pos[1],round(angle_lat,2)), file=report)
                scoliosis_type.append(False)
            else:
                # angle1, image1 = cal_angles_with_pos(pts_all[1], ori_image_points_all[1], ptss[i])
                # angle2, image2 = cal_angles_with_pos(pts_all[2], ori_image_points_all[2], ptss[i])

                # image = bl_image if bl_angle < br_angle else br_image
                # i_lr = 1 if bl_angle < br_angle else 2

                if angle_bending > 25:
                    print("#####{}为结构弯：脊椎{}-{}：{}，bending位：{}，矢状面：{}-{}：{}".format(scoliosis_seg[i],j1, j2, round(angles[i],2), round(angle_bending,2), pos[0],pos[1],round(angle_lat,2)), file=report)
                    scoliosis_type.append(True)
                else:

                    # angle_lat, image = cal_angles_with_pos(pts_all[3], ori_image_points_all[3], pos)
                    # cv2.imwrite(os.path.join('./results/', patient_id, scoliosis_seg[i]+'_'+index[3]+'.jpg'), image)

                    if angle_lat < 20:
                        print("#####{}为非结构弯：脊椎{}-{}：{}，bending位：{}，矢状面：{}-{}：{}".format(scoliosis_seg[i],j1, j2, round(angles[i],2), round(angle_bending,2), pos[0],pos[1],round(angle_lat,2)), file=report)
                        scoliosis_type.append(False)
                    else:
                        print("#####{}为结构弯：脊椎{}-{}：{}，bending位：{}，矢状面：{}-{}：{}".format(scoliosis_seg[i],j1, j2, round(angles[i],2), round(angle_bending,2), pos[0],pos[1],round(angle_lat,2)), file=report)
                        scoliosis_type.append(True)

def lat_report():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    heads = {'hm': 1,
                'reg': 2*1,
                'wh': 2*4,}

    model = spinal_net.SpineNet(heads=heads,
                                        pretrained=True,
                                        down_ratio=4,
                                        final_kernel=1,
                                        head_conv=256)
    # num_classes = args.num_classes
    decoder = DecDecoder(K=17, conf_thresh=0.2)
    # dataset = BaseDataset
    lat_save_path = './models/pretrain/lat_model_best.pth'
    model = load_model(model, lat_save_path)
    model = model.to(device)
    model.eval()

    # dataset_module = dataset
    dsets = BaseDataset_lat(data_dir='./data/DATAEXPORT_TEST',
                            phase='test',
                            input_h=1024,
                            input_w=512,
                            down_ratio=4
                            )

    data_loader = torch.utils.data.DataLoader(dsets,
                                                batch_size=1,
                                                shuffle=False,
                                                num_workers=1,
                                                pin_memory=True)

    for cnt, data_dict in enumerate(data_loader):
        image = data_dict['images'][0]
        patient_id = data_dict['patient_id'][0]
        image = image.to(device)
        print("###############################################")
        print('processing {}/{} patient ... {}'.format(cnt, len(data_loader), patient_id))

        # print(image.shape)
        with torch.no_grad():
            output = model(image)
            hm = output['hm']
            wh = output['wh']
            reg = output['reg']

        # torch.cuda.synchronize(device)
        pts2 = decoder.ctdet_decode(hm, wh, reg)   # 17, 11
        pts0 = pts2.copy()
        pts0[:,:10] *= 4

        # print('totol pts num is {}'.format(len(pts2)))

        ori_image = dsets.load_images(dsets.patient_ids.index(patient_id))
        # ori_image = ori_images[3]
        # print(ori_image.shape)
        k = 256
        ori_image = ori_image * k

        ori_image_regress = cv2.resize(ori_image, (512, 1024))
        ori_image_points = ori_image_regress.copy()

        h,w,c = ori_image.shape
        pts0 = np.asarray(pts0, np.float32)
        # pts0[:,0::2] = pts0[:,0::2]/args.input_w*w
        # pts0[:,1::2] = pts0[:,1::2]/args.input_h*h
        sort_ind = np.argsort(pts0[:,1])
        pts0 = pts0[sort_ind]

        # ori_image_regress, ori_image_points = draw_points.draw_landmarks_regress_test(pts0,
        #                                                                             ori_image_regress,
        #                                                                             ori_image_points)
        if not os.path.exists(os.path.join('./results', patient_id)):
            os.mkdir(os.path.join('./results', patient_id))
        report = open(os.path.join('./results', patient_id, 'report.txt'), 'a')

        pts3 = []
        for p in pts0:
            point1 = [p[2], p[3]]
            point2 = [p[4], p[5]]
            point3 = [p[6], p[7]]
            point4 = [p[8], p[9]]
            pts3.extend([point1,point2,point3,point4])
        pts3 = np.array(pts3)
        print(pts3.shape)
        pos = (4, 11)
        angles, image = cal_angles_with_pos(pts3, ori_image_points, pos)
        # print(image.shape)
        if not os.path.exists('./results'):
            os.mkdir('./results/')

        cv2.imwrite(os.path.join('./results/', patient_id, 'lat.jpg'), image)
        print('T5-T12:{}'.format(round(angles,2)), file=report)
        if angles < 10:
            print('矢状面： -类', file=report)
        elif angles < 40:
            print('矢状面： N类', file=report)
        else:
            print('矢状面： +类', file=report)

def csvl_report():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    heads = {'hm': 1,
            'reg': 2*1,
            'wh': 2*4,}

    model_ap = spinal_net.SpineNet(heads=heads,
                                    pretrained=True,
                                    down_ratio=4,
                                    final_kernel=1,
                                    head_conv=256)

    # model_lat = spinal_net.SpineNet(heads=heads,
    #                                 pretrained=True,
    #                                 down_ratio=4,
    #                                 final_kernel=1,
    #                                 head_conv=256)

    # self.num_classes = args.num_classes
    decoder = DecDecoder(K=18, conf_thresh=0.2)
    # dataset = BaseDataset

    ap_save_path = './models/pretrain/s1.pth'
    model_ap = load_model(model_ap, ap_save_path)
    model_ap = model_ap.to(device)
    model_ap.eval()

    model_csvl = DeepLab(num_classes=2,
                            backbone='resnet',
                            output_stride=16,
                            sync_bn=None,
                            freeze_bn=False)
    csvl_save_path = './models/pretrain/csvl_model_best.pth.tar'
    model_csvl = load_model(model_csvl, csvl_save_path)
    model_csvl = model_csvl.to(device)
    model_csvl.eval()

    dsets = BaseDataset_csvl(data_dir='./data/DATAEXPORT_TEST/',
                            phase='test',
                            input_h=1024,
                            input_w=512,
                            down_ratio=4
                            )

    data_loader = torch.utils.data.DataLoader(dsets,
                                            batch_size=1,
                                            shuffle=False,
                                            num_workers=1,
                                            pin_memory=True)

    print("--load data done--")
    for cnt, data_dict in enumerate(data_loader):
        images = data_dict['images'][0]
        patient_id = data_dict['patient_id'][0]
        oimage = data_dict['oimage']
        # print(images.shape, oimage.shape)
        oimage = oimage.to(device)
        images = images.to(device)
        print('processing {}/{} image ... {}'.format(cnt, len(data_loader), patient_id))

        if not os.path.exists(os.path.join('./results', patient_id)):
            os.mkdir(os.path.join('./results', patient_id))
        report = open(os.path.join('./results', patient_id, 'report.txt'), 'a')
        path = os.path.join('./results', patient_id, 'report.txt')
        with torch.no_grad():
            output = model_ap(images)
            hm = output['hm']
            wh = output['wh']
            reg = output['reg']

        # torch.cuda.synchronize(device)
        pts2 = decoder.ctdet_decode(hm, wh, reg)   # 17, 11
        pts0 = pts2.copy()
        pts0[:,:10] *= 4

        print('totol pts num is {}'.format(len(pts2)))

        ori_image = dsets.load_image(dsets.patient_ids.index(patient_id))
        k = 256 / 1024
        ori_image = ori_image * k

        ori_image_regress = cv2.resize(ori_image, (512, 1024))
        ori_image_points = ori_image_regress.copy()

        h,w,c = ori_image.shape
        pts0 = np.asarray(pts0, np.float32)
        # pts0[:,0::2] = pts0[:,0::2]/args.input_w*w
        # pts0[:,1::2] = pts0[:,1::2]/args.input_h*h
        sort_ind = np.argsort(pts0[:,1])
        pts0 = pts0[sort_ind]

        pts3 = []
        for p in pts0:
            point1 = [p[2], p[3]]
            point2 = [p[4], p[5]]
            point3 = [p[6], p[7]]
            point4 = [p[8], p[9]]
            pts3.extend([point1,point2,point3,point4])
        pts3 = np.array(pts3)

        jieduan = int(min(pts3[52][1], pts3[53][1]))
        sly = (pts3[-4][1] - jieduan) *1000/ (images.shape[2] - jieduan)
        sry = (pts3[-3][1] - jieduan) *1000/ (images.shape[2] - jieduan)
        slx =  (pts3[-4][0]/images.shape[3]) * 1000
        srx =  (pts3[-3][0]/images.shape[3]) * 1000

        jieduan = int((jieduan/images.shape[2]) * oimage.shape[2])
        csvl = int((slx + srx) / 2)

        oimage = oimage[:,:,jieduan:,:]
        oimage = oimage[0]
        simage = oimage.data
        oimage = oimage.data.cpu().numpy().transpose(1,2,0)

        composed_transforms = transforms.Compose([
            tr.FixScaleCrop(crop_size=1000),
            tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            tr.ToTensor()])
        sample = {'image': Image.fromarray(oimage.astype(np.uint8)), 'label': Image.fromarray(oimage.astype(np.uint8))}
        sample = composed_transforms(sample)
        oimage = sample['image'].to(device)

        oimage = oimage.unsqueeze(0)
        # print(oimage)

        with torch.no_grad():
            output = model_csvl(oimage)
        result = output.data.cpu().numpy()
        result = np.argmax(result, axis=1)
        result = result.squeeze()
        result = result * 255
        result = np.clip(result, 0, 255)
        result = np.array(result,np.uint8)

        kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, (5, 5))
        result = cv.morphologyEx(result, cv.MORPH_OPEN, kernel, iterations=5)

        # result = result[0]
        # png_file = patient_id + '.png'
        png_file = os.path.join('./results', patient_id, 'csvl_seg.png')
        output = Image.fromarray(result.astype(np.uint8)).convert('P')

        simage = oimage.data.cpu().squeeze()
        std = torch.tensor((0.229, 0.224, 0.225)).view(3,1,1).repeat(1,1000,1000)
        mu = torch.tensor((0.485, 0.456, 0.406)).view(3,1,1).repeat(1,1000,1000)
        simage = simage * std + mu
        simage = simage.numpy()
        simage = np.transpose(simage, (1,2,0))
        simage = simage * 255
        simage = np.clip(simage, 0, 255)
        simage = np.array(simage,np.uint8)

        ret,thresh=cv2.threshold(result,127,255,0)

        # temp = np.ones(result.shape,np.uint8)*255
        contours,hierarchy=cv2.findContours(thresh,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
        # print(len(contours))
        cv2.drawContours(simage,contours,-1,(0,255,0),5)

        # simage = np.expand_dims(simage,2).repeat(3,2)

        cv2.line(simage,
                 (int(csvl), int(0)),
                 (int(csvl), int(simage.shape[0])),
                 color=(0, 255, 255), thickness=2, lineType=2)
        cv2.line(simage,
                 (int(slx), int(sly)),
                 (int(srx), int(sry)),
                 color=(0, 255, 255), thickness=2, lineType=2)

        left, right = get_csvl_ans(result, csvl)

        cv2.line(simage,
                 (left[1], left[0]),
                 (right[1],right[0]),
                 color=(255, 0, 0), thickness=2, lineType=2)
        if left[1] < csvl:   # 1, <
            if right[1] < csvl:  # 1, <
                print('csvl类型： C type', file=report)
            else:
                print('csvl类型： B type', file=report)
        else:
            print('csvl类型： A type',file=report)

        cv2.imwrite(png_file, simage)
        report.close()
        final_lenke = lenke_judge(path)
        with open(r'' + path, 'a') as f:
            f.write(final_lenke)
            f.close()

def get_csvl_ans(mask, csvl):
    pts = np.argwhere(mask ==  255)
    pts = [[x[0], x[1]-csvl] for x in pts]
    pts = sorted(pts, key=lambda x: -abs(x[1]))
    max_pt = pts[0]

    max_pt[1] += csvl

    right = max_pt[:]
    left = max_pt[:]
    que = []
    que.append(max_pt)
    dir = [[0,1],[0,-1],[1,0],[-1,0]]
    while len(que) > 0:
        p = que.pop(0)
        for i in range(4):
            new_p = [p[0]+dir[i][0], p[1]+dir[i][1]]
            if new_p[0] < mask.shape[0] and new_p[0] >= 0 and new_p[1] < mask.shape[1] and new_p[1] >= 0 and mask[new_p[0],new_p[1]] == 255:
                que.append(new_p)
                if new_p[1] > right[1]:
                    right = new_p[:]
                if new_p[1] < left[1]:
                    left = new_p[:]
                mask[new_p[0],new_p[1]] = 0

    return left, right

def post_report():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    heads = {'hm': 1,
             'reg': 2 * 1,
             'wh': 2 * 4, }

    model_ap = spinal_net.SpineNet(heads=heads,
                                   pretrained=True,
                                   down_ratio=4,
                                   final_kernel=1,
                                   head_conv=256)

    # self.num_classes = args.num_classes
    decoder = DecDecoder(K=17, conf_thresh=0.2)
    # dataset = BaseDataset

    ap_save_path = './models/pretrain/ap_model_best.pth'
    model_ap = load_model(model_ap, ap_save_path)
    model_ap = model_ap.to(device)
    model_ap.eval()

    # dataset_module = dataset
    dsets = BaseDataset_ap(data_dir='./data/DATAEXPORT_TEST/',
                           phase='test',
                           input_h=1024,
                           input_w=512,
                           down_ratio=4
                           )

    data_loader = torch.utils.data.DataLoader(dsets,
                                              batch_size=1,
                                              shuffle=False,
                                              num_workers=1,
                                              pin_memory=True)

    model = regression.regression(input_dim=34).to(device)
    checkpoint = torch.load('./models/pretrain/post_model_best.tar', map_location=device)
    model.load_state_dict(checkpoint['state_dict'])

    model.eval()

    for cnt, data_dict in enumerate(data_loader):
        images = data_dict['images'][0]
        patient_id = data_dict['patient_id'][0]
        images = images.to(device)
        print("###############################################")
        print('processing {}/{} patient ... {}'.format(cnt, len(data_loader), patient_id))

        #####################处理 ap#################
        with torch.no_grad():
            output = model_ap(images[0])
            hm = output['hm']
            wh = output['wh']
            reg = output['reg']

        # torch.cuda.synchronize(device)
        pts2 = decoder.ctdet_decode(hm, wh, reg)  # 17, 11
        pts0 = pts2.copy()
        pts0[:, :10] *= 4

        # h,w,c = ori_image.shape
        pts0 = np.asarray(pts0, np.float32)
        # pts0[:,0::2] = pts0[:,0::2]/args.input_w*w
        # pts0[:,1::2] = pts0[:,1::2]/args.input_h*h
        sort_ind = np.argsort(pts0[:, 1])
        pts0 = pts0[sort_ind]

        pts3 = []
        for p in pts0:
            point1 = [p[2], p[3]]
            point2 = [p[4], p[5]]
            point3 = [p[6], p[7]]
            point4 = [p[8], p[9]]
            pts3.extend([point1, point2, point3, point4])
        pts3 = np.array(pts3)  # 68 * 2

        pts3[:, 0] = (pts3[:, 0] - min(pts3[:, 0])) / (max(pts3[:, 0]) - min(pts3[:, 0]))
        pts3[:, 1] = (pts3[:, 1] - min(pts3[:, 1])) / (max(pts3[:, 1]) - min(pts3[:, 1]))

        center1, hm1 = generate_gt(pts3)
        center1 = torch.tensor(center1)
        center1 = center1.to(device)

        center1 = Variable(center1)
        center1 = center1.unsqueeze(0)
        # print(center1)
        pre_center = model(center1)

        result_dir = os.path.join('./results', patient_id)
        pre = center1[0].cpu().data.view(17, 2).numpy()
        post = pre_center[0].cpu().view(17, 2).data.numpy()

        pre[:, 1] = (1 - pre[:, 1]) * 4.5  # 调整脊椎的长宽比，影响侧弯的角度
        # print(pre)
        post[:, 1] = (1 - post[:, 1]) * 6

        # pre[:,0] = -pre[:,0]
        # post[:,0] = -post[:,0]

        plt.figure()
        plt.xlim(0, 1)
        plt.ylim(0, 4)
        plt.subplot(131)
        plt.axis('equal')
        # plt.xticks([]),plt.yticks([])
        # 拟合曲线
        z1 = np.polyfit(pre[:, 1], pre[:, 0], 40)
        p1 = np.poly1d(z1)
        yvals = p1(pre[:, 1])
        plt.plot(yvals, pre[:, 1], 'grey', linewidth=15, marker="o", markersize=5, markerfacecolor="white")

        plt.subplot(133)
        plt.xlim(-0.5, 1.5)
        # plt.ylim(0,2)
        # plt.xticks([]),plt.yticks([])
        # print(post.shape)

        z1 = np.polyfit(post[:, 1], post[:, 0], 40)
        p1 = np.poly1d(z1)
        yvals = p1(post[:, 1])
        plt.plot(yvals, post[:,1], 'grey', linewidth=15, marker="o", markersize=5, markerfacecolor="white")

        plt.savefig(os.path.join(result_dir, 'post_operation.jpg'))

def generate_gt(pts):
    boxes = []
    centers = []
    hm = []
    # print(pts.shape, '###')
    for k in range(0, len(pts), 4):
        pts_4 = pts[k:k + 4, :]
        x_inds = np.argsort(pts_4[:, 0])
        pt_l = np.asarray(pts_4[x_inds[:2], :])
        pt_r = np.asarray(pts_4[x_inds[2:], :])
        y_inds_l = np.argsort(pt_l[:, 1])
        y_inds_r = np.argsort(pt_r[:, 1])
        tl = pt_l[y_inds_l[0], :]
        bl = pt_l[y_inds_l[1], :]
        tr = pt_r[y_inds_r[0], :]
        br = pt_r[y_inds_r[1], :]
        # boxes.append([tl, tr, bl, br])
        boxes.append(tl)
        boxes.append(tr)
        boxes.append(bl)
        boxes.append(br)
        c = np.mean(pts_4, axis=0)
        centers.append(c)
        hm.append(tl - c)
        hm.append(tr - c)
        hm.append(bl - c)
        hm.append(br - c)
    # bboxes = np.asarray(boxes, np.float32) #每个脊椎4个点顺序排好
    # rearrange top to bottom sequence
    # print('########')
    centers = np.asarray(centers, np.float32)
    hm = np.asarray(hm, np.float32)
    return centers, hm

def segType_Trans(num):
    segType = dict(
        zip(["T1", "T2", "T3", "T4", "T5", "T6", "T7", "T8", "T9", "T10", "T11", "T12", "L1", "L2", "L3", "L4",
             "L5"], list(range(1, 18))))
    for j in list(segType.values()):
        if num == j:
            ptss_sup=list(segType.keys())[list(segType.values()).index(j)]
        continue
    return ptss_sup
def lenke_judge(path):
    lenke_type = ["Lenke1", "Lenke2", "Lenke3", "Lenke4", "Lenke5", "Lenke6"]
    report=open(r''+path)
    a_list = report.readlines()
    pt = re.search(pattern=r"(\u975e?)\u7ed3\u6784\u5f2f", string=a_list[5].encode('utf8').decode('utf8')).group()
    mt = re.search(pattern=r"(\u975e?)\u7ed3\u6784\u5f2f", string=a_list[7].encode('utf8').decode('utf8')).group()
    tl = re.search(pattern=r"(\u975e?)\u7ed3\u6784\u5f2f", string=a_list[9].encode('utf8').decode('utf8')).group()
    pt_angle = float(re.search(pattern="([TL][0-9]+-[TL][0-9]+)：(\d+.\d+)", string=a_list[5]).group(2)) #.strip('，bending'))
    mt_angle = float(re.search(pattern="([TL][0-9]+-[TL][0-9]+)：(\d+.\d+)", string=a_list[7]).group(2)) #.strip('，bending'))
    tl_angle = float(re.search(pattern="([TL][0-9]+-[TL][0-9]+)：(\d+.\d+)", string=a_list[9]).group(2)) #.strip('，bending'))
    sag_mod = re.search(pattern=r"(.)\u7c7b", string=a_list[1].encode('utf8').decode('utf8')).group().strip('类')
    csvl_mod = re.search(pattern="([A-Z])", string=a_list[10]).group()
    if pt=="非结构弯" and mt == "结构弯" and tl == "非结构弯" :
        label_type= lenke_type[0]
    elif pt=="结构弯" and mt == "结构弯" and tl == "非结构弯" :
        label_type = lenke_type[1]
    elif pt=="非结构弯" and mt == "结构弯" and tl == "结构弯" and mt_angle > tl_angle:
        label_type = lenke_type[2]
    elif pt=="结构弯" and mt == "结构弯" and tl == "结构弯" :
        label_type = lenke_type[3]
    elif pt=="非结构弯" and mt == "非结构弯" and tl == "结构弯" :
        label_type = lenke_type[4]
    elif pt=="非结构弯" and mt == "结构弯" and tl == "结构弯" and mt_angle < tl_angle:
        label_type = lenke_type[5]
    else:
        label_type="None lenke type"
    return label_type+csvl_mod+sag_mod

if __name__ == '__main__':

    print("################################generate lat report####################################")
    lat_report()
    print("################################generate ap report####################################")
    ap_report_v2()
    print("################################generate csvl report##### ###############################")
    csvl_report()
    print("################################generate post operation report####################################")
    post_report()


import functools
import copy
import math
import time
import logging
import json

#from itertools import product
from io import StringIO
import numpy as np
import pandas as pd

#Grow tf as needed
import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)

import keras
from keras.models import load_model

from rdkit import Chem, DataStructs
from rdkit.Chem import AllChem, Draw
import networkx as nx

class Log(object):
    """Logger object, for when we need to store special things like the Nodes where actions fail"""
    __shared_state = {} #Borg Pattern, class level properties shared amoung instances

    def __init__(self):
        self.__dict__ = self.__shared_state
        if not hasattr(self, "initialized"):
            self.reinitialize()

    def reinitialize(self):
        self.log_stream = StringIO()
        logging.basicConfig(stream=self.log_stream, level=logging.INFO)
        self.logger = logging.getLogger()
        self.runaway_nodes = []
        self.initialized = True

    def log(self):
        print(self.log_stream.getvalue())

    def set_level(self, level="DEBUG"):
        self.logger.setLevel(level)

class Stock(object):
    __shared_state = {} #Borg Pattern, class level properties shared amoung instances

    def __init__(self):
        self.__dict__ = self.__shared_state

    def load_stock(self, stock_name, overwrite=False):
        if overwrite or not hasattr(self, "stock_df"):
            self.stock_df = pd.read_hdf(stock_name, 'table')
            self.stock_inchikeys = set(self.stock_df.inchi_key.values)
        else:
            df = pd.read_hdf(stock_name, 'table')
            self.stock_df = pd.concat([self.stock_df, df], sort=False)
            self.stock_inchikeys = set(self.stock_df.inchi_key.values)

    def reinitialize(self):
        del(self.stock_df)
        del(self.stock_inchikeys)

    #Simple check if compound is in stock
    def mol_in_stock(self, mol):
        inchi_key = Chem.MolToInchiKey( mol )
        return inchi_key in self.stock_inchikeys

    def smiles_in_stock(self, smiles):
        return self.mol_in_stock( Chem.MolFromSmiles(smiles))

class Policy(object):
    __shared_state = {} #Borg, class level properties shared among all instances

    def __init__(self):
        self.__dict__ = self.__shared_state
        if not hasattr(self, "initialized"):
            self.initialize()

    def initialize(self):
        self.cutoff_cumulative = 0.95
        self.cutoff_number = 10
        self.C = 1.4 #Exploitation/exploration balance, used by the Nodes
        self.max_transforms = 5 #Max direct transforms, used by states
        self.max_depth = 20 #Max depth in tree search, used by states
        self.default_prior = 0.5 #Default prior set by nodes, This should more probably be set for the expected average score of the terminal solutions in the tree
        self.use_prior = True #If true, will use the predictions from the NN as priors
        self.initialized = True

    def reinitialize(self):
        del(self.policy)
        del(self.templates_df)
        del(self.templates)
        self.initialize()

    def load_policy(self, modelfile):
        """
        Loads the policy network
        """
        top10_acc = functools.partial(keras.metrics.top_k_categorical_accuracy, k=10)
        top10_acc.__name__ = 'top10_acc'

        top50_acc = functools.partial(keras.metrics.top_k_categorical_accuracy, k=50)
        top50_acc.__name__ = 'top50_acc'

        self.policy = load_model(modelfile, custom_objects={'top10_acc': top10_acc, 'top50_acc': top50_acc})

        self.input_dims = int(self.policy.input.shape[1])

    def load_templates(self, templatefile, template_col="retro_template"):
        """ Load the template file"""
        self.templates_df = pd.read_hdf(templatefile, "table")
        self.templates = self.templates_df[template_col]

    def smiles_to_fp(self, smiles):
        """Convert a SMILES into a fingerprint usable for prediction"""
        return self.mol_to_fp( Chem.MolFromSmiles(smiles) )

    def mol_to_fp(self, mol):
        """Convert a RDKit Mol object into a fingerprint usable for prediction"""
        AllChem.SanitizeMol(mol)
        fp = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=self.input_dims)
        arr = np.zeros((1,))
        DataStructs.ConvertToNumpyArray(fp, arr) #Sub ms conversion
        return  arr.reshape([1,self.input_dims])

    def predict_policy(self, mol):
        """Get predictions of the policy given a SMILES string of the product"""
        fp = self.mol_to_fp(mol)
        return np.array(self.policy.predict(fp)).flatten()

    def cutoff_prediction(self, predictions):
        """Get the idx of the top predictions

        cutoff using cumulated prediction or a max number

        self.cutoff_cumulative
        self.cutoff_number
        """
        #Get the top transforms, truncated at cummulated 99 or at most 10
        sortidx = np.argsort(predictions)[::-1]
        cumsum = np.cumsum(predictions[sortidx])
        maxidx = min(np.argmin(cumsum < self.cutoff_cumulative), self.cutoff_number)
        return sortidx[:maxidx + 1]

    def get_actions(self, state):
        """Get all the probable actions of a state, using the policy and given cutoff's"""
        possible_actions = []
        priors = []

        #Itereate through the molecules and find the possible moves
        for i, mol in enumerate(state.mols):
            #Only try to break down molecules that are not in stock
            if not state.stock.mol_in_stock(mol):
                all_transforms_prop = self.predict_policy(mol)
                probable_transforms_idx = self.cutoff_prediction(all_transforms_prop)
                possible_moves = self.templates[probable_transforms_idx]
                probs = all_transforms_prop[probable_transforms_idx]

                for move,prob in zip(possible_moves, probs):
                    possible_actions.append( (i, move, 0 ) ) #Using tuple makes it hashable and thus usable as key, 0 is default as assumes only one set of reactants
                    priors.append( prob )
        return possible_actions, priors

    def display_action(self, action_key):
        reaction = AllChem.ReactionFromSmarts(action_key[1])
        legend = "Selected molecule number %i"%(action_key[0]+1)
        print(legend)
        display(Draw.ReactionToImage(reaction, subImgSize=(100, 100), useSVG=False))

    def display_transformation(self, transformation):
        """Display transformation in jupyter notebook"""
        display(AllChem.ReactionFromSmarts(transformation))

class Tree(object):
    __shared_state = {} #Borg Pattern, class level properties shared amoung instances

    def __init__(self):
        self.__dict__ = self.__shared_state
        if not hasattr(self, "initialized"):
            self.initialize()

    def initialize(self):
        self.G = nx.DiGraph()
        #self.clear_route()
        self.initialized = True

    def add_root(self, node):
        self.root = node
        self.G.add_node(node)

    def add_child(self, parent, child, action_key):
        self.G.add_nodes_from([parent, child])
        self.G.add_edge(parent, child, action=action_key)

    def get_node_actions(self, node):
        return [child['action'] for child in self.G[node].values()]

    def is_action_instantiated(self, node, action_key):
        #Looks in the dictionaries of the children, if action is there
        return action_key in [child["action"] for child in self.G[node].values()]

    def get_parent(self, node):
        for parent, _ in self.G.in_edges(node):
            return parent #Early return only first

    def get_child(self, node, action_key):
        return  [child for child, edge in self.G[node].items() if edge["action"] == action_key][0]

    def get_action_from_nodes(self, parent, child):
        return self.G[parent][child]["action"]

    def is_root(self, node):
        return self.G.in_degree(node) == 0

    def find_route_to_node(self, node):
        """Find the route (list of states) going from the root to the query node"""
        route = nx.shortest_path(self.G, self.root, node)
        return route

class MolNode(object):
    """MolNode class for reaction networks in NetworkX

    As the __hash__ is default based on the InchiKey, there can be only one molecule node in the network
    can be switched off at instantiation time with the molhash variable or by deleting the .hash property

    """
    def __init__(self, rdmol, molhash=True):
        self.mol = rdmol
        if molhash:
            self.hash = hash(AllChem.MolToInchi(self.mol))

        self.type = "mol"

    def __hash__(self):
        """if custom mol hash, return that, otherwise use default python hash"""
        if hasattr(self, "hash"):
            return self.hash
        else:
            return super(MolNode,self).__hash__()

    def __eq__(self, other):
        return self.__hash__() == other.__hash__()

    def __str__(self):
        return Chem.MolToSmiles(self.mol)

    def __repr__(self):
        return ("MolNode(Chem.MolFromSmiles(%s))"%Chem.MolToSmiles(self.mol))

class ReactionNode(object):
    """ReactionNode object for reaction networks in NetworkX

    The hash is the standard object hash, so there can be several of the same reaction in the tree

    """

    def __init__(self, reaction):
        self.reaction = reaction
        self.type = "reaction"
        #Chem.ReactionFromSmarts(reactionSmarts)

    def __str__(self):
        reactants = self.reaction.GetReactants()
        products = self.reaction.GetProducts()
        forstring = []
        for i,mol in enumerate(reactants):
            if i > 0: forstring.append("+")
            forstring.append(Chem.MolToSmiles(mol))

        forstring.append("==>")
        for i, mol in enumerate(products):
            if i > 0: forstring.append("+")
            forstring.append(Chem.MolToSmiles(mol))

        return " ".join(forstring)

    def __repr__(self):
        return "ReactionNode(AllChem.ReactionFromSmarts('%s'))"%AllChem.ReactionToSmarts(self.reaction)


import logging
import argparse
import torch
import numpy as np
import pandas as pd
import math
import os, fnmatch, shutil
import matplotlib.pyplot as plt
import cv2
import itertools
import pickle
import yaml
from time import time

logger = logging.getLogger(__name__)
logging.basicConfig(format='%(asctime)s - %(message)s', level=logging.INFO)

class FailException:
    def __init__(self, func2failnum):
        self.func2failnum = func2failnum

    def __call__(self, func):
        def inner_function(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            except Exception as e:
                e.args = (f'{self.func2failnum[func.__name__]} in {func.__name__} => ' + e.args[0] ,)
                raise

        return inner_function

def load_model_configs(args):
    with open(args.config_path) as f:
        configs = yaml.load(f, Loader=yaml.FullLoader)
    return configs

def split_pooled(current_method, unique_labels):

    tot_db = {}
    for _label in unique_labels:
        tmp_db = [a for a in current_method if a['label'] == _label]
        tot_db[_label] = tmp_db

    train_db = {}
    for _label in unique_labels:
        tmp_db = [a for a in current_method if (a['label'] == _label) and (a['split'] in ('train', 'val'))]
        train_db[_label] = tmp_db

    test_db = {}
    for _label in unique_labels:
        tmp_db = [a for a in current_method if (a['label'] == _label) and (a['split'] == 'test')]
        test_db[_label] = tmp_db

    return tot_db, train_db, test_db

def retrieve_from_db(query, db, k, print_rate = False):
    '''
    query: dictionary form {id, feature, pair_id}
    '''

    db_same_label_num = len(db)
    if db_same_label_num == 0:
        return None, None, None, None

    else:
        query_feat = torch.from_numpy(query['feature'])
        query_feat_norm = query_feat / query_feat.norm()

        db_feats_list = [a['feature'] for a in db]

        db_feats = torch.from_numpy(np.array(db_feats_list))
        db_feats_norm = db_feats.t() / db_feats.norm(dim = 1)

        # feature similarity
        feats_sim = torch.mm(query_feat_norm.unsqueeze(0), db_feats_norm).squeeze(0)

        # similarity index(descending order)
        feats_sim_index = feats_sim.sort(descending = True)[1].tolist()

        # sorted db
        sorted_db = [db[a] for a in feats_sim_index]

        recall_rate_list, hit_rate_list, rr_list = [], [], []
        for target_k in k:

            k = min(db_same_label_num, target_k)

            recom_db = sorted_db[:k]
            recom_ids = [a['id'] for a in recom_db]

            # db에 있는 동일한 pair_id 개수
            query_pair_id = query['pair_id']
            same_id_dbs = [a for a in db if a['pair_id'] == query_pair_id]
            query_id_len = len(same_id_dbs)

            if not same_id_dbs:
                continue

            else:
                # -----------------
                # RECALL
                # -----------------
                max_recall = query_id_len
                recall_success = len([rec for rec in recom_db if rec['pair_id'] == query_pair_id])
                recall_rate = recall_success/max_recall

                # -----------------
                # HIT: 관련 항목(pair_id) 중 하나라도 추천(k) 내에 속하면 1, 아니면 0
                # -----------------
                hit_success = len([rec for rec in recom_db if rec['pair_id'] == query_pair_id])
                hit_rate = min(1, hit_success)

                # -----------------
                # MRR: k 개중 관련 항목(pair_id)가 처음으로 등장하는 rank를 r이라고 하고, 이것의 역수 = RR, 전체 평균 = MRR
                # -----------------
                same_query_id_index = [i for i, a in enumerate(recom_db) if a['pair_id'] == query_pair_id]
                if same_query_id_index:
                    _rr = min([i for i, a in enumerate(recom_db) if a['pair_id'] == query_pair_id]) + 1
                    rr = 1/_rr
                else:
                    rr = 0

            recall_rate_list.append(recall_rate)
            hit_rate_list.append(hit_rate)
            rr_list.append(rr)

        # return
        return recom_ids, recall_rate_list, hit_rate_list, rr_list

def calc_metrics(unique_labels, set_k, result_table, *DB):

    train_db, test_db = DB[0], DB[1]

    real_fin_result = []
    for target_label in unique_labels:
        logger.info(f'================ Test on [{target_label}] ==================')
        tot_num_per_target_label = len(test_db[target_label])

        M_recall_rate_list, M_hit_rate_list, M_rr_list = [], [], []

        for idx in range(tot_num_per_target_label):

            query = test_db[target_label][idx]
            db    = train_db[target_label]

            _, recall_rate_list, hit_rate_list, rr_list = retrieve_from_db(query, db, k = set_k, print_rate = False)

            if recall_rate_list and hit_rate_list and rr_list:
                M_recall_rate_list.append(recall_rate_list)
                M_hit_rate_list.append(hit_rate_list)
                M_rr_list.append(rr_list)

        fin_result = []
        for i, k in enumerate(set_k):

            RC = np.mean(np.array([a[i] for a in M_recall_rate_list]))
            HT = np.mean(np.array([b[i] for b in M_hit_rate_list]))
            RR = np.mean(np.array([c[i] for c in M_rr_list]))

            logger.info(f'Tested on: {target_label} with k: {k} || Recall: {RC:.4f}, Hit: {HT:.4f}, MRR: {RR:.4f}')
            fin_result.append((RC, HT, RR))

        real_fin_result.append(fin_result)

    # ---------------------------
    #     Calculate(Average)
    # ---------------------------
    RESULT = np.mean(np.array(real_fin_result), axis = 0)
    for i, k in enumerate(set_k):
        logger.info(f"For k = {k} RECALL: {RESULT[i][0]:.4f}  HIT: {RESULT[i][1]:.4f}  MRR: {RESULT[i][2]:.4f}")

    return RESULT

if __name__ == "__main__":

    t0 = time()

    parser = argparse.ArgumentParser()

    parser.add_argument("--extractor_path", type = str, default = './dataset/feature_extraction')
    parser.add_argument("--extractor_type", type = str, default = 'cgd_pca')
    parser.add_argument("--set_ks", type = str, default = "1,2,5,10,20", help = "Set of k's used for testing")
    parser.add_argument("--config_path", type = str, default = './src/configs.yaml')

    args = parser.parse_args()

    configs = load_model_configs(args)

    # ---------------------
    #      EXCEPTION
    # ---------------------
    fail_exception = FailException(configs['exception']['rec_test'])
    # ---------------------
    #       FAIL-001
    # ---------------------
    @fail_exception
    def check_cgd_load(cgd_dir):
        if not os.path.exists(cgd_dir):
            raise Exception('Check directory of the selected extractor')

    selected_method_dir = os.path.join(args.extractor_path, ''.join((args.extractor_type, '.pkl')))
    check_cgd_load(selected_method_dir)

    with open(selected_method_dir, 'rb') as tmp_file:
        current_method = pickle.load(tmp_file)

    # label split
    unique_labels = list(set([a['label'] for a in current_method]))
    tot_db, train_db, test_db = split_pooled(current_method, unique_labels)

    set_ks = list(map(int, args.set_ks.split(',')))
    result_table = pd.DataFrame(columns = ['Label', 'k', 'Recall@', 'Hit@', 'MRR@'])

    result_metrics = calc_metrics(unique_labels, set_ks, result_table, train_db, test_db)

    print("[DONE] Total time spent : {:.4f} seconds.".format(time()-t0))


'''
# Reconstruct image and evalute the performance by Generalization
# Author: Yuki Saeki
'''

import os
import argparse
import numpy as np
import SimpleITK as sitk
import csv
import dataIO as io
import matplotlib.pyplot as plt
import os
import torch
import torch.utils.data
from torch import nn, optim
from torch.nn import functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image
from torchsummary import summary
from torch.utils.tensorboard import SummaryWriter
import cloudpickle
from utils import *
from tqdm import trange
import pickle
from importlib import machinery

parser = argparse.ArgumentParser(description='VAE test')
parser.add_argument('--input', type=str, default="E:/git/TFRecord_example/input/CT/patch/size9/cc1/filename.txt",
                    help='File path of input images')
parser.add_argument('--model', type=str, default="E:/git/pytorch/vae/results/z_24/B_0.1/L_0.005/model.pkl",
                    help='File path of model')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='enables CUDA')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--constrain', '-c', type=bool, default=False, help='topo con')
parser.add_argument('--outdir', type=str, default="E:/git/pytorch/vae/results/z_24/B_0.1/L_0.005/",
                    help='File path of output images')
parser.add_argument('--mode', type=int, default=1,
                    help='[mode: process] = [0: artificial], [1: real]')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()

patch_side = 9
latent_dim = 24

if args.mode==0:
    num_of_data = 10000
    num_of_test = 2000
    num_of_val = 2000
elif args.constrain == True:
    num_of_data = 1978
    num_of_test = 467
    num_of_val = 425
else:
    num_of_data = 3039
    num_of_test = 607
    num_of_val = 607

torch.manual_seed(args.seed)

device = torch.device("cuda" if args.cuda else "cpu")

# check folder
if not (os.path.exists(args.outdir + 'spe/')):
    os.makedirs(args.outdir + 'spe/')

# get data
data_set = get_dataset(args.input, patch_side, num_of_data)
data = data_set.reshape(num_of_data, patch_side * patch_side * patch_side)
data = min_max(data, axis=1)
# divide data
test_data = torch.from_numpy(data[:num_of_test]).float()
train_data = torch.from_numpy(data[num_of_test+num_of_val:]).float().to(device)
train_loader = torch.utils.data.DataLoader(train_data,
                          batch_size=1,
                          shuffle=False,
                          num_workers=0,
                          pin_memory=False,
                          drop_last=False)
test_loader = torch.utils.data.DataLoader(test_data,
                          batch_size=1,
                          shuffle=False,
                          num_workers=0,
                          pin_memory=False,
                          drop_last=False)
# load model
with open(args.model, 'rb') as f:
    model = cloudpickle.load(f)

# testing
specificity = []
mean = []
std = []
mean_single = []
std_single = []
num_of_gen = 5000
patch_side = 9

def spe_test(model):
    #  calculate mu and sigma
    for i in enumerate(train_loader):
        with torch.no_grad():
            train_data_cuda = train_data.to(device)
            recon_batch, mean_batch, logvar_batch = model(train_data_cuda)
            mean_single = mean_batch[0, :]
            std_single = torch.exp(0.5 * logvar_batch)[0, :]
            mean.append(mean_single.cpu().numpy())
            std.append(std_single.cpu().numpy())

    mu = np.mean(mean)
    sigma = np.mean(std)

    model.eval()
    with torch.no_grad():
        gen = np.zeros(num_of_gen, dtype="float")
        ori = np.reshape(test_data.cpu().numpy(), [num_of_test, patch_side, patch_side, patch_side])
        for j in trange(num_of_gen):
            # sample_z = np.random.normal(mu, sigma, (1, latent_dim))
            sample_z = torch.normal(mu, sigma, (1, latent_dim)).to(device)
            gen_batch = model.decode(sample_z)
            gen_single = gen_batch.cpu().numpy()
            gen = np.reshape(gen_single, [patch_side, patch_side, patch_side])
            # EUDT
            eudt_image = sitk.GetImageFromArray(gen)
            eudt_image.SetSpacing([0.885, 0.885, 1])
            eudt_image.SetOrigin([0, 0, 0])

            # calculate spe
            case_min_specificity = 1.0
            for image_index in range(num_of_test):
                specificity_tmp = L1norm(ori[image_index] ,gen)
                if specificity_tmp < case_min_specificity:
                    case_min_specificity = specificity_tmp
            specificity.append([case_min_specificity])

            # output image
            io.write_mhd_and_raw(eudt_image, '{}.mhd'.format(os.path.join(args.outdir, 'spe', '{}'.format(str(j).zfill(4)))))

    return specificity

specificity = spe_test(model)

print('specificity = %f' % np.mean(specificity))
np.savetxt(os.path.join(args.outdir, 'specificity.csv'), specificity, delimiter=",")


# -*- coding: utf-8 -*-
"""
Created on Wed Mar 24 16:25:27 2021

Snapshot data visualization script.
Uses snapshot_visualization_mod and the spreadsheet stored at \snapshot_pilot\SS_data_collector.xlsx,
which contains: list of mice/slides/slices obtained through snapshot pilot experiment, raw cell counts by slice, basic histology info,
and the 4OHT dosage the mouse received.

Green = gfp (eGFP)
Red = tom (tdTomato)

@author: dpaynter
"""
## Imports:
import snapshot_visualization_mod as ssv
import matplotlib.pyplot as plt
import numpy as np
import sys, os

date = str(input("What is today's date?"))

ss_dir = os.path.dirname(os.path.dirname(os.getcwd()))
datapath = ss_dir + r'\SS_data_collector.xlsx'
savefigpath = ss_dir + r'\reports\figures'

# Create an object for each mouse
DP_210202A = ssv.snapshot_mouse("DP_210202A", datapath)
DP_210202B = ssv.snapshot_mouse("DP_210202B", datapath)
DP_210202C = ssv.snapshot_mouse("DP_210202C", datapath)
DP_210202D = ssv.snapshot_mouse("DP_210202D", datapath)
DP_210203A = ssv.snapshot_mouse("DP_210203A", datapath)
DP_210203B = ssv.snapshot_mouse("DP_210203B", datapath)
DP_210308A = ssv.snapshot_mouse("DP_210308A", datapath)
DP_210308B = ssv.snapshot_mouse("DP_210308B", datapath)
DP_210308C = ssv.snapshot_mouse("DP_210308C", datapath)
DP_210416A = ssv.snapshot_mouse("DP_210416A", datapath)
DP_210416B = ssv.snapshot_mouse("DP_210416B", datapath)
DP_210417 = ssv.snapshot_mouse("DP_210417", datapath)

# Create a concatenated dataframe for all mice loaded:
all_mice_df = ssv.concat_dfs([DP_210202A.mouse_df, DP_210202B.mouse_df, DP_210202C.mouse_df,
                              DP_210202D.mouse_df, DP_210203A.mouse_df, DP_210203B.mouse_df,
                              DP_210308A.mouse_df, DP_210308B.mouse_df, DP_210308C.mouse_df,
                              DP_210416A.mouse_df, DP_210416B.mouse_df, DP_210417.mouse_df])

#Get single-labeled cell proportions for each mouse
DP_210202A_green_prop = DP_210202A.green_only_prop
DP_210202B_green_prop = DP_210202B.green_only_prop
DP_210202C_green_prop = DP_210202C.green_only_prop
DP_210202D_green_prop = DP_210202D.green_only_prop
DP_210203A_green_prop = DP_210203A.green_only_prop
DP_210203B_green_prop = DP_210203B.green_only_prop
DP_210308A_green_prop = DP_210308A.green_only_prop
DP_210308B_green_prop = DP_210308B.green_only_prop
DP_210308C_green_prop = DP_210308C.green_only_prop
DP_210416A_green_prop = DP_210416A.green_only_prop
DP_210416B_green_prop = DP_210416B.green_only_prop
DP_210417_green_prop = DP_210417.green_only_prop

DP_210202A_red_prop = DP_210202A.red_only_prop
DP_210202B_red_prop = DP_210202B.red_only_prop
DP_210202C_red_prop = DP_210202C.red_only_prop
DP_210202D_red_prop = DP_210202D.red_only_prop
DP_210203A_red_prop = DP_210203A.red_only_prop
DP_210203B_red_prop = DP_210203B.red_only_prop
DP_210308A_red_prop = DP_210308A.red_only_prop
DP_210308B_red_prop = DP_210308B.red_only_prop
DP_210308C_red_prop = DP_210308C.red_only_prop
DP_210416A_red_prop = DP_210416A.red_only_prop
DP_210416B_red_prop = DP_210416B.red_only_prop
DP_210417_red_prop = DP_210417.red_only_prop

##### MAKE FIGURE 1 #####
# Plot a bar graph with these proportions by mouse
labels = ['DP_210202A \n {}'.format(int(DP_210202A.total_cells)), 'DP_210202B \n {}'.format(int(DP_210202B.total_cells)),
          'DP_210202C \n {}'.format(int(DP_210202C.total_cells)), 'DP_210202D \n {}'.format(int(DP_210202D.total_cells)),
          'DP_210203A \n {}'.format(int(DP_210203A.total_cells)), 'DP_210203B \n {}'.format(int(DP_210203B.total_cells)),
          'DP_210308A \n {}'.format(int(DP_210308A.total_cells)), 'DP_210308B \n {}'.format(int(DP_210308B.total_cells)),
          'DP_210308C \n {}'.format(int(DP_210308C.total_cells)), 'DP_210416A \n {}'.format(int(DP_210416A.total_cells)),
          'DP_210416B \n {}'.format(int(DP_210416B.total_cells)), 'DP_210417 \n {}'.format(int(DP_210417.total_cells))]
gfp_props = [DP_210202A_green_prop, DP_210202B_green_prop, DP_210202C_green_prop,
             DP_210202D_green_prop, DP_210203A_green_prop, DP_210203B_green_prop,
             DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop,
             DP_210416A_green_prop, DP_210416B_green_prop, DP_210417_green_prop]
tom_props = [DP_210202A_red_prop, DP_210202B_red_prop, DP_210202C_red_prop,
             DP_210202D_red_prop, DP_210203A_red_prop, DP_210203B_red_prop,
             DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop,
             DP_210416A_red_prop, DP_210416B_red_prop, DP_210417_red_prop]

gfp_err = [mouse.green_std for mouse in [DP_210202A, DP_210202B, DP_210202C, DP_210202D, DP_210203A, DP_210203B,
                                         DP_210308A, DP_210308B, DP_210308C, DP_210416A, DP_210416B, DP_210417]]
tom_err = [mouse.red_std for mouse in [DP_210202A, DP_210202B, DP_210202C, DP_210202D, DP_210203A, DP_210203B,
                                         DP_210308A, DP_210308B, DP_210308C, DP_210416A, DP_210416B, DP_210417]]

x = np.arange(len(labels))
width = 0.3

fig, ax = plt.subplots()
fig.set_size_inches(11,6)
set1 =  ax.bar(x - width/2, gfp_props, width, yerr=gfp_err, label='GFP-only', color='limegreen')
set2 = ax.bar(x + width/2, tom_props, width, yerr=tom_err, label='TOM-only', color='red')

ax.set_ylabel('Proportion')
ax.set_xlabel('Mouse')
ax.set_title('Proportion single-labeled cells by mouse')
ax.set_xticks(x)
ax.set_xticklabels(labels, fontdict ={'fontsize': 8.5})
ax.legend()

fig.tight_layout()

plt.savefig(savefigpath + r'\fig1_' + date + '.png')

##### MAKE FIGURE 1 WITHOUT ERROR BARS #####
# Plot a bar graph with these proportions by mouse without error bars
labels = ['DP_210202A', 'DP_210202B', 'DP_210202C', 'DP_210202D', 'DP_210203A', 'DP_210203B', 'DP_210308A',
          'DP_210308B', 'DP_210308C', 'DP_210416A', 'DP_210416B', 'DP_210417']
gfp_props = [DP_210202A_green_prop, DP_210202B_green_prop, DP_210202C_green_prop,
             DP_210202D_green_prop, DP_210203A_green_prop, DP_210203B_green_prop,
             DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop,
             DP_210416A_green_prop, DP_210416B_green_prop, DP_210417_green_prop]
tom_props = [DP_210202A_red_prop, DP_210202B_red_prop, DP_210202C_red_prop,
             DP_210202D_red_prop, DP_210203A_red_prop, DP_210203B_red_prop,
             DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop,
             DP_210416A_red_prop, DP_210416B_red_prop, DP_210417_red_prop]

x = np.arange(len(labels))
width = 0.35

fig, ax = plt.subplots()
fig.set_size_inches(10,6)
set1 =  ax.bar(x - width/2, gfp_props, width, label='GFP-only', color='limegreen')
set2 = ax.bar(x + width/2, tom_props, width, label='TOM-only', color='red')

ax.set_ylabel('Proportion')
ax.set_xlabel('Mouse')
ax.set_title('Proportion single-labeled cells by mouse')
ax.set_xticks(x)
ax.set_xticklabels(labels, fontdict ={'fontsize': 7.5})
ax.legend()

fig.tight_layout()

plt.savefig(savefigpath + r'\fig1_noerr_' + date + '.png')

##### MAKE FIGURE 2 #####
# Plot a bar graph with these proportions by condition or dosage; average across mice:
labels = ['50mg/kg 2x', '100mg/kg 1x', '100mg/kg 2x', '200mg/kg 1x']
gfp_props = [np.average([DP_210202A_green_prop, DP_210202B_green_prop, DP_210202C_green_prop]),
             np.average([DP_210202D_green_prop, DP_210203A_green_prop, DP_210203B_green_prop]),
             np.average([DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop,
                         DP_210417_green_prop]),
             np.average([DP_210416A_green_prop, DP_210416B_green_prop])]

gfp_scatter_props = [[DP_210202A_green_prop, DP_210202B_green_prop, DP_210202C_green_prop],
             [DP_210202D_green_prop, DP_210203A_green_prop, DP_210203B_green_prop],
             [DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop,
              DP_210417_green_prop],[DP_210416A_green_prop, DP_210416B_green_prop]]

tom_props = [np.average([DP_210202A_red_prop, DP_210202B_red_prop, DP_210202C_red_prop]),
             np.average([DP_210202D_red_prop, DP_210203A_red_prop, DP_210203B_red_prop]),
             np.average([DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop,
                         DP_210417_red_prop]),
             np.average([DP_210416A_red_prop, DP_210416B_red_prop])]

tom_scatter_props = [[DP_210202A_red_prop, DP_210202B_red_prop, DP_210202C_red_prop],
             [DP_210202D_red_prop, DP_210203A_red_prop, DP_210203B_red_prop],
             [DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop,
              DP_210417_red_prop],
             [DP_210416A_red_prop, DP_210416B_red_prop]]

x = np.arange(len(labels))
width = 0.35

gfp_err = [np.std([DP_210202A_green_prop, DP_210202B_green_prop, DP_210202C_green_prop]),
           np.std([DP_210202D_green_prop, DP_210203A_green_prop, DP_210203B_green_prop]),
           np.std([DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop,
                    DP_210417_green_prop]),
           np.std([DP_210416A_green_prop, DP_210416B_green_prop])]

tom_err = [np.std([DP_210202A_red_prop, DP_210202B_red_prop, DP_210202C_red_prop]),
           np.std([DP_210202D_red_prop, DP_210203A_red_prop, DP_210203B_red_prop]),
           np.std([DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop,
                    DP_210417_red_prop]),
           np.std([DP_210416A_red_prop, DP_210416B_red_prop])]

fig, ax = plt.subplots()

for condition in range(4):
    for i in range(len(gfp_scatter_props[condition])):
    # distribute scatter randomly across whole width of bar
        ax.scatter(x[condition] - width/2, gfp_scatter_props[condition][i],
                   zorder=10, color='black', s=14)
        ax.scatter(x[condition] + width/2, tom_scatter_props[condition][i],
                   zorder=10, color='black', s=14)

set1 =  ax.bar(x - width/2, gfp_props, width, yerr=gfp_err, label='GFP-only', color='limegreen', zorder=0)
set2 = ax.bar(x + width/2, tom_props, width, yerr=tom_err, label='TOM-only', color='red', zorder=0)

ax.set_ylabel('Proportion')
ax.set_xlabel('4OHT condition')
ax.set_title('Proportion single-labeled cells by condition \n Avg over mice')
ax.set_xticks(x)
ax.set_xticklabels(labels, fontdict ={'fontsize':8})
ax.legend()

fig.tight_layout()

plt.savefig(savefigpath + r'\fig2_' + date + '.png')

print ('50x2 GFP-only avg + std: {} {}'.format(str(gfp_props[0]) + ',', gfp_err[0]))
print ('100x1 GFP-only avg + std: {} {}'.format(str(gfp_props[1]) + ',', gfp_err[1]))
print ('100x2 GFP-only avg + std: {} {}'.format(str(gfp_props[2]) + ',', gfp_err[2]))
print ('200x1 GFP-only avg + std: {} {}'.format(str(gfp_props[3]) + ',', gfp_err[3]))

##### MAKE FIGURE 3 #####
# Plot a bar graph with these proportions by condition or dosage; use the whole cell population from the group:

gfp_50x2 = np.sum([DP_210202A.green_cells, DP_210202B.green_cells, DP_210202C.green_cells]) / np.sum([DP_210202A.total_cells, DP_210202B.total_cells, DP_210202C.total_cells])
gfp_100x1 = np.sum([DP_210202D.green_cells, DP_210203A.green_cells, DP_210203B.green_cells]) /np.sum([DP_210202D.total_cells, DP_210203A.total_cells, DP_210203B.total_cells])
gfp_100x2 = np.sum([DP_210308A.green_cells, DP_210308B.green_cells, DP_210308C.green_cells, DP_210417.green_cells]) / np.sum([DP_210308A.total_cells, DP_210308B.total_cells, DP_210308C.total_cells, DP_210417.total_cells])
gfp_200x1 = np.sum([DP_210416A.green_cells, DP_210416B.green_cells]) /  np.sum([DP_210416A.total_cells, DP_210416B.total_cells])

tom_50x2 = np.sum([DP_210202A.red_cells, DP_210202B.red_cells, DP_210202C.red_cells]) / np.sum([DP_210202A.total_cells, DP_210202B.total_cells, DP_210202C.total_cells])
tom_100x1 = np.sum([DP_210202D.red_cells, DP_210203A.red_cells, DP_210203B.red_cells]) /np.sum([DP_210202D.total_cells, DP_210203A.total_cells, DP_210203B.total_cells])
tom_100x2 = np.sum([DP_210308A.red_cells, DP_210308B.red_cells, DP_210308C.red_cells, DP_210417.red_cells]) / np.sum([DP_210308A.total_cells, DP_210308B.total_cells, DP_210308C.total_cells, DP_210417.total_cells])
tom_200x1 = np.sum([DP_210416A.red_cells, DP_210416B.red_cells]) /  np.sum([DP_210416A.total_cells, DP_210417.total_cells])

labels = ['50mg/kg 2x', '100mg/kg 1x', '100mg/kg 2x', '200mg/kg 1x']
gfp_props = [gfp_50x2, gfp_100x1, gfp_100x2, gfp_200x1]

tom_props = [tom_50x2, tom_100x1, tom_100x2, tom_200x1]

x = np.arange(len(labels))
width = 0.35

gfp_err = [np.std([DP_210202A_green_prop, DP_210202B_green_prop, DP_210202C_green_prop]),
           np.std([DP_210202D_green_prop, DP_210203A_green_prop, DP_210203B_green_prop]),
           np.std([DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop, DP_210417_green_prop]),
           np.std([DP_210416A_green_prop, DP_210416B_green_prop])]

tom_err = [np.std([DP_210202A_red_prop, DP_210202B_red_prop, DP_210202C_red_prop]),
           np.std([DP_210202D_red_prop, DP_210203A_red_prop, DP_210203B_red_prop]),
           np.std([DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop, DP_210417_red_prop]),
           np.std([DP_210416A_red_prop, DP_210416B_red_prop])]

fig, ax = plt.subplots()
set1 =  ax.bar(x - width/2, gfp_props, width, yerr=gfp_err, label='GFP-only', color='limegreen')
set2 = ax.bar(x + width/2, tom_props, width, yerr=tom_err, label='TOM-only', color='red')

ax.set_ylabel('Proportion')
ax.set_xlabel('4OHT condition')
ax.set_title('Proportion single-labeled cells by condition \n Avg across all cells')
ax.set_xticks(x)
ax.set_xticklabels(labels, fontdict ={'fontsize':8})
ax.legend()

fig.tight_layout()

plt.savefig(savefigpath + r'\fig3_' + date + '.png')

print ('50x2 GFP-only avg + std: {} {}'.format(str(gfp_props[0]) + ',', gfp_err[0]))
print ('100x1 GFP-only avg + std: {} {}'.format(str(gfp_props[1]) + ',', gfp_err[1]))
print ('100x2 GFP-only avg + std: {} {}'.format(str(gfp_props[2]) + ',', gfp_err[2]))
print ('200x1 GFP-only avg + std: {} {}'.format(str(gfp_props[3]) + ',', gfp_err[3]))

##### MAKE FIGURE 4 #####
# Plot correlation between total number of cells in a slice, and proportion single-labeled cells
sub_df = all_mice_df[['total_cell_count','green_only_proportion','red_only_proportion']]
filtered_df = sub_df[sub_df['green_only_proportion'].notna()]

total_cells = list(filtered_df['total_cell_count'])
green_prop = list(filtered_df['green_only_proportion'])
red_prop = list(filtered_df['red_only_proportion'])
fig, ax = plt.subplots()
for fluo, color in zip(['green_only_proportion', 'red_only_proportion'], ['limegreen','red']):
   x, y = total_cells, filtered_df[fluo]
   ax.scatter(x ,y, label=fluo, c=color, s=6)
ax.legend(loc='upper right', fontsize='small')
for label in (ax.get_xticklabels() + ax.get_yticklabels()):
    label.set_fontsize(8)
ax.set_title('Proportion single-labeled cells by total cells in slice')
ax.set_xlabel('Total labeled cell count')
plt.savefig(savefigpath + r'\fig4_' + date + '.png')

##### MAKE FIGURE 5 #####
# Plot a bar graph; group by number of injections:
labels = ['One Day', 'Two Days']
gfp_props = [np.average([DP_210202D_green_prop, DP_210203A_green_prop, DP_210203B_green_prop, DP_210416A_green_prop, DP_210416B_green_prop]),
             np.average([DP_210202A_green_prop, DP_210202B_green_prop, DP_210202C_green_prop, DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop,
                         DP_210417_green_prop])]

tom_props = [np.average([DP_210202D_red_prop, DP_210203A_red_prop, DP_210203B_red_prop, DP_210416A_red_prop, DP_210416B_red_prop]),
             np.average([DP_210202A_red_prop, DP_210202B_red_prop, DP_210202C_red_prop, DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop,
                         DP_210417_red_prop])]
x = np.arange(len(labels))
width = 0.35

gfp_err = [np.std([DP_210202D_green_prop, DP_210203A_green_prop, DP_210203B_green_prop, DP_210416A_green_prop, DP_210416B_green_prop]),
             np.std([DP_210202A_green_prop, DP_210202B_green_prop, DP_210202C_green_prop, DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop,
                         DP_210417_green_prop])]

tom_err = [np.std([DP_210202D_red_prop, DP_210203A_red_prop, DP_210203B_red_prop, DP_210416A_red_prop, DP_210416B_red_prop]),
             np.std([DP_210202A_red_prop, DP_210202B_red_prop, DP_210202C_red_prop, DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop,
                         DP_210417_red_prop])]

fig, ax = plt.subplots()
set1 =  ax.bar(x - width/2, gfp_props, width, yerr=gfp_err, label='GFP-only', color='limegreen')
set2 = ax.bar(x + width/2, tom_props, width, yerr=tom_err, label='TOM-only', color='red')

ax.set_ylabel('Proportion')
ax.set_xlabel('Number of Days with Injections')
ax.set_title('Proportion single-labeled cells by number of treatment days')
ax.set_xticks(x)
ax.set_xticklabels(labels, fontdict ={'fontsize':8})
ax.legend()

fig.tight_layout()

plt.savefig(savefigpath + r'\fig5_' + date + '.png')

print ('One-inj GFP-only avg + std: {} {}'.format(str(gfp_props[0]) + ',', gfp_err[0]))
print ('Two-inj GFP-only avg + std: {} {}'.format(str(gfp_props[1]) + ',', gfp_err[1]))

##### MAKE FIGURE 6 #####
# Plot a bar graph; group by total dose (disregard number of injections):
labels = ['100mg/kg', '200mg/kg']

gfp_props = [np.average([DP_210202A_green_prop, DP_210202B_green_prop, DP_210202C_green_prop,
                         DP_210202D_green_prop, DP_210203A_green_prop, DP_210203B_green_prop]),
             np.average([DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop,
                         DP_210417_green_prop, DP_210416A_green_prop, DP_210416B_green_prop])]

tom_props = [np.average([DP_210202A_red_prop, DP_210202B_red_prop, DP_210202C_red_prop,
                         DP_210202D_red_prop, DP_210203A_red_prop, DP_210203B_red_prop]),
             np.average([DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop,
                         DP_210417_red_prop, DP_210416A_red_prop, DP_210416B_red_prop])]

x = np.arange(len(labels))

width = 0.35

gfp_err = [np.std([DP_210202A_green_prop, DP_210202B_green_prop, DP_210202C_green_prop,
                         DP_210202D_green_prop, DP_210203A_green_prop, DP_210203B_green_prop]),
           np.std([DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop,
                         DP_210417_green_prop, DP_210416A_green_prop, DP_210416B_green_prop])]

tom_err = [np.std([DP_210202A_red_prop, DP_210202B_red_prop, DP_210202C_red_prop,
                         DP_210202D_red_prop, DP_210203A_red_prop, DP_210203B_red_prop]),
           np.std([DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop,
                         DP_210417_red_prop, DP_210416A_red_prop, DP_210416B_red_prop])]

fig, ax = plt.subplots()
set1 =  ax.bar(x - width/2, gfp_props, width, yerr=gfp_err, label='GFP-only', color='limegreen')
set2 = ax.bar(x + width/2, tom_props, width, yerr=tom_err, label='TOM-only', color='red')

ax.set_ylabel('Proportion')
ax.set_xlabel('Total dose 4OHT')
ax.set_title('Proportion single-labeled cells by total dose of 4OHT')
ax.set_xticks(x)
ax.set_xticklabels(labels, fontdict ={'fontsize':8})
ax.legend()

fig.tight_layout()

plt.savefig(savefigpath + r'\fig6_' + date + '.png')

print ('100mg/kg GFP-only avg + std: {} {}'.format(str(gfp_props[0]) + ',', gfp_err[0]))
print ('200mg/kg GFP-only avg + std: {} {}'.format(str(gfp_props[1]) + ',', gfp_err[1]))

##### MAKE FIGURE 7 #####
# Plot a bar graph to show overall average across all mice:
labels = ['All Mice, All Conditions']
gfp_props = [np.average([DP_210202A_green_prop, DP_210202B_green_prop, DP_210202C_green_prop,
                         DP_210202D_green_prop, DP_210203A_green_prop, DP_210203B_green_prop,
                         DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop,
                         DP_210417_green_prop, DP_210416A_green_prop, DP_210416B_green_prop])]

tom_props = [np.average([DP_210202A_red_prop, DP_210202B_red_prop, DP_210202C_red_prop,
                         DP_210202D_red_prop, DP_210203A_red_prop, DP_210203B_red_prop,
                         DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop,
                         DP_210417_red_prop, DP_210416A_red_prop, DP_210416B_red_prop])]

x = np.arange(len(labels))

width = 0.35
gfp_err = [np.std([DP_210202A_green_prop, DP_210202B_green_prop, DP_210202C_green_prop,
                         DP_210202D_green_prop, DP_210203A_green_prop, DP_210203B_green_prop,
                         DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop,
                         DP_210417_green_prop, DP_210416A_green_prop, DP_210416B_green_prop])]

tom_err = [np.std([DP_210202A_red_prop, DP_210202B_red_prop, DP_210202C_red_prop,
                         DP_210202D_red_prop, DP_210203A_red_prop, DP_210203B_red_prop,
                         DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop,
                         DP_210417_red_prop, DP_210416A_red_prop, DP_210416B_red_prop])]

fig, ax = plt.subplots()
set1 =  ax.bar(x - width/2, gfp_props, width, yerr=gfp_err,  label='GFP-only', color='limegreen')
set2 = ax.bar(x + width/2, tom_props, width, yerr=tom_err, label='TOM-only', color='red')

ax.set_ylabel('Proportion')
ax.set_title('Proportion single-labeled cells')
ax.set_xticks(x)
ax.set_xticklabels(labels, fontdict ={'fontsize':8})
ax.legend()

fig.tight_layout()

plt.savefig(savefigpath + r'\fig7_' + date + '.png')
print ('GFP-only avg + std: {} {}'.format(str(gfp_props[0]) + ',', gfp_err[0]))

##### MAKE FIGURE 8 #####
### Plot males with 200mg/kg total dose against females with 200mg/kg total dose:
labels = ['males', 'females']

gfp_props = [np.average([DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop]),
             np.average([DP_210417_green_prop, DP_210416A_green_prop, DP_210416B_green_prop])]

tom_props = [np.average([DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop]),
             np.average([DP_210417_red_prop, DP_210416A_red_prop, DP_210416B_red_prop])]

x = np.arange(len(labels))

width = 0.35

gfp_err = [np.std([DP_210308A_green_prop, DP_210308B_green_prop, DP_210308C_green_prop]),
           np.std([DP_210417_green_prop, DP_210416A_green_prop, DP_210416B_green_prop])]

tom_err = [np.std([DP_210308A_red_prop, DP_210308B_red_prop, DP_210308C_red_prop]),
           np.std([DP_210417_red_prop, DP_210416A_red_prop, DP_210416B_red_prop])]

fig, ax = plt.subplots()
set1 =  ax.bar(x - width/2, gfp_props, width, yerr=gfp_err, label='GFP-only', color='limegreen')
set2 = ax.bar(x + width/2, tom_props, width, yerr=tom_err, label='TOM-only', color='red')

ax.set_ylabel('Proportion')
ax.set_title('Proportion single-labeled cells in high-dose mice by sex')
ax.set_xticks(x)
ax.set_xticklabels(labels, fontdict ={'fontsize':8})
ax.legend()

fig.tight_layout()

plt.savefig(savefigpath + r'\fig8_' + date + '.png')

print ('Males high-dose GFP-only avg + std: {} {}'.format(str(gfp_props[0]) + ',', gfp_err[0]))
print ('Females high-dose GFP-only avg + std: {} {}'.format(str(gfp_props[1]) + ',', gfp_err[1]))

##############################################################################################################
## Specific analysis for cortex and thalamus separately:

#Get single-labeled cell proportions for each mouse

#DP_210202A_green_cort = DP_210202A.green_prop_cort
#DP_210202B_green_cort = DP_210202B.green_prop_cort
#DP_210202C_green_cort = DP_210202C.green_prop_cort
DP_210202D_green_cort = DP_210202D.green_prop_cort
#DP_210203A_green_cort = DP_210203A.green_prop_cort
#DP_210203B_green_cort = DP_210203B.green_prop_cort
DP_210308A_green_cort = DP_210308A.green_prop_cort
DP_210308B_green_cort = DP_210308B.green_prop_cort
#DP_210308C_green_cort = DP_210308C.green_prop_cort
DP_210416A_green_cort = DP_210416A.green_prop_cort
#DP_210416B_green_cort = DP_210416B.green_prop_cort
DP_210417_green_cort = DP_210417.green_prop_cort

#DP_210202A_red_cort = DP_210202A.red_prop_cort
#DP_210202B_red_cort = DP_210202B.red_prop_cort
#DP_210202C_red_cort = DP_210202C.red_prop_cort
#DP_210202D_red_cort = DP_210202D.red_prop_cort
#DP_210203A_red_cort = DP_210203A.red_prop_cort
#DP_210203B_red_cort = DP_210203B.red_prop_cort
DP_210308A_red_cort = DP_210308A.red_prop_cort
DP_210308B_red_cort = DP_210308B.red_prop_cort
#DP_210308C_red_cort = DP_210308C.red_prop_cort
DP_210416A_red_cort = DP_210416A.red_prop_cort
#DP_210416B_red_cort = DP_210416B.red_prop_cort
DP_210417_red_cort = DP_210417.red_prop_cort

#DP_210202A_green_thal = DP_210202A.green_prop_thal
#DP_210202B_green_thal = DP_210202B.green_prop_thal
#DP_210202C_green_thal = DP_210202C.green_prop_thal
#DP_210202D_green_thal = DP_210202D.green_prop_thal
#DP_210203A_green_thal = DP_210203A.green_prop_thal
#DP_210203B_green_thal = DP_210203B.green_prop_thal
DP_210308A_green_thal = DP_210308A.green_prop_thal
DP_210308B_green_thal = DP_210308B.green_prop_thal
#DP_210308C_green_thal = DP_210308C.green_prop_thal
DP_210416A_green_thal = DP_210416A.green_prop_thal
#DP_210416B_green_thal = DP_210416B.green_prop_thal
DP_210417_green_thal = DP_210417.green_prop_thal

#DP_210202A_red_thal = DP_210202A.red_prop_thal
#DP_210202B_red_thal = DP_210202B.red_prop_thal
#DP_210202C_red_thal = DP_210202C.red_prop_thal
DP_210202D_red_thal = DP_210202D.red_prop_thal
#DP_210203A_red_thal = DP_210203A.red_prop_thal
#DP_210203B_red_thal = DP_210203B.red_prop_thal
DP_210308A_red_thal = DP_210308A.red_prop_thal
DP_210308B_red_thal = DP_210308B.red_prop_thal
#DP_210308C_red_thal = DP_210308C.red_prop_thal
DP_210416A_red_thal = DP_210416A.red_prop_thal
#DP_210416B_red_thal = DP_210416B.red_prop_thal
DP_210417_red_thal = DP_210417.red_prop_thal

## Plot the correlation between green-only prop in thal and green-only prop in cortex:
green_cort = [DP_210308A_green_cort, DP_210308B_green_cort, DP_210416A_green_cort, DP_210417_green_cort]
green_thal = [DP_210308A_green_thal, DP_210308B_green_thal, DP_210416A_green_thal, DP_210417_green_thal]

red_cort = [DP_210308A_red_cort, DP_210308B_red_cort, DP_210416A_red_cort, DP_210417_red_cort]
red_thal = [DP_210308A_red_thal, DP_210308B_red_thal, DP_210416A_red_thal, DP_210417_red_thal]

fig, ax = plt.subplots()
ax.scatter(green_cort, green_thal)
#ax.legend(loc='upper right', fontsize='small')
for label in (ax.get_xticklabels() + ax.get_yticklabels()):
    label.set_fontsize(8)
ax.set_title('Correlation of single-labeled cells in cortex vs thalamus')
ax.set_xlabel('prop single-labeled in cortex')
ax.set_ylabel('prop single-labeled in thalamus')


import cv2
from matplotlib import pyplot
from PIL import Image
import numpy as np
import pandas as pd
from scipy.spatial.distance import cosine
from mtcnn import MTCNN
from keras.models import load_model
from keras_vggface.utils import preprocess_input
# from flask import Flask, request, jsonify
from firebase_admin import credentials, firestore, storage, initialize_app, db

# extract a single face from a given photograph
def extract_face(filename, required_size=(224, 224)):
	# load image from file
	pixels = pyplot.imread(filename)
	# create the detector, using default weights
	detector = MTCNN()
	# detect faces in the image
	results = detector.detect_faces(pixels)
	# extract the bounding box from the first face
	x1, y1, width, height = results[0]['box']
	x2, y2 = x1 + width, y1 + height
	# extract the face
	face = pixels[y1:y2, x1:x2]
	# resize pixels to the model size
	image = Image.fromarray(face)
	image = image.resize(required_size)
	face_array = np.asarray(image)
	return face_array

# extract faces and calculate face embeddings for a list of photo files
def get_embeddings(filenames):
	# extract faces
	faces = [extract_face(f) for f in filenames]
	# convert into an array of samples
	samples = np.asarray(faces, 'float32')
	# prepare the face for the model, e.g. center pixels
	samples = preprocess_input(samples, version=2)
	# load the vggface model
	model = load_model('VGGFace.h5') # VGGFace(model='resnet50', include_top=False, input_shape=(224, 224, 3), pooling='avg')
	# perform prediction
	yhat = model.predict(samples)
	return yhat

# determine if a candidate face is a match for a known face
def is_match(known_embedding, candidate_embedding, thresh=0.5):
	# calculate distance between embeddings
	score = cosine(known_embedding, candidate_embedding)
	if score <= thresh:
		return True, score, thresh # print('>face is a Match (%.3f <= %.3f)' % (score, thresh))
	else:
		return False, score, thresh #print('>face is NOT a Match (%.3f > %.3f)' % (score, thresh))

size = 4

# load the model
model = load_model('VGGFace.h5')

# load the xml file
classifier = cv2.CascadeClassifier('../haarcascade_frontalface_alt.xml')

# load image database
image_db = pd.read_csv('image_db.csv')

webcam = cv2.VideoCapture(0) #Using default WebCam connected to the PC.

while True:
	(rval, im) = webcam.read()
	im=cv2.flip(im,1,0) #Flip to act as a mirror

	# Resize the image to speed up detection
	mini = cv2.resize(im, (int(im.shape[1]/size), int(im.shape[0]/size)))

	# detect MultiScale / faces
	faces = classifier.detectMultiScale(mini)
	if len(faces):
		# get the first face only
		f = faces[0]
		# Draw rectangles around the face
		(x, y, w, h) = [v * size for v in f] #Scale the shapesize backup
		cv2.rectangle(im, (x,y), (x+w,y+h), (0,255,0), 4)

	# Show the image
	cv2.imshow('Capture', im)
	key = cv2.waitKey(10)

	if key == 32 and len(faces):
		# get the first face only
		f = faces[0]
		(x, y, w, h) = [v * size for v in f] #Scale the shapesize backup

		#Save just the rectangle faces in SubRecFaces
		sub_face = im[y:y+h, x:x+w]
		sub_face = cv2.resize(sub_face, (224, 224))
		samples = np.asarray(sub_face, 'float32')
		arr4d = np.expand_dims(samples, 0)
		samples_array = preprocess_input(arr4d, version=2)
		embedding_main = model.predict(samples_array)
		embedding_main = embedding_main[0]
		grant = False
		for index, row in image_db.iterrows():
			# get embeddings file filenames
			embeddings = get_embeddings([row['image_path']])

			match, score, thresh = is_match(embedding_main, embeddings[0])
			if match:
				print('>face is a Match with guest %s (%.3f <= %.3f)' % (row['name'], score, thresh))
				grant = True
		if not grant: print('>face is NOT a Match')

	# if Esc key is press then break out of the loop
	elif key == 27: #The Esc key
		break


# %%
import torch
import numpy as np
import matplotlib.pyplot as plt
import logging
import MIA_torch
from datasets_torch import *
from utils import setup_logger

import argparse

parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')

# training setting ()
parser.add_argument('--arch', default="vgg11_bn", type=str, help='please type save_file name for the testing purpose')
parser.add_argument('--cutlayer', default=4, type=int, help='number of layers in local')
parser.add_argument('--batch_size', default=128, type=int, help='training batch size')
parser.add_argument('--filename', required=True, type=str, help='please type save_file name for the testing purpose')
parser.add_argument('--folder', default="saves", type=str, help='please type folder name for the testing purpose')
parser.add_argument('--num_client', default=1, type=int, help='number of client')
parser.add_argument('--num_epochs', default=200, type=int, help='number of epochs')
parser.add_argument('--learning_rate', default=0.01, type=float, help='Learning Rate for server-side model')
parser.add_argument('--dataset_portion', default=1.0, type=float, help='dataset_portion')
parser.add_argument('--client_sample_ratio', default=1.0, type=float, help='client_sample_ratio')
parser.add_argument('--noniid', default=1.0, type=float, help='noniid_ratio, if = 0.1, meaning 1 out of 10 class per client')
parser.add_argument('--local_lr', default=-1, type=float, help='Learning Rate for client-side model')
parser.add_argument('--dataset', default="cifar10", type=str, help='number of classes for the testing dataset')
parser.add_argument('--scheme', default="V2_epoch", type=str, help='the name of the scheme, either V3 or others')

#regularization setting ()
parser.add_argument('--regularization', default="None", type=str, help='apply regularization in multi-client training.')
parser.add_argument('--regularization_strength', default=0, type=float, help='regularization_strength of regularization in multi-client training.')
parser.add_argument('--ssim_threshold', default=0.0, type=float, help='regularization threshold for the SSIM score.')
parser.add_argument('--gan_AE_type', default="custom", type=str, help='the name of the AE used in GAN_adv, option: custom, simple, simplest')
parser.add_argument('--gan_loss_type', default="SSIM", type=str, help='loss type of training defensive decoder: SSIM or MSE')
parser.add_argument('--bottleneck_option', default="None", type=str, help='set bottleneck option')
parser.add_argument('--optimize_computation', default=1, type=int, help='set interval N to optimize_computation')
parser.add_argument('--decoder_sync', action='store_true', default=False, help='if True, we sync decoder')

#training dataset setting ()
parser.add_argument('--load_from_checkpoint', action='store_true', default=False, help='if True, we load_from_checkpoint')
parser.add_argument('--load_from_checkpoint_server', action='store_true', default=False, help='if True, we load_from_checkpoint for server-side model')
parser.add_argument('--transfer_source_task', default="cifar100", type=str, help='the name of the transfer_source_task, option: cifar10, cifar100')
parser.add_argument('--finetune_freeze_bn', action='store_true', default=False, help='if True, we finetune_freeze_bn')
parser.add_argument('--save_more_checkpoints', action='store_true', default=False, help='if True, we save_more_checkpoints')
parser.add_argument('--initialize_different', action='store_true', default=False, help='if True, we initialze differently for different client')

#training randomseed setting ()
parser.add_argument('--random_seed', default=123, type=int, help='random_seed')

args = parser.parse_args()

random_seed = args.random_seed
torch.manual_seed(random_seed)
np.random.seed(random_seed)
batch_size = args.batch_size
cutting_layer = args.cutlayer
num_client = args.num_client
save_dir_name = "./{}/{}".format(args.folder, args.filename)
mi = MIA_torch.MIA(args.arch, cutting_layer, batch_size, n_epochs = args.num_epochs, scheme = args.scheme,
                 num_client = num_client, dataset=args.dataset, save_dir=save_dir_name,
                 regularization_option=args.regularization, regularization_strength = args.regularization_strength,
                 initialize_different=args.initialize_different, learning_rate = args.learning_rate, local_lr = args.local_lr, gan_AE_type = args.gan_AE_type,
                 load_from_checkpoint = args.load_from_checkpoint, bottleneck_option = args.bottleneck_option,
                 optimize_computation = args.optimize_computation, decoder_sync = args.decoder_sync,
                 finetune_freeze_bn = args.finetune_freeze_bn, gan_loss_type=args.gan_loss_type, ssim_threshold = args.ssim_threshold,
                 source_task = args.transfer_source_task, load_from_checkpoint_server = args.load_from_checkpoint_server, save_more_checkpoints = args.save_more_checkpoints,
                 dataset_portion = args.dataset_portion, noniid = args.noniid, client_sample_ratio = args.client_sample_ratio)
mi.logger.debug(str(args))

log_frequency = 500

LOG = mi(verbose=True, progress_bar=True, log_frequency=log_frequency)

# %%


"""Object detector based on EfficientDet model.

This module supports inferenceing with the official EfficientDet model.
For more details: https://github.com/yinguobing/efficientdet-runner
"""
import cv2
import numpy as np
import tensorflow as tf

class Detector(object):
    """Mini module to run the EfficientDet model."""

    def __init__(self, saved_model):
        """Build an EfficientDet model runner.

        Args:
            saved_model: the string path to the SavedModel.
        """
        self.scale_width = 0
        self.scale_height = 0
        self.input_size = 512

        # Load the SavedModel object.
        imported = tf.saved_model.load('saved_model')
        self.__predict_fn = imported.signatures["serving_default"]

    def preprocess(self, image):
        """Preprocess the input image."""

        # Scale the image first.
        height, width, _ = image.shape
        self.ratio = self.input_size / max(height, width)
        image = cv2.resize(
            image, (int(self.ratio * width), int(self.ratio * height)))

        # Convert to RGB.
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # Then pad the image to input size.
        self.padding_h = self.input_size - int(self.ratio * width)
        self.padding_v = self.input_size - int(self.ratio * height)
        image = cv2.copyMakeBorder(
            image, 0, self.padding_v, 0, self.padding_h, cv2.BORDER_CONSTANT, (0, 0, 0))

        return image

    def __filter(self, detections, threshold):
        """Filter the detection results by score threshold."""
        # Get the detection results.
        boxes = detections['output_0'].numpy()[0]
        scores = detections['output_1'].numpy()[0]
        classes = detections['output_2'].numpy()[0]

        # Filter out the results by score threshold.
        mask = scores > threshold
        boxes = boxes[mask]
        scores = scores[mask]
        classes = classes[mask]

        return boxes, scores, classes

    def predict(self, image, threshold):
        """Run inference with image inputs.

        Args:
            image: a numpy array as an input image.

        Returns:
            predictions: result.
        """
        frame_tensor = tf.constant(image, dtype=tf.uint8)
        frame_tensor = tf.expand_dims(frame_tensor, axis=0)
        detections = self.__predict_fn(frame_tensor)
        boxes, scores, classes = self.__filter(detections, threshold)

        # Scale the box back to the original image size.
        boxes /= self.ratio

        # Crop out the padding area.
        boxes[:, 2] = np.minimum(
            boxes[:, 2], (self.input_size - self.padding_v)/self.ratio)
        boxes[:, 1] = np.minimum(
            boxes[:, 1], (self.input_size - self.padding_h)/self.ratio)

        return boxes, scores, classes

    def transform_to_square(self, boxes, scale=1.0, offset=(0, 0)):
        """Get the square bounding boxes.

        Args:
            boxes: input boxes [[ymin, xmin, ymax, xmax], ...]
            scale: ratio to scale the boxes
            offset: a tuple of offset to move the boxes (x, y)

        Returns:
            square boxes.
        """
        # Make them squares.
        ymins, xmins, ymaxs, xmaxs = np.split(boxes, 4, 1)
        center_x = np.floor_divide(xmins + xmaxs, 2)
        center_y = np.floor_divide(ymins + ymaxs, 2)
        margin = np.floor_divide(np.maximum(ymaxs - ymins,
                                            xmaxs - xmins) * scale, 2)
        boxes = np.concatenate((center_y-margin, center_x-margin,
                                center_y+margin, center_x+margin), axis=1)

        # Move
        boxes[:, 1::2] += offset[0]
        boxes[:, ::2] += offset[1]

        return boxes

    def clip_boxes(self, boxes, margins):
        """Clip the boxes to the safe margins.

        Args:
            boxes: input boxes [[ymin, xmin, ymax, xmax], ...].
            margins: a tuple of 4 int (top, left, bottom, right) as safe margins.

        Returns:
            boxes: clipped boxes.
            clip_mark: the mark of clipped sides.
        """
        top, left, bottom, right = margins

        clip_mark = (boxes[:, 0] < top, boxes[:, 1] < left,
                     boxes[:, 2] > bottom, boxes[:, 3] > right)

        boxes[:, 0] = np.maximum(boxes[:, 0], top)
        boxes[:, 1] = np.maximum(boxes[:, 1], left)
        boxes[:, 2] = np.minimum(boxes[:, 2], bottom)
        boxes[:, 3] = np.minimum(boxes[:, 3], right)

        return boxes, clip_mark