import os

from djitellopy import Tello
import cv2
import numpy as np
import argparse
from detect_person import detectPerson

# TODO: Add better distance values
sizes = [800, 700, 600, 500, 406]

class Drone:
    def __init__(self, speed, forward_backward_speed, steering_speed, up_down_speed, save_session, save_path, distance,
                 safety_x, safety_y, safety_z):
        # Initialize Drone
        self.tello = Tello()

        # Initialize Person Detector
        self.detector = detectPerson()

        self.speed = speed
        self.fb_speed = forward_backward_speed
        self.steering_speed = steering_speed
        self.ud_speed = up_down_speed
        self.save_session = save_session
        self.save_path = save_path
        self.distance = distance
        self.safety_x = safety_x
        self.safety_y = safety_y
        self.safety_z = safety_z

        self.for_back_velocity = 0
        self.left_right_velocity = 0
        self.up_down_velocity = 0
        self.yaw_velocity = 0
        self.flight_mode = False
        self.send_rc_control = False
        self.dimensions = (960, 720)

        if self.save_session:
            os.makedirs(self.save_path, exist_ok=True)

    def run(self):
        if not self.tello.connect():
            print("Tello not connected")
            return
        elif not self.tello.set_speed(self.speed):
            print('Not set speed to lowest possible')
            return
        elif not self.tello.streamoff():
            print("Could not stop video stream")
            return
        elif not self.tello.streamon():
            print("Could not start video stream")
            return

        self.tello.get_battery()

        frame_read = self.tello.get_frame_read()

        count = 0
        while True:
            if frame_read.stopped:
                frame_read.stop()
                break

            # Listen for key presses
            k = cv2.waitKey(20)

            if k == 8:
                self.flight_mode = not self.flight_mode
            elif k == 27:
                break
            elif k == ord('t'):
                self.tello.takeoff()
                self.send_rc_control = True
            elif k == ord('l'):
                self.tello.land()
                self.send_rc_control = False

            # read frame
            frameBGR = frame_read.frame

            if self.save_path:
                cv2.imwrite(f'{self.save_path}/{count}.jpg', frameBGR)
                count += 1

            x_min, y_min, x_max, y_max = self.detector.detectSingle(frameBGR)
            cv2.rectangle(frameBGR, (x_min, y_min), (x_max, y_max), (255, 0, 0), 2)

            if not self.flight_mode and self.send_rc_control and x_max != 0 and y_max != 0:
                # these are our target coordinates
                targ_cord_x = int((x_min + x_max) / 2)
                targ_cord_y = int((y_min + y_max) / 2)

                # This calculates the vector from your face to the center of the screen
                vTrue = np.array((int(self.dimensions[0]/2), int(self.dimensions[1]/2), sizes[self.distance]))
                vTarget = np.array((targ_cord_x, targ_cord_y, (x_max-x_min)*2))
                vDistance = vTrue - vTarget

                # turn drone
                if vDistance[0] < -self.safety_x:
                    self.yaw_velocity = self.steering_speed
                elif vDistance[0] > self.safety_x:
                    self.yaw_velocity = -self.steering_speed
                else:
                    self.yaw_velocity = 0

                # for up & down
                if vDistance[1] > self.safety_y:
                    self.up_down_velocity = self.ud_speed
                elif vDistance[1] < -self.safety_y:
                    self.up_down_velocity = -self.ud_speed
                else:
                    self.up_down_velocity = 0

                # forward & backward
                if vDistance[2] > self.safety_z:
                    self.for_back_velocity = self.fb_speed
                elif vDistance[2] < self.safety_z:
                    self.for_back_velocity = -self.fb_speed
                else:
                    self.for_back_velocity = 0

                # always set left_right_velocity to 0
                self.left_right_velocity = 0

                # Draw the target as a circle
                cv2.circle(frameBGR, (targ_cord_x, targ_cord_y), 10, (0, 255, 0), 2)

                # Draw the safety zone
                cv2.rectangle(frameBGR, (targ_cord_x - self.safety_x, targ_cord_y - self.safety_y), (targ_cord_x + self.safety_x, targ_cord_y + self.safety_y),
                              (0, 255, 0), 2)
            elif not self.flight_mode and self.send_rc_control and x_max==0 and y_max==0:
                self.for_back_velocity = 0
                self.left_right_velocity = 0
                self.yaw_velocity = 0
                self.up_down_velocity = 0
            elif self.flight_mode and self.send_rc_control:
                # fligh forward and back
                if k == ord('w'):
                    self.for_back_velocity = self.speed
                elif k == ord('s'):
                    self.for_back_velocity = -self.speed
                else:
                    self.for_back_velocity = 0

                # fligh left & right
                if k == ord('d'):
                    self.left_right_velocity = self.speed
                elif k == ord('a'):
                    self.left_right_velocity = -self.speed
                else:
                    self.left_right_velocity = 0

                # fligh up & down
                if k == ord('u'):
                    self.up_down_velocity = self.speed
                elif k == ord('j'):
                    self.up_down_velocity = -self.speed
                else:
                    self.up_down_velocity = 0

                # turn
                if k == ord('h'):
                    self.yaw_velocity = self.speed
                elif k == ord('k'):
                    self.yaw_velocity = -self.speed
                else:
                    self.yaw_velocity = 0

            if self.send_rc_control:
                # Send velocities to Drone
                self.tello.send_rc_control(self.left_right_velocity, self.for_back_velocity, self.up_down_velocity,
                                           self.yaw_velocity)

            cv2.imshow('Tello Drone', frameBGR)
        # Destroy cv2 windows and end drone connection
        cv2.destroyAllWindows()
        self.tello.end()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-os', '--overwrite_speed', type=int, default=30, help='Speed for manual use')
    parser.add_argument('-fbs', '--forward_backward_speed', type=int, default=20,
                        help='Speed of forward and backward movement drone')
    parser.add_argument('-ss', '--steering_speed', type=int, default=40,
                        help='Steering speed')
    parser.add_argument('-uds', '--up_down_speed', type=int, default=20,
                        help='Speed of up and down movement drone')
    parser.add_argument('-sa', '--save_session', action='store_true', help='Record flight')
    parser.add_argument('-sp', '--save_path', type=str, default="session/", help="Path where images will get saved")
    parser.add_argument('-d', '--distance', type=int, default=0,
                        help='use -d to change the distance of the drone. Range 0-2')
    parser.add_argument('-sx', '--safety_x', type=int, default=100,
                        help='use -sx to change the safety bound on the x axis . Range 0-480')
    parser.add_argument('-sy', '--safety_y', type=int, default=55,
                        help='use -sy to change the safety bound on the y axis . Range 0-360')
    parser.add_argument('-sz', '--safety_z', type=int, default=30,
                        help='use -sz to change the safety bound on the z axis.')

    args = parser.parse_args()

    drone = Drone(args.overwrite_speed, args.forward_backward_speed, args.steering_speed, args.up_down_speed,
                  args.save_session, args.save_path, args.distance, args.safety_x, args.safety_y, args.safety_z,)
    drone.run()


# --------------------------------------------------------
# Visual Detection: State-of-the-Art
# Copyright: Hanbo Zhang
# Licensed under The MIT License [see LICENSE for details]
# Written by Hanbo Zhang
# --------------------------------------------------------

import torch
import torch.nn as nn
from model.utils.config import cfg
from model.SSD import SSD
from model.Detectors import VMRN

import numpy as np

from model.rpn.bbox_transform import bbox_transform_inv, clip_boxes

