import mmcv
import numpy as np
import pycocotools.mask as mask_util
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.utils import _pair

from mmdet.core import auto_fp16, force_fp32, mask_target
from ..builder import build_loss
from ..registry import HEADS
from ..utils import ConvModule

@HEADS.register_module
class SemanticFCNMaskHead(nn.Module):

    def __init__(self,
                 num_convs=4,
                 roi_feat_size=14,
                 in_channels=256,
                 conv_kernel_size=3,
                 conv_out_channels=256,
                 upsample_method='deconv',
                 upsample_ratio=2,
                 num_classes=81,
                 semantic_dims=300,
                 seen_class=True,
                 gzsd=False,
                 share_semantic=False,
                 sync_bg=False,
                 voc_path=None,
                 vec_path=None,
                 with_learnable_kernel=True,
                 with_decoder=False,
                 class_agnostic=False,
                 conv_cfg=None,
                 norm_cfg=None,
                 loss_mask=dict(
                     type='CrossEntropyLoss', use_mask=True, loss_weight=1.0),
                 loss_ed=dict(type='MSELoss', loss_weight=0.5)):
        super(SemanticFCNMaskHead, self).__init__()
        if upsample_method not in [None, 'deconv', 'nearest', 'bilinear']:
            raise ValueError(
                'Invalid upsample method {}, accepted methods '
                'are "deconv", "nearest", "bilinear"'.format(upsample_method))

        self.seen_class = seen_class
        self.gzsd = gzsd
        self.share_semantic = share_semantic
        self.with_learnable_kernel = with_learnable_kernel
        self.with_decoder = with_decoder
        self.num_convs = num_convs
        # WARN: roi_feat_size is reserved and not used
        self.roi_feat_size = _pair(roi_feat_size)
        self.in_channels = in_channels
        self.conv_kernel_size = conv_kernel_size
        self.conv_out_channels = conv_out_channels
        self.upsample_method = upsample_method
        self.upsample_ratio = upsample_ratio
        self.num_classes = num_classes
        self.class_agnostic = class_agnostic
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.fp16_enabled = False
        self.loss_mask = build_loss(loss_mask)
        self.loss_ed = build_loss(loss_ed)
        self.sync_bg=sync_bg

        self.convs = nn.ModuleList()
        for i in range(self.num_convs):
            in_channels = (
                self.in_channels if i == 0 else self.conv_out_channels)
            padding = (self.conv_kernel_size - 1) // 2
            self.convs.append(
                ConvModule(
                    in_channels,
                    self.conv_out_channels,
                    self.conv_kernel_size,
                    padding=padding,
                    conv_cfg=conv_cfg,
                    norm_cfg=norm_cfg))
        self.convT = ConvModule(
                    in_channels,
                    300,
                    3,
                    padding=1,
                    conv_cfg=conv_cfg,
                    norm_cfg=norm_cfg)
        if self.with_decoder:
            self.dconvT = ConvModule(
                        300,
                        in_channels,
                        3,
                        padding=1,
                        conv_cfg=conv_cfg,
                        norm_cfg=norm_cfg)
        upsample_in_channels = (
            self.conv_out_channels if self.num_convs > 0 else in_channels)
        if self.upsample_method is None:
            self.upsample = None
        elif self.upsample_method == 'deconv':
            self.upsample = nn.ConvTranspose2d(
                upsample_in_channels,
                self.conv_out_channels,
                self.upsample_ratio,
                stride=self.upsample_ratio)
        else:
            self.upsample = nn.Upsample(
                scale_factor=self.upsample_ratio, mode=self.upsample_method)

        self.relu = nn.ReLU(inplace=True)
        self.debug_imgs = None

        if voc_path is not None:
            voc = np.loadtxt(voc_path, dtype='float32', delimiter=',')
        else:
            voc = None

        vec_load = np.loadtxt(vec_path, dtype='float32', delimiter=',')

        vec = vec_load[:, :num_classes]
        vec_unseen = np.concatenate([vec_load[:, 0:1], vec_load[:, num_classes:]], axis=1)
        vec = torch.tensor(vec, dtype=torch.float32)
        if voc is not None:
            voc = torch.tensor(voc, dtype=torch.float32)
        vec_unseen = torch.tensor(vec_unseen, dtype=torch.float32)
        self.vec_unseen = vec_unseen.cuda()
        self.vec = vec.cuda()  # 300*n
        self.conv_vec = nn.Conv2d(300, num_classes, 1, bias=False)

        self.conv_vec.weight.data = torch.unsqueeze(torch.unsqueeze(self.vec.t(), -1), -1)

        if not self.seen_class:
            self.con_vec_t = nn.Conv2d(num_classes, 300, 1, bias=False)
            self.con_vec_t.weight.data = torch.unsqueeze(torch.unsqueeze(self.vec, -1), -1)
            self.conv_vec_unseen = nn.Conv2d(300, vec_unseen.shape[1], 1, bias=False)
            self.conv_vec_unseen.weight.data = torch.unsqueeze(torch.unsqueeze(self.vec_unseen.t(), -1), -1)

        if voc is not None:
            self.voc = voc.cuda()  # 300*66
            self.conv_voc = nn.Conv2d(300, self.voc.size(1), 1, bias=False)
            self.conv_voc.weight.data = torch.unsqueeze(torch.unsqueeze(self.voc.t(), -1), -1)
        else:
            self.voc = None

        self.vec_unseen = vec_unseen.cuda()
        if self.with_learnable_kernel:
            if self.voc is not None:
                self.kernel_semantic = nn.Conv2d(self.voc.size(1), 300, kernel_size=3, padding=1)
                if self.with_decoder:
                    self.d_kernel_semantic = nn.Conv2d(300, self.voc.size(1), kernel_size=3, padding=1)
            else:
                self.kernel_semantic = nn.Conv2d(300, 300, kernel_size=3, padding=1)
                if self.with_decoder:
                    self.d_kernel_semantic = nn.Conv2d(300, 300, kernel_size=3, padding=1)

    def init_weights(self):
        for m in [self.upsample]:
            if m is None:
                continue
            nn.init.kaiming_normal_(
                m.weight, mode='fan_out', nonlinearity='relu')
            nn.init.constant_(m.bias, 0)
        if self.with_learnable_kernel:
            for m in [self.kernel_semantic]:
                if m is None:
                    continue
                nn.init.kaiming_normal_(
                    m.weight, mode='fan_out', nonlinearity='relu')
                nn.init.constant_(m.bias, 0)
        for m in [self.conv_vec]:
            for param in m.parameters():
                param.requires_grad = False
        if self.voc is not None:
            for m in [self.conv_voc]:
                for param in m.parameters():
                    param.requires_grad = False

    @auto_fp16()
    def forward(self, x, bg_vector=None):
        if bg_vector and self.sync_bg:
            with torch.no_grad():
                self.conv_vec.weight.data[0] = bg_vector[0]
                if not self.seen_class:
                    self.conv_vec_unseen.weight.data[0] = bg_vector[0]
        for conv in self.convs:
            conv4_x = conv(x)
        if self.upsample is not None:
            conv4_x = self.upsample(conv4_x)
            if self.upsample_method == 'deconv':
                conv4_x = self.relu(conv4_x)

        # encoder
        x = self.convT(conv4_x)
        if self.voc is not None:
            x = self.conv_voc(x)
        if self.with_learnable_kernel:
            x = self.kernel_semantic(x)

        # decoder
        if self.with_decoder:
            if self.with_learnable_kernel:
                d_x = self.d_kernel_semantic(x)
            d_x = self.dconvT(d_x)
        # classification module
        mask_pred_seen = self.conv_vec(x)
        if not self.seen_class and not self.gzsd:
            mask_pred = self.con_vec_t(mask_pred_seen)
            mask_pred_unseen = self.conv_vec_unseen(mask_pred)
            return mask_pred_unseen
        if self.gzsd:
            mask_pred = self.con_vec_t(mask_pred_seen)
            mask_pred_unseen = self.conv_vec_unseen(mask_pred)
            return mask_pred_seen, mask_pred_unseen
        if not self.with_decoder:
            return mask_pred_seen
        else:
            return mask_pred_seen, conv4_x, d_x

    def get_target(self, sampling_results, gt_masks, rcnn_train_cfg):
        pos_proposals = [res.pos_bboxes for res in sampling_results]
        pos_assigned_gt_inds = [
            res.pos_assigned_gt_inds for res in sampling_results
        ]
        mask_targets = mask_target(pos_proposals, pos_assigned_gt_inds,
                                   gt_masks, rcnn_train_cfg)
        return mask_targets

    @force_fp32(apply_to=('mask_pred', ))
    def loss(self, mask_pred, mask_targets, labels, conv4_x=None, d_x=None):
        loss = dict()
        if self.class_agnostic:
            loss_mask = self.loss_mask(mask_pred, mask_targets,
                                       torch.zeros_like(labels))
        else:
            loss_mask = self.loss_mask(mask_pred, mask_targets, labels)
        loss['loss_mask'] = loss_mask
        if self.with_decoder and conv4_x is not None and d_x is not None:
            loss_encoder_decoder = self.loss_ed(conv4_x, d_x)
            loss['loss_ed'] = loss_encoder_decoder
        return loss

    def get_seg_masks(self, mask_pred, det_bboxes, det_labels, rcnn_test_cfg,
                      ori_shape, scale_factor, rescale):
        """Get segmentation masks from mask_pred and bboxes.

        Args:
            mask_pred (Tensor or ndarray): shape (n, #class+1, h, w).
                For single-scale testing, mask_pred is the direct output of
                model, whose type is Tensor, while for multi-scale testing,
                it will be converted to numpy array outside of this method.
            det_bboxes (Tensor): shape (n, 4/5)
            det_labels (Tensor): shape (n, )
            img_shape (Tensor): shape (3, )
            rcnn_test_cfg (dict): rcnn testing config
            ori_shape: original image size

        Returns:
            list[list]: encoded masks
        """
        if isinstance(mask_pred, torch.Tensor):
            mask_pred = mask_pred.sigmoid().cpu().numpy()
        assert isinstance(mask_pred, np.ndarray)
        # when enabling mixed precision training, mask_pred may be float16
        # numpy array
        mask_pred = mask_pred.astype(np.float32)

        cls_segms = [[] for _ in range(self.num_classes - 1)]
        bboxes = det_bboxes.cpu().numpy()[:, :4]
        labels = det_labels.cpu().numpy() + 1
        if rescale:
            img_h, img_w = ori_shape[:2]
        else:
            img_h = np.round(ori_shape[0] * scale_factor).astype(np.int32)
            img_w = np.round(ori_shape[1] * scale_factor).astype(np.int32)
            scale_factor = 1.0

        if not self.gzsd:
            for i in range(bboxes.shape[0]):
                bbox = (bboxes[i, :] / scale_factor).astype(np.int32)
                label = labels[i]
                w = max(bbox[2] - bbox[0] + 1, 1)
                h = max(bbox[3] - bbox[1] + 1, 1)
                if not self.class_agnostic:
                    mask_pred_ = mask_pred[i, label, :, :]
                else:
                    mask_pred_ = mask_pred[i, 0, :, :]
                im_mask = np.zeros((img_h, img_w), dtype=np.uint8)

                bbox_mask = mmcv.imresize(mask_pred_, (w, h))
                bbox_mask = (bbox_mask > rcnn_test_cfg.mask_thr_binary).astype(
                    np.uint8)
                im_mask[bbox[1]:bbox[1] + h, bbox[0]:bbox[0] + w] = bbox_mask
                rle = mask_util.encode(
                    np.array(im_mask[:, :, np.newaxis], order='F'))[0]
                cls_segms[label - 1].append(rle)

        if self.gzsd:
            if self.num_classes == 49:
                cls_segms = [[] for _ in range(48+17)]

            elif self.num_classes == 66:
                cls_segms = [[] for _ in range(65 + 15)]

            for i in range(bboxes.shape[0]):
                bbox = (bboxes[i, :] / scale_factor).astype(np.int32)
                label = labels[i]
                w = max(bbox[2] - bbox[0] + 1, 1)
                h = max(bbox[3] - bbox[1] + 1, 1)
                temp_label = label
                if mask_pred.shape[1] == 66:
                    if label >= 66:
                        continue
                if mask_pred.shape[1] == 16:
                    if label < 66:
                        continue
                    label -= 65
                if mask_pred.shape[1] == 48:
                    if label >= 48:
                        continue
                if mask_pred.shape[1] == 18:
                    if label < 48:
                        continue
                    label -= 47
                if not self.class_agnostic:
                    mask_pred_ = mask_pred[i, label, :, :]
                else:
                    mask_pred_ = mask_pred[i, 0, :, :]
                im_mask = np.zeros((img_h, img_w), dtype=np.uint8)

                bbox_mask = mmcv.imresize(mask_pred_, (w, h))
                bbox_mask = (bbox_mask > rcnn_test_cfg.mask_thr_binary).astype(
                    np.uint8)
                im_mask[bbox[1]:bbox[1] + h, bbox[0]:bbox[0] + w] = bbox_mask
                rle = mask_util.encode(
                    np.array(im_mask[:, :, np.newaxis], order='F'))[0]
                cls_segms[temp_label - 1].append(rle)

        return cls_segms


