import datetime
import torch
from sklearn.metrics import precision_recall_fscore_support as prfs
from utils.parser import get_parser_with_args
from utils.helpers import (get_loaders, get_criterion,
                           load_model, initialize_metrics, get_mean_metrics,
                           set_metrics)
import os
import logging
import json
from tensorboardX import SummaryWriter
from tqdm import tqdm
import random
import numpy as np

"""
Initialize Parser and define arguments
"""
parser, metadata = get_parser_with_args()
opt = parser.parse_args()

"""
Initialize experiments log
"""
logging.basicConfig(level=logging.INFO)
writer = SummaryWriter(opt.log_dir + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/')

"""
Set up environment: define paths, download data, and set device
"""
os.environ["CUDA_VISIBLE_DEVICES"] = "4,5,6,7"
dev = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
logging.info('GPU AVAILABLE? ' + str(torch.cuda.is_available()))

def seed_torch(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    # torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

seed_torch(seed=777)

train_loader, val_loader = get_loaders(opt)

"""
Load Model or pretrained model
"""
path = opt.pretrain_path

if opt.pretrain_mode:
    try:
        logging.info('Loading pretrained dict model')
        model = load_model(opt, dev)
        model.load_state_dict(torch.load(path), strict=False)
    except:
        logging.info('Changed, Loading dict of full model')
        model = load_model(opt, dev)
        ptmodel = torch.load(path)
        ptmodel_dict = ptmodel.state_dict()
        model.load_state_dict(ptmodel_dict, strict=False)
else:
    logging.info('LOADING vanilla Model')
    model = load_model(opt, dev)

### load full model directly ###
#device_ids = list(range(opt.num_gpus))
#model = torch.load(path).to(dev)
#model = torch.nn.DataParallel(model, device_ids=device_ids)

"""
Activate loss, optimizer, scheduler function
"""
criterion = get_criterion(opt)
optimizer = torch.optim.AdamW(model.parameters(), lr=opt.learning_rate) # Be careful when you adjust learning rate, you can refer to the linear scaling rule
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=8, gamma=0.5)

"""
Set starting values
"""
best_metrics = {'cd_f1scores': -1, 'cd_recalls': -1, 'cd_precisions': -1}
logging.info('STARTING training')
total_step = -1

for epoch in range(opt.epochs):
    train_metrics = initialize_metrics()
    val_metrics = initialize_metrics()

    """
    Begin Training
    """
    model.train()
    logging.info('SET model mode to train!')
    batch_iter = 0
    tbar = tqdm(train_loader)
    for batch_img1, batch_img2, labels in tbar:
        tbar.set_description("epoch {} info ".format(epoch) + str(batch_iter) + " - " + str(batch_iter+opt.batch_size))
        batch_iter = batch_iter+opt.batch_size
        total_step += 1
        # Set variables for training
        batch_img1 = batch_img1.float().to(dev)
        batch_img2 = batch_img2.float().to(dev)
        labels = labels.long().to(dev)

        # Zero the gradient
        optimizer.zero_grad()

        # Get model predictions, calculate loss, backprop
        cd_preds = model(batch_img1, batch_img2)
        '''
        cd_shape = cd_preds[-1]
        print(" ")
        print('cd_shape: ' + str(cd_shape.shape))
        '''
        cd_loss = criterion(cd_preds, labels)
        loss = cd_loss
        loss.backward()
        optimizer.step()

        cd_preds = cd_preds[-1]
        _, cd_preds = torch.max(cd_preds, 1)

        # Calculate and log other batch metrics
        cd_corrects = (100 *
                       (cd_preds.squeeze().byte() == labels.squeeze().byte()).sum() /
                       (labels.size()[0] * (opt.patch_size**2)))

        cd_train_report = prfs(labels.data.cpu().numpy().flatten(),
                               cd_preds.data.cpu().numpy().flatten(),
                               average='binary',
                               pos_label=1)

        train_metrics = set_metrics(train_metrics,
                                    cd_loss,
                                    cd_corrects,
                                    cd_train_report,
                                    scheduler.get_last_lr())

        # log the batch mean metrics
        mean_train_metrics = get_mean_metrics(train_metrics)

        for k, v in mean_train_metrics.items():
            writer.add_scalars(str(k), {'train': v}, total_step)

        # clear batch variables from memory
        del batch_img1, batch_img2, labels

    scheduler.step()
    logging.info("EPOCH {} TRAIN METRICS".format(epoch) + str(mean_train_metrics))

    """
    Begin Validation
    """
    model.eval()
    with torch.no_grad():
        for batch_img1, batch_img2, labels in val_loader:
            # Set variables for training
            batch_img1 = batch_img1.float().to(dev)
            batch_img2 = batch_img2.float().to(dev)
            labels = labels.long().to(dev)

            # Get predictions and calculate loss
            cd_preds = model(batch_img1, batch_img2)

            cd_loss = criterion(cd_preds, labels)

            cd_preds = cd_preds[-1]
            _, cd_preds = torch.max(cd_preds, 1)

            # Calculate and log other batch metrics
            cd_corrects = (100 *
                           (cd_preds.squeeze().byte() == labels.squeeze().byte()).sum() /
                           (labels.size()[0] * (opt.patch_size**2)))

            cd_val_report = prfs(labels.data.cpu().numpy().flatten(),
                                 cd_preds.data.cpu().numpy().flatten(),
                                 average='binary',
                                 pos_label=1)

            val_metrics = set_metrics(val_metrics,
                                      cd_loss,
                                      cd_corrects,
                                      cd_val_report,
                                      scheduler.get_last_lr())

            # log the batch mean metrics
            mean_val_metrics = get_mean_metrics(val_metrics)

            for k, v in mean_train_metrics.items():
                writer.add_scalars(str(k), {'val': v}, total_step)

            # clear batch variables from memory
            del batch_img1, batch_img2, labels

        logging.info("EPOCH {} VALIDATION METRICS".format(epoch)+str(mean_val_metrics))

        """
        Store the weights of good epochs based on validation results
        """
        if ((mean_val_metrics['cd_precisions'] > best_metrics['cd_precisions'])
                or
                (mean_val_metrics['cd_recalls'] > best_metrics['cd_recalls'])
                or
                (mean_val_metrics['cd_f1scores'] > best_metrics['cd_f1scores'])):

            # Insert training and epoch information to metadata dictionary
            logging.info('updata the model')
            metadata['validation_metrics'] = mean_val_metrics

            # Save model(state_dict) and log
            if not os.path.exists('./tmp'):
                os.mkdir('./tmp')
            with open('./tmp/metadata_epoch_' + str(epoch) + '.json', 'w') as fout:
                json.dump(metadata, fout)

            torch.save(model.state_dict(), './tmp/checkpoint_epoch_'+str(epoch)+'.pt')

            # comet.log_asset(upload_metadata_file_path)
            best_metrics = mean_val_metrics

        print('An epoch finished.')
writer.close()  # close tensor board
print('Done!')


"""
datasets
--------

This utility loads data, inheriting from the PyTorch Dataset model.
This approach allows very large datasets---possibly larger than RAM---to inform models.

FPADataset
############
"""

import torch
from torch.utils.data import Dataset
from astropy.io import fits
import numpy as np
import kornia
import ccdproc
import astropy.units as u
from astropy.nddata import CCDData
from sklearn.cluster import KMeans
import warnings
import logging

logging.getLogger("ccdproc").setLevel(logging.ERROR)

# custom dataset loader
class FPADataset(Dataset):
    r"""Read in two AB nods of an Echellogram (experimental)

    Args:
        ybounds (tuple of ints): the :math:`(y_0, y_{max})` bounds for isolated echelle trace
        root_dir (str): path to the directory containing 2D echellograms as fits files
        inpaint_bad_pixels (bool): flag for whether or not to inpaint bad pixels

    """

    def __init__(
        self,
        ybounds=(425, 510),
        root_dir=None,
        inpaint_bad_pixels=False,
        inpaint_cosmic_rays=False,
    ):
        super().__init__()

        if root_dir is None:
            root_dir = "/home/gully/GitHub/ynot/test/data/2012-11-27/"
        self.root_dir = root_dir
        self.nirspec_collection = self.create_nirspec_collection()
        # self.unique_objects = self.get_unique_objects()
        # self.label_nirspec_nods()
        nodA_path = self.root_dir + "/NS.20121127.49332.fits"
        nodA_data = fits.open(nodA_path)[0].data.astype(np.float64)
        nodA = torch.tensor(nodA_data)

        nodB_path = self.root_dir + "/NS.20121127.50726.fits"
        nodB_data = fits.open(nodB_path)[0].data.astype(np.float64)
        nodB = torch.tensor(nodB_data)

        # Read in the Bad Pixel mask
        self.bpm = self.load_bad_pixel_mask()

        data_full = torch.stack([nodA, nodB])  # Creates NxHxW tensor
        # Inpaint bad pixels.  In the future we will simply neglect these pixels
        if inpaint_bad_pixels:
            data_full = self.inpaint_bad_pixels(data_full)

        self.n_images = len(data_full[:, 0, 0])
        self.gain = 5.8  # e/ADU, per NIRSPEC documentation

        if inpaint_cosmic_rays:
            for ii in range(self.n_images):
                nod_ccd = CCDData(data_full[ii].numpy(), unit="adu")
                out = ccdproc.cosmicray_lacosmic(
                    nod_ccd,
                    readnoise=23.0,
                    gain=self.gain,
                    verbose=False,
                    satlevel=1.0e7,
                    sigclip=7.0,
                    sepmed=False,
                    cleantype="medmask",
                    fsmode="median",
                )
                data_full[ii] = torch.tensor(out.data)
        else:
            data_full = data_full * self.gain

        data = data_full[:, ybounds[0] : ybounds[1], :]
        data = data.permute(0, 2, 1)

        self.pixels = data
        self.index = torch.tensor([0, 1])

    def __getitem__(self, index):
        return (self.index[index], self.pixels[index])

    def __len__(self):
        return len(self.pixels[:, 0, 0])

    def load_bad_pixel_mask(self):
        """Load the global bad pixel mask"""
        bad_pixel_mask_path = "../../zoja/ccdproc/reduced/static/bad_pixel_mask.fits"
        bpm_data = fits.open(bad_pixel_mask_path)[0].data.astype(np.bool)
        return torch.tensor(bpm_data)

    def inpaint_bad_pixels(self, data_tensor):
        """Inpaint the bad pixels

        Args:
            data_tensor (tensor of shape NHW): Tensor of data destined for masking
                in on the HW axis, and reapplied over all elements in batch axis.

        """
        smoothed_data = kornia.filters.median_blur(
            data_tensor.unsqueeze(1), (5, 5)
        ).squeeze()
        data_tensor[:, self.bpm] = smoothed_data[:, self.bpm]
        return data_tensor

    def create_nirspec_collection(self):
        """Create a collection of Keck NIRSPEC echelle spectra"""
        keywords = [
            "imagetyp",
            "filname",
            "slitname",
            "itime",
            "frameno",
            "object",
            "ut",
            "airmass",
            "dispers",
            "slitwidt",
            "slitlen",
            "ra",
            "dec",
        ]
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            ims = (
                ccdproc.ImageFileCollection(
                    self.root_dir, keywords=keywords, glob_include="NS*.fits",
                )
                .filter(dispers="high")
                .filter(regex_match=True, slitlen="12|24")
            )
        return ims

    def get_unique_objects(self):
        """Return the unique object names from a NIRSPEC collection"""

        objects = np.unique(
            self.nirspec_collection.filter(imagetyp="object")
            .summary["object"]
            .data.data
        )
        return objects

    def label_nods_from_coordinates(self, coords):
        """Label the ABBA nods given Telescope RA, Dec coordinates"""
        nod_dict = {0: "A", 1: "B"}
        kmeans = KMeans(n_clusters=2)
        kmeans.fit(coords)
        y_pred = kmeans.predict(coords)
        # Enforce that the first nod is usually "A"
        if y_pred[0] != 0:
            y_pred = 1 - y_pred
        return [nod_dict[key] for key in y_pred]

    # def label_nirspec_nods(self):
    #    """Label the ABBA nods in-place given a nirspec collection"""
    #    for object in self.unique_objects:
    #        target_subset = self.nirspec_collection.filter(
    #            imagetyp="object", object=object
    #        )
    #        target_table = target_subset.summary[["file", "object", "ra", "dec"]]
    #        coords = np.vstack([target_table[col].data.data for col in ["ra", "dec"]]).T
    #        target_subset.summary["nod"] = self.label_nods_from_coordinates(coords)

import numpy,import astropy,from astropy


import os
import sys
import numpy as np
import tensorflow as tf
import logging
tf.get_logger().setLevel(logging.ERROR)
from model import TrRefiner

n_chunk_default = 5
n_filter_default = 128
def parse_input_args():
    import argparse
    opt = argparse.ArgumentParser\
            (description='''AngPred: Backbone torsion angle prediction using ResNet''')
    #
    # Training config
    opt.add_argument('-train_param', dest='train_param', metavar='TRAIN_PARAM', default='param_train.json',\
                         help="json file containing all the training related parameters (e.g. learning rate, etc)")
    #
    # Model config
    opt.add_argument('-num_chunks', '--num_chunks', dest='num_chunks', metavar='num_chunks', type=int, default=n_chunk_default, \
                         help='The number of ResNet chunks [%d]'%n_chunk_default)
    opt.add_argument('-n2d_filters', '--n2d_filters', dest='n2d_filters', metavar='n2d_filters', type=int, default=n_filter_default, \
                         help='The number of 2D filters [%d]'%n_filter_default)
    #
    # Flag for training or not
    opt.add_argument('-train', '--train', dest='train', default=False, action='store_true',\
                         help='for training [False]')
    opt.add_argument('-att', '--att', dest='att', default=False, action='store_true',\
                         help='use pixel-attention')
    #
    # Inputs for prediction mode
    opt.add_argument('-a3m_fn', dest='a3m_fn', default=None, \
                        help="MSA file in a3m format")
    opt.add_argument('-pdb_fn_s', dest='pdb_fn_s', default=[None], nargs='+',\
                        help="Input structures in pdb format")
    opt.add_argument('-npz_fn_s', dest='npz_fn_s', default=[None], nargs='+',\
                        help="Previous 6D coords distribution prediction results in npz format. It must contain dist, omega, theta, phi")
    opt.add_argument('-SS_fn_s', dest='SS_fn_s', default=[None], nargs='+',\
                        help="predicted SS in npz file from Emin")
    opt.add_argument('-acc_fn_s', dest='acc_fn_s', default=[None], nargs='+',\
                        help="predicted acc in npz file for all input pdbs")
    opt.add_argument('-out_fn', dest='out_fn', default=None,\
                        help="Output file name")
    opt.add_argument('-model_dir', dest='model_dir', default="%s/model"%(os.path.dirname(__file__)),\
                        help="Directory for pre-trained ML model checkpoint files")

    if len(sys.argv) == 1:
        opt.print_help()
        sys.exit()
    #
    FLAGS = opt.parse_args()

    return FLAGS

def main():
    args = parse_input_args()
    #
    if not args.train: # prediction mode
        error = False
        if args.a3m_fn == None:
            sys.stderr.write("ERROR: Please provide MSA in a3m format\n")
            error = True
        if args.out_fn == None:
            sys.stderr.write("ERROR: Please provide output file name\n")
            error = True
        if args.pdb_fn_s[0] == None:
            sys.stderr.write("ERROR: Please provide structure model generated by trRosetta\n")
            error = True
        if args.npz_fn_s[0] == None:
            sys.stderr.write("ERROR: please provide 6d coords prediction results in npz format\n")
            error = True
        if not os.path.exists("%s_rep1"%args.model_dir):
            sys.stderr.write("ERROR: please check directory path for pre-trained checkpoint files, %s\n"%args.model_dir)
            error = True
        if error:
            return

    run_config = tf.ConfigProto(
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.92)
    )
    with tf.Session(config=run_config) as sess:

        if args.train:
            model = TrRefiner(sess, num_chunks=args.num_chunks, n2d_filters=args.n2d_filters, pdb_in=False, train=True, config_file=args.train_param, attention=args.att)
            model.train(args.train_param)
        else:
            model = TrRefiner(sess, num_chunks=args.num_chunks, n2d_filters=args.n2d_filters, pdb_in=True, attention=args.att)
            model.predict(args.model_dir, args.a3m_fn, args.npz_fn_s, args.pdb_fn_s, args.acc_fn_s, args.SS_fn_s, args.out_fn)