class SSD_VMRN(SSD, VMRN):
    """Single Shot Multibox Architecture
    The network is composed of a base VGG network followed by the
    added multibox conv layers.  Each multibox layer branches into
        1) conv2d for class conf scores
        2) conv2d for localization predictions
        3) associated priorbox layer to produce default bounding
           boxes specific to the layer's feature map size.
    See: https://arxiv.org/pdf/1512.02325.pdf for more details.
    Args:
        phase: (string) Can be "test" or "train"
        size: input image size
        base: VGG16 layers for input, size of either 300 or 500
        extras: extra layers that feed to multibox loc and conf layers
        head: "multibox head" consists of loc and conf conv layers
    """

    def __init__(self, classes, class_agnostic, feat_name, feat_list=('conv3', 'conv4'), pretrained = True):
        super(SSD_VMRN, self).__init__(classes, class_agnostic, feat_name, feat_list, pretrained)

    def forward(self, data_batch):
        """Applies network layers and ops on input image(s) x.
        Args:
            x: input image or batch of images. Shape: [batch,3,300,300].
        Return:
            Depending on phase:
            test:
                Variable(tensor) of output class label predictions,
                confidence score, and corresponding location predictions for
                each object detected. Shape: [batch,topk,7]
            train:
                list of concat outputs from:
                    1: confidence layers, Shape: [batch*num_priors,num_classes]
                    2: localization layers, Shape: [batch,num_priors*4]
                    3: priorbox layers, Shape: [2,num_priors*4]
        """
        x = data_batch[0]
        im_info = data_batch[1]
        gt_boxes = data_batch[2]
        num_boxes = data_batch[3]
        rel_mat = data_batch[4]

        if self.training:
            self.iter_counter += 1

        sources = []

        base_feat, x = self.FeatExt(x)
        base_feat = self.L2Norm(base_feat)
        sources.append(base_feat)

        for m in self.extra_conv:
            x = m(x)
            sources.append(x)

        loc, conf = self._get_obj_det_result(sources)
        SSD_loss_cls, SSD_loss_bbox = 0, 0
        if self.training:
            predictions = (
                loc,
                conf,
                self.priors.type_as(loc)
            )
            SSD_loss_bbox, SSD_loss_cls = self.criterion(predictions, gt_boxes, num_boxes)
        conf = self.softmax(conf)

        # generate object RoIs.
        obj_rois, obj_num = torch.Tensor([]).type_as(loc), torch.Tensor([]).type_as(num_boxes)
        # online data
        if not self.training or (cfg.TRAIN.VMRN.TRAINING_DATA == 'all' or 'online'):
            obj_rois, obj_num = self._object_detection(self.priors.type_as(loc), conf, loc, self.batch_size,
                                                       im_info.data)
        # offline data
        if self.training and (cfg.TRAIN.VMRN.TRAINING_DATA == 'all' or 'offline'):
            for i in range(self.batch_size):
                img_ind = (i * torch.ones(num_boxes[i].item(), 1)).type_as(gt_boxes)
                obj_rois = torch.cat([obj_rois, torch.cat([img_ind, (gt_boxes[i][:num_boxes[i]])], 1)])
            obj_num = torch.cat([obj_num, num_boxes])

        obj_labels = torch.Tensor([]).type_as(gt_boxes).long()
        if obj_rois.size(0) > 0:
            obj_labels = obj_rois[:, 5]
            obj_rois = obj_rois[:, :5]

        VMRN_rel_loss_cls, reg_loss = 0, 0
        if (obj_num > 1).sum().item() > 0:
            rel_cls_score, rel_cls_prob, reg_loss = self._get_rel_det_result(base_feat, obj_rois, obj_num, im_info)
            if self.training:
                obj_pair_rel_label = self._generate_rel_labels(obj_rois, gt_boxes, obj_num, rel_mat,
                                                               rel_cls_prob.size(0))
                VMRN_rel_loss_cls = self._rel_det_loss_comp(obj_pair_rel_label.type_as(gt_boxes).long(), rel_cls_score)
            else:
                rel_cls_prob = self._rel_cls_prob_post_process(rel_cls_prob)
        else:
            rel_cls_prob = torch.Tensor([]).type_as(conf)

        rel_result = None
        if not self.training:
            if obj_rois.numel() > 0:
                pred_boxes = obj_rois.data[:, 1:5]
                pred_boxes[:, 0::2] /= im_info[0][3].item()
                pred_boxes[:, 1::2] /= im_info[0][2].item()
                rel_result = (pred_boxes, obj_labels, rel_cls_prob.data)
            else:
                rel_result = (obj_rois.data, obj_labels, rel_cls_prob.data)

        return loc, conf, rel_result, SSD_loss_bbox, SSD_loss_cls, VMRN_rel_loss_cls, reg_loss

    def _init_modules(self):
        assert cfg.TRAIN.VMRN.USE_REL_CLS_GRADIENTS, "No gradients are applied to relationship convolutional layers."
        SSD._init_modules(self)
        VMRN._init_modules(self)

    def _init_weights(self):
        SSD._init_weights(self)
        VMRN._init_weights(self)

    def create_architecture(self, object_detector_path=''):
        # TODO: support pretrained object detector here
        pass


from db_control import *
import os
from model import *
import numpy as np
import time

import logging
logging.basicConfig(filename='predict.log', level=logging.DEBUG)

logging.debug("predict log")

TICK ="3"
TRAIND_MODEL_PATH = "./models/"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
IMG_SIZE = DATA_LENGTH = 128
CHANNELS_IMG = 5
NUM_CLASSES = 3
FEATURES_CRITIC = 16

# codes = get_company_code_for_trade()

files = os.listdir("./models/")
codes=list()
for f in files:
    codes.append(f.replace(".pt",""))

init_data = dict()
for code in codes:
    init_data[code] = Code_obj()
    try:
        company, ret, data = get_last_data(code, int(TICK), DATA_LENGTH)
        init_data[code].last_data = ret
        init_data[code].data = data

    except Exception as e:
        print(e)
        print(code)
        break
for code in codes:
    critic = Critic_128(CHANNELS_IMG, FEATURES_CRITIC, NUM_CLASSES, IMG_SIZE).to(device)
    file_path = TRAIND_MODEL_PATH+code+".pt"
    if os.path.isfile(file_path):
        critic.load_state_dict(torch.load(file_path, map_location=device))
        init_data[code].model=critic

check_list = list()
# 모델 있는 것만 체크
for code in codes:
    obj = init_data[code]
    if obj.model is None:
        continue
    else:
        check_list.append(code)

while True:
    for code in check_list:
        company, ret, data = get_last_data(code, int(TICK), DATA_LENGTH)
        obj = init_data[code]
        if not obj.check(ret):
            # if code == "000070":
            #     print(obj.last_data)
            #     print(ret)
            #     print(obj.last_data==ret)
            res = obj.update(ret)
            action, rate = obj.forward()
            print("company",company, "action",action.item(), "rate",rate.item())
            if action.item() != 0:
                c_time = time.strftime('%X', time.localtime(time.time()))
                order_time = c_time.replace(":","")
                company,order_time, trade_price, order_action =code,order_time,ret[0],action.item()
                insert_action(company, order_time, trade_price, order_action, rate)


import sys
import math
import random
from collections import namedtuple
import time
from pyrf.util import (compute_usable_bins, adjust_usable_fstart_fstop,
    trim_to_usable_fstart_fstop, find_saturation)

import numpy as np
from twisted.internet import defer