#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""

Maxwell Macros
v0.7.0

by Sam Neurohack
from /team/laser

Launchpad set a "current path"

"""

from OSC3 import OSCServer, OSCClient, OSCMessage
import time
import numpy as np
import rtmidi
from rtmidi.midiutil import open_midiinput
from threading import Thread
from rtmidi.midiconstants import (CHANNEL_PRESSURE, CONTROLLER_CHANGE, NOTE_ON, NOTE_OFF,
                                  PITCH_BEND, POLY_PRESSURE, PROGRAM_CHANGE)

import os, json
import midi3

if os.uname()[1]=='raspberrypi':
    pass

port = 8090
ip = "127.0.0.1"
mididest = 'Session 1'
djdest = 'Port'

midichannel = 1
computerIP = ['127.0.0.1','192.168.2.95','192.168.2.52','127.0.0.1',
              '127.0.0.1','127.0.0.1','127.0.0.1','127.0.0.1']
computer = 0

# store current value for computer 1
cc1 =[0]*140

current = {
    "patch":  0,
    "prefixLeft": "/osc/left/X",
    "prefixRight": "/osc/right/X",
    "suffix": "/amp",
    "path": "/osc/left/X/curvetype",
    "pathLeft": "/osc/left/X/curvetype",
    "pathRight": "/osc/left/X/curvetype",
    "previousmacro": -1,
    "LeftCurveType":  0,
    "lfo": 1,
    "rotator": 1,
    "translator": 1
    }

specificvalues = {

# Sine: 0-32, Tri: 33-64, Square: 65-96, Line: 96-127
"curvetype": {"sin": 0, "saw": 33, "squ": 95, "lin": 127},
"freqlimit": {"1": 0, "4": 26, "16": 52, "32": 80, "127": 127},
"amptype": {"constant": 0, "lfo1": 33, "lfo2": 95, "lfo3": 127},
"phasemodtype": {"linear": 0,"sin": 90},
"phaseoffsettype": {"manual": 0, "lfo1": 33, "lfo2": 95, "lfo3": 127},
"ampoffsettype": { "manual": 0, "lfo1": 33, "lfo2": 95, "lfo3": 127},
"inversion": {"off": 0, "on": 127},
"colortype": {"solid": 0, "lfo": 127},
"modtype": {"sin": 0,"linear": 127},
"switch": {"off": 0,"on": 127},
"operation": {"+": 0, "-": 50, "*": 127}
}

#
# Maxwell CCs
#

def FindCC(FunctionName):

    for Maxfunction in range(len(maxwell['ccs'])):
        if FunctionName == maxwell['ccs'][Maxfunction]['Function']:
            #print(FunctionName, "is CC", Maxfunction)
            return Maxfunction

def LoadCC():
    global maxwell

    print("Loading Maxwell CCs Functions...")

    if os.path.exists('maxwell.json'):
        #print('File maxwell.json exits')
        f=open("maxwell.json","r")

    else:
        if os.path.exists('../maxwell.json'):
            #print('File ../maxwell.json exits')
            f=open("../maxwell.json","r")

    s = f.read()
    maxwell = json.loads(s)
    print(len(maxwell['ccs']),"Functions")
    print("Loaded.")

# /cc cc number value
def cc(ccnumber, value, dest=mididest):

    #print('Output CC',[CONTROLLER_CHANGE+midichannel-1, ccnumber, value], dest)
    midi3.MidiMsg([CONTROLLER_CHANGE+midichannel-1,ccnumber,value], dest)

def NoteOn(note,velocity, dest=mididest):
    midi3.NoteOn(note,velocity, mididest)

def NoteOff(note, dest=mididest):
    midi3.NoteOn(note, mididest)

def Send(oscaddress,oscargs=''):

    oscmsg = OSCMessage()
    oscmsg.setAddress(oscaddress)
    oscmsg.append(oscargs)

    osclient = OSCClient()
    osclient.connect((ip, port))

    print("sending OSC message : ", oscmsg, "to", ip, ":",port)
    try:
        osclient.sendto(oscmsg, (ip, port))
        oscmsg.clearData()
        return True
    except:
        print ('Connection to', ip, 'refused : died ?')
        return False

def ssawtooth(samples,freq,phase):

    t = np.linspace(0+phase, 1+phase, samples)
    for ww in range(samples):
        samparray[ww] = signal.sawtooth(2 * np.pi * freq * t[ww])
    return samparray

def ssquare(samples,freq,phase):

    t = np.linspace(0+phase, 1+phase, samples)
    for ww in range(samples):
        samparray[ww] = signal.square(2 * np.pi * freq * t[ww])
    return samparray

def ssine(samples,freq,phase):

    t = np.linspace(0+phase, 1+phase, samples)
    for ww in range(samples):
        samparray[ww] = np.sin(2 * np.pi * freq  * t[ww])
    return samparray

def MixerLeft(value):

    if value == 127:
        Send("/mixer/value", 0)

def MixerRight(value):

    if value == 127:
        Send("/mixer/value", 127)

def MixerTempo(tempo):

    for counter in range(127):
        Send("/mixer/value", counter)

# Jog send 127 to left and 1 to right
# increase or decrease current CC defined in current path
def jogLeft(value):
    path = current["pathLeft"]
    print("jog : path =",path, "CC :", FindCC(path), "value", value)
    MaxwellCC = FindCC(current["pathLeft"])
    if value == 127:
        # decrease CC
        if cc1[MaxwellCC] > 0:
            cc1[MaxwellCC] -= 1
    else:
        if cc1[MaxwellCC] < 127:
            cc1[MaxwellCC] += 1
    #print("sending", cc1[MaxwellCC], "to CC", MaxwellCC )
    cc(MaxwellCC, cc1[MaxwellCC] , dest ='to Maxwell 1')
    #RotarySpecifics(MaxwellCC, path[path.rfind("/")+1:len(path)], value)

# Jog send 127 to left and 1 to right
# increase or decrease current CC defined in current path
def jogRight(value):
    path = current["pathRight"]
    print("jog : path =",path, "CC :", FindCC(path), "value", value)
    MaxwellCC = FindCC(current["pathRight"])
    if value == 127:
        # decrease CC
        if cc1[MaxwellCC] > 0:
            cc1[MaxwellCC] -= 1
    else:
        if cc1[MaxwellCC] < 127:
            cc1[MaxwellCC] += 1
    #print("sending", cc1[MaxwellCC], "to CC", MaxwellCC )
    cc(MaxwellCC, cc1[MaxwellCC] , dest ='to Maxwell 1')
    #RotarySpecifics(MaxwellCC, path[path.rfind("/")+1:len(path)], value)

# Parameter change  : to left 127 / to right 0 or 1
def RotarySpecifics( MaxwellCC, specificsname, value):
    global maxwell

    print("Maxwell CC :",MaxwellCC)
    print("Current :",maxwell['ccs'][MaxwellCC]['init'])
    print("Specifics :",specificvalues[specificsname])
    print("midi value :", value)

    elements = list(enumerate(specificvalues[specificsname]))
    print(elements)
    nextype = maxwell['ccs'][MaxwellCC]['init']

    for count,ele in elements:

        if ele == maxwell['ccs'][MaxwellCC]['init']:
            if count > 0 and value == 127:
                nextype = elements[count-1][1]

            if count < len(elements)-1 and value < 2:
                #print("next is :",elements[count+1][1])
                nextype = elements[count+1][1]

    print("result :", nextype, "new value :", specificvalues[specificsname][nextype], "Maxwell CC", MaxwellCC)
    maxwell['ccs'][MaxwellCC]['init'] = nextype
    cc(MaxwellCC, specificvalues[specificsname][nextype], dest ='to Maxwell 1')

# Change type : trig with only with midi value 127 on a CC event
def ButtonSpecifics127( MaxwellCC, specificsname, value):
    global maxwell

    print("Maxwell CC :",MaxwellCC)
    print("Current :",maxwell['ccs'][MaxwellCC]['init'])
    print("Specifics :",specificvalues[specificsname])
    print("midi value :", value)

    elements = list(enumerate(specificvalues[specificsname]))
    print(elements)
    nextype = maxwell['ccs'][MaxwellCC]['init']

    for count,ele in elements:

        if ele == maxwell['ccs'][MaxwellCC]['init']:
            if count >0 and value == 127:
                nextype = elements[count-1][1]

            if count < len(elements)-1 and value < 2:
                #print("next is :",elements[count+1][1])
                nextype = elements[count+1][1]

    print("result :", nextype, "new value :", specificvalues[specificsname][nextype], "Maxwell CC", MaxwellCC)
    maxwell['ccs'][MaxwellCC]['init'] = nextype
    cc(MaxwellCC, specificvalues[specificsname][nextype], dest ='to Maxwell 1')

# Left cue button 127 = on  0 = off
def PrevPatch(value):
    global current

    print('PrevPatch function')
    if value == 127 and current['patch'] - 1 > -1:
        cc(9, 127, dest=djdest)
        time.sleep(0.1)
        current['patch'] -= 1
        print("Current patch is now :",current['patch'])
        midi3.NoteOn(current['patch'], 127, 'to Maxwell 1')
        cc(9, 0, dest=djdest)

# Right cue button 127 = on  0 = off
def NextPatch(value):
    global current

    print('NextPatch function', current["patch"])
    if value == 127 and current["patch"] + 1 < 41:
        cc(3, 127, dest = djdest)
        current["patch"] += 1
        #ModeNote(current["patch"], 127, 'to Maxwell 1')
        midi3.NoteOn(current["patch"], 127, 'to Maxwell 1')
        print("Current patch is now :",current["patch"])
        time.sleep(0.1)
        cc(3, 0, dest = djdest)

# increase/decrease a CC
def changeCC(value, path):
    global current

    #path = current["pathLeft"]
    MaxwellCC = FindCC(path)
    cc1[MaxwellCC] += value
    print("Change Left CC : path =",path, "CC :", FindCC(path), "is now ", cc1[MaxwellCC])
    cc(MaxwellCC, cc1[MaxwellCC] , dest ='to Maxwell 1')

def PlusTenLeft(value):
    value = 10
    changeCC(value, current["pathLeft"])

def MinusTenLeft(value):
    value = -10
    changeCC(value, current["pathLeft"])

def PlusOneLeft(value):
    value = 1
    changeCC(value, current["pathLeft"])

def MinusOneLeft(value):
    value = -1
    changeCC(value, current["pathLeft"])

def PlusTenRight(value):
    value = 10
    changeCC(value, current["pathRight"])

def MinusTenRight(value):
    value = -10
    changeCC(value, current["pathRight"])

def PlusOneRight(value):
    value = 1
    changeCC(value, current["pathRight"])

def MinusOneRight(value):
    value = -1
    changeCC(value, current["pathRight"])

def ChangeCurveLeft(value):

    MaxwellCC = FindCC(current["prefixLeft"] + '/curvetype')
    RotarySpecifics(MaxwellCC, "curvetype", value)

def ChangeFreqLimitLeft(value):

    MaxwellCC = FindCC(current["prefixLeft"] + '/freqlimit')
    RotarySpecifics(MaxwellCC, "curvetype", value)

def ChangeATypeLeft(value):

    MaxwellCC = FindCC(current["prefixLeft"] + '/freqlimit')
    RotarySpecifics(MaxwellCC, "curvetype", value)

def ChangePMTypeLeft(value):

    MaxwellCC = FindCC(current["prefixLeft"] + '/phasemodtype')
    RotarySpecifics(MaxwellCC, "curvetype", value)

def ChangePOTypeLeft(value):

    MaxwellCC = FindCC(current["prefixLeft"] + '/phaseoffsettype')
    RotarySpecifics(MaxwellCC, "curvetype", value)

def ChangeAOTypeLeft(value):

    MaxwellCC = FindCC(current["prefixLeft"] + '/ampoffsettype')
    RotarySpecifics(MaxwellCC, "curvetype", value)

def ChangeCurveRight(value):

    MaxwellCC = FindCC(current["prefixRight"] + '/curvetype')
    RotarySpecifics(MaxwellCC, "curvetype", value)

def ChangeCurveLFO(value):

    MaxwellCC = FindCC('/lfo/'+ current["lfo"] +'/curvetype')
    RotarySpecifics(MaxwellCC, "curvetype", value)

def ChangeCurveRot(value):

    MaxwellCC = FindCC('/rotator/'+ current["rotator"] +'/curvetype')
    RotarySpecifics(MaxwellCC, "curvetype", value)

def ChangeCurveTrans(value):

    MaxwellCC = FindCC('/translator/'+ current["translator"] +'/curvetype')
    RotarySpecifics(MaxwellCC, "curvetype", value)


import os
import cv2
import numpy as np
import tensorflow as tf
import tensorflow.compat.v1 as tf
from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw
from flask_ngrok import run_with_ngrok
from flask import Flask,request,send_from_directory,render_template

# GLOBAL ACCESS
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
THIS_FOLDER = os.path.dirname(os.path.abspath(__file__))
# SETUP APPLICATION
tf.disable_eager_execution()
app = Flask(__name__, static_url_path='')
run_with_ngrok(app)

def sparse_tensor_to_strs(sparse_tensor):
    indices= sparse_tensor[0][0]
    values = sparse_tensor[0][1]
    dense_shape = sparse_tensor[0][2]
    strs = [ [] for i in range(dense_shape[0]) ]
    string = []
    ptr = 0
    b = 0
    for idx in range(len(indices)):
        if indices[idx][0] != b:
            strs[b] = string
            string = []
            b = indices[idx][0]
        string.append(values[ptr])
        ptr = ptr + 1
    strs[b] = string
    return strs

def normalize(image):
    return (255. - image)/255.

def resize(image, height):
    width = int(float(height * image.shape[1]) / image.shape[0])
    sample_img = cv2.resize(image, (width, height))
    return sample_img

voc_file = "vocabulary_semantic.txt"
model = "semantic_model/semantic_model.meta"
tf.reset_default_graph()
sess = tf.InteractiveSession()
# Read the dictionary
dict_file = open(voc_file,'r')
dict_list = dict_file.read().splitlines()
int2word = dict()
for word in dict_list:
    word_idx = len(int2word)
    int2word[word_idx] = word
dict_file.close()

# Restore weights
saver = tf.train.import_meta_graph(model)
saver.restore(sess,model[:-5])
graph = tf.get_default_graph()
input = graph.get_tensor_by_name("model_input:0")
seq_len = graph.get_tensor_by_name("seq_lengths:0")
rnn_keep_prob = graph.get_tensor_by_name("keep_prob:0")
height_tensor = graph.get_tensor_by_name("input_height:0")
width_reduction_tensor = graph.get_tensor_by_name("width_reduction:0")
logits = tf.get_collection("logits")[0]

# Constants that are saved inside the model itself
WIDTH_REDUCTION, HEIGHT = sess.run([width_reduction_tensor, height_tensor])

decoded, _ = tf.nn.ctc_greedy_decoder(logits, seq_len)
# HOME
@app.route("/")
def root():
    return render_template('index.html')
# IMAGE REQUEST
@app.route('/img/<filename>')
def send_img(filename):
    return send_from_directory('', filename)
# ANDROID REQUEST
@app.route('/android/predict', methods = ['GET', 'POST'])
def login():
    return 'Yeah, it works.'
# GET
@app.route('/users/<var>')
def hello_user(var):
    """
    this serves as a demo purpose
    :param user:
    :return: str
    """
    return "Wow, the GET works %s!" % var

# POST
@app.route('/api/post_some_data', methods=['POST'])
def get_text_prediction():
    """
    predicts requested text whether it is ham or spam
    :return: json
    """
    json = request.get_json()
    print(json)
    if len(json['text']) == 0:
        return jsonify({'error': 'invalid input'})

    return jsonify({'This is the KEY': json['This is the value?']})

#UPLOAD_FOLDER = 'static/upload'
#app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
#@app.route('/upload', methods=['GET','POST'])
#def upload():
#    if flask.request.method == "POST":
#        files = flask.request.files.getlist("file")
#        for file in files:
#            file.save(os.path.join(app.config['UPLOAD_FOLDER'], file.filename))

# MODEL PREDICTION
@app.route('/predict', methods = ['GET', 'POST'])
def predict():
    if request.method == 'POST':
        f = request.files['file']
        img = f
        image = Image.open(img).convert('L')
        image = np.array(image)
        image = resize(image, HEIGHT)
        image = normalize(image)
        image = np.asarray(image).reshape(1,image.shape[0],image.shape[1],1)
        seq_lengths = [ image.shape[2] / WIDTH_REDUCTION ]
        prediction = sess.run(decoded,
                          feed_dict={
                              input: image,
                              seq_len: seq_lengths,
                              rnn_keep_prob: 1.0,
                          })
        str_predictions = sparse_tensor_to_strs(prediction)
        array_of_notes = []

        for w in str_predictions[0]:
            array_of_notes.append(int2word[w])
        notes=[]
        for i in array_of_notes:
            if i[0:5]=="note-":
                if not i[6].isdigit():
                    notes.append(i[5:7])
                else:
                    notes.append(i[5])
        img = Image.open(img).convert('L')
        size = (img.size[0], int(img.size[1]*1.5))
        layer = Image.new('RGB', size, (255,255,255))
        layer.paste(img, box=None)
        img_arr = np.array(layer)
        height = int(img_arr.shape[0])
        width = int(img_arr.shape[1])
        print(img_arr.shape[0])
        draw = ImageDraw.Draw(layer)
        # font = ImageFont.truetype(<font-file>, <font-size>)
        font = ImageFont.truetype("Aaargh.ttf", 16)
        # draw.text((x, y),"Sample Text",(r,g,b))
        j = width / 9
        for i in notes:
            draw.text((j, height-40), i, (0,0,0), font=font)
            j+= (width / (len(notes) + 4))
        layer.save("img/annotated.png")
        return render_template('result.html')

if __name__=="__main__":
    app.run()


import os, sys
import numpy as np
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.autograd.profiler as profiler
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm, trange

# our own codes
from core.trainer import *
from core.utils.ray_utils import *
from core.utils.run_nerf_helpers import *
from core.utils.evaluation_helpers import evaluate_metric, evaluate_pampjpe_from_smpl_params
from core.utils.skeleton_utils import draw_skeletons_3d

from core.raycasters import create_raycaster
from core.pose_opt import create_popt
from core.load_data import load_data

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
np.random.seed(0)
DEBUG = True

@torch.no_grad()
def render_path(render_poses, hwf, chunk, render_kwargs,
                centers=None, kp=None, skts=None, cyls=None, bones=None,
                gt_imgs=None, bg_imgs=None, bg_indices=None,
                cams=None, subject_idxs=None, render_factor=0,
                white_bkgd=False, ret_acc=False,
                ext_scale=0.00035, base_bg=1.0):

    H, W, focal = hwf

    if render_factor!=0:
        # Render downsampled for speed
        H = H//render_factor
        W = W//render_factor
        if isinstance(focal, float):
            focal = focal / render_factor
            if centers is not None:
                centers = centers / render_factor
        else:
            focal = focal.copy() / render_factor
            if centers is not None:
                centers = centers.copy() / render_factor

    if kp is not None or cyls is not None:
        # only render part of the image
        rays, valid_idxs, cyls, bboxes = kp_to_valid_rays(render_poses, H, W,
                                                          focal, kps=kp, cylinder_params=cyls,
                                                          skts=skts, ext_scale=ext_scale,
                                                          centers=centers)
    else:
        rays, valid_idxs = None, None

    rgbs, disps, accs = [], [], []

    t = time.time()
    # reuse human pose if #render_poses > #human_poses and #human_poses > 0
    def reuse_input(x, expand=None):
        if x is None:
            return x

        if x.shape[0] > 1:
            y = x[i%x.shape[0]:i%x.shape[0]+1].clone()
        else:
            y = x.clone()

        if expand is not None:
            y = y.expand(expand, *x.shape[1:])
        return y
    #reuse_input = lambda x: x[i%x.shape[0]:i%x.shape[0]+1] if x is not None and x.shape[0] > 1 else x

    for i, c2w in enumerate(tqdm(render_poses)):
        print(i, time.time() - t)
        t = time.time()
        h = H if isinstance(H, int) else H[i]
        w = W if isinstance(W, int) else W[i]

        ray_input = rays[i] if rays is not None else None
        expand = len(ray_input[0])
        kp_input = reuse_input(kp, expand)
        skt_input = reuse_input(skts, expand)
        cyl_input = reuse_input(cyls, expand)
        bone_input = reuse_input(bones, expand)
        cam_input = reuse_input(cams, expand)
        subject_input = reuse_input(subject_idxs, expand)
        #center_input = reuse_input(centers, expand)

        if len(ray_input[0]) > 0:
            ret_dict = render(h, w, focal, rays=ray_input, chunk=chunk, c2w=c2w[:3,:4],
                              kp_batch=kp_input, skts=skt_input, cyls=cyl_input, cams=cam_input,
                              subject_idxs=subject_input,
                              bones=bone_input, **render_kwargs)
            rgb, disp, acc = ret_dict['rgb_map'], ret_dict['disp_map'], ret_dict['acc_map']

        if valid_idxs is not None:
            # in this case, we only render the rays that are within the foreground or cylinder
            valid_idx = valid_idxs[i]

            # initialize the images
            if bg_imgs is not None and not white_bkgd:

                if bg_indices is not None:
                    bg = torch.tensor(bg_imgs[bg_indices[i]]).permute(2, 0, 1)[None]
                else:
                    # TODO: HACK, FIX THIS
                    bg = torch.tensor(bg_imgs[0]).permute(2, 0, 1)[None]
                rgb_img = F.interpolate(bg, size=(h, w), mode='bilinear',
                                   align_corners=False)[0].permute(1, 2, 0).view(h * w, 3)
            else:
                rgb_img = torch.zeros(h * w, 3) if not white_bkgd else torch.ones(h * w, 3)
            disp_img = torch.zeros(h * w)

            if len(valid_idx) > 0:
                # check if we have rays for this image
                bg = (1. - acc[..., None]) * rgb_img[valid_idx]

                # assign values to the corresponding ray location
                rgb_img[valid_idx] = rgb + bg
                disp_img[valid_idx] = disp
                if ret_acc:
                    acc_img = torch.zeros(h * w)
                    acc_img[valid_idx] = acc
                    accs.append(acc_img.view(h, w, 1).detach().cpu().numpy())

            rgbs.append(rgb_img.view(h, w, 3).detach().cpu().numpy())
            disps.append(disp_img.view(h, w, 1).detach().cpu().numpy())

        else:
            # render the whole image in this case, can append directly
            rgbs.append(rgb.cpu().numpy())
            disps.append(disp.cpu().numpy())

    rgbs = np.stack(rgbs, 0)
    disps = np.stack(disps, 0)
    disps_nan = np.isnan(disps)
    disps[disps_nan] = 0.
    if ret_acc:
        accs = np.stack(accs, 0)

    return rgbs, disps, accs, valid_idxs, bboxes

@torch.no_grad()
def render_testset(poses, hwf, args, render_kwargs, kps=None, skts=None, cyls=None, cams=None,
                   bones=None, subject_idxs=None, gt_imgs=None, gt_masks=None, bg_imgs=None, bg_indices=None,
                   vid_base=None, rgb_vid="rgb.mp4", disp_vid="disp.mp4", eval_metrics=False,
                   render_factor=0, eval_postfix="", save_npy=False, eval_both=False, centers=None,
                   save_image=False):

    render_kwargs["ray_caster"].eval()
    rgbs, disps, _, valid_idxs, bboxes = render_path(poses, hwf, args.chunk//8, render_kwargs,
                                                     bg_imgs=bg_imgs, bg_indices=bg_indices,
                                                     centers=centers, kp=kps, skts=skts, cyls=cyls, bones=bones,
                                                     cams=cams, subject_idxs=subject_idxs, render_factor=args.render_factor,
                                                     ext_scale=args.ext_scale, white_bkgd=args.white_bkgd)
    render_kwargs["ray_caster"].train()
    if save_image:
        print("Warning: this is super hacky. Should terminate in the right way.")
        return

    if save_npy:
        import pickle as pkl
        np.save(vid_base + rgb_vid + ".npy", rgbs)
        np.save(vid_base + disp_vid + ".npy", disps)

    if eval_metrics:
        metrics = evaluate_metric(rgbs, gt_imgs, disps, gt_masks, valid_idxs, poses=poses,
                                  kps=kps, hwf=hwf, ext_scale=args.ext_scale, rgb_vid=rgb_vid,
                                  disp_vid=disp_vid, vid_base=vid_base, eval_postfix=eval_postfix,
                                  centers=centers, eval_both=eval_both, white_bkgd=False,
                                  render_factor = args.render_factor)
    else:
        metrics = {"psnr": None, "ssim": None}
    disps = disps / disps.max()

    return metrics, rgbs, disps

def config_parser():

    import configargparse
    parser = configargparse.ArgumentParser()
    parser.add_argument('--config', is_config_file=True,
                        help='config file path')
    parser.add_argument("--expname", type=str,
                        help='experiment name')
    parser.add_argument("--basedir", type=str, default='./logs/',
                        help='where to store ckpts and logs')
    parser.add_argument("--datadir", type=str, default='./data/llff/fern',
                        help='input data directory')

    # training options
    parser.add_argument("--lindisp", action='store_true',
                        help='sampling linearly in disparity rather than depth')
    parser.add_argument("--netdepth", type=int, default=8,
                        help='layers in network')
    parser.add_argument("--netwidth", type=int, default=256,
                        help='channels per layer')
    parser.add_argument("--netdepth_fine", type=int, default=8,
                        help='layers in fine network')
    parser.add_argument("--netwidth_fine", type=int, default=256,
                        help='channels per layer in fine network')

    parser.add_argument("--N_rand", type=int, default=32*32*4,
                        help='batch size (number of random rays per gradient step)')
    parser.add_argument("--lrate", type=float, default=5e-4,
                        help='learning rate')
    parser.add_argument("--lrate_decay", type=int, default=250,
                        help='exponential learning rate decay (in args.decay_unit (default=1000) steps)')
    parser.add_argument("--lrate_decay_rate", type=float, default=0.1,
                        help='learning rate decay')
    parser.add_argument("--decay_unit", type=int, default=1000,
                        help='number of steps until the next decay happen')
    parser.add_argument("--weight_decay", type=float, default=None,
                        help='weight decay to apply on the nerf model')
    parser.add_argument("--single_net", action='store_true',
                        help='use a single network for coarse and fine network')
    parser.add_argument("--coarse_weight", type=float, default=1.0,
                        help='loss weight for coarse samples.')

    parser.add_argument("--use_temp_loss", action='store_true',
                        help='use temporal smoothness loss')
    parser.add_argument("--temp_coef", type=float, default=0.05,
                        help='coefficient on temporal smoothness loss')

    parser.add_argument("--chunk", type=int, default=1024*32,
                        help='number of rays processed in parallel, decrease if running out of memory')
    parser.add_argument("--netchunk", type=int, default=1024*64,
                        help='number of pts sent through network in parallel, decrease if running out of memory')
    parser.add_argument("--no_reload", action='store_true',
                        help='do not reload weights from saved ckpt')
    parser.add_argument("--ft_path", type=str, default=None,
                        help='specific weights npy file to reload for coarse network')
    parser.add_argument("--n_iters", type=int, default=200000,
                        help="number of training iter")
    parser.add_argument("--loss_fn", type=str, default='MSE',
                        help='Type of loss to use')
    parser.add_argument("--loss_beta", type=float, default=0.1,
                        help='beta for smoothed l1 loss')
    parser.add_argument("--reg_fn", type=str, default=None,
                        help='to regularize occupancy')
    parser.add_argument("--reg_coef", type=float, default=0.1,
                        help='coefficient for regularization loss')
    parser.add_argument("--init_poseopt", type=str, default=None,
                        help='initial poseopt layer from a specific checkpoint')
    parser.add_argument("--no_poseopt_reload", action='store_true',
                        help='ignore the poseopt ckpt')
    parser.add_argument("--finetune", action='store_true',
                        help='flag to indicate fine tune stage, which will not load optimizer and global step from ckpt')
    parser.add_argument("--fix_layer", type=int, default=0,
                        help='flag to fix the pts_linears layer weights before the specified layer')
    parser.add_argument("--use_yuv", action='store_true',
                        help='loss in yuv space instead of rgb')

    # rendering options
    parser.add_argument("--density_scale", type=float, default=1.0,
                       help='to scale the density')
    parser.add_argument("--N_samples", type=int, default=64,
                        help='number of coarse samples per ray')
    parser.add_argument("--N_importance", type=int, default=0,
                        help='number of additional fine samples per ray')
    parser.add_argument("--perturb", type=float, default=1.,
                        help='set to 0. for no jitter, 1. for jitter')
    parser.add_argument("--P_nms", type=float, default=0.0,
                        help='percentage (in [0, 1] * 100%) of samples from out-of-mask area.')
    parser.add_argument("--use_viewdirs", action='store_true',
                        help='use full 5D input instead of 3D')
    parser.add_argument("--i_embed", type=int, default=0,
                        help='set 0 for default positional encoding, -1 for none')
    parser.add_argument("--multires", type=int, default=10,
                        help='log2 of max freq for positional encoding (kp encoding)')
    parser.add_argument("--multires_pts", type=int, default=5,
                        help='log2 of max freq for positional encoding (mapped location)')
    parser.add_argument("--multires_views", type=int, default=4,
                        help='log2 of max freq for positional encoding (2D direction)')
    parser.add_argument("--multires_bones", type=int, default=0,
                        help='log2 of max freq for positional encoding on joint rotations (24 x 3D)')
    parser.add_argument("--raw_noise_std", type=float, default=0.,
                        help='std dev of noise added to regularize sigma_a output, 1e0 recommended')
    parser.add_argument("--ray_noise_std", type=float, default=0.,
                        help='std dev of noises to add stochasticity for pose optimization')

    parser.add_argument("--render_factor", type=int, default=0,
                        help='downsampling factor to speed up rendering, set 4 or 8 for fast preview')
    parser.add_argument("--save_image", action='store_true',
                        help='save rendering outcomes as images instead of videos.')

    # training options
    parser.add_argument("--nerf_type", type=str, default="nerf",
                        help='type of NeRF model to use')
    parser.add_argument("--precrop_iters", type=int, default=0,
                        help='number of steps to train on central crops')
    parser.add_argument("--precrop_frac", type=float,
                        default=.5, help='fraction of img taken for central crops')
    parser.add_argument("--density_type", type=str, default='relu',
                        help='function for transforming raw density values to values >= 0')
    parser.add_argument("--softplus_shift", type=float, default=1.0,
                        help='shift the softplus activation input by subtracting this amount')

    # TODO: get this from dataloader?
    parser.add_argument("--n_subjects", type=int, default=2,
                        help='numbers of subjects to learn')

    # per-frame code optimization options
    parser.add_argument("--opt_framecode", action='store_true',
                        help='jointly optimize per-frame codes')
    parser.add_argument("--n_framecodes", type=int, default=None,
                        help="overwrite the number of framecode for the NeRF model")
    parser.add_argument("--framecode_size", type=int, default=16,
                        help='size of per-view frame code')

    # pose optimization options
    parser.add_argument("--opt_rot6d",action='store_true',
                        help='use continuous rotation matrix')
    parser.add_argument("--opt_posecode", action='store_true',
                        help='jointly optimize per-pose code')
    parser.add_argument("--opt_pose", action='store_true',
                        help='jointly optimize pose')
    parser.add_argument("--opt_pose_stop", type=int, default=None,
                        help='stop updating after this many steps')
    parser.add_argument("--opt_pose_coef", type=float, default=0.,
                        help='regularize so that the optimized pose wouldn\'t be too far from the original')
    parser.add_argument("--opt_pose_tol", type=float, default=0.,
                        help='tolerance of pose adjustment')
    parser.add_argument("--opt_pose_type", type=str, default="B",
                        help="type of objective to use for pose optimization")
    parser.add_argument("--opt_pose_step", type=int, default=1,
                        help='frequency of updating parameters')
    parser.add_argument("--opt_pose_lrate", type=float, default=5e-4,
                        help='learning rate for pose update')
    parser.add_argument("--opt_pose_lrate_decay", type=int, default=250,
                        help='exponential step size decay for pose optimizer (in opt_pose_decay steps)')
    parser.add_argument("--opt_pose_decay_rate", type=float, default=1.0,
                        help='exponential decay rate')
    parser.add_argument("--opt_pose_warmup", type=int, default=0,
                        help='wait this amount of iteration before we optimizing keypoints')
    parser.add_argument("--opt_pose_decay_unit", type=int, default=400,
                        help='unit of decay for opt pose')
    parser.add_argument("--opt_pose_cache", action='store_true',
                        help='use cache for slight speed up lol')
    parser.add_argument("--opt_pose_joint", action='store_true',
                        help='jointly learn NeRF and pose when pose_turn == True')
    parser.add_argument("--testopt", action='store_true',
                        help='doing test time optimization only, no NeRF update')

    # additional network
    parser.add_argument("--use_bgnet", action='store_true',
                        help='learn a background network')
    parser.add_argument("--use_uncertainty", action='store_true',
                        help='use uncertainty for density')
    parser.add_argument("--bgnet_stop", type=int, default=500000,
                        help='stop training bgnet after this amount of iterations')
    parser.add_argument("--bgnet_reg", type=float, default=0.01,
                        help='penalize bgnet for deviating from bg values')
    parser.add_argument("--use_bgfill", action='store_true',
                        help='fill in mean bg image at foreground, while keeping othe the same')

    parser.add_argument("--lbsnet_type", type=str, default="default",
                        help='type of LBSNET')
    parser.add_argument("--use_lbsnet", action='store_true',
                        help='use LBS blending')
    parser.add_argument("--n_lbs", type=int, default=1,
                        help='number of blended coordinates')
    parser.add_argument("--multires_lbs", type=int, default=10,
                        help='number of frequencies')
    parser.add_argument("--multires_lbsviews", type=int, default=4,
                        help='number of frequencies')

    parser.add_argument("--use_ckpt_anchor", action='store_true',
                        help='set ckpt as anchor')

    # dataset options
    parser.add_argument("--num_workers", type=int, default=16,
                        help='number of workers for dataloader (works only when --image_batching=True)')
    parser.add_argument("--dataset_type", type=str, default=['h36m'], nargs='+',
                        help='options: h36m / surreal / perfcap / mixamo')
    parser.add_argument("--subject", type=str, default=["S9"], nargs='+',
                        help='subject to train with on mixamo')
    parser.add_argument("--camera", type=int, default=None,
                        help='camera to use')
    parser.add_argument("--split", type=str, default='train',
                        help='use validation / train set')
    parser.add_argument("--white_bkgd", action='store_true',
                        help='set to render synthetic data on a white bkgd (always use for dvoxels)')
    parser.add_argument("--ext_scale", type=float, default=0.001,
                        help="scaling the extrinsic matrix (world coordinate)")
    parser.add_argument("--use_background", action='store_true',
                        help="use static background during reconstructions.")
    parser.add_argument("--fg_ratio", type=float, default=None,
                        help="sample at least this amount of rays from the mask")
    parser.add_argument("--kp_dist_type", type=str, default="reldist",
                        help="specify the way we calculate kp input")
    parser.add_argument("--view_type", type=str, default="relray",
                        help="type of view as input")
    parser.add_argument("--bone_type", type=str, default="reldir",
                        help="type of bone as input")
    parser.add_argument("--pts_tr_type", type=str, default="local",
                        help='type of transformation to apply on 3D world points')

    # surreal dataset option
    parser.add_argument("--train_skip", type=int, default=1,
                        help="skip animated frames in surreal dataset")
    parser.add_argument("--view_skip", type=int, default=1,
                        help="skip camera view in surreal dataset")
    parser.add_argument("--N_cams", type=int, default=None,
                        help='number of cams to use for surreal dataset')
    parser.add_argument("--use_cutoff", action='store_true',
                        help='use cutoff embedder')
    parser.add_argument("--normalize_cutoff", action='store_true',
                        help='normalize cutoff so that we have a complete wave inside the cutoff range')
    parser.add_argument("--cutoff_mm", type=float, default=500,
                        help='distance to apply cutoff (default: 0.5m = 500mm)')
    parser.add_argument("--cutoff_inputs", action='store_true',
                        help='apply cutoff to the input as well')
    parser.add_argument("--cut_to_dist", action='store_true',
                        help='cut input distance to cutoff_mm')
    parser.add_argument("--cutoff_shift", action='store_true',
                        help='shift the input distance by (dist * 2 - 1), so that distance span a complete period')
    parser.add_argument("--cutoff_viewdir", action='store_true',
                        help='apply cutoff to view direction input using the initial inputs!')
    parser.add_argument("--opt_cutoff", action='store_true',
                        help='optimize cutoff threshold')
    parser.add_argument("--cutoff_step", type=int, default=250,
                        help='steps (in 1000) to increase the temperature parameter of cutoff by cutoff_rate')
    parser.add_argument("--cutoff_rate", type=float, default=10.,
                        help='exponential cutoff temperature increase')

    parser.add_argument("--cutoff_bones", action='store_true',
                        help='apply cutoff to bone rotations')
    parser.add_argument("--cutoff_ancestors", type=int, default=5,
                        help='numbers of ancestors to keep along the kinematic chain when doing bone cutoff')

    parser.add_argument("--freq_schedule", action='store_true',
                         help='schedule frequencies as in BARF')
    parser.add_argument("--freq_schedule_step", type=int, default=5,
                         help='target step to enable all frequencies (in 1000 steps)')
    parser.add_argument("--init_freq", type=float, default=0.,
                        help='initial frequencies that are enabled')

    # h36m dataset
    # TODO: enable this
    parser.add_argument("--multiview", action='store_true',
                        help='use multiview optimization')
    # TODO: REMOVE
    parser.add_argument("--training_res", type=float, default=1.0,
                        help='resize training images by this amount')

    parser.add_argument("--val_seq", nargs="+", type=int, default=[6, 18],
                        help='list of sequence number for validation')
    #parser.add_argument("--rand_train_kps", type=int, default=None,
    #                    help='randomly select a number of kps for training (note: for fixed random seq for 180 and 420)')
    parser.add_argument("--rand_train_kps", type=str, default=None,
                        help='randomly select a number of kps for training (note: for fixed random seq for 180 and 420)')

    parser.add_argument("--N_sample_images", type=int, default=8,
                        help='number of images to sample rays from')
    parser.add_argument("--image_batching", action='store_true',
                        help='sample rays from N_sample images')
    parser.add_argument("--mask_image", action='store_true',
                        help='mask out pixels that are not in the foreground when providing image target')
    parser.add_argument("--patch_size", type=int, default=1,
                        help='sample patches of rays from the image')
    parser.add_argument("--load_refined", action='store_true',
                        help='load refined poses for training')

    # logging/saving options
    parser.add_argument("--i_print",   type=int, default=100,
                        help='frequency of console printout and metric loggin')
    parser.add_argument("--i_weights", type=int, default=10000,
                        help='frequency of weight ckpt saving')
    parser.add_argument("--i_pose_weights", type=int, default=2000,
                        help='frequency of saving pose weights')
    parser.add_argument("--i_testset", type=int, default=50000,
                        help='frequency of testset saving')
    parser.add_argument("--i_video", type=int, default=10000,
                        help='frequency of render_poses video saving')

    # debug
    parser.add_argument("--debug", action='store_true',
                        help='debug flag')

    return parser

def train():

    parser = config_parser()
    args = parser.parse_args()

    train_loader, render_data, data_attrs = load_data(args)
    skel_type = data_attrs['skel_type']
    hwf = data_attrs["hwf"]
    H, W, focal = hwf
    print("Loader initialized.")

    # Create log dir and copy the config file
    basedir = args.basedir
    expname = args.expname
    os.makedirs(os.path.join(basedir, expname), exist_ok=True)
    f = os.path.join(basedir, expname, 'args.txt')
    with open(f, 'w') as file:
        for arg in sorted(vars(args)):
            attr = getattr(args, arg)
            file.write('{} = {}\n'.format(arg, attr))
    if args.config is not None:
        f = os.path.join(basedir, expname, 'config.txt')
        with open(f, 'w') as file:
            file.write(open(args.config, 'r').read())

    # Create nerf model
    render_kwargs_train, render_kwargs_test, start, grad_vars, optimizer, loaded_ckpt = create_raycaster(args, data_attrs, device=device)
    global_step = start

    popt_kwargs, pose_optimizer = None, None
    if args.opt_pose:
        gt_kps = data_attrs['gt_kp3d']
        pose_optimizer, popt_kwargs = create_popt(args, data_attrs, ckpt=loaded_ckpt, device=device)
    print('done creating popt')

    N_iters = args.n_iters + 1

    # tensorboard summary writer
    writer = SummaryWriter(os.path.join(basedir, expname))

    start = start + 1
    ray_caster = render_kwargs_train['ray_caster']
    trainer = Trainer(args, data_attrs, optimizer, pose_optimizer,
                      render_kwargs_train, render_kwargs_test,
                      popt_kwargs, device=device)

    train_iter = iter(train_loader)
    for i in trange(start, N_iters):
        time0 = time.time()
        batch = next(train_iter)
        loss_dict, stats = trainer.train_batch(batch, i, global_step)

        # Rest is logging
        if i % args.i_weights == 0:
            path = os.path.join(basedir, expname, '{:06d}.tar'.format(i))
            trainer.save_nerf(path, global_step)

        save_opt_pose = args.opt_pose and ((args.opt_pose_stop is None) or (args.opt_pose_stop > i))
        if (i % args.i_pose_weights == 0) and save_opt_pose:
            path = os.path.join(basedir, expname, 'pose_weights_{:06d}.tar'.format(i))
            trainer.save_popt(path, i)

        # TODO: deal with this
        if i % args.i_testset==0 and i > 0:
            if args.opt_pose and render_data["kp_idxs"] is not None:
                popt_layer = popt_kwargs['popt_layer']
                with torch.no_grad():
                    kp_val, bone_val, skt_val, _, _ = popt_layer(render_data["kp_idxs"])
            else:
                kp_val = torch.tensor(render_data["kp3d"]).to(device)
                skt_val = torch.tensor(render_data["skts"]).to(device)
                bone_val = torch.tensor(render_data["bones"]).to(device)

            pose_val = torch.tensor(render_data["c2ws"]).to(device)

            gt_imgs = render_data["imgs"]
            gt_masks = render_data["fgs"]
            bg_imgs = render_data["bgs"]
            bg_indices = render_data.get("bg_idxs", None)
            centers = render_data['center']
            cams_val = torch.tensor(render_data["cam_idxs"]) if args.opt_framecode else None
            subject_val = render_data.get("subject_idxs", None)
            subject_val = torch.tensor(subject_val) if subject_val is not None else None

            moviebase = os.path.join(basedir, expname, '{}_val_{:06d}_'.format(expname, i))
            # Intentionally not feeding cyl here to speed ikt up
            if bg_indices is not None:
                masked_gts = gt_imgs * gt_masks + (1 - gt_masks) * bg_imgs[bg_indices]
            else:
                masked_gts = gt_imgs * gt_masks + (1 - gt_masks) * bg_imgs

            metrics, rgbs, disps = render_testset(pose_val, render_data["hwf"], args, render_kwargs_test, cams=cams_val,
                                                  kps=kp_val, skts=skt_val, bones=bone_val, subject_idxs=subject_val,
                                                  gt_imgs=masked_gts, gt_masks=gt_masks, vid_base=moviebase, centers=centers,
                                                  bg_imgs=bg_imgs, bg_indices=bg_indices, eval_metrics=True, eval_both=True)

            fps  = 5
            writer.add_video("Val/ValRGB", torch.tensor(rgbs).permute(0, 3, 1, 2)[None], i, fps=fps)
            writer.add_video("Val/ValDIPS", torch.tensor(disps).permute(0, 3, 1, 2)[None], i, fps=fps)
            if args.opt_pose:
                RH, RW, Rfocals = render_data["hwf"]
                # set the resolution right!
                if args.render_factor != 0:
                    RH, RW = RH // args.render_factor, RW // args.render_factor
                    Rfocals = Rfocals / args.render_factor
                skel_imgs = draw_skeletons_3d((rgbs * 255).astype(np.uint8), kp_val.cpu().numpy(), pose_val.cpu().numpy(),
                                              RH, RW, Rfocals)
                writer.add_video("Val/Skeleton", torch.tensor(skel_imgs).permute(0, 3, 1, 2)[None], i, fps=fps)
            writer.add_scalar("Val/PSNR", metrics["psnr"], i)
            writer.add_scalar("Val/SSIM", metrics["ssim"], i)

        if i % args.i_print == 0:
            mem = torch.cuda.max_memory_allocated() / 1024. / 1024.
            loss, psnr, alpha, total_norm = loss_dict['total_loss'], stats['psnr'], stats['alpha'], stats['total_norm']
            tqdm.write(f"[TRAIN] Iter: {i} Loss: {loss.item()}  PSNR: {psnr}, Alpha: {alpha}, GradNorm {total_norm}, Mem: {mem}")
            for k in loss_dict:
                if k == 'total_loss':
                    writer.add_scalar(f'Loss/{k}_{args.loss_fn}', loss_dict[k], i)
                else:
                    writer.add_scalar(f'Loss/{k}', loss_dict[k], i)

            for k in stats:
                writer.add_scalar(f'Stats/{k}', stats[k], i)

        global_step += 1

if __name__=='__main__':
    torch.set_default_tensor_type('torch.cuda.FloatTensor')
    torch.multiprocessing.set_start_method('spawn')

    train()


#!/usr/bin/env python3

import os
import base64

import tensorflow as tf
import numpy as np
import cv2

tf.flags.DEFINE_string(
    'input', '', 'Input (TSV or directory with JPEG/PNG images).')

tf.flags.DEFINE_string(
    'output_file', 'output.csv', 'Output CSV file.')

tf.flags.DEFINE_string(
    'meta_graph', '', 'Exported MetaGraph file.')

tf.flags.DEFINE_string(
    'model', '', 'Exported model file.')

tf.flags.DEFINE_integer(
    'batch_size', 32, 'How many images process at one time.')

FLAGS = tf.flags.FLAGS

def aligned_tsv_loader(input_tsv):
    with open(input_tsv) as fin:
        for line in fin:
            freebase_mid, search_rank, image_url, page_url, \
                face_id, box_b64, image_b64 = line.strip().split("\t")

            fullpath = "{},{},{}".format(freebase_mid, search_rank, face_id)
            buf = base64.decodebytes(image_b64.encode())
            buf = np.frombuffer(buf, dtype=np.uint8)
            raw = cv2.imdecode(buf, cv2.IMREAD_COLOR)

            yield fullpath, raw

def directory_image_loader(input_dir):
    for fname in os.listdir(input_dir):
        if os.path.splitext(fname)[-1].lower() not in [".jpg", ".png"]:
            continue

        fullpath = os.path.join(input_dir, fname)
        raw = cv2.imread(fullpath, cv2.IMREAD_COLOR)
        yield fullpath, raw

def load_images(input_path, target_size, batch_size):
    if os.path.isdir(input_path):
        loader = directory_image_loader(input_path)
    elif os.path.isfile(input_path) and input_path.lower().endswith(".tsv"):
        loader = aligned_tsv_loader(input_path)
    else:
        raise Exception("Unknown input!")

    batch_shape = [batch_size] + list(target_size) + [3]
    images = np.zeros(batch_shape, dtype=np.uint8)
    filenames = []
    idx = 0

    for fullpath, raw in loader:
        resized = cv2.resize(raw, target_size[::-1])

        images[idx, ...] = resized
        filenames.append(fullpath)

        idx += 1
        if idx == batch_size:
            yield filenames, images
            filenames = []
            images = np.zeros(batch_shape, dtype=np.uint8)
            idx = 0

    if idx > 0:
        yield filenames, images

def main(_):
    tf.logging.set_verbosity(tf.logging.INFO)

    saver = tf.train.import_meta_graph(FLAGS.meta_graph, clear_devices=True)
    graph = tf.get_default_graph()
    input_tf = graph.get_tensor_by_name("input:0")
    output_tf = graph.get_tensor_by_name("towerp0/linear/output:0")
    prob_tf = tf.nn.softmax(output_tf)

    image_height = input_tf.shape[1].value
    image_width = input_tf.shape[2].value

    with tf.Session() as sess:
        saver.restore(sess, FLAGS.model)
        with open(FLAGS.output_file, "w") as fout:
            for filenames, images in load_images(FLAGS.input,
                                                 (image_height, image_width),
                                                 FLAGS.batch_size):
                prob = sess.run(prob_tf, feed_dict={input_tf: images})
                for idx, path in enumerate(filenames):
                    label = prob[idx, :].argmax()
                    confidence = prob[idx, label]
                    print(path, label, confidence, sep=",", file=fout)
                    fout.flush()

if __name__ == "__main__":
    tf.app.run()


import torch
import torch.nn.functional as F

import numpy as np
import pdb, os, argparse
from skimage import io
from tqdm import trange
# from model.CPD_models import CPD_VGG
from model.CPD_ResNet_models import CPD_ResNet
from model.Sal_CNN import Sal_CNN
from evaluateSOD.main import evalateSOD
from data import test_dataset

def eval_data(dataset_path, test_datasets, ckpt_name):

    parser = argparse.ArgumentParser()
    parser.add_argument('--testsize', type=int, default=352, help='testing size')
    parser.add_argument('--snapshot', type=str, default=ckpt_name, help='checkpoint name')
    cfg = parser.parse_args()

    model_rgb = CPD_ResNet()
    model_depth = CPD_ResNet()
    model = Sal_CNN()
    model_rgb.load_state_dict(torch.load('./ckpt/'+'DSU_rgb.pth' +cfg.snapshot))
    model_depth.load_state_dict(torch.load('./ckpt/' + 'DSU_depth.pth' +cfg.snapshot))
    model.load_state_dict(torch.load('./ckpt/' +'DSU.pth' + cfg.snapshot))

    cuda = torch.cuda.is_available()
    if cuda:
        model_rgb.cuda()
        model_depth.cuda()
        model.cuda()
    model_rgb.eval()
    model_depth.eval()
    model.eval()

    for dataset in test_datasets:
        save_path = './results/' + dataset + '/'
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        image_root = dataset_path + dataset + '/test_images/'
        gt_root = dataset_path + dataset + '/test_masks/'
        depth_root = dataset_path + dataset + '/test_depth/'
        test_loader = test_dataset(image_root, gt_root, depth_root, cfg.testsize)
        print('Evaluating dataset: %s' %(dataset))

        '''~~~ Our DSU FRAMEWORK~~~'''
        for i in trange(test_loader.size):
            image, gt, depth, name = test_loader.load_data()

            if cuda:
                image = image.cuda()
                depth = depth.cuda()

            # The inference stage involves only RGB stream, i.e., only RGB image is used for predicting.
            _, res_r, _ = model_rgb(image)

            res = res_r
            res = res.sigmoid().data.cpu().numpy().squeeze()

            res = (res - res.min()) / (res.max() - res.min() + 1e-8)
            io.imsave(save_path+name, np.uint8(res * 255))
        _ = evalateSOD(save_path, gt_root, dataset,ckpt_name,switch=False) # This eval code is for reference.
        # All reported results in this paper are evaluated by the same code
        # as in 'https://github.com/jiwei0921/Saliency-Evaluation-Toolbox'.
    return

if __name__ == '__main__':
    dataset_path = '../Dataset/test_data/'
    #test_datasets=['NLPR']
    test_datasets = ['DUT','NJUD', 'NLPR','STERE1000', 'SIP','LFSD', 'RGBD135','SSD']

    ckpt_name = '.24'
    eval_data(dataset_path,test_datasets,ckpt_name)


import click
import numpy as np

from pathlib import Path

from .. import deepcell
from ... import io

if deepcell.deepcell_available:
    import yaml
else:
    yaml = None

deepcell_cli_available = deepcell.deepcell_available

_applications = {
    "mesmer": deepcell.Application.MESMER,
}

@click.command(name="deepcell", help="Run a object segmentation batch using DeepCell")
@click.option(
    "--app",
    "application_name",
    type=click.Choice(list(_applications.keys()), case_sensitive=True),
    show_choices=True,
    default="mesmer",
    show_default=True,
    help="DeepCell application name",
)
@click.option(
    "--model",
    "model_path_or_name",
    type=click.STRING,
    default="MultiplexSegmentation",
    show_default=True,
    help="Path/name of custom Keras model",
)
@click.option(
    "--modeldir",
    "keras_model_dir",
    type=click.Path(file_okay=False),
    default="/opt/keras/models",
    show_default=True,
    help="Path to Keras model directory",
)
@click.option(
    "--img",
    "img_dir",
    type=click.Path(exists=True, file_okay=False),
    default="img",
    show_default=True,
    help="Path to the image directory",
)
@click.option(
    "--minmax/--no-minmax",
    "channelwise_minmax",
    default=False,
    show_default=True,
    help="Channel-wise min-max normalization",
)
@click.option(
    "--zscore/--no-zscore",
    "channelwise_zscore",
    default=False,
    show_default=True,
    help="Channel-wise z-score normalization",
)
@click.option(
    "--panel",
    "panel_file",
    type=click.Path(exists=True, dir_okay=False),
    default="panel.csv",
    show_default=True,
    help="Path to the panel file",
)
@click.option(
    "--aggr",
    "aggr_func_name",
    type=click.STRING,
    default="mean",
    show_default=True,
    help="Numpy function for aggregating channel pixels",
)
@click.option(
    "--pixelsize",
    "pixel_size_um",
    type=click.FLOAT,
    default=1.0,
    show_default=True,
    help="[Mesmer] Pixel size in micrometers",
)
@click.option(
    "--type",
    "segmentation_type",
    type=click.Choice(["whole-cell", "nuclear"]),
    default="whole-cell",
    show_default=True,
    show_choices=True,
    help="[Mesmer] Segmentation type",
)
@click.option(
    "--preprocess",
    "preprocess_file",
    type=click.Path(exists=True, dir_okay=False),
    help="[Mesmer] Preprocessing parameters (YAML file)",
)
@click.option(
    "--postprocess",
    "postprocess_file",
    type=click.Path(exists=True, dir_okay=False),
    help="[Mesmer] Postprocessing parameters (YAML file)",
)
@click.option(
    "-o",
    "mask_dir",
    type=click.Path(file_okay=False),
    default="masks",
    show_default=True,
    help="Path to the mask output directory",
)
def deepcell_cmd(
    application_name,
    model_path_or_name,
    keras_model_dir,
    img_dir,
    channelwise_minmax,
    channelwise_zscore,
    panel_file,
    aggr_func_name,
    pixel_size_um,
    segmentation_type,
    preprocess_file,
    postprocess_file,
    mask_dir,
):
    channel_groups = None
    if Path(panel_file).exists():
        panel = io.read_panel(panel_file)
        if "deepcell" in panel and panel["deepcell"].notna().any():
            channel_groups = panel["deepcell"].values
    aggr_func = getattr(np, aggr_func_name)
    img_files = io.list_image_files(img_dir)
    model = None
    if model_path_or_name is not None:
        from tensorflow.keras.models import load_model

        if Path(model_path_or_name).exists():
            model = load_model(model_path_or_name, compile=False)
        elif Path(keras_model_dir).joinpath(model_path_or_name).exists():
            model = load_model(
                Path(keras_model_dir).joinpath(model_path_or_name),
                compile=False,
            )
    preprocess_kwargs = None
    if preprocess_file is not None:
        with Path(preprocess_file).open() as f:
            preprocess_kwargs = yaml.load(f, yaml.Loader)
    postprocess_kwargs = None
    if postprocess_file is not None:
        with Path(postprocess_file).open() as f:
            postprocess_kwargs = yaml.load(f, yaml.Loader)
    Path(mask_dir).mkdir(exist_ok=True)
    for img_file, mask in deepcell.try_segment_objects(
        img_files,
        _applications[application_name],
        model=model,
        channelwise_minmax=channelwise_minmax,
        channelwise_zscore=channelwise_zscore,
        channel_groups=channel_groups,
        aggr_func=aggr_func,
        pixel_size_um=pixel_size_um,
        segmentation_type=segmentation_type,
        preprocess_kwargs=preprocess_kwargs,
        postprocess_kwargs=postprocess_kwargs,
    ):
        mask_file = io._as_path_with_suffix(Path(mask_dir) / img_file.name, ".tiff")
        io.write_mask(mask, mask_file)
        click.echo(mask_file)


"""
Generate sequences based on trained models. Ensure you first have 'save/'
in your current working directory, containing saved model files. If not,
first run