main()


import time
import argparse
import os
import numpy as np
import glob
from sklearn.metrics import accuracy_score, confusion_matrix
from sklearn.externals import joblib

import config
import util

def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('--path', dest='path', help='Path to image', default=None, type=str)
    parser.add_argument('--accuracy', action='store_true', help='To print accuracy score')
    parser.add_argument('--plot_confusion_matrix', action='store_true')
    parser.add_argument('--execution_time', action='store_true')
    parser.add_argument('--store_activations', action='store_true')
    parser.add_argument('--novelty_detection', action='store_true')
    parser.add_argument('--model', type=str, required=True, help='Base model architecture',
                        choices=[config.MODEL_RESNET50, config.MODEL_RESNET152, config.MODEL_INCEPTION_V3,
                                 config.MODEL_VGG16])
    parser.add_argument('--data_dir', help='Path to data train directory')
    parser.add_argument('--batch_size', default=500, type=int, help='How many files to predict on at once')
    args = parser.parse_args()
    return args

def get_files(path):
    if os.path.isdir(path):
        files = glob.glob(path + '*.jpg')
    elif path.find('*') > 0:
        files = glob.glob(path)
    else:
        files = [path]

    if not len(files):
        print('No images found by the given path')
        exit(1)

    return files

def get_inputs_and_trues(files):
    inputs = []
    y_true = []

    for i in files:
        x = model_module.load_img(i)
        try:
            image_class = i.split(os.sep)[-2]
            keras_class = int(classes_in_keras_format[image_class])
            y_true.append(keras_class)
        except Exception:
            y_true.append(os.path.split(i)[1])

        inputs.append(x)

    return y_true, inputs

def predict(path):
    files = get_files(path)
    n_files = len(files)
    print('Found {} files'.format(n_files))

    if args.novelty_detection:
        activation_function = util.get_activation_function(model, model_module.noveltyDetectionLayerName)
        novelty_detection_clf = joblib.load(config.get_novelty_detection_model_path())

    y_trues = []
    predictions = np.zeros(shape=(n_files,))
    nb_batch = int(np.ceil(n_files / float(args.batch_size)))
    for n in range(0, nb_batch):
        print('Batch {}'.format(n))
        n_from = n * args.batch_size
        n_to = min(args.batch_size * (n + 1), n_files)

        y_true, inputs = get_inputs_and_trues(files[n_from:n_to])
        y_trues += y_true

        if args.store_activations:
            util.save_activations(model, inputs, files[n_from:n_to], model_module.noveltyDetectionLayerName, n)

        if args.novelty_detection:
            activations = util.get_activations(activation_function, [inputs[0]])
            nd_preds = novelty_detection_clf.predict(activations)[0]
            print(novelty_detection_clf.__classes[nd_preds])

        if not args.store_activations:
            # Warm up the model
            if n == 0:
                print('Warming up the model')
                start = time.clock()
                model.predict(np.array([inputs[0]]))
                end = time.clock()
                print('Warming up took {} s'.format(end - start))

            # Make predictions
            start = time.clock()
            ###############################
            # global out
            ###############################
            out = model.predict(np.array(inputs))
            print('out shape : ', out.shape)
            end = time.clock()
            predictions[n_from:n_to] = np.argmax(out, axis=1)
            print('Prediction on batch {} took: {}'.format(n, end - start))

    if not args.store_activations:
        #######################
        store_label = []
        ######################
        for i, p in enumerate(predictions):
            recognized_class = list(classes_in_keras_format.keys())[list(classes_in_keras_format.values()).index(p)]
            print('| should be {} ({}) -> predicted as {} ({}) -> score : {}'.format(y_trues[i],
                                                                                files[i].split(os.sep)[-2], p,
                                                                                recognized_class, np.max(out[i])))
            ###########################################
            if y_trues[i] == int(p):
                store_current_label = 1
            else:
                store_current_label = 0
            store_label.append([np.max(out[i]), store_current_label])
        # print(np.array(store_label).shape)
        # print(store_label)
        ###########################################
        if args.accuracy:
            print('Accuracy {}'.format(accuracy_score(y_true=y_trues, y_pred=predictions)))

        if args.plot_confusion_matrix:
            cnf_matrix = confusion_matrix(y_trues, predictions)
            util.plot_confusion_matrix(cnf_matrix, config.classes, normalize=False)
            util.plot_confusion_matrix(cnf_matrix, config.classes, normalize=True)
        #############
        return store_label