from pyrf.numpy_util import compute_fft
import struct
MAXIMUM_SPP = 32768

class correction_vector_acquire(object):
    data_buffer = ""
    v_type = "SIGNAL"
    dut = None
    complete_buffer = False
    d = None
    offset = 0
    size = 0
    transfer_size = 16*1024

    def get_vector_loop(self, data):
        self.data_buffer = b"".join([self.data_buffer, data])
        self.offset += len(data)
        if self.offset >= self.size:
            # we have gotten all out data, return this object
            if self.d is not None:
                self.d.callback(self)
        else:
            # more data, grab another set of data
            data1 = self.dut.correction_data(self.v_type, self.offset,
                                             self.transfer_size)
            # and add this function to the call back
            data1.addCallback(self.get_vector_loop)

    def get_vector_data(self, size):
        # We got out size
        if size is None:
            # size is return None threw our created deffered in get_vector
            if self.d is not None:
                self.d.callback(None)
            return

        self.size = int(size)
        if self.size == 0:
            if self.d is not None:
                self.d.callback(None)
            return

        if self.size < self.transfer_size:
            self.transfer_size = self.size
        # Grab our first set of data (deffered)
        data = self.dut.correction_data(self.v_type, self.offset,
                                        self.transfer_size)
        # add the self.get_vector_loop call back
        data.addCallback(self.get_vector_loop)
        # what happens to error back here?

    def error_b(self, failure):
        if self.d is not None:
            self.d.callback(None)
        return None

    def get_vector(self, v_type=None):
        #
        self.v_type = v_type
        self.offset = 0
        self.data_buffer = ""
        # Create a defered
        d = defer.Deferred()
        self.d = d

        # get our size (deffered)
        size = self.dut.correction_size(self.v_type)
        size.addCallback(self.get_vector_data)
        size.addErrback(self.error_b)
        # return our deferred
        return d

class correction_vector(object):
    correction_vectors = None
    frequency_index = None
    digest = None

    def __init__(self):
        self.frequency_index = []
        self.dy = np.dtype(np.int32)
        self.dy = self.dy.newbyteorder('>')
        self.correction_vectors = {}

    def _binary_search(self, freq):
        # Simple binary search, modified to work the object's datastructure
        lo = 0
        hi = len(self.frequency_index)
        while lo < hi:
            mid = (lo + hi) // 2
            if self.frequency_index[mid][0] * 1e3 < freq:
                lo = mid + 1
            else:
                hi = mid
        return lo

    def _interp(self, in_array, number_of_points):
        # array index of our orignal from 0 to size of vector - 1
        x = np.arange(0.0, self.vector_size, 1.0)
        # our new index
        z = np.linspace(0.0, self.vector_size - 1, number_of_points)
        # interpolate to get our new vector array
        out_array = np.interp(z, x, in_array)
        return out_array

    def get_correction_vector(self, freq, number_of_points):
        # binary search, retunrs our index
        index = self._binary_search(freq)
        # get the case where we go off the end
        if index == len(self.frequency_index):
            index = index - 1

        # get our vector
        vector = self.correction_vectors[self.frequency_index[index][1]]
        # convert from micro db to db
        vector = vector / 1000000.0
        # interpolate our vector to the wanted size
        resampled_vector = self._interp(vector, number_of_points)
        return resampled_vector

    def buffer_to_vector(self, buffer_in):
        if buffer_in is None:
            raise ValueError

        if len(buffer_in) < 8 + 40:
            raise ValueError

        # Get the first 8 bytes
        offset = 0
        size = 8
        input_buffer = buffer_in[offset:offset + size]
        version, freq_num, vector_num, self.vector_size = struct.unpack("!HHHH", input_buffer)
        offset = size

        # Ignore the next 40 bytes, as not used know
        offset += 40

        # grab our frequency list
        size = 6 * freq_num
        input_buffer = buffer_in[offset:offset + size]
        offset += size

        if len(input_buffer) < size:
            raise ValueError
        # loop over our buffer, adding a frequency pair to the array
        for i in range(freq_num):
            freq, index = struct.unpack("!LH", input_buffer[i*6:i*6+6])
            self.frequency_index.append([freq, index])

        # grab our correction vectors
        for i in range(vector_num):

            # Grab out index
            size = 2
            input_buffer = buffer_in[offset:offset + size]
            index = struct.unpack(">H", input_buffer)[0]
            offset += size

            # get our correction vector
            size = 4 * self.vector_size
            input_buffer = buffer_in[offset:offset + size]
            micro_db = np.frombuffer(input_buffer, dtype=self.dy,
                                     count=self.vector_size)
            self.correction_vectors[index] = micro_db
            offset += size

class SweepDeviceError(Exception):
    """
    Exception for the sweep device to state an error() has occured
    """
    pass

class SweepSettings(object):
    """
    An object used to keep track of the sweep settings
    """

    def __init__(self):
        # start frequency of the results we will eventually return
        self.bandstart = 0.0

        # stop frequency of the results we will eventually return
        self.bandstop = 0.0

        # sweep entry's start frequency
        self.fstart = 0.0

        # sweep entry's stop frequency
        self.fstop = 0.0

        # sweep entry frequency step
        self.fstep = 0.0

        # sweep entry's RFE mode
        self.rfe_mode = None

        # determine if a second entry is required
        self.dd_mode = False

        # determines if a non dd entry is needed
        self.beyond_dd = True

        # entry attenuation
        self.attenuation = 0

        # entry ppb
        self.ppb = 1

        # sweep entry's spp
        self.spp = 0.0

        # sweep capture iterations
        self.iterations = 0

        # expected spectral points
        self.spectral_points = 0

        # determines if a sweep entry is required at the end
        self.make_end_entry = False

        # determine the frequency of the end entry
        self.end_entry_freq = 0.0

        # how many steps are in this sweep
        self.step_count = 0

        # what's the actual RBW of what we're capturing
        self.rbw = 0

    def __str__(self):
        return "SweepSettings[ bandstart = %d, bandstop = %d, fstart = %d, fstop = %d, fstep = %d, step_count = %d, rfe_mode = %s, dd_mode = %s, beyond_dd = %s, attenuation = %s, ppb = %d, spp = %d, iterations = %d, spectral_points = %d, make_end_entry = %s, end_entry_freq = %d, rbw = %f ]" % (self.bandstart, self.bandstop, self.fstart, self.fstop, self.fstep, self.step_count, self.rfe_mode, self.dd_mode, self.beyond_dd, self.attenuation, self.ppb, self.spp, self.iterations, self.spectral_points, self.make_end_entry, self.end_entry_freq, self.rbw)