$ python -m bayz.train

then

$ python -m bayz.generate

to generate a short sample of music. In order to display the music, ensure
you have a working musicxml renderer configured with music21. For instructions
on how, see: https://web.mit.edu/music21/doc/usersGuide/usersGuide_08_installingMusicXML.html

author: William Tong (wlt2115@columbia.edu)
"""

import pickle
from pathlib import Path

from music21 import bar
from music21 import chord
from music21 import note
from music21 import stream
from music21 import tie

import numpy as np
import scipy

from .common import BayesianGaussianTypeModel

REST_WORD = 'REST'
START_WORD = 'START'
END_WORD = 'END'

norm_prob_cache = {}

def load_model(path: Path) -> 'Any':
    """
    Load a model file from the given path

    path: a pathlib.Path object
    return: the loaded model
    """

    with path.open('rb') as fp:
        return pickle.load(fp)

def to_token(types: list, mixture: BayesianGaussianTypeModel, embedding: 'Word2Vec') -> list:
    """
    Samples a set of tokens from the provided mixture model

    param types: a list of types (mixture ids) corresponding to the music to be
                 generated
    param mixture: a mixture model fitted to a training corpus
    param embedding: a Word2Vec model trained on the corpus

    return: a list of abstract note names, ready to be sampled into real music
    """

    words = []
    for symbol in types:
        new_vec = mixture.emit(symbol, words) # instantiate and use previous words
        words.append(_decode(embedding, new_vec))
    token_seq = [word.split('_') for word in words]

    return token_seq

def _decode(embedding, vector):
    return embedding.similar_by_vector(vector, topn=1)[0][0]

def to_score(token_seq: list, texture: 'function', **texture_args) -> 'Score':
    """
    Samples a music21.Score object from a list of tokens

    param token_seq: list of tokens
    param texture: a function that determines the texture with which to give
                   the sampled music
    param **texture_args: arguments passed to the texture function

    return: a generated music21.Score object
    """

    score = stream.Stream()
    for note in texture(token_seq, **texture_args):
        score.append(note)

    return score

def chord_texture(token_seq, duration=1):
    """
    Defines a texture in which tokens are rendered as chords

    param token_seq: tokens to be rendered
    param duration: the duration of each chord
    """

    for token in token_seq:
        if token[0] == REST_WORD:
            yield note.Rest()
        elif token[0] in (START_WORD, END_WORD):
            yield bar.Barline('double')
        else:
            print('processing token', token)
            yield chord.Chord(set(token), quarterLength=duration)

def chord_with_ties_texture(token_seq, duration=1):
    """
    Defines a texture in which tokens are rendered as chords. Adjacent chords
    with shared notes have ties that connect the notes, creating a continuous
    sound.

    param token_seq: tokens to be rendered
    param duration: the duration of each chord
    """

    chords = chord_texture(token_seq, duration)
    chords_with_ties = []

    prev = next(chords)
    for stack in chords:
        if type(stack) is chord.Chord:
            for elem in prev:
                if elem in stack:
                    elem.tie = tie.Tie()
            chords_with_ties.append(prev)
            prev = stack

    chords_with_ties.append(prev)
    return chords_with_ties

def melody_texture(token_seq, duration=0.25, use_last=False):
    """
    Defines a texture in which tokens are rendered as individual notes

    param token_seq: tokens to be rendered
    param duration: the duration of each chord
    param use_last: whether to sample notes from the last note of each word, or
                    to use all notes in a word
    """

    for token in token_seq:
        if token[0] == REST_WORD:
            yield note.Rest()
        elif token[0] in (START_WORD, END_WORD):
            yield bar.Barline('double')
        else:
            if use_last:
                yield(note.Note(token[-1], quarterLength=duration))
            else:
                for elem in token:
                    yield(note.Note(elem, quarterLength=duration))

def to_midi(token_seq: list, use_last=False) -> list:
    """
    Renders a set of tokens directly into midi values, rather than a Score. It
    uses the same scheme as the melody_texture.

    param token_seq: tokens to be rendered
    param use_last: whether to sample the last note in a word, or to use the
                    entire word

    return: list of midi note values
    """

    for token in token_seq:
        if token[0] != REST_WORD or token[0] not in (START_WORD, END_WORD):
            candidate_notes = []
            if use_last:
                candidate_notes.append(note.Note(token[-1]))
            else:
                for elem in token:
                    candidate_notes.append(note.Note(elem))

            for n in candidate_notes:
                yield n.pitch.midi

if __name__ == '__main__':
    embedding = load_model(Path('save/embedding.wv'))
    mixture = BayesianGaussianTypeModel(embedding)
    mixture.load_model(Path('save/mixture.pk'))
    hmm = load_model(Path('save/hmm.pk'))

    types, _ = hmm.sample(5)
    types = types.flatten()
    print('Sampled', types)

    tokens = to_token(types, mixture, embedding)
    score = to_score(tokens, melody_texture)
    score.show('musicxml')


"""
Call variant based on a list of SAMMPileUpRecord where list[i] is the record of i-th position