###################
# calc AP
###################
def metrix(data):
    import copy

    # data = [[0.23, 0], [0.76, 1], [0.01, 0], [0.91, 1], [0.13, 0],
    #         [0.45, 0], [0.12, 1], [0.03, 0], [0.38, 1], [0.11, 0],
    #         [0.03, 0], [0.09, 0], [0.65, 0], [0.07, 0], [0.12, 0],
    #         [0.24, 1], [0.1, 0], [0.23,  0], [0.46, 0], [0.08, 1]]

    # # list to numpy
    data = np.array(data)

    # according to the first column to sort
    # array([2, 7, 10, 13, 19, 11, 16, 9, 14, 6, 4, 0, 17, 15, 8, 5, 18, 12, 1, 3], dtype=int64)
    # means that 2 -> 0.01  7 -> 0.03  ....  3 -> 0.91
    data1 = np.lexsort(data[:, ::-1].T)
    # print(data[data1])

    # reverse order
    # array([ 3  1  8 15  6 19 12 18  5 17  0  4 14  9 16 11 13 10  7  2], dtype=int64)
    data1_reverse = data1[::-1]
    # print(data1_reverse)

    # sorted data
    # [[0.91 1] [0.76 1] ... [0.01 0]]
    data_result = data[data1_reverse]
    # print(data_result)

    # sorted label
    data_result_label = data_result[:, -1]
    # print(data_result_label)

    # calc true positive + false positive, that is those samples whose label is 1
    all_label_is_1_sum = sum(data_result_label == 1)
    # print(label_is_1_sum)

    precision = []
    recall = []

    # print(data_result_label[0:1])

    for i in range(1, len(data_result_label) + 1):
        # i starting from 1
        # while loop to the current
        current_label_is_1_sum = np.sum(data_result_label[:i] == 1)
        precision_temp = current_label_is_1_sum / len(data_result_label[:i])
        recall_temp = current_label_is_1_sum / all_label_is_1_sum
        precision.append(round(precision_temp, 10))
        recall.append(round(recall_temp, 10))

    # print('precision:',  precision)
    # print('recall:', recall)

    assert len(recall) == len(precision) and len(recall) > 0

    temp = []
    for i in range(len(recall)):
        current_recall_value_sum = recall.count(recall[i])
        if current_recall_value_sum > 1:
            # current recall index  first
            current_recall_index = recall.index(recall[i])
            current_recall_info = [current_recall_index, current_recall_value_sum]
            if current_recall_info not in temp:
                temp.append(current_recall_info)

    # print(temp)

    max_precision = copy.deepcopy(precision)

    for temp_every in temp:
        precision_start_index = temp_every[0]
        precision_end_index = temp_every[1] + temp_every[0]  # not include

        # get max
        temp_max = np.max(max_precision[temp_every[0]:(temp_every[1] + temp_every[0])])

        for i in range(precision_start_index, precision_end_index):
            max_precision[i] = temp_max

    # print('max_presicion : ', max_precision)

    total = sum(np.array(max_precision))
    print(total / len(max_precision))
    return (total / len(max_precision))

if __name__ == '__main__':
    tic = time.clock()

    args = parse_args()
    print('=' * 50)
    print('Called with args:')
    print(args)

    if args.data_dir:
        config.data_dir = args.data_dir
        config.set_paths()
    if args.model:
        config.model = args.model

    util.set_img_format()
    model_module = util.get_model_class_instance()
    model = model_module.load()

    classes_in_keras_format = util.get_classes_in_keras_format()

    all_metrix = []
    print('args.path:', os.listdir(args.path))
    for cow_dir in os.listdir(args.path):
        root = args.path + '/' + str(cow_dir) + '/'
        store_label = predict(root)
        # print(store_label)
        all_metrix.append(metrix(store_label))

    total = 0
    for i in range(len(all_metrix)):
        total += all_metrix[i]
    print('mAP:', (total / len(all_metrix)))

    if args.execution_time:
        toc = time.clock()
        print('Time: %s' % (toc - tic))


"""
Copyright (R) @huawei.com, all rights reserved
-*- coding:utf-8 -*-
"""
import sys
import cv2 as cv
import numpy as np
import os
import re
import time
path = os.path.dirname(os.path.abspath(__file__))

sys.path.append(os.path.join(path, ".."))
sys.path.append(os.path.join(path, "../../../../common/"))
sys.path.append(os.path.join(path, "../../../../common/acllite"))

from acllite_model import AclLiteModel
from acllite_image import AclLiteImage
from acllite_resource import AclLiteResource
from acllite import presenteragent
import presenteragent.presenter_channel as presenter_channel

MODEL_WIDTH = 224
MODEL_HEIGHT = 224
model_path = '../model/colorization.om'
COLORIZATION_CONF ="../scripts/param.conf"

def preprocess(frame):
    """
    Preprocess the images before they are sent into the model
    """
    #Read the image
    bgr_img = frame.astype(np.float32)
    #Opencv reads the picture as (N) HWC to get the HW value
    orig_shape = bgr_img.shape[:2]
    #Normalize the picture
    bgr_img = bgr_img / 255.0
    #Convert the picture to Lab space
    lab_img = cv.cvtColor(bgr_img, cv.COLOR_BGR2Lab)
    #Gets the L
    orig_l = lab_img[:, :, 0]

    if not orig_l.flags['C_CONTIGUOUS']:
        orig_l = np.ascontiguousarray(orig_l)

    #resize
    lab_img = cv.resize(lab_img, (MODEL_WIDTH, MODEL_HEIGHT)).astype(np.float32)

    l_data = lab_img[:, :, 0]
    if not l_data.flags['C_CONTIGUOUS']:
        l_data = np.ascontiguousarray(l_data)

    #The L-part minus the average
    l_data = l_data - 50

    return orig_shape, orig_l, l_data

def postprocess(result_list, orig_shape, orig_l):
    """
    Post process the images after model reasoning
    """
    result_list[0] = result_list[0].reshape(1, 2, 56, 56).transpose(0, 2, 3, 1)
    result_array = result_list[0][0]
    #Scales the predicted ab channel value to the original picture size
    ab_data = cv.resize(result_array, orig_shape[::-1])
    #get the Lab value of the predicted image
    result_lab = np.concatenate((orig_l[:, :, np.newaxis], ab_data), axis = 2)
    #Convert Lab to RGBU888
    result_bgr = (255 * np.clip(cv.cvtColor(result_lab, cv.COLOR_Lab2BGR), 0, 1)).astype('uint8')
    return cv.imencode('.jpg', result_bgr)[1]

def main():
    """
    acl resource initialization
    """
    acl_resource = AclLiteResource()
    acl_resource.init()
    #load model
    model = AclLiteModel(model_path)
    chan = presenter_channel.open_channel(COLORIZATION_CONF)
    if chan is None:
        print("Open presenter channel failed")
        return

    lenofUrl = len(sys.argv)

    if lenofUrl <= 1:
        print("[ERROR] Please input mp4/Rtsp URL")
        exit()
    elif lenofUrl >= 3:
        print("[ERROR] param input Error")
        exit()
    URL = sys.argv[1]
    URL1 = re.match('rtsp://', URL)
    URL2 = re.search('.mp4', URL)

    if URL1 is None and URL2 is None:
        print("[ERROR] should input correct URL")
        exit()

    cap = cv.VideoCapture(URL)
    #Gets the total frames
    frames_num = cap.get(7)
    currentFrames = 0

    while True:
        #read image
        ret, frame = cap.read()
        if ret is not True:
            print("read None image, break")
            break

        if currentFrames == frames_num - 1:
            currentFrames = 0
            cap.set(1, 0)

        currentFrames += 1
        #Gets the L channel value
        orig_shape, orig_l, l_data = preprocess(frame)
        result_list = model.execute([l_data,])
        result_jpeg = postprocess(result_list, orig_shape, orig_l)
        chan.send_image(orig_shape[0], orig_shape[1], result_jpeg)

if __name__ == '__main__':
    main()


import mlprocessors as mlpr
import os
import time
import numpy as np
from os.path import join
from subprocess import Popen, PIPE
import shlex
import random
import string
import shutil
from spikeforest import mdaio
import spikeextractors as se
from spikeforest import SFMdaRecordingExtractor, SFMdaSortingExtractor
from mountaintools import client as mt
import json
import traceback
from .install_jrclust import install_jrclust

class JRClust(mlpr.Processor):
    """
    JRClust v4 wrapper
      written by J. James Jun, May 8, 2019
      modified from `spikeforest/spikesorters/ironclust/ironclust.py`

    [Installation instruction in SpikeForest environment]
    1. Run `git clone https://github.com/JaneliaSciComp/JRCLUST`
    2. Activate conda environment for SpikeForest
    3. Create `JRCLUST_PATH` and `MDAIO_PATH`

    See: James Jun, et al. Real-time spike sorting platform for
    high-density extracellular probes with ground-truth validation and drift correction
    https://github.com/JaneliaSciComp/JRCLUST
    """

    NAME = 'JRClust'
    VERSION = '0.1.0'
    ENVIRONMENT_VARIABLES = [
        'NUM_WORKERS', 'MKL_NUM_THREADS', 'NUMEXPR_NUM_THREADS', 'OMP_NUM_THREADS', 'TEMPDIR']
    ADDITIONAL_FILES = ['*.m', '*.prm']
    CONTAINER = None

    recording_dir = mlpr.Input('Directory of recording', directory=True)
    channels = mlpr.IntegerListParameter(
        description='List of channels to use.', optional=True, default=[])
    firings_out = mlpr.Output('Output firings file')

    detect_sign = mlpr.IntegerParameter(
        optional=True, default=-1, description='Use -1, 0, or 1, depending on the sign of the spikes in the recording')
    adjacency_radius = mlpr.FloatParameter(
        optional=True, default=50, description='')
    detect_threshold = mlpr.FloatParameter(
        optional=True, default=4.5, description='detection threshold')
    freq_min = mlpr.FloatParameter(
        optional=True, default=300, description='Use 0 for no bandpass filtering')
    freq_max = mlpr.FloatParameter(
        optional=True, default=3000, description='Use 0 for no bandpass filtering')
    merge_thresh = mlpr.FloatParameter(
        optional=True, default=0.98, description='Threshold for automated merging')
    pc_per_chan = mlpr.IntegerParameter(
        optional=True, default=1, description='Number of principal components per channel')

    # added in version 0.2.4
    filter_type = mlpr.StringParameter(
        optional=True, default='bandpass', description='{none, bandpass, wiener, fftdiff, ndiff}')
    nDiffOrder = mlpr.FloatParameter(optional=True, default=2, description='')
    common_ref_type = mlpr.StringParameter(
        optional=True, default='none', description='{none, mean, median}')
    min_count = mlpr.IntegerParameter(
        optional=True, default=30, description='Minimum cluster size')
    fGpu = mlpr.BoolParameter(
        optional=True, default=False, description='Use GPU if available')
    fParfor = mlpr.BoolParameter(
        optional=True, default=True, description='Use parfor if available')
    feature_type = mlpr.StringParameter(
        optional=True, default='gpca', description='{gpca, pca, vpp, vmin, vminmax, cov, energy, xcov}')

    def run(self):
        tmpdir = _get_tmpdir('jrclust')

        try:
            recording = SFMdaRecordingExtractor(self.recording_dir)
            params = read_dataset_params(self.recording_dir)
            if len(self.channels) > 0:
                recording = se.SubRecordingExtractor(
                    parent_recording=recording, channel_ids=self.channels)
            if not os.path.exists(tmpdir):
                os.mkdir(tmpdir)

            all_params = dict()
            for param0 in self.PARAMETERS:
                all_params[param0.name] = getattr(self, param0.name)
            sorting = jrclust_helper(
                recording=recording,
                tmpdir=tmpdir,
                params=params,
                **all_params,
            )
            SFMdaSortingExtractor.write_sorting(
                sorting=sorting, save_path=self.firings_out)
        except:
            if os.path.exists(tmpdir):
                if not getattr(self, '_keep_temp_files', False):
                    shutil.rmtree(tmpdir)
            raise
        if not getattr(self, '_keep_temp_files', False):
            shutil.rmtree(tmpdir)