class SweepPlanner(object):
    """
    An object that plans a sweep based on  given paramaters.

    :param dev_prop: the sweep device properties
    :type dev_prop: dict
    """

    def __init__(self, dev_prop):
        self.dev_properties = dev_prop
        self._prev_settings = SweepSettings()

    def plan_sweep(self, fstart, fstop, rbw, mode, dev_settings = {}):
        """
        Plan the sweep given the inputs
        """

        # initialize the sweep settings variable
        sweep_settings = SweepSettings()

        # assign the sweep mode and start/stop
        sweep_settings.rfe_mode = mode
        sweep_settings.bandstart = fstart
        sweep_settings.bandstop = fstop

        if 'attenuator' in dev_settings:
            sweep_settings.attenuation = dev_settings['attenuator']

        # grab the usable bw of the current mode
        usable_bw = self.dev_properties.USABLE_BW[mode]

        # calculate the required SPP to get the RBW desired
        sweep_settings.spp = self.dev_properties.FULL_BW[mode] / rbw

        # find closest multiple of 32 because hardware
        sweep_settings.spp = int(32 * round(float(sweep_settings.spp) / 32))

        # double the points for SH/SHN mode
        if mode in ['SH', 'SHN']:
            sweep_settings.spp = sweep_settings.spp * 2

        # if we're using zif mode, but we have a DD entry, we have half the SPP avaible, since DD is I-only and ZIF is IQ
        if (mode == 'ZIF') and sweep_settings.dd_mode:
            maxspp = self.dev_properties.MAX_SPP / 2
        else:
            maxspp = self.dev_properties.MAX_SPP

        # adjust SPP if it's too big
        sweep_settings.spp = min(maxspp, sweep_settings.spp)

        # figure out our actual RBW (account for real vs complex data)
        sweep_settings.rbw = self.dev_properties.FULL_BW[mode] / sweep_settings.spp
        if not (mode == 'ZIF'):
            sweep_settings.rbw = sweep_settings.rbw * 2

        # make sure our result is atleast 1 RBW big
        if (sweep_settings.bandstop - sweep_settings.bandstart) < sweep_settings.rbw:
            fstop = sweep_settings.bandstart + sweep_settings.rbw
            sweep_settings.bandstop = fstop

        # change fstart and stop by a bit to account for floating point errors
        # TODO: make this take into account tuning resolution
        fstart -= sweep_settings.rbw * 4
        fstop += sweep_settings.rbw * 4

        # calculate fstart frequency
        if fstart < self.dev_properties.MIN_TUNABLE[mode]:
            sweep_settings.dd_mode = True
            sweep_settings.fstart = self.dev_properties.MIN_TUNABLE[mode] + (usable_bw / 2)
            sweep_settings.step_count += 1

        # make sure we don't accidentally make an fstart that's beyond our tuning range
        elif (fstart + (usable_bw / 2)) > self.dev_properties.MAX_TUNABLE[mode]:
            sweep_settings.dd_mode = False
            sweep_settings.fstart = self.dev_properties.MAX_TUNABLE[mode] - (usable_bw / 2)

        else:
            sweep_settings.dd_mode = False
            sweep_settings.fstart = fstart + (usable_bw / 2)

        # check if non-dd mode is required
        if fstop <= self.dev_properties.MIN_TUNABLE[mode]:
            sweep_settings.beyond_dd = False
        else:
            sweep_settings.beyond_dd = True
            sweep_settings.step_count += 1

            # assign the sweep entry's step frequency reducing by a couple rbw to account for floating point errors
            # TODO: make this take into account tuning resolution
            sweep_settings.fstep = usable_bw - (sweep_settings.rbw * 4)

            # calculate the fstop of the sweep entry from fstart and how many usable_bw's we need
            fspan = fstop - sweep_settings.fstart - sweep_settings.rbw
            required_steps = round(fspan / sweep_settings.fstep)
            sweep_settings.fstop = sweep_settings.fstart + (required_steps * sweep_settings.fstep)
            sweep_settings.step_count += required_steps

        # make sure fstop is lower than max tunable
        # - it can sometimes be higher if an fstart is chosen, such that our
        #   fstep causes our fstop to go beyond fmax to cover all the band required
        sweep_settings.make_end_entry = False
        sweep_settings.end_entry_freq = 0
        if sweep_settings.fstop > self.dev_properties.MAX_TUNABLE[mode]:
            # go back one step
            sweep_settings.fstop -= sweep_settings.fstep

            # add an entry for fmax
            sweep_settings.make_end_entry = True
            sweep_settings.end_entry_freq = self.dev_properties.MAX_TUNABLE[mode] - (usable_bw / 2)

        # calculate the expected number of spectral bins required for the SweepEntry
        sweep_settings.spectral_points = int(round((sweep_settings.bandstop - sweep_settings.bandstart) / sweep_settings.rbw))

        # return the sweep_settings
        return sweep_settings