Most of the code follows Juliet's code at
https://github.com/PacificBiosciences/minorseq/blob/develop/src/AminoAcidCaller.cpp

"""
import os, sys
import scipy.stats as stats
from collections import Counter

class MPileUPVariant(object):
    def __init__(self, record_list, min_cov, err_sub, expected_strand, pval_cutoff=0.01):
        """
        :param record_list: list of SAMMPileUpRecord
        :param min_cov: minimum coverage to call variant
        :param err_sub: substitution error, right now a fixed float
        :param expected_strand: expected strand of the transcript (+ or -)
        """
        self.record_by_pos = dict((r.pos, r) for r in record_list)
        self.min_cov = min_cov
        self.err_sub = err_sub

        self.pval_cutoff = pval_cutoff
        self.expected_strand = expected_strand

        self.prep_records()
        self.positions_to_call = self.get_positions_to_call()

        # must first call positions to call, then prep records, then number of tests
        self.number_of_tests = sum(self.record_by_pos[pos].clean_type for pos in self.positions_to_call)

        self.variant = {} # position --> in sorted order, (base, count)
        self.ref_base = {} # position --> ref base

        self.call_variant()

    def is_in_or_near_HP(self, pos, hp_size=4):
        """
        We define a HP region as stretches of 4 or more same nucleotides
        :return: True/False for in/hear HP region
        """
        def find_hp_region_size(cur):
            if cur not in self.record_by_pos: return 0
            end = cur+1
            while end in self.record_by_pos and self.record_by_pos[end].ref == self.record_by_pos[cur].ref:
                end += 1
            start = cur-1
            while start in self.record_by_pos and self.record_by_pos[start].ref == self.record_by_pos[cur].ref:
                start -= 1
            # hp region is from start+1 to end
            return end-(start+1)

        return (find_hp_region_size(pos) >= hp_size) or \
               (find_hp_region_size(pos-1) >= hp_size) or \
               (find_hp_region_size(pos+1) >= hp_size)

    def get_positions_to_call(self):
        """
        Identify list of positions to try to call SNPs. Must have:
        1. minimum coverage >= min_cov
        2. the first and second most frequent base are NOT an indel
        3. not next to or inside a homopolymer region
        4. has at least two or more keys
        """
        positions_to_call = []
        for pos in self.record_by_pos:
            if self.record_by_pos[pos].clean_type < 2: continue # only one base at this position, skip
            elif self.record_by_pos[pos].clean_cov < self.min_cov: continue # insufficient cov, skip
            else:
                # find the first and second most freq base in the "non-clean" counts
                m = self.record_by_pos[pos].clean_counts.most_common()
                # ex: m = [('a', 10), ('-ct', 20), ('+t', 10)]
                if m[0][0][0]in ('+','-') or m[1][0][0] in ('+','-') or self.is_in_or_near_HP(pos): continue
                else:
                    positions_to_call.append(pos)
        return positions_to_call

    def prep_records(self):
        """
        Prepare the records by:
        1. remove all 'N' bases
        2. remove all bases that were not on the expected strand
        3. remove all indels

        Creates three new vars: clean_counts, clean_cov, clean_type
        DOES NOT ALTER the original counts or other variables!!!

        If + strand, then ATCG
        If - strand, then atcg
        """
        for pos in self.record_by_pos:
            r = self.record_by_pos[pos]
            if self.expected_strand == '+':
                bases = 'ATCG'
            else:
                bases = 'atcg'

            r.clean_counts = Counter(r.counts)
            keys = r.counts.keys()
            for k in keys:
                if k not in bases:
                    del r.clean_counts[k]
            r.clean_cov = sum(r.clean_counts.itervalues())
            r.clean_type = len(r.clean_counts)

    def call_variant(self):
        """
        mirrors AminoAcidCaller::CallVariants() in
        https://github.com/PacificBiosciences/minorseq/blob/develop/src/AminoAcidCaller.cpp

        For each position (that has sufficient coverage),
         do Fisher exact test w/ correction
         if p-val < threshold, then store it.

        Stores results in self.variant as:

        self.variant[position] = desc list of (base, count).
        NOTE: base must be either all in lower case (which means - strand)
              or call upper case (+ strand).
              If - strand and ('a', 10), it means the ref base in A on the + strand,
              and the transcript should be T on the - strand.

        Only positions with more than the ref base is stored.
        """
        for pos in self.positions_to_call:
            r = self.record_by_pos[pos]
            alt_variant = []
            for base, count in r.clean_counts.most_common()[1:]:
                assert not base.startswith('+') and not base.startswith('-') # clean counts should NOT have indels
                exp = r.clean_cov * self.err_sub
                odds, pval = stats.fisher_exact([[count, r.clean_cov-count], [exp, r.clean_cov-exp]], alternative='greater')
                pval *= self.number_of_tests
                if pval < self.pval_cutoff: # store variant if below cutoff
                    alt_variant.append((base, count))
            if len(alt_variant) > 0: # only record this variant if there's at least two haps
                self.variant[pos] = [r.clean_counts.most_common()[0]] + alt_variant
                self.ref_base[pos] = r.ref


import numpy
import pyautogui

from common import (
    UserError,
    ROOM_HEIGHT_IN_TILES,
    ROOM_WIDTH_IN_TILES,
    TILE_SIZE,
)
from .consts import ROOM_ORIGIN_X, ROOM_ORIGIN_Y
from util import find_color

_ROOM_UPPER_EDGE_COLOR = (32, 60, 74)
_ROOM_UPPER_EDGE_LENGTH = 838
_ROOM_UPPER_EDGE_START_X = 162
_ROOM_UPPER_EDGE_START_Y = 39
_DROD_WINDOW_WIDTH = 1024
_DROD_WINDOW_HEIGHT = 768

_OVERLAY_COLOR = (0, 255, 0)
_OVERLAY_WIDTH = 5

_MINIMAP_ROOM_ORIGIN_X = 61
_MINIMAP_ROOM_ORIGIN_Y = 631

async def get_drod_window(return_debug_images=False):
    """Take a screenshot and find the DROD window.

    Only the inside of the window is included, not the window
    bar which is OS-dependent.

    Parameters
    ----------
    return_debug_images
        If this is True, return an additional list of (name, debug_image).

    Returns
    -------
    origin_x : int
        The X coordinate of upper left corner of the window.
    origin_y : int
        The Y coordinate of upper left corner of the window.
    image : numpy.ndarray
        The DROD window.
    debug_images
        Debug images. Only returned if `return_debug_images` is True.
    """
    if return_debug_images:
        debug_images = []
    raw_image = numpy.array(pyautogui.screenshot())
    if return_debug_images:
        debug_images.append(("Screenshot", raw_image))

    # Try finding the upper edge of the room, which is a long line of constant color
    correct_color = find_color(raw_image, _ROOM_UPPER_EDGE_COLOR)
    if return_debug_images:
        debug_images.append(("Find upper edge color", correct_color))

    lines = _find_horizontal_lines(correct_color, _ROOM_UPPER_EDGE_LENGTH)
    if return_debug_images:
        # We can't show the line coordinates directly, so we'll overlay the line on
        # the screenshot
        with_line = raw_image.copy()
        for (start_x, start_y, end_x, _) in lines:
            # Since we're only dealing with horizontal lines, we can do the overlay
            # by indexing the array directly
            with_line[
                start_y : start_y + _OVERLAY_WIDTH, start_x:end_x, :
            ] = _OVERLAY_COLOR
        debug_images.append(("Find upper edge line", with_line))

    if len(lines) > 1:
        raise UserError("Cannot identify DROD window, too many candidate lines")
    elif len(lines) == 0:
        raise UserError("Cannot identify DROD window, is it open and unblocked?")

    # Extract the window
    line_start_x = lines[0][0]
    line_start_y = lines[0][1]
    window_start_x = line_start_x - _ROOM_UPPER_EDGE_START_X
    window_start_y = line_start_y - _ROOM_UPPER_EDGE_START_Y
    window_end_x = window_start_x + _DROD_WINDOW_WIDTH
    window_end_y = window_start_y + _DROD_WINDOW_HEIGHT
    drod_window = raw_image[
        window_start_y:window_end_y,
        window_start_x:window_end_x,
        :,
    ]
    if return_debug_images:
        return window_start_x, window_start_y, drod_window, debug_images
    return window_start_x, window_start_y, drod_window

def extract_room(window_image):
    """Extract the room from a window image.

    Parameters
    ----------
    window_image
        The window to extract the room from.

    Returns
    -------
    An image of the room.
    """
    room_end_x = ROOM_ORIGIN_X + ROOM_WIDTH_IN_TILES * TILE_SIZE
    room_end_y = ROOM_ORIGIN_Y + ROOM_HEIGHT_IN_TILES * TILE_SIZE
    room_image = window_image[ROOM_ORIGIN_Y:room_end_y, ROOM_ORIGIN_X:room_end_x, :]
    return room_image

def extract_minimap(window_image):
    """Extract the room minimap from a window image.

    Parameters
    ----------
    window_image
        The window to extract the room minimap from.

    Returns
    -------
    An image of the room minimap.
    """
    room_end_x = _MINIMAP_ROOM_ORIGIN_X + ROOM_WIDTH_IN_TILES
    room_end_y = _MINIMAP_ROOM_ORIGIN_Y + ROOM_HEIGHT_IN_TILES
    minimap_room_image = window_image[
        _MINIMAP_ROOM_ORIGIN_Y:room_end_y, _MINIMAP_ROOM_ORIGIN_X:room_end_x, :
    ]
    return minimap_room_image

def reconstruct_from_tiles(tiles):
    """Reconstruct a room image given the tiles.

    Parameters
    ----------
    tiles
        A dict mapping (x, y) coordinates to tile images. The images
        can be in color or grayscale.

    Returns
    -------
    A room image.
    """
    tile = list(tiles.values())[0]
    if len(tile.shape) == 3:
        room_image = numpy.zeros(
            (
                ROOM_HEIGHT_IN_TILES * TILE_SIZE,
                ROOM_WIDTH_IN_TILES * TILE_SIZE,
                tile.shape[2],
            ),
            dtype=numpy.uint8,
        )
    else:
        room_image = numpy.zeros(
            (ROOM_HEIGHT_IN_TILES * TILE_SIZE, ROOM_WIDTH_IN_TILES * TILE_SIZE),
            dtype=numpy.uint8,
        )
    for x in range(ROOM_WIDTH_IN_TILES):
        for y in range(ROOM_HEIGHT_IN_TILES):
            start_x = x * TILE_SIZE
            end_x = (x + 1) * TILE_SIZE
            start_y = y * TILE_SIZE
            end_y = (y + 1) * TILE_SIZE
            room_image[start_y:end_y, start_x:end_x] = tiles[(x, y)]
    return room_image

def _find_horizontal_lines(boolean_array, length):
    # Pad, to find lines that start or end at the edges
    int_array = numpy.pad(boolean_array.astype(int), (0, 1))
    # Coords of the first pixel in each line
    starts = numpy.argwhere(numpy.diff(int_array) == 1) + (0, 1)
    # Coords of the pixel after the last in each line
    ends = numpy.argwhere(numpy.diff(int_array) == -1) + (0, 1)
    lines = []
    for coords in starts:
        start_x = coords[1]
        start_y = coords[0]
        end_x = [e[1] for e in ends if e[0] == start_y]
        if start_x + length in end_x and not any(
            [x > start_x and x < start_x + length for x in end_x]
        ):
            # We have an uninterrupted line of the specified length
            lines.append((start_x, start_y, start_x + length, start_y))
    return lines