def jrclust_helper(
        *,
        recording,  # Recording object
        tmpdir,  # Temporary working directory
        params=dict(),
        **kwargs):

    jrclust_path = os.environ.get('JRCLUST_PATH_DEV', None)
    if jrclust_path:
        print('Using jrclust from JRCLUST_PATH_DEV directory: {}'.format(jrclust_path))
    else:
        try:
            print('Auto-installing jrclust.')
            jrclust_path = install_jrclust(
                repo='https://github.com/JaneliaSciComp/JRCLUST.git',
                commit='68ffb3ef064f97aca7043b7faac49c34a58997d9'
            )
        except:
            traceback.print_exc()
            raise Exception('Problem installing jrclust. You can set the JRCLUST_PATH_DEV to force to use a particular path.')
    print('Using jrclust from: {}'.format(jrclust_path))

    dataset_dir = os.path.join(tmpdir, 'jrclust_dataset')
    # Generate three files in the dataset directory: raw.mda, geom.csv, params.json
    SFMdaRecordingExtractor.write_recording(
        recording=recording, save_path=dataset_dir, params=params)

    samplerate = recording.get_sampling_frequency()

    print('Reading timeseries header...')
    raw_mda = os.path.join(dataset_dir, 'raw.mda')
    HH = mdaio.readmda_header(raw_mda)
    num_channels = HH.dims[0]
    num_timepoints = HH.dims[1]
    duration_minutes = num_timepoints / samplerate / 60
    print('Num. channels = {}, Num. timepoints = {}, duration = {} minutes'.format(
        num_channels, num_timepoints, duration_minutes))

    print('Creating argfile.txt...')
    txt = ''
    for key0, val0 in kwargs.items():
        txt += '{}={}\n'.format(key0, val0)
    if 'scale_factor' in params:
        txt += 'bitScaling={}\n'.format(params["scale_factor"])
    txt += 'sampleRate={}\n'.format(samplerate)
    _write_text_file(dataset_dir + '/argfile.txt', txt)

    # new method
    source_path = os.path.dirname(os.path.realpath(__file__))
    print('Running jrclust in {tmpdir}...'.format(tmpdir=tmpdir))
    cmd = '''
        addpath('{jrclust_path}', '{source_path}', '{source_path}/mdaio');
        try
            p_jrclust('{tmpdir}', '{dataset_dir}/raw.mda', '{dataset_dir}/geom.csv', '{tmpdir}/firings.mda', '{dataset_dir}/argfile.txt');
        catch
            fprintf('----------------------------------------');
            fprintf(lasterr());
            quit(1);
        end
        quit(0);
    '''
    cmd = cmd.format(jrclust_path=jrclust_path, tmpdir=tmpdir, dataset_dir=dataset_dir, source_path=source_path)

    matlab_cmd = mlpr.ShellScript(cmd, script_path=tmpdir + '/run_jrclust.m', keep_temp_files=True)
    matlab_cmd.write()

    shell_cmd = '''
        #!/bin/bash
        cd {tmpdir}
        matlab -nosplash -nodisplay -r run_jrclust
    '''.format(tmpdir=tmpdir)
    shell_cmd = mlpr.ShellScript(shell_cmd, script_path=tmpdir + '/run_jrclust.sh', keep_temp_files=True)
    shell_cmd.write(tmpdir + '/run_jrclust.sh')
    shell_cmd.start()

    retcode = shell_cmd.wait()

    if retcode != 0:
        raise Exception('jrclust returned a non-zero exit code')

    # parse output
    result_fname = tmpdir + '/firings.mda'
    if not os.path.exists(result_fname):
        raise Exception('Result file does not exist: ' + result_fname)

    firings = mdaio.readmda(result_fname)
    sorting = se.NumpySortingExtractor()
    sorting.set_times_labels(firings[1, :], firings[2, :])
    return sorting

def _read_text_file(fname):
    with open(fname) as f:
        return f.read()

def _write_text_file(fname, str):
    with open(fname, 'w') as f:
        f.write(str)

def _run_command_and_print_output(command):
    with Popen(shlex.split(command), stdout=PIPE, stderr=PIPE) as process:
        while True:
            output_stdout = process.stdout.readline()
            output_stderr = process.stderr.readline()
            if (not output_stdout) and (not output_stderr) and (process.poll() is not None):
                break
            if output_stdout:
                print(output_stdout.decode())
            if output_stderr:
                print(output_stderr.decode())
        rc = process.poll()
        return rc

def read_dataset_params(dsdir):
    # ca = _load_required_modules()
    fname1 = dsdir + '/params.json'
    fname2 = mt.realizeFile(path=fname1)
    if not fname2:
        raise Exception('Unable to find file: ' + fname1)
    if not os.path.exists(fname2):
        raise Exception('Dataset parameter file does not exist: ' + fname2)
    with open(fname2) as f:
        return json.load(f)

# To be shared across sorters (2019.05.05)
def _get_tmpdir(sorter_name):
    code = ''.join(random.choice(string.ascii_uppercase) for x in range(10))
    tmpdir0 = os.environ.get('TEMPDIR', '/tmp')
    tmpdir = os.path.join(tmpdir0, '{}-tmp-{}'.format(sorter_name, code))
    # reset the output folder
    if os.path.exists(tmpdir):
        shutil.rmtree(str(tmpdir))
    else:
        os.makedirs(tmpdir)
    return tmpdir


# Copyright 2022 The T5X Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Provides model subclasses with Mixture of Experts support."""

import dataclasses
from typing import Callable, Mapping, Optional, Sequence, Tuple, Union

import clu.metrics as clu_metrics
from flax import core as flax_core
from flax import linen as nn
from flax import optim
from flax import traverse_util
from flax.core import scope as flax_scope
import jax.numpy as jnp
import seqio
from t5x import decoding
from t5x import losses
from t5x import metrics as metrics_lib
from t5x import models

AveragePerStep = metrics_lib.AveragePerStep
DecodeFnCallable = models.DecodeFnCallable
FrozenVariableDict = flax_scope.FrozenVariableDict
MetricsMap = metrics_lib.MetricsMap
PyTreeDef = models.PyTreeDef
Sum = metrics_lib.Sum

@dataclasses.dataclass()
class ExpertMetrics:
  """Metrics for analyzing diversity among experts in mixture of experts models.

  Attributes:
    auxiliary_loss: Auxiliary load balancing loss.
    router_z_loss: Router z-loss. Encourages router logits to remain small in an
      effort to improve stability.
    fraction_tokens_left_behind: Fraction of tokens NOT processed by any expert.
    expert_usage: Fraction of total capacity, across all experts, used to
      process tokens. Larger expert capacities or non-uniform token routing will
      result in smaller expert usage values.
    router_confidence: How confident the router is about the tokens that it has
      routed.
  """
  auxiliary_loss: float
  router_z_loss: float

  fraction_tokens_left_behind: float
  expert_usage: float
  router_confidence: float

class MoeEncoderDecoderModel(models.EncoderDecoderModel):
  """Subclass which propagates MoE auxiliary loss and metrics."""

  def __init__(
      self,
      module: nn.Module,
      input_vocabulary: seqio.Vocabulary,
      output_vocabulary: seqio.Vocabulary,
      optimizer_def: optim.OptimizerDef,
      decode_fn: DecodeFnCallable = decoding.beam_search,
      feature_converter_cls: Optional[Callable[...,
                                               seqio.FeatureConverter]] = None,
      label_smoothing: float = 0.0,
      z_loss: float = 0.0,
      loss_normalizing_factor: Optional[float] = None,
      aux_loss_factor: float = 0.,
      router_z_loss_factor: float = 0.):
    super().__init__(
        module=module,
        input_vocabulary=input_vocabulary,
        output_vocabulary=output_vocabulary,
        optimizer_def=optimizer_def,
        decode_fn=decode_fn,
        feature_converter_cls=feature_converter_cls,
        label_smoothing=label_smoothing,
        z_loss=z_loss,
        loss_normalizing_factor=loss_normalizing_factor)
    self.aux_loss_factor = aux_loss_factor
    self.router_z_loss_factor = router_z_loss_factor

  def loss_fn(
      self, params: models.PyTreeDef, batch: Mapping[str, jnp.ndarray],
      dropout_rng: Optional[jnp.ndarray]) -> Tuple[jnp.ndarray, MetricsMap]:
    """Cross-entropy loss function with auxiliary MoE load balancing loss.

    Args:
      params: Model parameters.
      batch: Batch of training examples.
      dropout_rng: Random number generator key for dropout.

    Returns:
      - Model loss.
      - Metrics.
    """
    logits, state = self._compute_logits(
        params, batch, dropout_rng, mutable=['intermediates'])
    loss_normalizing_factor: Optional[Union[
        float, int, str, losses.SpecialLossNormalizingFactor]]
    (loss_normalizing_factor,
     weights) = losses.get_loss_normalizing_factor_and_weights(
         self._loss_normalizing_factor, batch)

    targets = batch['decoder_target_tokens']
    total_loss, z_loss, _ = losses.compute_weighted_cross_entropy(
        logits,
        targets=targets,
        weights=weights,
        label_smoothing=self._label_smoothing,
        z_loss=self._z_loss,
        loss_normalizing_factor=loss_normalizing_factor)

    # Extract and add MoE losses to total loss.
    diversity_metrics = _extract_diversity_metrics(state)
    aux_loss, router_z_loss = _expert_losses(diversity_metrics,
                                             self.aux_loss_factor,
                                             self.router_z_loss_factor)
    total_loss += aux_loss + router_z_loss

    metrics = self._compute_metrics(
        logits=logits,
        targets=targets,
        mask=weights,
        loss=total_loss,
        z_loss=z_loss)
    metrics.update(
        _expert_metrics(
            diversity_metrics,
            total_loss,
            z_loss,
            aux_loss,
            router_z_loss,
            num_tokens=targets.size))

    return total_loss, metrics

def _extract_diversity_metrics(
    state: flax_scope.FrozenVariableDict) -> Sequence[ExpertMetrics]:
  """Extract expert diversity metrics from sown state intermediates.

  Args:
    state: Model state holding sown intermediate metrics.

  Returns:
    Single diversity metrics instance per MoE layer.

  Raises:
    ValueError if unable to extract any diversity metrics from model state.
  """
  state_dict = traverse_util.flatten_dict(flax_core.unfreeze(state))
  diversity_metrics = [
      metric for path, metric in state_dict.items()
      if path[-1] == 'diversity_metrics'
  ]
  if not diversity_metrics:
    raise ValueError(
        'Unable to find any expert diversity metrics. Please check that MoE '
        'metrics and losses are correctly sown.')
  # Convert modeling library DiversityMetrics objects to local ExpertMetrics
  # objects to avoid modeling library dependencies.
  return [
      ExpertMetrics(metric.auxiliary_loss, metric.router_z_loss,
                    metric.fraction_tokens_left_behind, metric.expert_usage,
                    metric.router_confidence) for metric in diversity_metrics
  ]