class SweepDevice(object):
    """
    Virtual device that generates power spectrum from a given frequency range
    by sweeping the frequencies with a real device and piecing together the FFT results.

    :param real_device: the RF device that will be used for capturing data,
                        typically a :class:`pyrf.devices.thinkrf.WSA` instance.
    :param async_callback: a callback to use for async operation (not used if
                     *real_device* is using a blocking :class:`PlainSocketConnector`)
    """
    # keep track of the mode
    rfe_mode = None

    # keep track of the fstart/fstop and rbw
    fstart = None
    fstop = None
    rbw = None

    # keep track of non-standard device settings
    device_settings = None

    # keep track of whether DD mode is needed
    dd_mode = False

    # keep track of the sweep settings
    _sweep_settings = None

    # keep track of the packet count
    packet_count = 0

    # determine if a new entry is required
    _new_entry = True

    # array to place spectral data
    spectral_data = []

    capture_count = 0

    sp_corr_obj = None
    nf_corr_obj = None
    _flattening_enabled = True

    def __init__(self, real_device, async_callback=None):

        # init log string
        self.logstr = ''
        self.logtype = 'NONE'

        # initialize the real device
        self.real_device = real_device

        # request read permission from device
        self.real_device.request_read_perm()

        # keep track of the device properties
        self.dev_properties = self.real_device.properties

        # initialize the geolocation callback
        self._geo_callback_func = None
        self._geo_callback_data = None

        # initialize the sweep planner
        self._sweep_planner = SweepPlanner(self.dev_properties)

        # make sure user passes async callback if the device has async connector
        if real_device.async_connector():
            if not async_callback:
                raise SweepDeviceError(
                    "async_callback required for async operation")

            # disable receiving data until we are expecting it
            real_device.set_async_callback(None)

            # Function to be called when async data is done capturing
            def _save_correction_vector(data_buffer):
                if data_buffer is None:
                    return None
                try:
                    if data_buffer.v_type == "SIGNAL":
                        self.sp_corr_obj = correction_vector()
                        self.sp_corr_obj.buffer_to_vector(data_buffer.data_buffer)
                    elif data_buffer.v_type == "NOISE":
                        self.nf_corr_obj = correction_vector()
                        self.nf_corr_obj.buffer_to_vector(data_buffer.data_buffer)
                except AttributeError:
                    if data_buffer.v_type == "SIGNAL":
                        self.sp_corr_obj = None
                    elif data_buffer.v_type == "NOISE":
                        self.nf_corr_obj = None

            # function to catch the errback of the async code. Used to handle
            # the case when we can get the correction vectors.
            def _catch_timeout(failure):
                failure.trap(IOError)
                return None

            vector_obj = correction_vector_acquire()
            vector_obj.dut = real_device

            vector_obj1 = correction_vector_acquire()
            vector_obj1.dut = real_device

            d1 = vector_obj.get_vector("NOISE")
            d1.addCallback(_save_correction_vector)
            d1.addErrback(_catch_timeout)

            d2 = vector_obj1.get_vector("SIGNAL")
            d2.addCallback(_save_correction_vector)
            d2.addErrback(_catch_timeout)

        else:

            # make sure user doesnt pass async callback if the connector uses blocking sockets
            if async_callback:
                raise SweepDeviceError(
                    "async_callback not applicable for sync operation")

            def _get_correction(dut, v_type=None):
                if v_type.upper() == "SIGNAL" or v_type.upper() == "NOISE":
                    v_type = v_type.upper()
                else:
                    raise ValueError

                max_buf_size = 16*1024
                offset = 0
                bin_data = ""
                try:
                    signal_size = dut.correction_size(v_type)
                except (IOError, OSError):  # this will handle socket.error's
                    raise ValueError

                # We have nothing to transfer
                if signal_size == 0:
                    return None

                # check to see if tere is more data than can be transfer in one
                # go
                if signal_size > max_buf_size:
                    # if so transfer our max buffer size
                    transfer_size = max_buf_size
                else:
                    # if not grab only what we need
                    transfer_size = signal_size

                # While we still have data remaining
                while offset < signal_size:
                    # get the data
                    data_buffer = dut.correction_data(v_type, offset,
                                                      transfer_size)
                    # figure out how many bytes were transfered
                    transfered = len(data_buffer)
                    # append the data to the buffer of what we have allready
                    # got
                    bin_data = b"".join([bin_data, data_buffer])
                    # increase the offset
                    offset = offset + transfered
                return bin_data

            self.sp_corr_obj = correction_vector()
            try:
                self.sp_corr_obj.buffer_to_vector(_get_correction(self.real_device, "SIGNAL"))
            except ValueError:
                self.sp_corr_obj = None

            self.nf_corr_obj = correction_vector()
            try:
                self.nf_corr_obj.buffer_to_vector(_get_correction(self.real_device, "NOISE"))
            except ValueError:
                self.nf_corr_obj = None

        self.async_callback = async_callback
        self.continuous = False

        # init the sweep id
        self._next_sweep_id = 0

        # init last finished (technically, it hasn't finished, but for our purposes, it has)
        self._last_finished = True

    # Private function
    def log(self, firstmsg, *msgs):
        if self.logtype == 'LOG':
            self.logstr += firstmsg.__str__()
            for msg in msgs:
                self.logstr += ", "
                self.logstr += msg.__str__()
            self.logstr += "\n"
        elif self.logtype == 'PRINT':
            sys.stdout.write(firstmsg.__str__())
            for msg in msgs:
                sys.stdout.write(", ")
                sys.stdout.write(msg.__str__())
            sys.stdout.write("\n")

    def enable_flattening(self, enable=None):
        """
        :param enable: enable or disable spectral flattening
        :type enable: bool or None
        """
        if enable is None:
            return self._flattening_enabled
        else:
            self._flattening_enabled = enable

    def set_geolocation_callback(self, func, data = None):
        """
        set a callback that will get called whenever the geolocation information
        of the device is updated.
        The callback function should accept two parameters.  The first parameter
        will be the callback data that was passed in this function
        set_geolocation_callback(func, data, geolocation_dictionary).
        The geolocation_dictionary will have the following properties:
        - oui
        - seconds
        - altitude
        - longitude
        - speedoverground
        - secondsfractional
        - track
        - latitude
        - magneticvariation
        - heading
        See the programmer's guide for usage on each of these properties.

        :param func: the function to be called
        :param data: the data to be passed to the function
        :returns: None
        """

        self._geo_callback_func = func
        self._geo_callback_data = data

    def capture_power_spectrum(self,
                               fstart,
                               fstop,
                               rbw,
                               device_settings=None,
                               mode='SH',
                               continuous=False):
        """
        Initiate a data capture from the *real_device* by setting up a sweep list
        and starting a single sweep, and then return power spectral density data
        along with the **actual** sweep start and stop frequencies set (which
        might not be exactly the same as the requested *fstart* and *fstop*).

        .. note:: This function does not pipeline, and if the last sweep isn't received before starting a new one, it will generate a failure.

        :param int fstart: sweep starting frequency in Hz
        :param int fstop: sweep ending frequency in Hz
        :param float rbw: the resolution bandwidth (RBW) in Hz of the data to be captured (output RBW may be smaller than requested)
        :param device_settings: attenuation and other device settings
        :type device_settings: dict
        :param str mode: sweep mode, 'ZIF', 'SH', or 'SHN'
        :param bool continuous: set sweep to be continuously or not (once only)

        :returns: fstart, fstop, power_data
        """

        self.log("- capture_power_spectrum", fstart, fstop, rbw, device_settings, mode, continuous)

        if continuous and not self.async_callback:
            raise SweepDeviceError(
                "continuous mode only applies to async operation")

        # see if the last sweep has finished
        if not self._last_finished:
            raise SweepDeviceError(
                "previous sweep must have finished before starting a new one")
        self._last_finished = False

        # increment the sweep id
        if self._next_sweep_id < 0x00000000ffffffff:
            self._next_sweep_id += 1
        else:
            self._next_sweep_id = 0

        # keep track if this is a continuous sweep
        self.continuous = continuous

        # plan the sweep
        self._sweep_planner = SweepPlanner(self.dev_properties)
        self._sweep_settings = self._sweep_planner.plan_sweep(fstart, fstop, rbw, mode, device_settings)
        self.log("self._sweep_settings = %s" % self._sweep_settings)

        # remember our last sweep for optimization purposes
        self._last_sweep = (fstart, fstop, rbw, mode, device_settings, continuous)

        # configure the device with the sweep_settings
        self.real_device.sweep_clear()
        self.real_device.sweep_add(self._sweep_settings)

        # configure the iteration
        self.real_device.sweep_iterations(1)

        # capture the sweep data
        return self._perform_full_sweep()

    def _perform_full_sweep(self):

        # perform the sweep using async socket
        if self.async_callback:

            # set the async callback
            self.real_device.set_async_callback(self._vrt_receive)

            # start the sweep sequence
            self._start_sweep()
            return

        # perform sweep using blocking sockets
        self._start_sweep()
        result = None
        while result is None:
            result = self._vrt_receive(self.real_device.read())

        return result

    def _start_sweep(self):

        self._vrt_context = {}

        # initialize the array we'll use to hold results
        self.spectral_data = np.zeros(self._sweep_settings.spectral_points)

        # keep track of packets recieved
        self.packet_count = 0

        self.real_device.sweep_start(self._next_sweep_id)

    def _vrt_receive(self, packet):

        # context packet just update our context dictionary
        if packet.is_context_packet():
            # look for any geolocation info
            geo = { }
            for field in [ 'latitude', 'longitude', 'altitude', 'speedoverground', 'heading', 'track', 'magneticvariation' ]:
                if field in packet.fields:
                    geo[field] = packet.fields[field]
            if geo and self._geo_callback_func:
                # execute callback
                func = self._geo_callback_func
                func(self._geo_callback_data, geo)

            self._vrt_context.update(packet.fields)
            self.log(packet)
            return

        # check to see if we recieved our sweep ID
        if not ('sweepid' in self._vrt_context):
            return

        # make sure we are receiving packets for the right sweep
        if not (self._vrt_context['sweepid'] == self._next_sweep_id):
            raise SweepDeviceError("data packets received before start of sweep received!  cur = %d, next = %d" % (self._vrt_context['sweepid'], self._next_sweep_id))

        # increment the packet count
        self.packet_count += 1
        self.log("#%d of %d - %s" % (self.packet_count, self._sweep_settings.step_count, packet))

        # retrieve the frequency and usable BW of the packet
        packet_freq = self._vrt_context['rffreq']
        usable_bw = self.dev_properties.USABLE_BW[self._sweep_settings.rfe_mode]

        # compute the fft
        pow_data = compute_fft(self.real_device, packet, self._vrt_context)

        # calc rbw for this packet
        rbw = float(self.dev_properties.FULL_BW[self._sweep_settings.rfe_mode]) / len(pow_data)
        self.log("rbw = %f, %f" % (rbw, self._sweep_settings.rbw))
        if self._flattening_enabled:
            # Check if we are above 50 MHz and in SH mode
            if packet_freq >= 50e6 and self._sweep_settings.rfe_mode == "SH":
                number_of_points = len(pow_data)
                # check if we have correction vectors (Noise)
                if self.nf_corr_obj is not None:
                    # if so grab them
                    nf_cal = \
                            self.nf_corr_obj.get_correction_vector(packet_freq,
                                                                   number_of_points)
                else:
                    # if no set it to 0
                    nf_cal = np.zeros(number_of_points)

                # check if we have corrrection vectors (Spectrum)
                if self.sp_corr_obj is not None:
                    # if so grab them
                    sp_cal = \
                            self.sp_corr_obj.get_correction_vector(packet_freq,
                                                                   number_of_points)
                else:
                    # if not set it to 0
                    sp_cal = np.zeros(number_of_points)

                # if the data is spectraly inverted, invert the vectors
                if packet.spec_inv:
                    nf_cal = np.flipud(nf_cal)
                    sp_cal = np.flipud(sp_cal)

                # calculate the correction threshold
                correction_thresh = (-135.0 + ((10.0 * packet_freq / 1e6)
                                               / 27000.0) + 10.0
                                     * np.log10(rbw)
                                     + self._sweep_settings.attenuation)
                # creat the spectrum. per bin, if the ampltitude is above
                # correction threshold do pow_data - sp_cal else do pow_data -
                # nf_cal
                pow_data = np.where(pow_data < correction_thresh,
                                    pow_data - nf_cal, pow_data - sp_cal)

        # check if DD mode was used in this sweep
        if self.packet_count == 1 and self._sweep_settings.dd_mode:
            # copy the data into the result array
            self._copy_data(0, self.dev_properties.FULL_BW['DD'], pow_data, self._sweep_settings.bandstart, self._sweep_settings.bandstop, self.spectral_data);

            if self._sweep_settings.beyond_dd:
                return
            else:
                return self._emit_data()

        # determine the usable bins in this config
        self.log("===> compute_usable_bins()", self._sweep_settings.rfe_mode, self._sweep_settings.spp, 1, 0)
        usable_bins = compute_usable_bins(self.dev_properties,
                                          self._sweep_settings.rfe_mode,
                                          self._sweep_settings.spp,
                                          1,
                                          0)
        self.log("<--- usable_bins", usable_bins)

        # adjust the usable range based on spectral inversion
        self.log("===> adjust_usable_fstart_fstop()", "self.dev_properties", self._sweep_settings.rfe_mode, len(pow_data) * 2, 1, packet_freq, packet.spec_inv, usable_bins)
        usable_bins, packet_start, packet_stop = adjust_usable_fstart_fstop(self.dev_properties,
                                                              self._sweep_settings.rfe_mode,
                                                              len(pow_data) * 2,
                                                              1,
                                                              packet_freq,
                                                              packet.spec_inv,
                                                              usable_bins)
        self.log("<--- adjust_usable_fstart_fstop", packet_start, packet_stop, usable_bins)
        #
        # WARNING: the start and stop returned from this function are HIGHLY sketchy
        #

        # calculate packet frequency range
        #packet_start = packet_freq - (self.dev_properties.FULL_BW[self._sweep_settings.rfe_mode] / 2)
        #packet_stop = packet_freq + (self.dev_properties.FULL_BW[self._sweep_settings.rfe_mode] / 2)
        #print "packet start/stop", packet_start, packet_stop

        #trim the FFT data, note decimation is 1, fshift is 0
        self.log("===> trim_to_usable_fstart_fstop()", "pow_data", usable_bins, packet_start, packet_stop)
        trimmed_spectrum, edge_data, usable_start, usable_stop = trim_to_usable_fstart_fstop(pow_data,
                                                                                usable_bins,
                                                                                packet_start,
                                                                                packet_stop)
        self.log("<--- trim_to_usable_fstart_fstop", usable_start, usable_stop, "trimmed_spectrum", edge_data)

        # copy the data
        self._copy_data(usable_start, usable_stop, trimmed_spectrum, self._sweep_settings.bandstart, self._sweep_settings.bandstop, self.spectral_data);

        # if there's no more packets, emit result
        if self.packet_count == self._sweep_settings.step_count:
            return self._emit_data()

        # all done
        return

    def _emit_data(self):

        # note that we finished this sweep
        self._last_finished = True

        # if async callback is available, emit the data
        if self.async_callback:

            self.async_callback(self._sweep_settings.bandstart, self._sweep_settings.bandstop, self.spectral_data)
            return
        # return the values if using blocking sockets
        else:
            return (self._sweep_settings.bandstart, self._sweep_settings.bandstop, self.spectral_data)

    def _copy_data(self, src_fstart, src_fstop, src_psd, dst_fstart, dst_fstop, dst_psd):
        self.log("_copy_data(%d, %d, src_psd, %d, %d, dst_psd)" % (src_fstart, src_fstop, dst_fstart, dst_fstop))

        # calc src len and dst len
        srclen = len(src_psd)
        dstlen = len(dst_psd)
        self.log("len -- src = %d, dst = %d" % (srclen, dstlen))

        # calc src and dest rbw
        srcrbw = float(src_fstop - src_fstart) / srclen
        dstrbw = float(dst_fstop - dst_fstart) / dstlen
        self.log("rbw = %f, %f, %f" % (srcrbw, dstrbw, self._sweep_settings.rbw))

        # check if packet start is before sweep start.  shouldn't happen, but check anyway
        self.log("boundary(start) = %f / %f" % (src_fstart, dst_fstart))
        if src_fstart < dst_fstart:
            self.log("foo")
            src_start_bin = int(float(dst_fstart - src_fstart) / srcrbw)
        else:
            self.log("bar")
            src_start_bin = 0

        # check if packet stop is after sweep stop.  this means we don't need the whole packet
        self.log("boundary(stop) = %f / %f" % (src_fstop, dst_fstop))
        if src_fstop > dst_fstop:
            self.log("foo")
            src_stop_bin = srclen - int(float(src_fstop - dst_fstop) / srcrbw)
        else:
            self.log("bar")
            src_stop_bin = srclen

        # how many values are we copying?
        tocopy = src_stop_bin - src_start_bin

        # calculate dest start index
        if src_fstart < dst_fstart:
            dst_start_bin = 0
        else:
            dst_start_bin = int(round(float(src_fstart - dst_fstart) / dstrbw))

        # calculate dest stop index
        dst_stop_bin = dst_start_bin + tocopy
        if dst_stop_bin > dstlen:
            dst_stop_bin = dstlen

            # adjust tocopy
            tocopy = dst_stop_bin - dst_start_bin

            # adjust src stop bin because we adjusted tocopy
            src_stop_bin = src_start_bin + tocopy

        # copy the data, if there's data that needs copying
        if ((dst_stop_bin - dst_start_bin) > 0) and ((src_stop_bin - src_start_bin) > 0):
            self.log("dst_psd[%d:%d] = src_psd[%d:%d]" % (dst_start_bin, dst_stop_bin, src_start_bin, src_stop_bin))
            dst_psd[dst_start_bin:dst_stop_bin] = src_psd[src_start_bin:src_stop_bin]