def _expert_losses(diversity_metrics: Sequence[ExpertMetrics],
                   auxiliary_loss_factor: float,
                   router_z_loss_factor: float) -> Tuple[float, float]:
  """Summarizes per-layer MoE auxiliary losses.

  For auxiliary losses, we take the mean across MoE layers.

  Args:
    diversity_metrics: Per-layer mixture of expert metrics.
    auxiliary_loss_factor: Factor by which to scale auxiliary load balancing
      loss for mixture of experts models. The raw auxiliary losses will be
      summed and then scaled by this factor.
    router_z_loss_factor: Factor by which to scale router z-loss for mixture of
      experts models.

  Returns:
    - Load balancing loss.
    - Router z-loss.
  """
  aux_loss = auxiliary_loss_factor * jnp.array(
      [m.auxiliary_loss for m in diversity_metrics], dtype=jnp.float32).mean()
  router_z_loss = router_z_loss_factor * jnp.array(
      [m.router_z_loss for m in diversity_metrics], dtype=jnp.float32).mean()
  return aux_loss, router_z_loss

def _expert_metrics(diversity_metrics: Sequence[ExpertMetrics],
                    total_loss: float, z_loss: float, auxiliary_loss: float,
                    router_z_loss: float, num_tokens: int) -> MetricsMap:
  """Summarizes per-layer expert metrics for the entire model.

  The return metrics map will also contain overrides for the cross entropy loss
  metrics to account for the MoE losses.

  Args:
    diversity_metrics: Per-layer mixture of expert metrics.
    total_loss: Total model loss.
    z_loss: Output logits z-loss (not MoE specific).
    auxiliary_loss: Auxiliary load balancing loss for MoE models.
    router_z_loss: Router z-loss for MoE models.
    num_tokens: Total number of target tokens.

  Returns:
    Expert diversity metrics.
  """
  cross_ent_loss = total_loss - z_loss - auxiliary_loss - router_z_loss
  return {
      'experts/auxiliary_loss':
          AveragePerStep.from_model_output(auxiliary_loss),
      'experts/router_z_loss':
          AveragePerStep.from_model_output(router_z_loss),
      'experts/fraction_tokens_left_behind':
          AveragePerStep.from_model_output(
              jnp.array(
                  [m.fraction_tokens_left_behind for m in diversity_metrics],
                  dtype=jnp.float32).mean()),
      'experts/expert_usage':
          AveragePerStep.from_model_output(
              jnp.array([m.expert_usage for m in diversity_metrics],
                        dtype=jnp.float32).mean()),
      'experts/router_confidence':
          AveragePerStep.from_model_output(
              jnp.array([m.router_confidence for m in diversity_metrics],
                        dtype=jnp.float32).mean()),
      # Override vanilla T5 cross entropy loss metrics with corrected loss that
      # accounts for MoE losses.
      'cross_ent_loss':
          metrics_lib.AveragePerStep(total=cross_ent_loss),
      'cross_ent_loss_per_all_target_tokens':
          clu_metrics.Average(total=jnp.sum(cross_ent_loss), count=num_tokens)
  }


import sys
import numpy as np
from panda3d.core import WindowProperties, FrameBufferProperties
from panda3d.core import GraphicsPipe, GraphicsEngine, GraphicsOutput
from panda3d.core import Texture

class Panda3dCameraSensor(object):
    def __init__(self, base, color=True, depth=False, size=None, near_far=None, hfov=None, title=None):
        if size is None:
            size = (640, 480)
        if near_far is None:
            near_far = (1.0, 1000.0)
        if hfov is None:
            hfov = 60
        winprops = WindowProperties.size(*size)
        winprops.setTitle(title or 'Camera Sensor')
        fbprops = FrameBufferProperties()
        # Request 8 RGB bits, 8 alpha bits, and a depth buffer.
        fbprops.setRgbColor(True)
        fbprops.setRgbaBits(8, 8, 8, 8)
        fbprops.setDepthBits(24)
        self.graphics_engine = GraphicsEngine(base.pipe)

        window_type = base.config.GetString('window-type', 'onscreen')
        flags = GraphicsPipe.BFFbPropsOptional
        if window_type == 'onscreen':
            flags = flags | GraphicsPipe.BFRequireWindow
        elif window_type == 'offscreen':
            flags = flags | GraphicsPipe.BFRefuseWindow

        self.buffer = self.graphics_engine.makeOutput(
            base.pipe, "camera sensor buffer", 0,
            fbprops, winprops, flags)

        if not color and not depth:
            raise ValueError("At least one of color or depth should be True")
        if color:
            self.color_tex = Texture("color_texture")
            self.buffer.addRenderTexture(self.color_tex, GraphicsOutput.RTMCopyRam,
                                         GraphicsOutput.RTPColor)
        else:
            self.color_tex = None
        if depth:
            self.depth_tex = Texture("depth_texture")
            self.buffer.addRenderTexture(self.depth_tex, GraphicsOutput.RTMCopyRam,
                                         GraphicsOutput.RTPDepth)
        else:
            self.depth_tex = None

        self.cam = base.makeCamera(self.buffer, scene=base.render, camName='camera_sensor')
        self.lens = self.cam.node().getLens()
        self.lens.setFov(hfov)
        self.lens.setFilmSize(*size)  # this also defines the units of the focal length
        self.lens.setNearFar(*near_far)

    def observe(self):
        for _ in range(self.graphics_engine.getNumWindows()):
            self.graphics_engine.renderFrame()
        self.graphics_engine.syncFrame()

        images = []

        if self.color_tex:
            data = self.color_tex.getRamImageAs('RGBA')
            if sys.version_info < (3, 0):
                data = data.get_data()
            image = np.frombuffer(data, np.uint8)
            image = image.reshape((self.color_tex.getYSize(), self.color_tex.getXSize(), self.color_tex.getNumComponents()))
            image = np.flipud(image)
            image = image[..., :-1]  # remove alpha channel; if alpha values are needed, set alpha bits to 8
            images.append(image)

        if self.depth_tex:
            depth_data = self.depth_tex.getRamImage()
            if sys.version_info < (3, 0):
                depth_data = depth_data.get_data()
            depth_image = np.frombuffer(depth_data, np.float32)
            depth_image = depth_image.reshape((self.depth_tex.getYSize(), self.depth_tex.getXSize(), self.depth_tex.getNumComponents()))
            depth_image = np.flipud(depth_image)
            depth_image = (255 ** depth_image).astype(np.uint8)
            images.append(depth_image)
        return tuple(images)


import os

os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.30
set_session(tf.Session(config=config))

import torch
torch.cuda.empty_cache()
# from Rec_text import rec_txt
# from crnn_torch.model import predict
# from ctpn.ctpn_blstm_test import text_predict
# from densent_ocr.model import predict
from direction_detection.direction_correct import fourier_demo
from pan.predict import text_predict
from crnn_torch.model import predict

from select_infomodel import select
# from warp_templates.get_text import get_texts

from wce_data import WCE
from viterbi import calculate
# from red_method import seal_eliminate

import numpy as np
from PIL import Image
import fitz
from full_ocr_local import single_ocr
# from template_warp import get_warp
import json
import pickle
#from sanic import Sanic, response

a = Image.open('test.png').convert('RGB')
ori_w, ori_h = a.size
b = a.resize((1000, 1000))
scale_w, scale_h = b.size
b = np.array(b)
scale_w, scale_h = ori_w / scale_w, ori_h / scale_h
text_predict(b, scale_w, scale_h, a)

#app = Sanic(__name__)

# with open('info.pkl', 'rb') as pkl:
#     info = pickle.load(pkl)
# info = pickle.load(open('info.pkl', 'rb'))

large_FT = ['FT001001001002', 'FT001001001003', 'FT001001001005', 'FT001003002001', 'FT001001002001', 'FT001001003001']

# @app.route('/fre', methods=['POST'])
# def get_text(request):
#     img_path = request.form.get('input')
#     print(request.form.get('input'))
#     par = request.form.get('par')
#     page = request.form.get('page')
#     print(img_path)
#     try:
#         if img_path.lower().endswith('.pdf'):
#             pdf = fitz.open(img_path)
#             page_num = pdf[int(page) - 1]
#             trans = fitz.Matrix(3, 3).preRotate(0)
#             pm = page_num.getPixmap(matrix=trans, alpha=False)
#             input_img = fourier_demo(Image.frombytes("RGB", [pm.width, pm.height], pm.samples), 'FT001')
#         else:
#             input_img = fourier_demo(Image.open(img_path).convert('RGB'), 'FT001')
#         image = input_img.copy()
#         image = np.array(image)
#         image = corp_margin(image)
#         Image.fromarray(image).save('ttt.jpg')
#         image = Image.fromarray(image).resize((299, 299))
#         FT = predict_all(image, par)
#         return response.json({'result': 'true', 'Images': [{'path': img_path, 'FT': i[0], 'proportion': str(i[1])} for i in FT]})
#     except Exception as e:
#         print(e)
#         return response.json({'result': 'false', 'Images': []})

#@app.route('/fce', methods=['POST'])
def get_text(request):
    img_path = request.get('img_path')
    par = request.get('par')
    task_id = request.get('task_id')
    FT = request.get('FT')
    page = request.get('page')
    print(img_path)
    try:
        if img_path.lower().endswith('.pdf'):
            pdf = fitz.open(img_path)
            page_num = pdf[int(page) - 1]
            trans = fitz.Matrix(3, 3).preRotate(0)
            pm = page_num.getPixmap(matrix=trans, alpha=False)
            ori_img = fourier_demo(Image.frombytes("RGB", [pm.width, pm.height], pm.samples), 'FT001')
        else:
            ori_img = fourier_demo(Image.open(img_path).convert('RGB'), 'FT001')
        f = select(FT[:11] + '001')
        print('FT:', FT[:11] + '001')

        # input_img = input_img.resize((2000, 2000), Image.ANTIALIAS)
        ori_w, ori_h = ori_img.size
        input_img = ori_img.copy()
        input_img.thumbnail((2000, 2000), Image.ANTIALIAS)
        # input_img = input_img.resize((2000, 2000), Image.ANTIALIAS)
        scale_w, scale_h = input_img.size
        scale_w, scale_h = ori_w / scale_w, ori_h / scale_h
        input_img = input_img.convert('RGB')
        data_image = str(os.path.splitext(img_path)[0].split('/')[-1]) + '_' + str(page)
        data_image = '/home/ddwork/wce_data/ori_images/{}_{}.jpg'.format(data_image, task_id)
        input_img.save(data_image)
        input_img = np.array(input_img)
        # input_img = seal_eliminate(input_img)
        import time
        start = time.time()
        print("text_predict zhiqian")
        images = text_predict(input_img, scale_w, scale_h, ori_img)
        print("text_predict zhihou")
        torch.cuda.empty_cache()

        print(111111111111111111111111111,"SAD", time.time() - start,'HAPPY')
        start = time.time()
        # image_positions = [[i[0].tolist(), rec_txt(i[1]).replace('“', '').replace('"', '')] for i
        #                    in
        #                    images]
        image_positions = []
        for j in images:
            try:
                print("predict front!!!!!!!!!!!!!!")
                content = predict(Image.fromarray(j[1]).convert('L'))
                print("predict back!!!!!!!!!!!!!!")
                for index, i in enumerate(content[1]):
                    if i[0] > 0.9:
                        content[0][index] = content[0][index][0]
                        content[1].pop(index)
                # if i[0] < 0.9:
                #     img = Image.fromarray(j[1]).convert('L')
                #     width, height = img.size[0], img.size[1]
                #     scale = height * 1.0 / 32
                #     width = int(width / scale)
                #
                #     img = img.resize([width, 32], Image.ANTIALIAS)
                #     img = np.array(img)
                #     new_img = img[:, (content[2][index] - 1) * 8:(content[2][index] + 2) * 8]
                #     word, prob = attention(new_img)
                #     if prob > 0.9:
                #         content[0][index] = word[0]
                #         content[1].pop(index)
                # else:
                #     content[0][index] = content[0][index][0]
                #     content[1].pop(index)
                content = calculate(content)
                image_positions.append([j[0], content.replace('“', '').replace('‘', '')])
            except Exception as e:
                print(e)
                continue
        # torch.cuda.empty_cache()
        # data_json[task_id] = [par, data_image, FT, image_positions]
        data_json = WCE.create(field_id=int(task_id), par=str(par), image_path=data_image, FT=FT, file_type=FT[:11],
                               image_positions=str(image_positions), edited=False, trained=False)
        data_json.save()
        print(222222222222222222222222222, time.time() - start)
        text = single_ocr(image_positions)
        print(text)
        # with open(img_path + '.txt', 'w', encoding='utf-8') as fd:
        #     fd.write(text)
        texts = f.extract_info(img_path, page, FT[:11] + '001', text)
        print(texts)
        # try:
        #     found = get_warp(input_img, image_positions, FT)
        #     found_texts = get_texts('warp_templates/{}/template.xml'.format(FT), found, img_path, task_id)
        # except Exception as e:
        #     print(e)
        found_texts = ''
        print('==================================================================')
        print(texts, found_texts)
        torch.cuda.empty_cache()
        # 资质证书判断
        if FT[:11] == 'FT001003110':
            FT = FT[:8] + texts.get('version')
        # 路径中取日期
        try:
            if texts.get('发证日期') == '' or not texts.get('发证日期'):
                import re
                date_path = re.search('([0-9]{4}[-/年][0-9]{1,2}[-/月][0-9]{1,2}日?)', os.path.split(img_path)[1])
                if date_path:
                    texts['发证日期'] = date_path.groups()[0]
        except:
            pass

        if texts == 'FT999' and found_texts:
            return {'result': 'true', 'message': '请求成功', 'taskid': task_id, 'fields': found_texts, 'FT': FT}
        if texts != 'FT999' and found_texts == '':
            return {'result': 'true', 'message': '请求成功', 'taskid': task_id, 'fields': texts, 'FT': FT}
        if found_texts:
            for key, value in texts.items():
                try:
                    if value == '':
                        texts[key] = found_texts[key]
                except:
                    continue
        blank = 0
        for key, value in texts.items():
            if value == '':
                blank += 1
        if blank == len(texts) - 1:
            return {'result': 'false', 'message': '请求失败', 'taskid': task_id, 'fields': {}, 'FT': 'FT999999999'}
        else:
            return {'result': 'true', 'message': '请求成功', 'taskid': task_id, 'fields': texts, 'FT': FT}
    except Exception as e:
        print(e)
        return {'result': 'false', 'message': '请求失败', 'taskid': task_id, 'fields': {}, 'FT': 'FT999999999'}

#@app.route('/check')
#def health_check(request):
#    return response.text('ok')

#if __name__ == '__main__':
#    app.config.KEEP_ALIVE = False
#    app.config.REQUEST_TIMEOUT = 900
#    app.config.RESPONSE_TIMEOUT = 900
#    app.run(host='0.0.0.0', port=8005)
from celery import Celery
import requests
import json
import os
import time
import socket
import urllib

ddwork = Celery('sh_tasks', broker='redis://172.30.81.208:30443/0', backend='redis://172.30.81.208:30445/1')
ddwork.conf.CELERYD_FORCE_EXECV = True

ddwork.conf.BROKER_TRANSPORT_OPTIONS = {'visibility_timeout': 36000}

requests.adapters.DEFAULT_RETRIES = 3 #重连次数

@ddwork.task
def fce_callback(url, data, task_id, work_type, token, dbname, state, project_id, timeout=36000):
    print ('URL:{},参数:{}'.format(url,data),task_id)
    odoo_data = {'access_token':token,'dbname':dbname}
    odoo_url = 'http://172.30.81.208:32621/restapi/1.0/extension_object/sh.work.task/ML_action_to_fce_callback'
    headers = {'User-Agent': 'User-Agent:Mozilla/5.0'}
    start_time = time.time()
    try:
        res = get_text(data)
        res = json.dumps(res)
        #res = requests.post(url, data=data, headers=headers, timeout=timeout).text
        print ('URL:{},返回值:{},类型:{}'.format(url,res,type(res)),task_id)
        #回调
        if work_type == 'fce':
            odoo_url = 'http://172.30.81.208:32621/restapi/1.0/extension_object/sh.work.task/ML_action_to_fce_callback'
            odoo_data.update({'result':res,'task_id':task_id,'all_time':str(time.time()-start_time),'state':state})
        odoo_res = requests.post(odoo_url, data=odoo_data, headers=headers, timeout=timeout).text
        print('回调返回值：{}'.format(json.loads(odoo_res)))
        return odoo_data
    except Exception as e:
        if not os.path.exists(r'/home/ddwork/projects/compound_log'):
            os.makedirs('/home/ddwork/projects/compound_log')
        if not os.path.exists(r'/home/ddwork/projects/compound_log/{}'.format(project_id)):
            os.makedirs('/home/ddwork/projects/compound_log/{}'.format(project_id))
        if not os.path.exists(r'/home/ddwork/projects/compound_log/{}/FCE'.format(project_id)):
            os.makedirs('/home/ddwork/projects/compound_log/{}/FCE'.format(project_id))
        with open(r'/home/ddwork/projects/compound_log/{}/FCE/{}_log.txt'.format(project_id,str(time.strftime('%Y-%m-%d'))), 'a', encoding='utf-8') as gap:
            gap.write(str(time.strftime('%Y-%m-%d %H:%M:%S'))+'\n'+ str(url) +'\n'+ str(data) + '\n'+str(task_id) +'\n'+str(e)+'\n\n')
            if odoo_url:
                gap.write(str(time.strftime('%Y-%m-%d %H:%M:%S'))+'\n'+ str(odoo_url) +'\n'+ str(odoo_data) + '\n'+str(task_id) +'\n'+str(e)+'\n\n')
            gap.close()
        odoo_data.update({'result':{"res":str(e)},'task_id':task_id,'all_time':str(time.time()-start_time),'state':state})
        odoo_res = requests.post(odoo_url, data=odoo_data, headers=headers, timeout=timeout).text
        print('回调异常返回值：{}'.format(json.loads(odoo_res)))
        return {'result':str(e),'task_id':task_id,'start_time':start_time}


import os
import sys
import json
import random
import numpy as np
import torch
from torch.utils.data import DataLoader
from torch import optim
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from absl import flags

from modules.data import datasets
from modules.models import (PPG2ECG_BASELINE_LSTM, PPG2ECG)
from modules.loss import QRSLoss
from modules.util import signal2waveform, dict2table

FLAGS = flags.FLAGS