# Mofidied work:
# --------------------------------------------------------
# Copyright (c) 2017 Preferred Networks, Inc.
# --------------------------------------------------------
#
# Original works by:
# --------------------------------------------------------
# Faster R-CNN implementation by Chainer
# Copyright (c) 2016 Shunta Saito
# Licensed under The MIT License [see LICENSE for details]
# https://github.com/mitmul/chainer-faster-rcnn
# --------------------------------------------------------
# Faster R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick and Sean Bell
# https://github.com/rbgirshick/py-faster-rcnn
# --------------------------------------------------------

from __future__ import division

import cv2
import numpy as np

import chainer
from chainer import cuda
import chainer.functions as F
from chainercv.links.model.faster_rcnn.utils.loc2bbox import loc2bbox
from chainercv.utils import non_maximum_suppression
import six

class MaskRCNN(chainer.Chain):

    mask_losses = [
        'softmax',
        'softmax_x2',
        'sigmoid_softmax',
        'sigmoid_sigmoid',
        'sigmoid_sigmoid+',
        'softmax_relook_softmax',
        'softmax_relook_softmax+',
        'softmax_relook_softmax+_res',
        'softmax_relook_softmax_cls',
        'softmax_relook_softmax+_cls',
        'softmax_relook_softmax_tt',
        'softmax_relook_softmax+_tt',
        'softmax_relook_softmax+_tt2',
        'softmax_relook_softmax_cls_tt',
        'softmax_relook_softmax+_cls_tt',
        'softmax_relook_softmax_bbox',
        'softmax_relook_softmax+_bbox',
    ]

    def __init__(
            self, extractor, rpn, head,
            mean,
            min_size=600,
            max_size=1000,
            loc_normalize_mean=(0., 0., 0., 0.),
            loc_normalize_std=(0.1, 0.1, 0.2, 0.2),
            detections_per_im=100,
            mask_loss='softmax'):
        super(MaskRCNN, self).__init__()
        with self.init_scope():
            self.extractor = extractor
            self.rpn = rpn
            self.head = head

        self.mean = mean
        self.min_size = min_size
        self.max_size = max_size
        self.loc_normalize_mean = loc_normalize_mean
        self.loc_normalize_std = loc_normalize_std

        self._detections_per_im = detections_per_im

        assert mask_loss in self.mask_losses
        self.mask_loss = mask_loss

        self.context = None
        self.nms_thresh = 0.3
        self.score_thresh = 0.05

    @property
    def n_class(self):
        # Total number of classes including the background.
        return self.head.n_class

    def prepare(self, imgs):
        prepared_imgs = []
        sizes = []
        scales = []
        for img in imgs:
            _, H, W = img.shape

            scale = 1.

            if self.min_size:
                scale = self.min_size / min(H, W)

            if self.max_size and scale * max(H, W) > self.max_size:
                scale = self.max_size / max(H, W)

            img = img.transpose(1, 2, 0)
            img = cv2.resize(img, None, fx=scale, fy=scale)
            img = img.transpose(2, 0, 1)

            img = (img - self.mean).astype(np.float32, copy=False)

            prepared_imgs.append(img)
            sizes.append((H, W))
            scales.append(scale)
        return prepared_imgs, sizes, scales

    def _suppress(self, raw_cls_bbox, raw_prob):
        bbox = list()
        label = list()
        score = list()
        # skip cls_id = 0 because it is the background class
        for l in range(1, self.n_class):
            cls_bbox_l = raw_cls_bbox.reshape((-1, self.n_class, 4))[:, l, :]
            prob_l = raw_prob[:, l]

            # thresholding by score
            keep = prob_l > self.score_thresh
            cls_bbox_l = cls_bbox_l[keep]
            prob_l = prob_l[keep]

            # thresholding by nms
            keep = non_maximum_suppression(
                cls_bbox_l, self.nms_thresh, prob_l)
            bbox.append(cls_bbox_l[keep])
            # The labels are in [0, self.n_class - 2].
            label.append((l - 1) * np.ones((len(keep),)))
            score.append(prob_l[keep])
        bbox = np.concatenate(bbox, axis=0).astype(np.float32)
        label = np.concatenate(label, axis=0).astype(np.int32)
        score = np.concatenate(score, axis=0).astype(np.float32)
        return bbox, label, score

    def _to_bbox_label_score(self, roi_cls_locs, roi_scores, rois, roi_indices, scale, size):  # NOQA
        # We are assuming that batch size is 1.
        roi_cls_loc = roi_cls_locs.data
        roi_score = roi_scores.data
        roi = rois / scale
        roi_index = roi_indices

        # Convert predictions to bounding boxes in image coordinates.
        # Bounding boxes are scaled to the scale of the input images.
        mean = self.xp.tile(self.xp.asarray(self.loc_normalize_mean),
                            self.n_class)
        std = self.xp.tile(self.xp.asarray(self.loc_normalize_std),
                           self.n_class)
        roi_cls_loc = (roi_cls_loc * std + mean).astype(np.float32)
        roi_cls_loc = roi_cls_loc.reshape((-1, self.n_class, 4))
        roi_cls = self.xp.broadcast_to(roi[:, None], roi_cls_loc.shape)
        cls_bbox = loc2bbox(roi_cls.reshape((-1, 4)),
                            roi_cls_loc.reshape((-1, 4)))
        cls_bbox = cls_bbox.reshape((-1, self.n_class * 4))
        # clip bounding box
        cls_bbox[:, 0::2] = self.xp.clip(cls_bbox[:, 0::2], 0, size[0])
        cls_bbox[:, 1::2] = self.xp.clip(cls_bbox[:, 1::2], 0, size[1])
        # clip roi
        roi[:, 0::2] = self.xp.clip(roi[:, 0::2], 0, size[0])
        roi[:, 1::2] = self.xp.clip(roi[:, 1::2], 0, size[1])

        prob = F.softmax(roi_score).data

        roi_index = self.xp.broadcast_to(
            roi_index[:, None], roi_cls_loc.shape[:2])
        raw_cls_bbox = cuda.to_cpu(cls_bbox)
        raw_prob = cuda.to_cpu(prob)

        if self.context:
            n_fg_class = self.n_class - 1
            for l in range(n_fg_class):
                if l not in self.context:
                    raw_prob[:, l + 1] = 0
            raw_prob = raw_prob / raw_prob.sum(axis=0)
        bbox, label, score = self._suppress(raw_cls_bbox, raw_prob)

        bbox_int = np.round(bbox).astype(np.int32)
        bbox_sizes = ((bbox_int[:, 2] - bbox_int[:, 0]) *
                      (bbox_int[:, 3] - bbox_int[:, 1]))
        keep = bbox_sizes > 0
        bbox = bbox[keep]
        label = label[keep]
        score = score[keep]

        if self._detections_per_im > 0:
            indices = np.argsort(score)
            keep = indices >= (len(indices) - self._detections_per_im)
            bbox = bbox[keep]
            label = label[keep]
            score = score[keep]

        return bbox, label, score

    def predict_masks(self, imgs):
        return self.predict(imgs)

    def predict(self, imgs):
        imgs, sizes, scales = self.prepare(imgs)

        bboxes = list()
        masks = list()
        labels = list()
        scores = list()
        for img, size, scale in zip(imgs, sizes, scales):
            with chainer.using_config('train', False), \
                    chainer.function.no_backprop_mode():
                x = chainer.Variable(self.xp.asarray(img[None]))
                h = self.extractor(x)
                rpn_locs, rpn_scores, rois, roi_indices, anchor =\
                    self.rpn(h, x.shape[2:], scale)
                roi_cls_locs, roi_scores, _, = self.head(
                    h, rois, roi_indices,
                    pred_bbox=True, pred_mask=False,
                    pred_bbox2=False, pred_mask2=False)

            bbox, label, score = self._to_bbox_label_score(
                roi_cls_locs, roi_scores, rois, roi_indices, scale, size)

            if len(bbox) == 0:
                bboxes.append(bbox)
                labels.append(label)
                scores.append(score)
                masks.append(np.zeros((0, size[0], size[1]), dtype=bool))
                continue

            if self.mask_loss in [
                'softmax_relook_softmax_bbox',
                'softmax_relook_softmax+_bbox',
            ]:
                with chainer.using_config('train', False), \
                        chainer.function.no_backprop_mode():
                    rois = self.xp.asarray(bbox) * scale
                    roi_indices = self.xp.zeros(
                        (len(bbox),), dtype=np.int32)
                    roi_cls_locs, roi_scores, _, = self.head(
                        h, rois, roi_indices,
                        pred_bbox=False, pred_mask=True,
                        pred_bbox2=True, pred_mask2=False,
                        labels=label + 1)
                    roi_cls_locs = roi_cls_locs[1]
                    roi_scores = roi_scores[1]
                bbox, label, score = self._to_bbox_label_score(
                    roi_cls_locs, roi_scores, rois, roi_indices, scale, size)

            bboxes.append(bbox)
            labels.append(label)
            scores.append(score)

            if len(bbox) == 0:
                masks.append(np.zeros((0, size[0], size[1]), dtype=bool))
                continue

            with chainer.using_config('train', False), \
                    chainer.function.no_backprop_mode():
                rois = self.xp.asarray(bbox) * scale
                roi_indices = self.xp.zeros(
                    (len(bbox),), dtype=np.int32)
                _, _, roi_masks = self.head(
                    x=h, rois=rois, roi_indices=roi_indices,
                    pred_bbox=False, pred_mask=True,
                    pred_bbox2=False, pred_mask2=True,
                    labels=label + 1)
            if self.mask_loss in ['softmax', 'softmax_x2']:
                _, n_fg_class_x3, roi_H, roi_W = roi_masks.shape
                n_fg_class = n_fg_class_x3 // 3
                roi_masks = roi_masks.reshape(
                    (roi_masks.shape[0], 3, n_fg_class, roi_H, roi_W))
                roi_masks = roi_masks[np.arange(len(label)), :, label, :, :]
                roi_masks = roi_masks.transpose(0, 2, 3, 1)
                roi_mask = cuda.to_cpu(roi_masks.data)
            elif self.mask_loss == 'sigmoid_softmax':
                roi_masks, roi_masks_bginv = roi_masks
                # roi_mask
                roi_masks = F.sigmoid(roi_masks)
                roi_masks = roi_masks[np.arange(len(label)), label, :, :]
                roi_mask = cuda.to_cpu(roi_masks.data)
                # roi_mask_bginv
                _, n_fg_class_x2, roi_H, roi_W = roi_masks_bginv.shape
                n_fg_class = n_fg_class_x2 // 2
                roi_masks_bginv = roi_masks_bginv.reshape(
                    (roi_masks_bginv.shape[0], 2, n_fg_class, roi_H, roi_W))
                roi_masks_bginv = roi_masks_bginv[
                    np.arange(len(label)), :, label, :, :]
                roi_masks_bginv = F.softmax(roi_masks_bginv)[:, 1, :, :]
                roi_mask_bginv = cuda.to_cpu(roi_masks_bginv.data)
            elif self.mask_loss in ['sigmoid_sigmoid', 'sigmoid_sigmoid+']:
                roi_masks, roi_masks_bginv = roi_masks
                # roi_mask
                roi_masks = F.sigmoid(roi_masks)
                roi_masks = roi_masks[np.arange(len(label)), label, :, :]
                roi_mask = cuda.to_cpu(roi_masks.data)
                # roi_mask_bginv
                roi_masks_bginv = F.sigmoid(roi_masks_bginv)
                roi_masks_bginv = roi_masks_bginv[
                    np.arange(len(label)), label, :, :]
                roi_mask_bginv = cuda.to_cpu(roi_masks_bginv.data)
            elif self.mask_loss in [
                'softmax_relook_softmax',
                'softmax_relook_softmax+',
                'softmax_relook_softmax+_res',
                'softmax_relook_softmax_cls',
                'softmax_relook_softmax+_cls',
                'softmax_relook_softmax_bbox',
                'softmax_relook_softmax+_bbox',
                'softmax_relook_softmax_tt',
                'softmax_relook_softmax+_tt',
                'softmax_relook_softmax+_tt2',
                'softmax_relook_softmax_cls_tt',
                'softmax_relook_softmax+_cls_tt',
            ]:
                roi_masks1, roi_masks2 = roi_masks
                roi_masks1 = roi_masks1.transpose(0, 2, 3, 1)
                roi_masks2 = roi_masks2.transpose(0, 2, 3, 1)
                roi_mask = cuda.to_cpu(roi_masks2.data)  # mainly used
                roi_mask_sub = cuda.to_cpu(roi_masks1.data)
            else:
                raise ValueError

            roi = np.round(bbox).astype(np.int32)
            n_roi = len(roi)
            mask = - np.ones((n_roi, size[0], size[1]), dtype=np.int32)
            for i in six.moves.range(n_roi):
                y1, x1, y2, x2 = roi[i]
                roi_H = y2 - y1
                roi_W = x2 - x1
                roi_mask_i = cv2.resize(roi_mask[i], (roi_W, roi_H))
                if self.mask_loss in ['softmax', 'softmax_x2']:
                    roi_mask_i_out = np.argmax(roi_mask_i, axis=2)
                elif self.mask_loss in [
                    'softmax_relook_softmax',
                    'softmax_relook_softmax+',
                    'softmax_relook_softmax+_res',
                    'softmax_relook_softmax_cls',
                    'softmax_relook_softmax+_cls',
                    'softmax_relook_softmax_bbox',
                    'softmax_relook_softmax+_bbox',
                    'softmax_relook_softmax_tt',
                    'softmax_relook_softmax+_tt',
                    'softmax_relook_softmax+_tt2',
                    'softmax_relook_softmax_cls_tt',
                    'softmax_relook_softmax+_cls_tt',
                ]:
                    if '+' in self.mask_loss:
                        roi_mask_sub_i = cv2.resize(
                            roi_mask_sub[i], (roi_W, roi_H))
                        roi_mask_i = roi_mask_i + roi_mask_sub_i
                    roi_mask_i_out = np.argmax(roi_mask_i, axis=2)
                elif self.mask_loss in [
                    'sigmoid_softmax',
                    'sigmoid_sigmoid',
                    'sigmoid_sigmoid+',
                ]:
                    roi_mask_bginv_i = cv2.resize(
                        roi_mask_bginv[i], (roi_W, roi_H))

                    roi_mask_i_out = - np.ones_like(roi_mask_i)
                    roi_mask_vis_i = roi_mask_i > 0.5
                    roi_mask_i_out[roi_mask_vis_i] = 1  # vis
                    roi_mask_i_out[~roi_mask_vis_i &
                                   (roi_mask_bginv_i <= 0.5)] = 0   # bg
                    roi_mask_i_out[~roi_mask_vis_i &
                                   (roi_mask_bginv_i > 0.5)] = 2   # inv
                else:
                    raise ValueError
                mask[i, y1:y2, x1:x2] = roi_mask_i_out
            masks.append(mask)

        return bboxes, masks, labels, scores