class Trainer(object):
    def __init__(self):
        self.device = torch.device(
            'cuda' if torch.cuda.is_available() else 'cpu')
        self.run_name = FLAGS.run_name
        self.log_path = os.path.join(FLAGS.logdir, self.run_name)
        self.save_path = os.path.join(FLAGS.logdir, self.run_name, 'model')
        os.makedirs(self.save_path, exist_ok=True)
        if FLAGS.seed != "":
            torch.manual_seed(int(FLAGS.seed))
            np.random.seed(int(FLAGS.seed))

        # tbwriter
        self.tbwriter = SummaryWriter(log_dir=self.log_path)
        flags_dict = dict(
            (item.name, item.value) for item
            in FLAGS.get_key_flags_for_module(sys.argv[0]))
        print(json.dumps(flags_dict, indent=4))
        self.tbwriter.add_text(self.run_name, dict2table(flags_dict), 0)

        # training parameters
        self.epoch = FLAGS.epoch
        self.eval_step = FLAGS.eval_step
        self.save_step = FLAGS.save_step

        # dataset
        dataset_class = datasets[FLAGS.dataset]
        train_set = dataset_class(sigma=FLAGS.qrs_sigma, use_aug=FLAGS.aug)
        test_set = dataset_class(sigma=FLAGS.qrs_sigma, is_train=False,
                                 use_aug=False)
        visual_set = random.choices(test_set, k=30)
        self.train_loader = DataLoader(
            dataset=train_set,
            batch_size=FLAGS.batch_size,
            shuffle=True,
            num_workers=3,
            pin_memory=True)
        self.test_loader = DataLoader(dataset=test_set, batch_size=1)
        self.visual_loader = DataLoader(dataset=visual_set, batch_size=1)

        # model
        if FLAGS.lstm:
            self.model = PPG2ECG_BASELINE_LSTM(50).to(self.device)
        else:
            self.model = PPG2ECG(FLAGS.input_size,
                                 use_stn=FLAGS.stn,
                                 use_attention=FLAGS.attn).to(self.device)
        # optim
        self.opt = optim.Adam(self.model.parameters(), lr=FLAGS.lr)
        # qrsloss
        if FLAGS.qrsloss:
            self.qrsloss = QRSLoss(beta=FLAGS.qrs_beta)
        else:
            self.qrsloss = None
        # learning rate scheduler
        if FLAGS.lr_sched:
            self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
                self.opt, FLAGS.epoch//10)
        else:
            self.scheduler = None

    def train(self, batch, iter):
        ppg, ecg = batch['ppg'].to(self.device), batch['ecg'].to(self.device)
        exp_rpeaks = batch['exp_rpeaks'].to(self.device)
        loss = {}
        self.opt.zero_grad()
        out_dict = self.model(ppg)
        output = out_dict['output']
        l1_loss = F.l1_loss(output, ecg)
        loss['train/L1'] = l1_loss.item()
        if self.qrsloss is not None:
            qrs_loss = self.qrsloss(output, ecg, exp_rpeaks)
            loss['train/QRS_L1'] = qrs_loss.item()
            qrs_loss.backward()
        else:
            l1_loss.backward()
        self.opt.step()
        # write scalar
        for name, value in loss.items():
            self.tbwriter.add_scalar(name, value, iter)

    def test(self, best, epoch, iter):
        self.model.eval()
        running_loss = 0
        running_qrs_loss = 0
        len_test_set = len(self.test_loader.dataset)
        pbar = tqdm(self.test_loader,
                    dynamic_ncols=True,
                    desc='Testing')
        for batch in pbar:
            ppg = batch['ppg'].to(self.device)
            ecg = batch['ecg'].to(self.device)
            exp_rpeaks = batch['exp_rpeaks'].to(self.device)
            loss = {}
            with torch.no_grad():
                out_dict = self.model(ppg)
                output = out_dict['output']
                l1_loss = F.l1_loss(output, ecg)
                loss['test/L1'] = l1_loss.item()
                if self.qrsloss is not None:
                    qrs_loss = self.qrsloss(output, ecg, exp_rpeaks)
                    loss['test/QRS_L1'] = qrs_loss.item()
            running_loss += loss['test/L1']
            running_qrs_loss += loss['test/QRS_L1']
        pbar.close()
        # write scalar
        running_loss /= len_test_set
        running_qrs_loss /= len_test_set
        self.tbwriter.add_scalar('test/L1', running_loss, iter)
        self.tbwriter.add_scalar('test/QRS_L1', running_qrs_loss, iter)
        # save best_checkpoint
        if running_loss < best["L1"]:
            best["L1"] = running_loss
            self.tbwriter.add_scalar('test/best_L1', running_loss, iter)
            if self.qrsloss is None:
                self.best_checkpoint = {
                    'epoch': epoch+1,
                    'net': self.model.state_dict(),
                }
        if running_qrs_loss < best["QRS_L1"]:
            best["QRS_L1"] = running_qrs_loss
            self.tbwriter.add_scalar(
                'test/best_QRS_L1', running_qrs_loss, iter)
            if self.qrsloss is not None:
                self.best_checkpoint = {
                    'epoch': epoch+1,
                    'net': self.model.state_dict(),
                }
        self.model.train()

    def visualize(self, epoch, iter):
        self.model.eval()
        pbar = tqdm(self.visual_loader,
                    dynamic_ncols=True,
                    desc='Visualizing')
        for batch in pbar:
            with torch.no_grad():
                ppg = batch['ppg'].to(self.device)
                ecg = batch['ecg'].to(self.device)
                path = os.path.basename(batch['path'][0])
                out_dict = self.model(ppg)
                output = out_dict['output']
                output_stn = out_dict['output_stn']
                ppg_waveform = signal2waveform(ppg.squeeze().cpu().numpy())
                ecg_waveform = signal2waveform(ecg.squeeze().cpu().numpy())
                out_waveform = signal2waveform(output.squeeze().cpu().numpy())
                out_stn_waveform = signal2waveform(
                    output_stn.squeeze().cpu().numpy())
                concat_img = np.concatenate([
                    ppg_waveform, np.ones((ppg_waveform.shape[0], 25)),
                    out_stn_waveform, np.ones((ppg_waveform.shape[0], 25)),
                    out_waveform, np.ones((ppg_waveform.shape[0], 25)),
                    ecg_waveform
                ], axis=1)
                self.tbwriter.add_image(
                    'waveform/{}'.format(path), concat_img, iter,
                    dataformats='HW'
                )
        pbar.close()

    def save_checkpoint(self, checkpoint, save_path, save_name):
        save_path = os.path.join(save_path, 'model_{}.pth'.format(save_name))
        torch.save(checkpoint, save_path)

    def run(self):
        best = {'L1': 100000, 'QRS_L1': 100000}
        self.best_checkpoint = None
        iters = 1
        pbar = tqdm(total=self.epoch*len(self.train_loader),
                    dynamic_ncols=True)
        for epoch in range(self.epoch):
            self.tbwriter.add_scalar('train/epoch', epoch+1, iters)
            self.model.train()
            for batch in self.train_loader:
                self.train(batch, iters)
                iters += 1
                pbar.update(1)
            if epoch % self.eval_step == 0 or epoch == self.epoch-1:
                self.test(best, epoch, iters)
                self.visualize(epoch, iters)
            if epoch % self.save_step == 0 or epoch == self.epoch-1:
                self.save_checkpoint(
                    {
                        'epoch': epoch+1,
                        'net': self.model.state_dict()
                    },
                    self.save_path,
                    '{}'.format(epoch+1)
                )
            if self.scheduler:
                self.scheduler.step()
        self.tbwriter.close()
        self.save_checkpoint(self.best_checkpoint, self.save_path, 'best')


import torch
from core_train import core_train
from fc_classification import FCNet
from torch import optim
from dataset import Refuge2, Resize2_640, RandomRotation, RandomFlip
from torchvision.transforms import Compose
from tqdm import tqdm
from torch.autograd import Variable
import torch.nn as nn
from dedicated_Resnet50 import ResNet50_Mod
from utils import DataLoaderX, collate_fn
from sklearn.metrics import roc_auc_score
import numpy as np

def getModel(base_lr=1e-4, cuda=False):
    model = FCNet()
    if cuda:
        model = model.cuda()
    optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=base_lr,
                          momentum=0.9, weight_decay=0.0001)
    lr_scheduler = optim.lr_scheduler.StepLR(optimizer, 30)
    return model, optimizer, lr_scheduler

def getResNet(size, base_lr=1e-4, cuda=False):
    model = ResNet50_Mod(input_size=size)
    if cuda:
        model = model.cuda()
    optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=0.0001)
    lr_scheduler = optim.lr_scheduler.StepLR(optimizer, 30)
    return model, optimizer, lr_scheduler

def test(data, model, batch_size, cuda):
    dataloader = DataLoaderX(data, batch_size=batch_size, shuffle=False, num_workers=1, collate_fn=collate_fn)
    iterator = tqdm(dataloader)
    classification = list()
    gts = list()
    for sample in iterator:
        img, gt_label = sample
        if cuda:
            img = Variable(img).cuda
        classification.append(model(img).cpu().detach().numpy())
        gts.append(gt_label.numpy())
    classification = np.stack(classification).flatten()
    gts = np.stack(gts).flatten()
    auc = roc_auc_score(classification, gts, average=None)
    return auc

def train_fcnet(data, gt_labels, gt_segmentations, batch_size=1, cuda=False):
    res = core_train(data, None, gt_segmentations, cuda=cuda)
    model, optimizer, lr_scheduler = getModel(cuda=cuda)
    transform = Compose(
        [
            RandomRotation(),
            RandomFlip(),
            Resize2_640()
        ]
    )
    dataset = Refuge2(data, gt_labels, None, transform=transform)
    epoch = 0
    best_auc = 0.
    while True:
        if epoch > 0 and epoch % 2 == 0:
            model.eval()
            auc = test(dataset, model, batch_size, cuda)
            best_auc = max(best_auc, auc)
            model.train()
        if epoch >= 10:
            break
        dataloader = DataLoaderX(dataset, batch_size=batch_size, shuffle=True, num_workers=1, collate_fn=collate_fn)
        iterator = tqdm(dataloader)
        for sample in iterator:
            optimizer.zero_grad()
            img, gt_label = sample
            if cuda:
                img = Variable(img).cuda()
                classification = model(img).cpu()
            else:
                classification = model(img)
            # loss = nn.BCELoss()(classification.view(classification.shape[0], -1), gt_label.view(gt_label.shape[0], -1))
            loss = nn.L1Loss()(classification.view(classification.shape[0], -1), gt_label.view(gt_label.shape[0], -1))
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 10.0)
            optimizer.step()
            lr_scheduler.step()
            epoch += 1
    return res, model

def train_resnet(data, gt_labels, batch_size=1, cuda=False):
    print("**********************************Start training resnet*********************************")
    model, optimizer, lr_scheduler = getResNet(size=data.shape[1], cuda=cuda)
    transform = Compose(
        [
            RandomRotation(),
            RandomFlip()
        ]
    )
    dataset = Refuge2(data, gt_labels, segmentations=None, transform=transform)
    epoch = 0
    best_auc = 0.
    while True:
        if epoch > 0 and epoch % 2 == 0:
            model.eval()
            auc = test(dataset, model, batch_size=batch_size, cuda=cuda)
            best_auc = max(best_auc, auc)
            model.train()
        if epoch >= 10:
            break
        dataloader = DataLoaderX(dataset, batch_size=batch_size, shuffle=True, num_workers=1, collate_fn=collate_fn)
        iterator = tqdm(dataloader)
        for sample in iterator:
            optimizer.zero_grad()
            img, gt_label = sample
            if cuda:
                img = Variable(img).cuda()
                classification = model(img).cpu()
            else:
                classification = model(img)
            loss = nn.BCELoss()(classification.view(classification.shape[0], -1), gt_label.view(gt_label.shape[0], -1))
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 10.0)
            optimizer.step()
            lr_scheduler.step()
            epoch += 1
    return model

if __name__ == "__main__":
    a = [0 for _ in range(10)]
    a = torch.Tensor(a)
    a = a.numpy()
    b = [a, a]
    b = np.stack(b).flatten()
    print(b.shape)


# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import argparse

import paddle
import paddlehub as hub
from paddlehub.module.module import moduleinfo, runnable, serving
import numpy as np
import rawpy
import cv2

from .util import base64_to_cv2

def pack_raw(raw):
    # pack Bayer image to 4 channels
    im = raw
    if not isinstance(raw, np.ndarray):
        im = raw.raw_image_visible.astype(np.float32)
    im = np.maximum(im - 512, 0) / (16383 - 512)  # subtract the black level

    im = np.expand_dims(im, axis=2)
    img_shape = im.shape
    H = img_shape[0]
    W = img_shape[1]

    out = np.concatenate((im[0:H:2, 0:W:2, :], im[0:H:2, 1:W:2, :], im[1:H:2, 1:W:2, :], im[1:H:2, 0:W:2, :]), axis=2)
    return out

@moduleinfo(
    name="seeinthedark", type="CV/denoising", author="paddlepaddle", author_email="", summary="", version="1.0.0")
class LearningToSeeInDark:
    def __init__(self):
        self.pretrained_model = os.path.join(self.directory, "pd_model/inference_model")
        self.cpu_have_loaded = False
        self.gpu_have_loaded = False

    def set_device(self, use_gpu=False):
        if use_gpu == False:
            if not self.cpu_have_loaded:
                exe = paddle.static.Executor(paddle.CPUPlace())
                [prog, inputs, outputs] = paddle.static.load_inference_model(
                    path_prefix=self.pretrained_model,
                    executor=exe,
                    model_filename="model.pdmodel",
                    params_filename="model.pdiparams")
                self.cpuexec, self.cpuprog, self.cpuinputs, self.cpuoutputs = exe, prog, inputs, outputs
                self.cpu_have_loaded = True

            return self.cpuexec, self.cpuprog, self.cpuinputs, self.cpuoutputs

        else:
            if not self.gpu_have_loaded:
                exe = paddle.static.Executor(paddle.CUDAPlace(0))
                [prog, inputs, outputs] = paddle.static.load_inference_model(
                    path_prefix=self.pretrained_model,
                    executor=exe,
                    model_filename="model.pdmodel",
                    params_filename="model.pdiparams")
                self.gpuexec, self.gpuprog, self.gpuinputs, self.gpuoutputs = exe, prog, inputs, outputs
                self.gpu_have_loaded = True

            return self.gpuexec, self.gpuprog, self.gpuinputs, self.gpuoutputs

    def denoising(self,
                  images: list = None,
                  paths: list = None,
                  output_dir: str = './enlightening_result/',
                  use_gpu: bool = False,
                  visualization: bool = True):
        '''
        Denoise a raw image in the low-light scene.

        images (list[numpy.ndarray]): data of images, shape of each is [H, W], must be sing-channel image captured by camera.
        paths (list[str]): paths to images
        output_dir: the dir to save the results
        use_gpu: if True, use gpu to perform the computation, otherwise cpu.
        visualization: if True, save results in output_dir.
        '''
        results = []
        paddle.enable_static()
        exe, prog, inputs, outputs = self.set_device(use_gpu)

        if images != None:
            for raw in images:
                input_full = np.expand_dims(pack_raw(raw), axis=0) * 300
                px = input_full.shape[1] // 512
                py = input_full.shape[2] // 512
                rx, ry = px * 512, py * 512
                input_full = input_full[:, :rx, :ry, :]
                output = np.random.randn(rx * 2, ry * 2, 3)
                input_full = np.minimum(input_full, 1.0)
                for i in range(px):
                    for j in range(py):
                        input_patch = input_full[:, i * 512:i * 512 + 512, j * 512:j * 512 + 512, :]
                        result = exe.run(prog, feed={inputs[0]: input_patch}, fetch_list=outputs)
                        output[i * 512 * 2:i * 512 * 2 + 512 * 2, j * 512 * 2:j * 512 * 2 + 512 * 2, :] = result[0][0]
                output = np.minimum(np.maximum(output, 0), 1)
                output = output * 255
                output = np.clip(output, 0, 255)
                output = output.astype('uint8')
                results.append(output)
        if paths != None:
            for path in paths:
                raw = rawpy.imread(path)
                input_full = np.expand_dims(pack_raw(raw), axis=0) * 300
                px = input_full.shape[1] // 512
                py = input_full.shape[2] // 512
                rx, ry = px * 512, py * 512
                input_full = input_full[:, :rx, :ry, :]
                output = np.random.randn(rx * 2, ry * 2, 3)
                input_full = np.minimum(input_full, 1.0)
                for i in range(px):
                    for j in range(py):
                        input_patch = input_full[:, i * 512:i * 512 + 512, j * 512:j * 512 + 512, :]
                        result = exe.run(prog, feed={inputs[0]: input_patch}, fetch_list=outputs)
                        output[i * 512 * 2:i * 512 * 2 + 512 * 2, j * 512 * 2:j * 512 * 2 + 512 * 2, :] = result[0][0]
                output = np.minimum(np.maximum(output, 0), 1)
                output = output * 255
                output = np.clip(output, 0, 255)
                output = output.astype('uint8')
                results.append(output)

        if visualization == True:
            if not os.path.exists(output_dir):
                os.makedirs(output_dir, exist_ok=True)
            for i, out in enumerate(results):
                cv2.imwrite(os.path.join(output_dir, 'output_{}.png'.format(i)), out[:, :, ::-1])

        return results

    @runnable
    def run_cmd(self, argvs: list):
        """
        Run as a command.
        """
        self.parser = argparse.ArgumentParser(
            description="Run the {} module.".format(self.name),
            prog='hub run {}'.format(self.name),
            usage='%(prog)s',
            add_help=True)

        self.arg_input_group = self.parser.add_argument_group(title="Input options", description="Input data. Required")
        self.arg_config_group = self.parser.add_argument_group(
            title="Config options", description="Run configuration for controlling module behavior, not required.")
        self.add_module_config_arg()
        self.add_module_input_arg()
        self.args = self.parser.parse_args(argvs)
        self.denoising(
            paths=[self.args.input_path],
            output_dir=self.args.output_dir,
            use_gpu=self.args.use_gpu,
            visualization=self.args.visualization)

    @serving
    def serving_method(self, images, **kwargs):
        """
        Run as a service.
        """
        images_decode = [base64_to_cv2(image) for image in images]
        results = self.denoising(images=images_decode, **kwargs)
        tolist = [result.tolist() for result in results]
        return tolist

    def add_module_config_arg(self):
        """
        Add the command config options.
        """
        self.arg_config_group.add_argument('--use_gpu', action='store_true', help="use GPU or not")

        self.arg_config_group.add_argument(
            '--output_dir', type=str, default='denoising_result', help='output directory for saving result.')
        self.arg_config_group.add_argument('--visualization', type=bool, default=False, help='save results or not.')

    def add_module_input_arg(self):
        """
        Add the command input options.
        """
        self.arg_input_group.add_argument(
            '--input_path', type=str, help="path to input raw image, should be raw file captured by camera.")


# -*-coding: utf-8 -*-
from __future__ import print_function
from __future__ import division
import argparse
from scipy import stats

def eval(file_source, file_gold, sense):
  fs = open(file_source)
  fg = open(file_gold)
  fs = fs.readlines()
  fg = fg.readlines()
  per_right_sense = []
  per_total_sense = []
  per_right_argument = []
  per_gold_argument = []
  per_retrieve_argument = []
  temp_right_sense = 0
  temp_total_sense = 0
  temp_right_argument = 0
  temp_gold_argument = 0
  temp_retrieve_argument = 0

  if len(fs) != len(fg):
    print("The length is not match, please check")
    print("source %d gold %d" % (len(fs), len(fg)))
  for line_num, (lines, lineg) in enumerate(zip(fs, fg)):
    try:
      lines = lines.strip().split()
      lineg = lineg.strip().split()
      if not lines and not lineg:
        per_right_sense.append(temp_right_sense)
        per_total_sense.append(temp_total_sense)
        per_right_argument.append(temp_right_argument)
        per_gold_argument.append(temp_gold_argument)
        per_retrieve_argument.append(temp_retrieve_argument)
        temp_right_sense = 0
        temp_total_sense = 0
        temp_right_argument = 0
        temp_gold_argument = 0
        temp_retrieve_argument = 0
        continue
      elif lines is not None and lineg is not None:
        if lineg[12] == 'Y':
          temp_total_sense += 1
          if lines[12] == 'Y' and lines[13].split('.')[1] == lineg[13].split('.')[1]:
            temp_right_sense += 1
        for p in xrange(14, len(lineg)):
          if lineg[p] != '_':
            temp_gold_argument += 1
          if lines[p] != '_':
            temp_retrieve_argument += 1
          if lineg[p] != '_' and lineg[p] == lines[p]:
            temp_right_argument += 1
      else:
        print("Error state in line %d" % (line_num))
    except:
      print("Error in line %d" % (line_num))
      print(lineg)
      print(lines)
      exit()

  right_sense = sum(per_right_sense)
  total_sense = sum(per_total_sense)
  right_argument = sum(per_right_argument)
  gold_arguent = sum(per_gold_argument)
  retrieve_argument = sum(per_retrieve_argument)
  if sense:
    print("## With sense")
    print("Precision: (%d + %d) / (%d + %d) = %f" %
          (right_argument, right_sense, retrieve_argument, total_sense,
           (right_argument + right_sense) / (retrieve_argument + total_sense)))
    print("Recall: (%d + %d) / (%d + %d) = %f" %
          (right_argument, right_sense, gold_arguent, total_sense,
           (right_argument + right_sense) / (gold_arguent + total_sense)))
    p = (right_argument + right_sense) / (retrieve_argument + total_sense)
    r = (right_argument + right_sense) / (gold_arguent + total_sense)
    print("F score: %f" % (2 * p * r / (p + r)))
    print("Sense Precision: %f " % (right_sense / total_sense))
  else:
    print("## without sense")
    print("Precision: %d / %d = %f" % (right_argument, retrieve_argument,
                                       right_argument / retrieve_argument))
    print("Recall: %d / %d = %f" % (right_argument, gold_arguent,
                                    right_argument / gold_arguent))
    p = right_argument / retrieve_argument
    r = right_argument / gold_arguent
    print("F score: %f" % (2 * p * r / (p + r)))

  file_name = "_per_sent_stat"
  if sense:
    file_name += "_with_sense"
  else:
    file_name += "_without_sense"
  file_name = file_source + file_name
  fout = open(file_name, "w")
  for item_right_sense, item_total_sense, \
    item_right_argument, item_retrieve_argument, item_gold_argument \
    in zip(per_right_sense, per_total_sense, per_right_argument, \
           per_retrieve_argument, per_gold_argument):
    if sense:
      if item_retrieve_argument + item_total_sense != 0:
        p = (item_right_argument + item_right_sense) / (item_retrieve_argument + item_total_sense)
      else:
        p = 0
      if item_gold_argument + item_total_sense != 0:
        r = (item_right_argument + item_right_argument) / (item_gold_argument + item_total_sense)
      else:
        r = 0
    else:
      if item_retrieve_argument != 0:
        p = (item_right_argument) / (item_retrieve_argument)
      else:
        p = 0
      if item_gold_argument != 0:
        r = (item_right_argument) / (item_gold_argument)
      else:
        r = 0
    if (p + r) != 0:
      f = 2 * p * r / (p + r)
    else:
      f = 0
    fout.write(str(f) + "\n")

def ttest(file_old, file_new):
  fold = open(file_old)
  fnew = open(file_new)
  base = [float(num.strip()) for num in fold.readlines()]
  test = [float(num.strip()) for num in fnew.readlines()]
  if len(base) != len(test):
    print("Length Error")
    exit()
    print("base %d test %d" % (len(base), len(test)))
  a, p = stats.ttest_rel(base, test)
  print("a = %f , p = %f" % (a, p))

if __name__=="__main__":
  argparser = argparse.ArgumentParser()
  argparser.add_argument('--source', help="The source file | The original file")
  argparser.add_argument('--gold', help="The gold file | The new file")
  argparser.add_argument('--sense', help="Include the sense", action="store_true")
  argparser.add_argument('--fscore', help="Evaluate the F score", action="store_true")
  argparser.add_argument('--ttest', help="T test", action="store_true")
  arg = argparser.parse_args()

  if arg.source == None or arg.gold == None:
    print("Source or Gold file is None")
  if arg.fscore:
    print(arg.sense)
    eval(arg.source, arg.gold, arg.sense)
  if arg.ttest:
    ttest(arg.source, arg.gold)