text:
import mmcv
import matplotlib.pyplot as plt
import os.path as osp
import numpy as np
from PIL import Image

# Check Pytorch installation
import torch, torchvision
print(torch.__version__, torch.cuda.is_available())

# Check MMSegmentation installation
import mmseg
print(mmseg.__version__)

# img = mmcv.imread('data/iccv09Data/images/6000124.jpg')
# plt.figure(figsize=(8, 6))
# plt.imshow(mmcv.bgr2rgb(img))
# plt.show()

if __name__ == '__main__':

    # # convert dataset annotation to semantic segmentation map
    data_root = 'data/iccv09Data'
    img_dir = 'images'
    ann_dir = 'labels'
    # # define class and plaette for better visualization
    classes = ('sky', 'tree', 'road', 'grass', 'water', 'bldg', 'mntn', 'fg obj')
    palette = [[128, 128, 128], [129, 127, 38], [120, 69, 125], [53, 125, 34],
               [0, 11, 123], [118, 20, 12], [122, 81, 25], [241, 134, 51]]
    # for file in mmcv.scandir(osp.join(data_root, ann_dir), suffix='.regions.txt'):
    #   seg_map = np.loadtxt(osp.join(data_root, ann_dir, file)).astype(np.uint8)
    #   seg_img = Image.fromarray(seg_map).convert('P')
    #   seg_img.putpalette(np.array(palette, dtype=np.uint8))
    #   seg_img.save(osp.join(data_root, ann_dir, file.replace('.regions.txt',
    #                                                          '.png')))

    # # Let's take a look at the segmentation map we got
    # import matplotlib.patches as mpatches
    # img = Image.open('data/iccv09Data/labels/6000124.png')
    # plt.figure(figsize=(8, 6))
    # im = plt.imshow(np.array(img.convert('RGB')))
    #
    # # create a patch (proxy artist) for every color
    # patches = [mpatches.Patch(color=np.array(palette[i])/255.,
    #                           label=classes[i]) for i in range(8)]
    # # put those patched as legend-handles into the legend
    # plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.,
    #            fontsize='large')
    #
    # plt.show()

    # # split train/val set randomly
    # split_dir = 'splits'
    # mmcv.mkdir_or_exist(osp.join(data_root, split_dir))
    # filename_list = [osp.splitext(filename)[0] for filename in mmcv.scandir(
    #     osp.join(data_root, ann_dir), suffix='.png')]
    # with open(osp.join(data_root, split_dir, 'train.txt'), 'w') as f:
    #   # select first 4/5 as train set
    #   train_length = int(len(filename_list)*4/5)
    #   f.writelines(line + '\n' for line in filename_list[:train_length])
    # with open(osp.join(data_root, split_dir, 'val.txt'), 'w') as f:
    #   # select last 1/5 as train set
    #   f.writelines(line + '\n' for line in filename_list[train_length:])

    from mmseg.datasets.builder import DATASETS
    from mmseg.datasets.custom import CustomDataset

    @DATASETS.register_module()
    class StandfordBackgroundDataset(CustomDataset):
      CLASSES = classes
      PALETTE = palette
      def __init__(self, split, **kwargs):
        super().__init__(img_suffix='.jpg', seg_map_suffix='.png',
                         split=split, **kwargs)
        assert osp.exists(self.img_dir) and self.split is not None

    from mmcv import Config
    cfg = Config.fromfile('configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py')

    from mmseg.apis import set_random_seed

    # Since we use ony one GPU, BN is used instead of SyncBN
    cfg.norm_cfg = dict(type='BN', requires_grad=True)
    cfg.model.backbone.norm_cfg = cfg.norm_cfg
    cfg.model.decode_head.norm_cfg = cfg.norm_cfg
    cfg.model.auxiliary_head.norm_cfg = cfg.norm_cfg
    # modify num classes of the model in decode/auxiliary head
    cfg.model.decode_head.num_classes = 8
    cfg.model.auxiliary_head.num_classes = 8

    # Modify dataset type and path
    cfg.dataset_type = 'StandfordBackgroundDataset'
    cfg.data_root = data_root

    cfg.data.samples_per_gpu = 8
    cfg.data.workers_per_gpu=8

    cfg.img_norm_cfg = dict(
        mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
    cfg.crop_size = (256, 256)
    cfg.train_pipeline = [
        dict(type='LoadImageFromFile'),
        dict(type='LoadAnnotations'),
        dict(type='Resize', img_scale=(320, 240), ratio_range=(0.5, 2.0)),
        dict(type='RandomCrop', crop_size=cfg.crop_size, cat_max_ratio=0.75),
        dict(type='RandomFlip', flip_ratio=0.5),
        dict(type='PhotoMetricDistortion'),
        dict(type='Normalize', **cfg.img_norm_cfg),
        dict(type='Pad', size=cfg.crop_size, pad_val=0, seg_pad_val=255),
        dict(type='DefaultFormatBundle'),
        dict(type='Collect', keys=['img', 'gt_semantic_seg']),
    ]

    cfg.test_pipeline = [
        dict(type='LoadImageFromFile'),
        dict(
            type='MultiScaleFlipAug',
            img_scale=(320, 240),
            # img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
            flip=False,
            transforms=[
                dict(type='Resize', keep_ratio=True),
                dict(type='RandomFlip'),
                dict(type='Normalize', **cfg.img_norm_cfg),
                dict(type='ImageToTensor', keys=['img']),
                dict(type='Collect', keys=['img']),
            ])
    ]

    cfg.data.train.type = cfg.dataset_type
    cfg.data.train.data_root = cfg.data_root
    cfg.data.train.img_dir = img_dir
    cfg.data.train.ann_dir = ann_dir
    cfg.data.train.pipeline = cfg.train_pipeline
    cfg.data.train.split = 'splits/train.txt'

    cfg.data.val.type = cfg.dataset_type
    cfg.data.val.data_root = cfg.data_root
    cfg.data.val.img_dir = img_dir
    cfg.data.val.ann_dir = ann_dir
    cfg.data.val.pipeline = cfg.test_pipeline
    cfg.data.val.split = 'splits/val.txt'

    cfg.data.test.type = cfg.dataset_type
    cfg.data.test.data_root = cfg.data_root
    cfg.data.test.img_dir = img_dir
    cfg.data.test.ann_dir = ann_dir
    cfg.data.test.pipeline = cfg.test_pipeline
    cfg.data.test.split = 'splits/val.txt'

    # We can still use the pre-trained Mask RCNN model though we do not need to
    # use the mask branch
    cfg.load_from = 'checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth'

    # Set up working dir to save files and logs.
    cfg.work_dir = './work_dirs/tutorial'

    cfg.total_iters = 200
    cfg.log_config.interval = 10
    cfg.evaluation.interval = 200
    cfg.checkpoint_config.interval = 200

    # Set seed to facitate reproducing the result
    cfg.seed = 0
    set_random_seed(0, deterministic=False)
    cfg.gpu_ids = range(1)

    # Let's have a look at the final config used for training
    print(f'Config:\n{cfg.pretty_text}')

    from mmseg.datasets import build_dataset
    from mmseg.models import build_segmentor
    from mmseg.apis import train_segmentor

    # Build the dataset
    datasets = [build_dataset(cfg.data.train)]

    # Build the detector
    # model = build_segmentor(cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg)
    train_cfg = cfg.get('train_cfg')
    test_cfg = cfg.get('test_cfg')
    model = build_segmentor(cfg.model, train_cfg=train_cfg, test_cfg=test_cfg)

    # Add an attribute for visualization convenience
    model.CLASSES = datasets[0].CLASSES

    # Create work_dir
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    train_segmentor(model, datasets, cfg, distributed=False, validate=True,
                    meta=dict())

    #
    # img = mmcv.imread('data/iccv09Data/images/6000124.jpg')
    #
    # model.cfg = cfg
    # result = inference_segmentor(model, img)
    # plt.figure(figsize=(8, 6))
    # show_result_pyplot(model, img, result, palette)

max_stars_repo_name: korhun/mmsegmentation

max_issues_repo_name: korhun/mmsegmentation

max_forks_repo_name: korhun/mmsegmentation


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

import h5py
import numpy as np
from imageio import imread

from src.config import parse_args

args = parse_args()

class MPII:
    """
    MPII dataset loader
    """

    def __init__(self):
        print("loading data...")
        tic = time.time()

        train_f = h5py.File(os.path.join(args.annot_dir, "train.h5"), "r")
        val_f = h5py.File(os.path.join(args.annot_dir, "valid.h5"), "r")

        self.t_center = train_f["center"][()]
        t_scale = train_f["scale"][()]
        t_part = train_f["part"][()]
        t_visible = train_f["visible"][()]
        t_normalize = train_f["normalize"][()]
        t_imgname = [None] * len(self.t_center)
        for i in range(len(self.t_center)):
            t_imgname[i] = train_f["imgname"][i].decode("UTF-8")

        self.v_center = val_f["center"][()]
        v_scale = val_f["scale"][()]
        v_part = val_f["part"][()]
        v_visible = val_f["visible"][()]
        v_normalize = val_f["normalize"][()]
        v_imgname = [None] * len(self.v_center)
        for i in range(len(self.v_center)):
            v_imgname[i] = val_f["imgname"][i].decode("UTF-8")

        self.center = np.append(self.t_center, self.v_center, axis=0)
        self.scale = np.append(t_scale, v_scale)
        self.part = np.append(t_part, v_part, axis=0)
        self.visible = np.append(t_visible, v_visible, axis=0)
        self.normalize = np.append(t_normalize, v_normalize)
        self.imgname = t_imgname + v_imgname

        print("Done (t={:0.2f}s)".format(time.time() - tic))

        self.num_examples_train, self.num_examples_val = self.getLength()

    def getLength(self):
        """
        get dataset length
        """
        return len(self.t_center), len(self.v_center)

    def setup_val_split(self):
        """
        get index for train and validation imgs
        index for validation images starts after that of train images
        so that loadImage can tell them apart
        """
        valid = [i + self.num_examples_train for i in range(self.num_examples_val)]
        train = [i for i in range(self.num_examples_train)]
        return np.array(train), np.array(valid)

    def get_img(self, idx):
        """
        get image
        """
        imgname = self.imgname[idx]
        path = os.path.join(args.img_dir, imgname)
        img = imread(path)
        return img

    def get_path(self, idx):
        """
        get image path
        """
        imgname = self.imgname[idx]
        path = os.path.join(args.img_dir, imgname)
        return path

    def get_kps(self, idx):
        """
        get key points
        """
        part = self.part[idx]
        visible = self.visible[idx]
        kp2 = np.insert(part, 2, visible, axis=1)
        kps = np.zeros((1, 16, 3))
        kps[0] = kp2
        return kps

    def get_normalized(self, idx):
        """
        get normalized value
        """
        n = self.normalize[idx]
        return n

    def get_center(self, idx):
        """
        get center of the person
        """
        c = self.center[idx]
        return c

    def get_scale(self, idx):
        """
        get scale of the person
        """
        s = self.scale[idx]
        return s

# Part reference
parts = {
    "mpii": [
        "rank",
        "rkne",
        "rhip",
        "lhip",
        "lkne",
        "lank",
        "pelv",
        "thrx",
        "neck",
        "head",
        "rwri",
        "relb",
        "rsho",
        "lsho",
        "lelb",
        "lwri",
    ]
}

flipped_parts = {"mpii": [5, 4, 3, 2, 1, 0, 6, 7, 8, 9, 15, 14, 13, 12, 11, 10]}

part_pairs = {"mpii": [[0, 5], [1, 4], [2, 3], [6], [7], [8], [9], [10, 15], [11, 14], [12, 13]]}

pair_names = {"mpii": ["ankle", "knee", "hip", "pelvis", "thorax", "neck", "head", "wrist", "elbow", "shoulder"]}

ext: py

max_stars_repo_name: leelige/mindspore

max_issues_repo_name: leelige/mindspore

max_forks_repo_name: leelige/mindspore

reason: import numpy


text:
import sys, os
import spacy
import pickle, h5py, json
import numpy as np
from tqdm import tqdm
from sklearn.metrics.pairwise import cosine_similarity

import argparse
sys.path.append(os.getcwd())
from utilities.dataset import Dictionary
from utilities import config

def filter_object_attributes(nlp, exp):
    doc = nlp(exp)
    object_attributes = []
    for token in doc:
        if token.pos_ == u'NUM' and token.dep_ == u'nummod' and token.head.pos_ == u'NOUN':
            obj = token.head.text
            attri = token.text
            key = str(obj) + ';' + str(attri)
            object_attributes.append(key)
        elif token.pos_ == u'ADJ' and token.dep_ == u'amod' and token.head.pos_ == u'NOUN':
            obj = token.head.text
            attri = token.text
            key = str(obj) + ';' + str(attri)
            object_attributes.append(key)
        elif token.pos_ == u'NOUN' and token.dep_ == u'compound' and token.head.pos_ == u'NOUN':
            obj = token.head.text
            attri = token.text
            key = str(obj) + ';' + str(attri)
            object_attributes.append(key)
    return object_attributes

def filter_objects(nlp, exp):
    obj_tokens = ""
    for k in ['question', 'answer']:
        doc = nlp(exp[k])
        for token in doc:
            if token.pos_ == u'NOUN':
                obj_tokens += " " + token.text
        print(exp[k])
    print(obj_tokens)
    return obj_tokens

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--exp', type=str, default='vqx', help='explanation from vqx-exp or vqa-qa')
    args = parser.parse_args()
    return args

def main():
    assert config.rcnn_output_size == 36
    args = parse_args()
    qid2hint = {}
    with h5py.File(config.rcnn_trainval_path, 'r') as hf:
        img_id2idx = {img_id: idx for idx, img_id in enumerate(hf.get('ids'))}
        cls_scores = np.array(hf.get('cls_score'))
        attr_scores = np.array(hf.get('attr_score'))

    obj_emb_dict = np.load(os.path.join(config.main_path, 'glove6b_init_objects_300d.npy'))
    attr_emb_dict = np.load(os.path.join(config.main_path, 'glove6b_init_attributes_300d.npy'))
    if args.exp == 'vqx':
        exp_emb_dict = np.load(os.path.join(config.main_path, 'glove6b_init_vqx_300d.npy'))
        exp_dict = Dictionary.load_from_file(os.path.join(config.main_path, 'exp_dictionary.json'))
        qid2exp = pickle.load(open(os.path.join(config.main_path, 'vqx_qid2exp.pkl'), 'rb'))
    else:
        exp_emb_dict = np.load(os.path.join(config.cache_root, 'glove6b_init_vqa_caption_300d.npy'))
        exp_dict = Dictionary.load_from_file(os.path.join(config.cache_root, 'vqa_caption_dictionary.json'))
        qid2exp = pickle.load(open(os.path.join(config.cache_root, 'vqa_qid2exp.pkl'), 'rb'))
    qids = list(qid2exp.keys())
    nlp = spacy.load('en_core_web_lg')

    mod = 1000 if config.version=='v2' else 10
    for i in tqdm(range(len(qids))):
        qid = qids[i]
        img_id = qid // mod
        img_idx = img_id2idx[img_id]
        clss = cls_scores[img_idx][:, 0].astype('int')
        attrs = attr_scores[img_idx][:, 0].astype('int')

        objs = obj_emb_dict[clss, :]
        atts = attr_emb_dict[attrs, :]

        exp = qid2exp[qid]
        if args.exp == 'vqx':
            tokens = exp_dict.tokenize(exp, False)
            object_attributes = filter_object_attributes(nlp, exp)
        else:
            tokens = exp_dict.tokenize(exp['question'], False)
            tokens += exp_dict.tokenize(exp['answer'], False)
            object_attributes = filter_object_attributes(nlp, exp['question'])
            object_attributes += filter_object_attributes(nlp, exp['answer'])
            # tokens = exp_dict.tokenize(filter_objects(nlp, exp), False)

        exp_emb = exp_emb_dict[tokens]
        obj_sim = cosine_similarity(objs, exp_emb)

        hint_score = np.zeros((36))
        # hint_score_attr = np.zeros((36))

        # mentioned_tokens = []
        # mentioned_attrs = []

        for j in range(36):
            for k in range(len(tokens)):
                if obj_sim[j, k] > 0.6:
                    if hint_score[j] <= obj_sim[j, k]:
                        hint_score[j] = obj_sim[j, k]
            for key in object_attributes:
                obj, attr = key.split(';')
                # print(key, exp)
                obj_token = exp_dict.tokenize(obj, False)[0]
                attr_token = exp_dict.tokenize(attr, False)[0]
                if cosine_similarity(exp_emb_dict[obj_token:obj_token + 1], objs[j:j + 1]) > 0.3:
                    if cosine_similarity(exp_emb_dict[attr_token:attr_token + 1], atts[j:j + 1]) > 0.3:
                        hint_score[j] = cosine_similarity(exp_emb_dict[attr_token:attr_token + 1], atts[j:j + 1])

        qid2hint[qid] = hint_score
    if args.exp == 'vqx':
        pickle.dump(qid2hint, open(os.path.join(config.main_path, 'vqx_hint.pkl'), 'wb'))
    else:
        pickle.dump(qid2hint, open(os.path.join(config.cache_root, 'qa_hint.pkl'), 'wb'))

if __name__ == '__main__':
    main()

max_stars_repo_name: BierOne/VQA-AttReg

max_issues_repo_name: BierOne/VQA-AttReg

max_forks_repo_name: BierOne/VQA-AttReg


text:
import openvr
import numpy as np
import time
import copy
from threading import Thread, Lock
from queue import Queue
import queue
import signal
import sys

from simple_pid import PID
import liegroups as lie
import transforms3d as tf3d

from manipulator_learning.sim.utils.general import convert_quat_tf_to_pb, convert_quat_pb_to_tf, TransformMat, q_convert
from manipulator_learning.learning.imitation.device_utils import Button

def openvr_pose_to_numpy(TrackedDevicePose_t):
    """ Convert a TrackedDevicePose_t to a 4x4 numpy array for its pose, a (3,) array for velocity,
    and another (3,) array for angular velocity"""
    pose_mat34_ovr = TrackedDevicePose_t.mDeviceToAbsoluteTracking
    vel_ovr = TrackedDevicePose_t.vVelocity
    avel_ovr = TrackedDevicePose_t.vAngularVelocity
    pose_mat = np.ndarray([4, 4])
    pose_mat[3, :4] = [0, 0, 0, 1]
    vel_vec = np.ndarray([3])
    avel_vec = np.ndarray([3])
    for i in range(3):
        pose_mat[i] = pose_mat34_ovr[i]
        vel_vec[i] = vel_ovr[i]
        avel_vec[i] = avel_ovr[i]
    return pose_mat, vel_vec, avel_vec

class VRTrackedDevice:
    def __init__(self, device_type, vr_obj, index, init_pose):
        self.device_type = device_type
        self.index = index
        self.vr = vr_obj
        self.ovr_pose = None  # set by update_pose
        self.np_pose = None  # set by update_pose
        self.np_vel = None  # set by update_pose
        self.np_avel = None  # set by update_pose
        self.update_pose(init_pose)

    def update_pose(self, TrackedDevicePose_t):
        """ Update openvr pose (TrackedDevicePose_t) as well as the numpy poses. This should be called
         instead of directly setting ovr_pose, np_pose, np_vel, or np_avel. """
        self.ovr_pose = TrackedDevicePose_t
        self.np_pose, self.np_vel, self.np_avel = openvr_pose_to_numpy(TrackedDevicePose_t)

class ViveController(VRTrackedDevice):
    def __init__(self, vr_obj, index, init_pose, button_hold_time=2.0):
        super().__init__(device_type=openvr.TrackedDeviceClass_Controller,
                         vr_obj=vr_obj, index=index, init_pose=init_pose)
        self.ovr_state = None  # set by update_buttons
        self.trackpad = [0, 0]
        self.trigger = 0
        self.button_state = dict(
            squeeze=False,
            trigger_bool=False,
            menu=False,
            trackpad_button=False,
            trackpad_up_click=False,  # these four are for a "D-pad" using the trackpad
            trackpad_down_click=False,
            trackpad_left_click=False,
            trackpad_right_click=False,
        )
        self.old_button_state = copy.deepcopy(self.button_state)
        self.button_edges_objs = dict()
        for k in self.button_state:
            self.button_edges_objs[k] = Button(hold_time_length=button_hold_time)
        self.thread_lock = Lock()

    def update_buttons(self, VRControllerState_t):
        """ Update openvr state as well as more readable states. Should only be used for devices that
        have buttons. """
        cur_time = time.time()
        self.ovr_state = VRControllerState_t
        self.trackpad[0] = VRControllerState_t.rAxis[0].x
        self.trackpad[1] = VRControllerState_t.rAxis[0].y
        self.trigger = VRControllerState_t.rAxis[1].x
        self.old_button_state = copy.deepcopy(self.button_state)
        button_str = format(VRControllerState_t.ulButtonPressed, '64b')
        self.button_state['menu'] = button_str[62] == '1'
        self.button_state['squeeze'] = button_str[61] == '1'
        self.button_state['trackpad_button'] = button_str[31] == '1'
        self.button_state['trigger_bool'] = button_str[30] == '1'
        min_trackpad_dpad = .5
        self.button_state['trackpad_right_click'] = self.button_state['trackpad_button'] \
                                                 and self.trackpad[0] > min_trackpad_dpad
        self.button_state['trackpad_left_click'] = self.button_state['trackpad_button'] \
                                                 and self.trackpad[0] < -min_trackpad_dpad
        self.button_state['trackpad_up_click'] = self.button_state['trackpad_button'] \
                                                 and self.trackpad[1] > min_trackpad_dpad
        self.button_state['trackpad_down_click'] = self.button_state['trackpad_button'] \
                                                 and self.trackpad[1] < -min_trackpad_dpad
        self.update_edges(cur_time)

    def update_edges(self, cur_time):
        for k in self.button_state:
            self.button_edges_objs[k].get_and_store_update(self.button_state[k], cur_time)

    def get_latest_button_data(self):
        ret_dict = dict()
        self.thread_lock.acquire()
        for k in self.button_state:
            if k == 'menu':
                continue  # only want in get_latest_menu_data
            ret_dict[k] = copy.deepcopy(self.button_edges_objs[k].stored_state)
            if any(ret_dict[k][edge] for edge in ret_dict[k].keys()):
                self.button_edges_objs[k].reset_state()
        self.thread_lock.release()
        return ret_dict

    def get_latest_menu_data(self):
        """ Called at a higher frequency by VRSteer object itself. """
        stored_menu_state = copy.deepcopy(self.button_edges_objs['menu'].stored_state)
        if any(stored_menu_state[edge] for edge in stored_menu_state.keys()):
            self.button_edges_objs['menu'].reset_state()
        return stored_menu_state

class VRSteer:
    def __init__(self, des_forward_axis, des_up_axis, headset_viewing=False, buttons_thread=True, vel_pid_p=5.0):
        self.vr = openvr.init(openvr.VRApplication_Other)

        # initialization borrowed from triad_openvr
        # Initializing object to hold indexes for various tracked objects
        self.object_names = {"Tracking Reference": [], "HMD": [], "Controller": [], "Tracker": []}
        self.devices = {}
        poses = self.__get_ovr_poses()
        # Iterate through the pose list to find the active devices and determine their type
        for i in range(openvr.k_unMaxTrackedDeviceCount):
            if poses[i].bPoseIsValid:
                device_class = self.vr.getTrackedDeviceClass(i)
                if device_class == openvr.TrackedDeviceClass_Controller:
                    device_name = "controller_" + str(len(self.object_names["Controller"]))
                    self.object_names["Controller"].append(device_name)
                    self.devices[device_name] = ViveController(self.vr, i, poses[i])
                elif device_class == openvr.TrackedDeviceClass_HMD:
                    device_name = "hmd_" + str(len(self.object_names["HMD"]))
                    self.object_names["HMD"].append(device_name)
                    self.devices[device_name] = VRTrackedDevice(device_class, self.vr, i, poses[i])
                elif device_class == openvr.TrackedDeviceClass_GenericTracker:
                    device_name = "tracker_" + str(len(self.object_names["Tracker"]))
                    self.object_names["Tracker"].append(device_name)
                    self.devices[device_name] = VRTrackedDevice(device_class, self.vr, i, poses[i])
                elif device_class == openvr.TrackedDeviceClass_TrackingReference:
                    device_name = "tracking_reference_" + str(len(self.object_names["Tracking Reference"]))
                    self.object_names["Tracking Reference"].append(device_name)
                    self.devices[device_name] = VRTrackedDevice(device_class, self.vr, i, poses[i])

        # various toggles for use in color blocks environment
        self.enter_toggle = False
        self.space_toggle = False
        self.d_pressed = False
        self.s_toggle = False
        self.menu_toggle = False
        self.d_pressed_start_time = None
        self.trackpad_down_click_toggle = False
        self.right_start_time = None
        self.enter_hold = False
        self.policy_play = False

        # needed for pose based control
        self.accept_control_toggle = False
        self._tool_ref_pose = None
        self._vr_ref_pose = None

        # set the desired forward and up axes
        self.pose = TransformMat(mat=np.eye(4))  # pose of the 0th controller
        self.vive_T_fix = np.eye(4)
        self.set_axes(des_forward_axis, des_up_axis)

        # for adding random uniform noise to the control outputs that manifests as a random walk
        self.random_walk_T = TransformMat(mat=np.eye(4))
        self.perturb_T = np.eye(4)
        self.last_control_noise = time.time()

        # for viewing in headset
        if headset_viewing:
            self.headset_pose = np.eye(4)

        # for giving user force feedback on hitting things
        self.force_feedback_dur = 0

        # for making a separate thread that polls the buttons faster and ensures presses aren't missed
        if buttons_thread:
            self.buttons_q = Queue()
            self.buttons_thread = Thread(target=self.buttons_worker, args=(self.buttons_q,))
            self.buttons_thread.daemon = True
            self.buttons_thread.start()
            signal.signal(signal.SIGINT, self.signal_handler)
        else:
            self.buttons_thread = None

        # for working the same as gamepadsteer and keyboardsteer
        self.btn_state = dict()

        self.vel_pid = PID(vel_pid_p, 0.0)

        self.last_grip_command = False

        self._last_des_pose = None

    def buttons_worker(self, q: Queue):
        while(True):
            try:
                q_data = q.get_nowait()
                if q_data == 'shutdown':
                    print('shutting down buttons thread in vr control')
                    break
            except queue.Empty:
                pass
            self.update_buttons()
            if self.force_feedback_dur > 0:
                self.trigger_haptic(self.force_feedback_dur)
            time.sleep(.01)  # 100Hz

    def signal_handler(self, sig, frame):
        self.shutdown_buttons_worker()
        sys.exit()

    def shutdown_buttons_worker(self):
        self.buttons_q.put('shutdown')
        self.buttons_thread.join()
        print('vr buttons thread shut down')

    def update_poses(self):
        """ Update all poses. Should be called once per iteration of a control loop. """
        poses = self.__get_ovr_poses()
        for dev in self.devices.keys():
            self.devices[dev].update_pose(poses[self.devices[dev].index])

        self.pose = TransformMat(mat=self.devices['controller_0'].np_pose)

    def update_buttons(self):
        """ Update all button states. Should be called once per iteration of a control loop if needed. """
        self.devices['controller_0'].thread_lock.acquire()
        for dev in self.devices.keys():
            if self.devices[dev].device_type == openvr.TrackedDeviceClass_Controller:
                index = self.devices[dev].index
                state = self.vr.getControllerState(index)
                assert (bool(state[0]) is True), "Cannot get state of %s. Is it off?" % dev
                self.devices[dev].update_buttons(state[1])

        self.btn_state = self.devices['controller_0'].button_state
        self.devices['controller_0'].thread_lock.release()

    def get_latest_button_state(self, dev='controller_0'):
        return self.devices[dev].get_latest_button_data()

    def process_events(self):
        """ Convenience function to give same functionality as KeyboardSteer and GamepadSteer"""
        self.update_poses()
        if self.buttons_thread is None:
            self.update_buttons()

    def __get_ovr_poses(self):
        """ Get the current openvr poses for all devices. """
        return self.vr.getDeviceToAbsoluteTrackingPose(openvr.TrackingUniverseStanding, 0,
                                                       openvr.k_unMaxTrackedDeviceCount)

    def set_axes(self, forward_axis, up_axis):
        """
        Set the desired forward axis and up axis of the frame to be controlled.

        rotation from the VR controller frame to the desired tool frame. I.e. the vr z axis points inward,
        but if the tool frame z axis points outward,  we need a 180 deg rotation about x to get that.

        If forward_axis is [0, 0, 1], and up_axis is [0, -1, 0], then this is the equivalent of
        tf3d.euler.euler2mat(np.pi, 0, 0, 'sxyz').

        :param forward_axis: The desired "forward movement" axis of the robot tool frame
        :param up_axis: The desired "upward movement" axis of the robot tool frame
        """
        # not 100% sure that this function is implemented correctly, but it seems okay.
        vive_forward_axis = [0, 0, -1]  # found from running various steam experiments
        # vive_forward_axis = [0, -1, -1]  # could use this as forward axis instead
        vive_up_axis = [0, 1, 0]
        # vive_up_axis = [0, 1, -1]

        vive_fa = np.array(vive_forward_axis) / np.linalg.norm(vive_forward_axis)
        vive_ua = np.array(vive_up_axis) / np.linalg.norm(vive_up_axis)
        vive_ra = np.cross(vive_fa, vive_ua)
        des_fa = np.array(forward_axis)
        des_ua = np.array(up_axis)
        des_ra = np.cross(des_fa, des_ua)
        vive_rot_mat = np.outer(des_fa, vive_fa) + np.outer(des_ua, vive_ua) + \
                       np.outer(des_ra, vive_ra)
        vive_T_fix = np.eye(4)
        vive_T_fix[:3, :3] = vive_rot_mat.T  # transpose to get the conversion in the direction we want

        # vive_T_fix = np.eye(4)
        # vive_T_fix[:3, :3] = tf3d.euler.euler2mat(np.pi, 0, 0, 'sxyz')

        self.vive_T_fix = vive_T_fix

    def reset_ref_poses(self):
        """ Allows an outside caller to cleanly reset ref poses. """
        self.devices['controller_0'].thread_lock.acquire()
        self._vr_ref_pose = None
        self._tool_ref_pose = None
        self._last_des_pose = None
        self.menu_toggle = False
        self.devices['controller_0'].thread_lock.release()

    def trigger_haptic(self, duration_micros=3999, axis=0):
        index = self.devices['controller_0'].index
        self.vr.triggerHapticPulse(index, axis, duration_micros)

    def move_robot(self, current_robot_pos,
                   current_robot_rot,
                   current_robot_base_pos=None,
                   current_robot_base_rot=None,
                   random_walk_noise=None,
                   drift_noise=None,
                   match_controller_orientation=True,
                   controller_index=0,
                   output_vel=False,
                   output_vel_p=10.0,
                   current_gripper_pos=None):
        """
         Move a robot using an HTC Vive Controller. Current robot pos and rot are only used for setting
         the initial frame by pressing the menu button.

        :param current_robot_pos: 3 element tuple, the current position of the robot in the robot control ref
            frame (probably either the base frame or the world frame).
        :param current_robot_rot: 4 element tuple, the current orientation (xyzw quat) of the robot in the
            robot control ref frame (probably either the base frame or the world frame).
        :param current_robot_base_pos: 3 element tuple, current base position, used for setting new ee relative
            positions. If None, all ee positions assumed to be in world frame.
        :param current_robot_base_rot: 4 element tuple, current base orietation (xyzw quat), used for setting new
        ee relative positions. If None, all ee positions assumed to be in world frame.
        :param controller_index: Controller index to use (if more than one detected)
        :param output_vel: Whether to output velocity instead of position.
        :return:
        """

        if self.buttons_thread is not None:
            self.devices['controller_0'].thread_lock.acquire()

        menu_state_dict = self.devices['controller_0'].get_latest_menu_data()
        if menu_state_dict['fe']:
            if self.menu_toggle:
                self.menu_toggle = False
            else:
                self.menu_toggle = True

        c_str = 'controller_' + str(controller_index)

        if self._tool_ref_pose is None:  # should only occur on first call
            self._tool_ref_pose = TransformMat(pb_pose=(current_robot_pos, current_robot_rot))
            self._last_des_pose = copy.deepcopy(self._tool_ref_pose)

        # Get tool pose in world frame
        if current_robot_base_pos is not None:
            current_robot_base_pose = TransformMat(pb_pose=(current_robot_base_pos, current_robot_base_rot))
            current_tool_world_pose = TransformMat(mat=np.dot(current_robot_base_pose(), self._tool_ref_pose()))
        else:
            current_tool_world_pose = self._tool_ref_pose

        if self.menu_toggle and not self.accept_control_toggle:
            # this ensures tool ref pose only gets set on the initial press of the menu button
            self.trigger_haptic()
            self.accept_control_toggle = True
            self._tool_ref_pose = TransformMat(pb_pose=(current_robot_pos, current_robot_rot))
            self._vr_ref_pose = TransformMat(mat=np.dot(self.pose(), self.vive_T_fix))
            self._tool_to_vr_ref_pose = TransformMat(mat=np.dot(self._tool_ref_pose(), self._vr_ref_pose()))
        elif not self.menu_toggle and self.accept_control_toggle:
            self.trigger_haptic()
            self.accept_control_toggle = False
            self._tool_ref_pose = TransformMat(pb_pose=(current_robot_pos, current_robot_rot))
            self.random_walk_T = TransformMat(mat=np.eye(4))

        if self.accept_control_toggle:
            vr_current_pose_mat = TransformMat(mat=self.pose())
            vr_fixed_pose = TransformMat(np.dot(vr_current_pose_mat(), self.vive_T_fix))
            if match_controller_orientation:
                # the desired pose should be set to the tool_ref + (vr_current - vr_ref)
                T_vr_ref_to_cur = TransformMat(mat=np.dot(np.linalg.inv(self._vr_ref_pose()), vr_fixed_pose()))
                if current_robot_base_pos is not None:
                    T_world_to_tool = TransformMat(mat=np.dot(current_robot_base_pose(), self._tool_ref_pose()))
                else:
                    T_world_to_tool = self._tool_ref_pose
                T_rob_ref_to_des = TransformMat(mat=np.dot(T_world_to_tool(), T_vr_ref_to_cur()))

                # add in random walk uniform noise if squeeze is being pressed
                if random_walk_noise is not None and self.devices[c_str].button_state['trackpad_down_click']:
                    c_noise_flat = np.array(random_walk_noise).flatten()
                    perturb_xi = np.random.uniform(low=-c_noise_flat, high=c_noise_flat, size=6)
                    perturb_T = lie.SE3.exp(perturb_xi).as_matrix()
                    self.random_walk_T = TransformMat(mat=np.dot(perturb_T, self.random_walk_T.pose_mat))
                    self.last_control_noise = time.time()

                # try with smooth drift added instead of random noise
                if drift_noise is not None and self.devices[c_str].button_state['trackpad_down_click']:
                    if self.trackpad_down_click_toggle:  # only set the drift values once
                        self.trackpad_down_click_toggle = False
                        c_noise_flat = np.array(drift_noise).flatten()
                        perturb_xi = np.random.uniform(low=-c_noise_flat, high=c_noise_flat, size=6)
                        self.perturb_T = lie.SE3.exp(perturb_xi).as_matrix()
                    self.random_walk_T = TransformMat(mat=np.dot(self.perturb_T, self.random_walk_T.pose_mat))
                    self.last_control_noise = time.time()

                # add in noise if it has changed, otherwise keep it the same as the last noise added output
                T_rob_ref_to_des = TransformMat(mat=np.dot(T_rob_ref_to_des.pose_mat, self.random_walk_T.pose_mat))

                pos, rot = T_rob_ref_to_des.to_pb()
            else:
                # get position by subtracting world coordinates of v0 from vi directly
                subtracted_pos = vr_fixed_pose()[:3, 3] - self._vr_ref_pose()[:3, 3]
                pos = current_tool_world_pose()[:3, 3] + subtracted_pos

                # rotations gotten based on C_w_ti = C_w_t0 * C_t0_v0 * C_v0_vi * C_v0_t0
                # equation can be simplified to C_w_vi * C_v0_w * C_w_t0
                C_w_vi = vr_fixed_pose()[:3, :3]
                C_v0_w = self._vr_ref_pose()[:3, :3].T
                C_w_t0 = current_tool_world_pose()[:3, :3]
                C_w_ti = C_w_vi @ C_v0_w @ C_w_t0
                quat = tf3d.quaternions.mat2quat(C_w_ti)
                rot = [quat[1], quat[2], quat[3], quat[0]]  # convert from wxyz to xyzw

                # uncomment to use position only
                # _, rot = TransformMat(mat=np.dot(current_robot_base_pose(), self._tool_ref_pose())).to_pb()
        else:
            # pos, rot = self._tool_ref_pose.to_pb()
            if current_robot_base_pos is not None:
                pos, rot = TransformMat(mat=np.dot(current_robot_base_pose(), self._tool_ref_pose())).to_pb()
            else:
                pos, rot = self._tool_ref_pose.to_pb()

        # safety checking
        last_pos, last_rot = self._last_des_pose.to_pb()
        if np.linalg.norm(np.array(pos) - np.array(last_pos)) > .08:
            print("Des pose greater than 8cm from cur pose, ignoring for safety.. is vr obscured?")
            print("Press menu to reenable")
            self.menu_toggle = False
            self.trigger_haptic()
            pos, rot = self._last_des_pose.to_pb()
        rot_wxyz = [rot[3], rot[0], rot[1], rot[2]]  # convert from xyzw to wxyz
        last_rot_wxyz = [last_rot[3], last_rot[0], last_rot[1], last_rot[2]]  # convert from xyzw to wxyz
        q_diff = tf3d.quaternions.qmult(tf3d.quaternions.qinverse(rot_wxyz), last_rot_wxyz)
        if q_diff[0] < 0:
            q_diff = -q_diff
        ax, ang = tf3d.quaternions.quat2axangle(q_diff)
        if abs(ang) > .7:
            print("Des orientation greater than .7 rad from cur, ignoring for safety.. is vr obscured?")
            print("Press menu to reenable")
            self.trigger_haptic()
            self.menu_toggle = False
            pos, rot = self._last_des_pose.to_pb()

        self._last_des_pose = TransformMat(pb_pose=(pos, rot))

        if self.accept_control_toggle:
            g_command = self.devices[c_str].button_state['trigger_bool']
        else:
            if current_gripper_pos is None:
                g_command = self.last_grip_command
            else:
                g_command = current_gripper_pos
        self.last_grip_command = g_command

        if output_vel:
            if rot[3] < 0:
                rot = -rot
            if current_robot_rot[3] < 0:
                current_robot_rot = -current_robot_rot
            rot = q_convert(rot, 'xyzw', 'wxyz')
            current_robot_rot = q_convert(current_robot_rot, 'xyzw', 'wxyz')

            self.vel_pid.setpoint = pos
            trans_vel = self.vel_pid(current_robot_pos)

            # if velocity is being given in tool frame
            # trans_vel = np.dot(tf3d.quaternions.quat2mat(current_robot_rot), trans_vel)

            q_diff = tf3d.quaternions.qmult(tf3d.quaternions.qinverse(current_robot_rot), rot)
            if q_diff[0] < 0:
                q_diff = -q_diff
            ax, ang = tf3d.quaternions.quat2axangle(q_diff)
            rot_vel = np.dot(tf3d.quaternions.quat2mat(current_robot_rot), ang * ax * output_vel_p)

            if self.buttons_thread is not None:
                self.devices['controller_0'].thread_lock.release()

            return trans_vel, rot_vel, g_command

        if self.buttons_thread is not None:
            self.devices['controller_0'].thread_lock.release()

        return pos, rot, g_command

    def _handle_button_toggles(self, controller_index):
        """ Handle button toggles the same as was done for KeyMouseSteer. Not a long term sol'n. """
        cur_time = time.time()
        hold_time = 2
        c_str = 'controller_' + str(controller_index)
        btn_state = self.devices[c_str].button_state
        old_btn_state = self.devices[c_str].old_button_state
        if btn_state['trackpad_right_click'] and not old_btn_state['trackpad_right_click']:
            self.right_start_time = cur_time
        if not btn_state['trackpad_right_click'] and old_btn_state['trackpad_right_click']:  # falling edge
            if not self.enter_toggle:
                self.enter_toggle = True
            else:
                self.enter_toggle = False
        if btn_state['trackpad_right_click'] and (cur_time - self.right_start_time > 2):
            self.enter_hold = True
        else:
            self.enter_hold = False
        if btn_state['trackpad_left_click'] and not old_btn_state['trackpad_left_click']:
            self.space_toggle = True
        if btn_state['trackpad_up_click'] and self.d_pressed_start_time is None and not \
                old_btn_state['trackpad_up_click']:
            self.d_pressed_start_time = cur_time
        if btn_state['trackpad_up_click'] and self.d_pressed_start_time is not None and \
                (cur_time - self.d_pressed_start_time) > 2:
            self.d_pressed = True
            self.d_pressed_start_time = None
        if not btn_state['trackpad_up_click']:
            self.d_pressed_start_time = None
        if btn_state['trackpad_up_click'] and not old_btn_state['trackpad_up_click']:
            if not self.policy_play:
                self.policy_play = True
            else:
                self.policy_play = False
        if btn_state['menu'] and not old_btn_state['menu']:
            if not self.menu_toggle:
                self.menu_toggle = True
            else:
                self.menu_toggle = False
        if btn_state['trackpad_down_click'] and not old_btn_state['trackpad_down_click']:
            if not self.trackpad_down_click_toggle:
                self.trackpad_down_click_toggle = True

if __name__ == '__main__':
    vrs = VRSteer((0, 0, 1), (0, -1, 0))

    test_tool_ref_rot = convert_quat_tf_to_pb(tf3d.euler.euler2quat(0, 0, 0, 'rxyz'))

    while 1:
        vrs.update_poses()
        vrs.update_buttons()
        des_pose = vrs.move_robot([1, 2, 3], test_tool_ref_rot)
        des_rot = tf3d.euler.quat2euler(convert_quat_pb_to_tf(des_pose[1]), 'rxyz')
        print("pos %.3f, %.3f, %.3f, rot %.3f, %.3f, %.3f, grip %s" % (*des_pose[0], *des_rot, str(des_pose[2])))
        time.sleep(.05)


text:
import numpy as np

######################################################
##Universal roation curve Persic1996
##with parameters for the MW from Reid2014
R0 = 8.31
V0 = 241
a1 = 241.
a2 = 0.9
a3 = 1.46

A = 1.97 * a2 ** -1.22 / (a2 ** -2 + 0.78 ** 2) ** 1.43
B = (1 + a3 ** 2) * a2 ** -2 / (a2 ** -2 + a3 ** 2)
beta = (V0 ** 2 / a1 ** 2 - B) / (A - B)

def V(R):
    x = R / (a2 * R0)
    val = beta * 1.97 * x ** 1.22 / (x ** 2 + 0.78 ** 2) ** 1.43
    val += (1 - beta) * (1 + a3 ** 2) * x ** 2 / (x ** 2 + a3 ** 2)
    val *= a1 ** 2
    val = np.sqrt(val)
    return val

########################################################

########################################################
##spiral arm model, model from Wainscoat 1992
##with rescuplting by Cordes2002
##translated from NE2001 fortran routine
Narm = 5
arms = [[4.25, 3.48, 0., 7.],
        [4.25, 3.48, 3.141, 6.],
        [4.89, 4.90, 2.525, 6.],
        [4.89, 3.76, 4.24, 8.],
        [4.57, 8.10, 5.847, 0.55]]

def rt2xy(rad, theta):
    x = -rad * np.sin(theta)
    y = rad * np.cos(theta)
    return x, y

def rt2lonvel(rad,theta):
    x, y = rt2xy(rad, theta)
    vel = V(rad)
    lon = (np.abs(y) < R0 + (np.abs(y) > R0) * (y < 0)) * np.arctan(x / (R0 - y)) + (np.abs(y) > R0) * (
        y > 0) * (np.sign(x) * np.pi - np.arctan(x / (y - R0)))
    vlsr = R0 * np.sin(lon) * (vel / rad - V0 / R0)
    lon = np.rad2deg(lon)
    return lon, vlsr

def rt2lond0(rad,theta):
    x, y = rt2xy(rad, theta)
    lon = (np.abs(y) < R0 + (np.abs(y) > R0) * (y < 0)) * np.arctan(x / (R0 - y)) + (np.abs(y) > R0) * (
        y > 0) * (np.sign(x) * np.pi - np.arctan(x / (y - R0)))
    lon = np.rad2deg(lon)
    dist = np.sqrt(np.power(x,2)+np.power(y-R0,2))
    return lon, dist

def lbd2vlsr(l,b,d):
    x, y = lbd2xy(l,b,d)
    rad = np.sqrt(x**2+y**2)
    vel = V(rad)
    vlsr = R0 * np.sin(np.deg2rad(l)) * (vel / rad - V0 / R0)
    return vlsr

def lbd2xy(l,b,d):
    d0 = d*np.cos(np.deg2rad(b))
    r = np.sqrt(R0**2+d0**2-2*R0*d0*np.cos(np.deg2rad(l)))
    x = d0*np.sin(np.deg2rad(l))
    y = np.sqrt(r**2-x**2)
    return x, y

def arm_polar(s):
    params = arms[s]
    rr = np.logspace(np.log10(params[1]),np.log10(30), 200)
    tt = params[0] * np.log(rr / params[1]) + params[2]
    rr = rr[tt < params[2] + params[3]]
    tt = tt[tt < params[2] + params[3]]
    if s == 1:
        rr[(tt > np.deg2rad(370)) & (tt < np.deg2rad(410))] *= 1. + 0.04 * np.cos(
            (tt[(tt > np.deg2rad(370)) & (tt < np.deg2rad(410))] - np.deg2rad(390.)) * 180. / 40.)
        rr[(tt > np.deg2rad(315)) & (tt < np.deg2rad(370))] *= 1. - 0.07 * np.cos(
            (tt[(tt > np.deg2rad(315)) & (tt < np.deg2rad(370))] - np.deg2rad(345.)) * 180. / 55.)
        rr[(tt > np.deg2rad(180)) & (tt < np.deg2rad(315))] *= 1. + 0.16 * np.cos(
            (tt[(tt > np.deg2rad(180)) & (tt < np.deg2rad(315))] - np.deg2rad(260.)) * 180. / 135.)
    elif s == 3:
        rr[(tt > np.deg2rad(290)) & (tt < np.deg2rad(395))] *= 1. -0.11 * np.cos(
            (tt[(tt > np.deg2rad(290)) & (tt < np.deg2rad(395))] - np.deg2rad(350.)) * 180. / 105.)
    #rr *= R0/8.5  # Wainscoat 1992 uses Rsun=8.5, why don't I need to correct if I change this value?
    return rr, tt

def arm_xy(s):
    rr, tt = arm_polar(s)
    x, y = rt2xy(rr, tt)
    return x, y

def arm_lv(s):
    rr, tt = arm_polar(s)
    lon,vlsr = rt2lonvel(rr,tt)

    return lon, vlsr

def arm_ld0(s):
    rr, tt = arm_polar(s)
    lon,dist = rt2lond0(rr,tt)

    return lon, dist


text:
from numpy import zeros
from numpy.random import random
from numpy import column_stack
from numpy import sort
from numpy import cos
from numpy import sin
from numpy import pi
from modules.helpers import random_points_in_circle

TWOPI = 2*pi

def get_rnd_circ(noise):
  def f(self):
    a = sort((random() + random(self.pnum))*TWOPI)
    rnd = self.xy + column_stack((cos(a), sin(a)))*self.size

    b = random(self.pnum)*TWOPI
    rad = random(size=(self.pnum,1))
    disp = column_stack((cos(b), sin(b)))*noise*rad
    return rnd, rnd+disp
  return f

def get_displaced_single(noise):
  def f(self):
    rnd = random_points_in_circle(
        self.pnum,
        self.xy[0],
        self.xy[1],
        self.size
        )
    a = random(self.pnum)*TWOPI
    rad = random(size=(self.pnum,1))
    disp = column_stack((cos(a), sin(a)))*noise*rad
    return rnd, rnd+disp
  return f

def get_connected():
  def f(self):
    rnd = random_points_in_circle(
        2*self.pnum-1,
        self.xy[0],
        self.xy[1],
        self.size
        )
    xy1 = rnd[:self.pnum,:]
    xy2 = zeros((self.pnum,2), 'float')
    xy2[0,:] = rnd[0,:]
    xy2[-1,:] = rnd[self.pnum-1,:]
    xy2[1:-1,:] = rnd[self.pnum:-1,:]
    return xy1, xy2
  return f

def get_displaced_multi(noise, num_offsets):
  from numpy import linspace

  def f(self):
    rnd = random_points_in_circle(
        self.pnum,
        self.xy[0],
        self.xy[1],
        self.size
        )

    columns = [rnd]

    aa = random(self.pnum)*TWOPI
    for s in linspace(noise/num_offsets, noise, num_offsets-1):
      # rad = random(size=(self.pnum,1))
      disp = column_stack((cos(aa), sin(aa)))*s
      columns.append(rnd+disp)
    return columns

  return f

max_stars_repo_name: bjornasm/sand-creatures

max_issues_repo_name: bjornasm/sand-creatures

max_forks_repo_name: bjornasm/sand-creatures

reason: from numpy


text:
"""
Python script for training a model version
"""
# Core Packages
import os
import json

# Third Party
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.model_selection import StratifiedKFold, cross_validate, cross_val_score
from sklearn.calibration import calibration_curve
from sklearn import metrics
import utils.credit as utils

# Bedrock and Boxkite
import bdrk
from bdrk.model_analyzer import ModelAnalyzer, ModelTypes
from boxkite.monitoring.service import ModelMonitoringService
import pickle
import logging

# ---------------------------------
# Constants
# ---------------------------------

ARTEFACT_PATH = "/artefact/model.pkl"

FAI_CONFIG = {
    'SEX': {
        'privileged_attribute_values': [1],
        'privileged_group_name': 'Male',  # privileged group name corresponding to values=[1]
        'unprivileged_attribute_values': [2],
        'unprivileged_group_name': 'Female',  # unprivileged group name corresponding to values=[0]
    }
}

# Retrieving environmental params from Bedrock HCL
DATA_DIR_LOCAL = os.getenv("DATA_DIR_LOCAL")
SEED = int(os.getenv("SEED"))
TH = float(os.getenv("TH"))
LR_REGULARIZER = float(os.getenv("LR_REGULARIZER"))
RF_N_ESTIMATORS = float(os.getenv("RF_N_ESTIMATORS"))
CB_ITERATIONS = float(os.getenv("CB_ITERATIONS"))
CLASSIFIER = os.getenv("CLASSIFIER")

# ---------------------------------
# Bedrock functions
# ---------------------------------

def compute_log_metrics(model, x_train,
                        x_test, y_test,
                        best_th=0.5,
                        model_name="tree_model",
                        model_type=ModelTypes.TREE,
                        fai_config=None,
                        artefact_path="/artefact/model.pkl"):
    """Compute and log metrics."""
    test_prob = model.predict_proba(x_test)[:, 1]
    test_pred = np.where(test_prob > best_th, 1, 0)

    acc = metrics.accuracy_score(y_test, test_pred)
    precision = metrics.precision_score(y_test, test_pred)
    recall = metrics.recall_score(y_test, test_pred)
    f1_score = metrics.f1_score(y_test, test_pred)
    roc_auc = metrics.roc_auc_score(y_test, test_prob)
    avg_prc = metrics.average_precision_score(y_test, test_prob)
    print("Evaluation\n"
          f"  Accuracy          = {acc:.4f}\n"
          f"  Precision         = {precision:.4f}\n"
          f"  Recall            = {recall:.4f}\n"
          f"  F1 score          = {f1_score:.4f}\n"
          f"  ROC AUC           = {roc_auc:.4f}\n"
          f"  Average precision = {avg_prc:.4f}")

    # --- Bedrock-native Integrations ---
    # Bedrock client library (bdrk): captures model metrics
    bdrk.init()
    with bdrk.start_run():
        # Log binary classifier metrics into a chart
        bdrk.log_binary_classifier_metrics(y_test.astype(int).tolist(),
                                              test_prob.flatten().tolist())

        # Log key-value pairs
        bdrk.log_metric("Accuracy", acc)
        bdrk.log_metric("Precision", precision)
        bdrk.log_metric("Recall", recall)
        bdrk.log_metric("F1 score", f1_score)
        bdrk.log_metric("ROC AUC", roc_auc)
        bdrk.log_metric("Avg precision", avg_prc)

        # Bedrock Model Analyzer: generates model explainability and fairness metrics
        # Analyzer (optional): generate explainability metrics
        analyzer = ModelAnalyzer(model[1],
                                 model_name=model_name,
                                 model_type=model_type)\
                        .train_features(x_train)\
                        .test_features(x_test)

        # Analyzer (optional): generate fairness metrics
        analyzer.fairness_config(fai_config)\
            .test_labels(y_test)\
            .test_inference(test_pred)

        # Run Model Analyzer
        # Returns a tuple of (shap_values, base_shap_values, global_explainability, fairness_metrics)
        analyzer.analyze()

        # IMPORTANT: Saving the Model Artefact to Bedrock
        with open(artefact_path, "wb") as model_file:
            pickle.dump(model, model_file)
        bdrk.log_model(artefact_path)

    return True

def main():
    # Extraneous columns (as might be determined through feature selection)
    drop_cols = ['ID']

    # --- Data ETL ---
    # Load into pandas dataframes
    # x_<name> : features
    # y_<name> : labels
    x_train, y_train = utils.load_dataset(os.path.join(DATA_DIR_LOCAL, 'creditdata_train_v2.csv'), drop_columns=drop_cols)
    x_test, y_test = utils.load_dataset(os.path.join(DATA_DIR_LOCAL, 'creditdata_test_v2.csv'), drop_columns=drop_cols)

    # --- Candidate Binary Classification Algos ---
    assert(CLASSIFIER in ["LR","RF","CB"])
    if CLASSIFIER == "LR":
        # MODEL 1: LOGISTIC REGRESSION
        # Use best parameters from a model selection and threshold tuning process
        model = utils.train_log_reg_model(x_train, y_train, seed=SEED, C=LR_REGULARIZER, upsample=True, verbose=True)
        model_name = "logreg_model"
        model_type = ModelTypes.LINEAR
    elif CLASSIFIER == "RF":
        # MODEL 2: RANDOM FOREST
        # Uses default threshold of 0.5 and model parameters
        model = utils.train_rf_model(x_train, y_train, seed=SEED, upsample=True, verbose=True)
        model_name = "randomforest_model"
        model_type = ModelTypes.TREE
    elif CLASSIFIER == "CB":
        # MODEL 3: CATBOOST
        # Uses default threshold of 0.5 and model parameters
        model = utils.train_catboost_model(x_train, y_train, seed=SEED, upsample=True, verbose=True)
        model_name = "catboost_model"
        model_type = ModelTypes.TREE

    # Compute and log metrics
    # --- Includes Bedrock-native Integrations ---
    compute_log_metrics(model=model,
                        x_train=x_train,
                        x_test=x_test,
                        y_test=y_test,
                        best_th=TH,
                        model_name=model_name,
                        model_type=model_type,
                        fai_config=FAI_CONFIG,
                        artefact_path=ARTEFACT_PATH)

    # Bedrock Model Monitoring: pre-requisite for monitoring concept drift
    # Prepare the inference probabilities
    train_prob = model.predict_proba(x_train)[:, 1]
    train_pred = np.where(train_prob > TH, 1, 0)

    # This step initialises the distribution from model training
    ModelMonitoringService.export_text(
        features=x_train.iteritems(),
        inference=train_prob.tolist(),
    )
    # --- End of Bedrock-native Integrations ---
    print("Done!")

if __name__ == "__main__":
    try:
        print("Hello Bedrock!")
        main()
    except Exception as e:
        print(e)
        print("Something went wrong...")

max_stars_repo_name: victor-basisai/credit-scoring

max_issues_repo_name: victor-basisai/credit-scoring

max_forks_repo_name: victor-basisai/credit-scoring


text:
import torch
import numpy as np
from onpolicy.utils.util import get_shape_from_obs_space, get_shape_from_act_space

def _flatten(T, N, x):
    return x.reshape(T * N, *x.shape[2:])

def _cast(x):
    return x.transpose(1, 2, 0, 3).reshape(-1, *x.shape[3:])

class SharedReplayBuffer(object):
    """
    Buffer to store training data.
    :param args: (argparse.Namespace) arguments containing relevant model, policy, and env information.
    :param num_agents: (int) number of agents in the env.
    :param obs_space: (gym.Space) observation space of agents.
    :param cent_obs_space: (gym.Space) centralized observation space of agents.
    :param act_space: (gym.Space) action space for agents.
    """

    def __init__(self, rspo_iter, args, num_agents, obs_space, cent_obs_space,
                 act_space):
        self.episode_length = args.episode_length
        self.n_rollout_threads = args.n_rollout_threads
        self.hidden_size = args.hidden_size
        self.recurrent_N = args.recurrent_N
        self.gamma = args.gamma
        self.gae_lambda = args.gae_lambda
        self._use_gae = args.use_gae
        self._use_popart = args.use_popart
        self._use_valuenorm = args.use_valuenorm
        self._use_proper_time_limits = args.use_proper_time_limits

        obs_shape = get_shape_from_obs_space(obs_space)
        share_obs_shape = get_shape_from_obs_space(cent_obs_space)

        if type(obs_shape[-1]) == list:
            obs_shape = obs_shape[:1]

        if type(share_obs_shape[-1]) == list:
            share_obs_shape = share_obs_shape[:1]

        self.share_obs = np.zeros(
            (self.episode_length + 1, self.n_rollout_threads, num_agents,
             *share_obs_shape),
            dtype=np.float32)
        self.obs = np.zeros((self.episode_length + 1, self.n_rollout_threads,
                             num_agents, *obs_shape),
                            dtype=np.float32)

        self.rnn_states = np.zeros(
            (self.episode_length + 1, self.n_rollout_threads, num_agents,
             self.recurrent_N, self.hidden_size),
            dtype=np.float32)
        self.critic_dim = rspo_iter + 1
        self.rnn_states_critic = np.zeros(
            (self.episode_length + 1, self.n_rollout_threads, num_agents,
             self.recurrent_N, self.hidden_size * self.critic_dim),
            dtype=np.float32)

        self.value_preds = np.zeros(
            (self.episode_length + 1, self.n_rollout_threads, num_agents,
             1 + rspo_iter),
            dtype=np.float32)
        self.returns = np.zeros_like(self.value_preds)

        if act_space.__class__.__name__ == 'Discrete':
            self.available_actions = np.ones(
                (self.episode_length + 1, self.n_rollout_threads, num_agents,
                 act_space.n),
                dtype=np.float32)
        else:
            self.available_actions = None

        act_shape = get_shape_from_act_space(act_space)

        self.actions = np.zeros((self.episode_length, self.n_rollout_threads,
                                 num_agents, act_shape),
                                dtype=np.float32)
        self.action_log_probs = np.zeros(
            (self.episode_length, self.n_rollout_threads, num_agents,
             act_shape),
            dtype=np.float32)
        self.rewards = np.zeros((self.episode_length, self.n_rollout_threads,
                                 num_agents, 1 + rspo_iter),
                                dtype=np.float32)

        self.masks = np.ones(
            (self.episode_length + 1, self.n_rollout_threads, num_agents, 1),
            dtype=np.float32)
        self.bad_masks = np.ones_like(self.masks)
        self.active_masks = np.ones_like(self.masks)

        self.step = 0

    def insert(self,
               share_obs,
               obs,
               rnn_states_actor,
               rnn_states_critic,
               actions,
               action_log_probs,
               value_preds,
               rewards,
               masks,
               bad_masks=None,
               active_masks=None,
               available_actions=None):
        """
        Insert data into the buffer.
        :param share_obs: (argparse.Namespace) arguments containing relevant model, policy, and env information.
        :param obs: (np.ndarray) local agent observations.
        :param rnn_states_actor: (np.ndarray) RNN states for actor network.
        :param rnn_states_critic: (np.ndarray) RNN states for critic network.
        :param actions:(np.ndarray) actions taken by agents.
        :param action_log_probs:(np.ndarray) log probs of actions taken by agents
        :param value_preds: (np.ndarray) value function prediction at each step.
        :param rewards: (np.ndarray) reward collected at each step.
        :param masks: (np.ndarray) denotes whether the environment has terminated or not.
        :param bad_masks: (np.ndarray) action space for agents.
        :param active_masks: (np.ndarray) denotes whether an agent is active or dead in the env.
        :param available_actions: (np.ndarray) actions available to each agent. If None, all actions are available.
        """
        self.share_obs[self.step + 1] = share_obs.copy()
        self.obs[self.step + 1] = obs.copy()
        self.rnn_states[self.step + 1] = rnn_states_actor.copy()
        self.rnn_states_critic[self.step + 1] = rnn_states_critic.copy()
        self.actions[self.step] = actions.copy()
        self.action_log_probs[self.step] = action_log_probs.copy()
        self.value_preds[self.step] = value_preds.copy()
        self.rewards[self.step] = rewards.copy()
        self.masks[self.step + 1] = masks.copy()
        if bad_masks is not None:
            self.bad_masks[self.step + 1] = bad_masks.copy()
        if active_masks is not None:
            self.active_masks[self.step + 1] = active_masks.copy()
        if available_actions is not None:
            self.available_actions[self.step + 1] = available_actions.copy()

        self.step = (self.step + 1) % self.episode_length

    def after_update(self):
        """Copy last timestep data to first index. Called after update to model."""
        self.share_obs[0] = self.share_obs[-1].copy()
        self.obs[0] = self.obs[-1].copy()
        self.rnn_states[0] = self.rnn_states[-1].copy()
        self.rnn_states_critic[0] = self.rnn_states_critic[-1].copy()
        self.masks[0] = self.masks[-1].copy()
        self.bad_masks[0] = self.bad_masks[-1].copy()
        self.active_masks[0] = self.active_masks[-1].copy()
        if self.available_actions is not None:
            self.available_actions[0] = self.available_actions[-1].copy()

    def compute_returns(self, next_value, value_normalizers=None):
        """
        Compute returns either as discounted sum of rewards, or using GAE.
        :param next_value: (np.ndarray) value predictions for the step after the last episode step.
        :param value_normalizer: (PopArt) If not None, PopArt value normalizer instance.
        """
        if value_normalizers is not None:
            denormalized_values = []
            value_splits = np.split(self.value_preds,
                                    self.value_preds.shape[-1],
                                    axis=-1)
            for v, vn in zip(value_splits, value_normalizers):
                denormalized_values.append(vn.denormalize(v))
            denormalized_values = np.concatenate(denormalized_values, -1)
        if self._use_proper_time_limits:
            if self._use_gae:
                self.value_preds[-1] = next_value
                gae = 0
                for step in reversed(range(self.rewards.shape[0])):
                    if self._use_popart or self._use_valuenorm:
                        # step + 1
                        delta = self.rewards[step] + self.gamma * denormalized_values[step + 1] * self.masks[step + 1] \
                                - denormalized_values[step]
                        gae = delta + self.gamma * self.gae_lambda * gae * self.masks[
                            step + 1]
                        gae = gae * self.bad_masks[step + 1]
                        self.returns[step] = gae + denormalized_values[step]
                    else:
                        delta = self.rewards[step] + self.gamma * self.value_preds[step + 1] * self.masks[step + 1] - \
                                self.value_preds[step]
                        gae = delta + self.gamma * self.gae_lambda * self.masks[
                            step + 1] * gae
                        gae = gae * self.bad_masks[step + 1]
                        self.returns[step] = gae + self.value_preds[step]
            else:
                self.returns[-1] = next_value
                for step in reversed(range(self.rewards.shape[0])):
                    if self._use_popart or self._use_valuenorm:
                        self.returns[step] = (self.returns[step + 1] * self.gamma * self.masks[step + 1] + self.rewards[
                            step]) * self.bad_masks[step + 1] \
                                             + (1 - self.bad_masks[step + 1]) * denormalized_values[step]
                    else:
                        self.returns[step] = (self.returns[step + 1] * self.gamma * self.masks[step + 1] + self.rewards[
                            step]) * self.bad_masks[step + 1] \
                                             + (1 - self.bad_masks[step + 1]) * self.value_preds[step]
        else:
            if self._use_gae:
                self.value_preds[-1] = next_value
                gae = 0
                for step in reversed(range(self.rewards.shape[0])):
                    if self._use_popart or self._use_valuenorm:
                        delta = self.rewards[step] + self.gamma * denormalized_values[step+1] * self.masks[step + 1] \
                                - denormalized_values[step]
                        gae = delta + self.gamma * self.gae_lambda * self.masks[
                            step + 1] * gae
                        self.returns[step] = gae + denormalized_values[step]
                    else:
                        delta = self.rewards[step] + self.gamma * self.value_preds[step + 1] * self.masks[step + 1] - \
                                self.value_preds[step]
                        gae = delta + self.gamma * self.gae_lambda * self.masks[
                            step + 1] * gae
                        self.returns[step] = gae + self.value_preds[step]
            else:
                self.returns[-1] = next_value
                for step in reversed(range(self.rewards.shape[0])):
                    self.returns[step] = self.returns[
                        step + 1] * self.gamma * self.masks[
                            step + 1] + self.rewards[step]

    def feed_forward_generator(self,
                               advantages,
                               interpolation_mask_=None,
                               num_mini_batch=None,
                               mini_batch_size=None):
        """
        Yield training data for MLP policies.
        :param advantages: (np.ndarray) advantage estimates.
        :param num_mini_batch: (int) number of minibatches to split the batch into.
        :param mini_batch_size: (int) number of samples in each minibatch.
        """
        episode_length, n_rollout_threads, num_agents = self.rewards.shape[0:3]
        batch_size = n_rollout_threads * episode_length * num_agents

        if mini_batch_size is None:
            assert batch_size >= num_mini_batch, (
                "PPO requires the number of processes ({}) "
                "* number of steps ({}) * number of agents ({}) = {} "
                "to be greater than or equal to the number of PPO mini batches ({})."
                "".format(n_rollout_threads, episode_length, num_agents,
                          n_rollout_threads * episode_length * num_agents,
                          num_mini_batch))
            mini_batch_size = batch_size // num_mini_batch

        rand = torch.randperm(batch_size).numpy()
        sampler = [
            rand[i * mini_batch_size:(i + 1) * mini_batch_size]
            for i in range(num_mini_batch)
        ]

        share_obs = self.share_obs[:-1].reshape(-1, *self.share_obs.shape[3:])
        obs = self.obs[:-1].reshape(-1, *self.obs.shape[3:])
        rnn_states = self.rnn_states[:-1].reshape(-1,
                                                  *self.rnn_states.shape[3:])
        rnn_states_critic = self.rnn_states_critic[:-1].reshape(
            -1, *self.rnn_states_critic.shape[3:])
        actions = self.actions.reshape(-1, self.actions.shape[-1])
        if self.available_actions is not None:
            available_actions = self.available_actions[:-1].reshape(
                -1, self.available_actions.shape[-1])
        value_preds = self.value_preds[:-1].reshape(-1,
                                                    self.value_preds.shape[-1])
        returns = self.returns[:-1].reshape(-1, self.returns.shape[-1])
        masks = self.masks[:-1].reshape(-1, 1)
        active_masks = self.active_masks[:-1].reshape(-1, 1)
        action_log_probs = self.action_log_probs.reshape(
            -1, self.action_log_probs.shape[-1])
        advantages = advantages.reshape(-1, 1)
        if interpolation_mask_ is not None:
            interpolation_mask = interpolation_mask_.reshape(
                -1, interpolation_mask_.shape[-1])

        for indices in sampler:
            # obs size [T+1 N M Dim]-->[T N M Dim]-->[T*N*M,Dim]-->[index,Dim]
            share_obs_batch = share_obs[indices]
            obs_batch = obs[indices]
            rnn_states_batch = rnn_states[indices]
            rnn_states_critic_batch = rnn_states_critic[indices]
            actions_batch = actions[indices]
            if self.available_actions is not None:
                available_actions_batch = available_actions[indices]
            else:
                available_actions_batch = None
            value_preds_batch = value_preds[indices]
            return_batch = returns[indices]
            masks_batch = masks[indices]
            active_masks_batch = active_masks[indices]
            old_action_log_probs_batch = action_log_probs[indices]
            if advantages is None:
                adv_targ = None
            else:
                adv_targ = advantages[indices]
            if interpolation_mask_ is not None:
                interpolation_mask_batch = interpolation_mask[indices]
            else:
                interpolation_mask_batch = None

            yield share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch,\
                  value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch,\
                  adv_targ, available_actions_batch, interpolation_mask_batch

    def recurrent_generator(self, advantages, interpolation_mask_,
                            num_mini_batch, data_chunk_length):
        """
        Yield training data for chunked RNN training.
        :param advantages: (np.ndarray) advantage estimates.
        :param num_mini_batch: (int) number of minibatches to split the batch into.
        :param data_chunk_length: (int) length of sequence chunks with which to train RNN.
        """
        episode_length, n_rollout_threads, num_agents = self.rewards.shape[0:3]
        batch_size = n_rollout_threads * episode_length * num_agents
        data_chunks = batch_size // data_chunk_length  # [C=r*T*M/L]
        mini_batch_size = data_chunks // num_mini_batch

        rand = torch.randperm(data_chunks).numpy()
        sampler = [
            rand[i * mini_batch_size:(i + 1) * mini_batch_size]
            for i in range(num_mini_batch)
        ]

        if len(self.share_obs.shape) > 4:
            share_obs = self.share_obs[:-1].transpose(
                1, 2, 0, 3, 4, 5).reshape(-1, *self.share_obs.shape[3:])
            obs = self.obs[:-1].transpose(1, 2, 0, 3, 4,
                                          5).reshape(-1, *self.obs.shape[3:])
        else:
            share_obs = _cast(self.share_obs[:-1])
            obs = _cast(self.obs[:-1])

        actions = _cast(self.actions)
        action_log_probs = _cast(self.action_log_probs)
        advantages = _cast(advantages)
        value_preds = _cast(self.value_preds[:-1])
        returns = _cast(self.returns[:-1])
        masks = _cast(self.masks[:-1])
        active_masks = _cast(self.active_masks[:-1])
        # rnn_states = _cast(self.rnn_states[:-1])
        # rnn_states_critic = _cast(self.rnn_states_critic[:-1])
        rnn_states = self.rnn_states[:-1].transpose(1, 2, 0, 3, 4).reshape(
            -1, *self.rnn_states.shape[3:])
        rnn_states_critic = self.rnn_states_critic[:-1].transpose(
            1, 2, 0, 3, 4).reshape(-1, *self.rnn_states_critic.shape[3:])

        if self.available_actions is not None:
            available_actions = _cast(self.available_actions[:-1])

        if interpolation_mask_ is not None:
            interpolation_mask = _cast(interpolation_mask_)

        for indices in sampler:
            share_obs_batch = []
            obs_batch = []
            rnn_states_batch = []
            rnn_states_critic_batch = []
            actions_batch = []
            available_actions_batch = []
            value_preds_batch = []
            return_batch = []
            interpolation_mask_batch = []
            masks_batch = []
            active_masks_batch = []
            old_action_log_probs_batch = []
            adv_targ = []

            for index in indices:

                ind = index * data_chunk_length
                # size [T+1 N M Dim]-->[T N M Dim]-->[N,M,T,Dim]-->[N*M*T,Dim]-->[L,Dim]
                share_obs_batch.append(share_obs[ind:ind + data_chunk_length])
                obs_batch.append(obs[ind:ind + data_chunk_length])
                actions_batch.append(actions[ind:ind + data_chunk_length])
                if self.available_actions is not None:
                    available_actions_batch.append(
                        available_actions[ind:ind + data_chunk_length])
                value_preds_batch.append(value_preds[ind:ind +
                                                     data_chunk_length])
                return_batch.append(returns[ind:ind + data_chunk_length])
                if interpolation_mask_ is not None:
                    interpolation_mask_batch.append(
                        interpolation_mask[ind:ind + data_chunk_length])
                masks_batch.append(masks[ind:ind + data_chunk_length])
                active_masks_batch.append(active_masks[ind:ind +
                                                       data_chunk_length])
                old_action_log_probs_batch.append(
                    action_log_probs[ind:ind + data_chunk_length])
                adv_targ.append(advantages[ind:ind + data_chunk_length])
                # size [T+1 N M Dim]-->[T N M Dim]-->[N M T Dim]-->[N*M*T,Dim]-->[1,Dim]
                rnn_states_batch.append(rnn_states[ind])
                rnn_states_critic_batch.append(rnn_states_critic[ind])

            L, N = data_chunk_length, mini_batch_size

            # These are all from_numpys of size (L, N, Dim)
            share_obs_batch = np.stack(share_obs_batch, axis=1)
            obs_batch = np.stack(obs_batch, axis=1)

            actions_batch = np.stack(actions_batch, axis=1)
            if self.available_actions is not None:
                available_actions_batch = np.stack(available_actions_batch,
                                                   axis=1)
            value_preds_batch = np.stack(value_preds_batch, axis=1)
            return_batch = np.stack(return_batch, axis=1)
            if interpolation_mask_ is not None:
                interpolation_mask_batch = np.stack(interpolation_mask_batch,
                                                    axis=1)
            masks_batch = np.stack(masks_batch, axis=1)
            active_masks_batch = np.stack(active_masks_batch, axis=1)
            old_action_log_probs_batch = np.stack(old_action_log_probs_batch,
                                                  axis=1)
            adv_targ = np.stack(adv_targ, axis=1)

            # States is just a (N, -1) from_numpy
            rnn_states_batch = np.stack(rnn_states_batch).reshape(
                N, *self.rnn_states.shape[3:])
            rnn_states_critic_batch = np.stack(
                rnn_states_critic_batch).reshape(
                    N, *self.rnn_states_critic.shape[3:])

            # Flatten the (L, N, ...) from_numpys to (L * N, ...)
            share_obs_batch = _flatten(L, N, share_obs_batch)
            obs_batch = _flatten(L, N, obs_batch)
            actions_batch = _flatten(L, N, actions_batch)
            if self.available_actions is not None:
                available_actions_batch = _flatten(L, N,
                                                   available_actions_batch)
            else:
                available_actions_batch = None
            value_preds_batch = _flatten(L, N, value_preds_batch)
            return_batch = _flatten(L, N, return_batch)
            if interpolation_mask_ is not None:
                interpolation_mask_batch = _flatten(L, N,
                                                    interpolation_mask_batch)
            else:
                interpolation_mask_batch = None
            masks_batch = _flatten(L, N, masks_batch)
            active_masks_batch = _flatten(L, N, active_masks_batch)
            old_action_log_probs_batch = _flatten(L, N,
                                                  old_action_log_probs_batch)
            adv_targ = _flatten(L, N, adv_targ)

            yield share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch,\
                  value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch,\
                  adv_targ, available_actions_batch, interpolation_mask_batch


text:
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

class BCELoss(nn.Module):
    def __init__(self,  weight=None, ignore_index=-100, reduce=None, reduction='mean'):
        super(BCELoss, self).__init__()
        self.reduction = reduction
    def forward(self, inputs, targets, mask):
        pos_id = (mask==1.0).float()
        neg_id = (mask==0.0).float()
        pos_loss = -pos_id * (targets * torch.log(inputs + 1e-14) + (1 - targets) * torch.log(1.0 - inputs + 1e-14))
        neg_loss = -neg_id * torch.log(1.0 - inputs + 1e-14)
        if self.reduction == 'mean':
            pos_loss = torch.mean(torch.sum(pos_loss, 1))
            neg_loss = torch.mean(torch.sum(neg_loss, 1))
            return pos_loss, neg_loss
        else:
            return pos_loss, neg_loss

class MSELoss(nn.Module):
    def __init__(self,  weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean'):
        super(MSELoss, self).__init__()
        self.reduction = reduction
    def forward(self, inputs, targets, mask):
        # We ignore those whose tarhets == -1.0.
        pos_id = (mask==1.0).float()
        neg_id = (mask==0.0).float()
        pos_loss = pos_id * (inputs - targets)**2
        neg_loss = neg_id * (inputs)**2
        if self.reduction == 'mean':
            pos_loss = torch.mean(torch.sum(pos_loss, 1))
            neg_loss = torch.mean(torch.sum(neg_loss, 1))
            return pos_loss, neg_loss
        else:
            return pos_loss, neg_loss

def generate_anchor(input_size, stride, anchor_scale, anchor_aspect):
    """
        The function is used to design anchor boxes by ourselves as long as you provide the scale and aspect of anchor boxes.
        Input:
            input_size : list -> the image resolution used in training stage and testing stage.
            stride : int -> the downSample of the CNN, such as 32, 64 and so on.
            anchor_scale : list -> it contains the area ratio of anchor boxes. For example, anchor_scale = [0.1, 0.5]
            anchor_aspect : list -> it contains the aspect ratios of anchor boxes for various anchor area.
                            For example, anchor_aspect = [[1.0, 2.0], [3.0, 1/3]]. And len(anchor_aspect) must
                            be equal to len(anchor_scale).
        Output:
            total_anchor_size : list -> [[h_1, w_1], [h_2, w_2], ..., [h_n, w_n]].
    """
    assert len(anchor_scale) == len(anchor_aspect)
    h, w = input_size
    hs, ws = h // stride, w // stride
    S_fmap = hs * ws
    total_anchor_size = []
    for ab_scale, aspect_ratio in zip(anchor_scale, anchor_aspect):
        for a in aspect_ratio:
            S_ab = S_fmap * ab_scale
            ab_w = np.floor(np.sqrt(S_ab))
            ab_h =ab_w * a
            total_anchor_size.append([ab_w, ab_h])
    return total_anchor_size

def compute_iou(anchor_boxes, gt_box):
    """
    Input:
        anchor_boxes : ndarray -> [[c_x_s, c_y_s, anchor_w, anchor_h], ..., [c_x_s, c_y_s, anchor_w, anchor_h]].
        gt_box : ndarray -> [c_x_s, c_y_s, anchor_w, anchor_h].
    Output:
        iou : ndarray -> [iou_1, iou_2, ..., iou_m], and m is equal to the number of anchor boxes.
    """
    # compute the iou between anchor box and gt box
    # First, change [c_x_s, c_y_s, anchor_w, anchor_h] ->  [xmin, ymin, xmax, ymax]
    # anchor box :
    ab_x1y1_x2y2 = np.zeros([len(anchor_boxes), 4])
    ab_x1y1_x2y2[:, 0] = anchor_boxes[:, 0] - anchor_boxes[:, 2] / 2  # xmin
    ab_x1y1_x2y2[:, 1] = anchor_boxes[:, 1] - anchor_boxes[:, 3] / 2  # ymin
    ab_x1y1_x2y2[:, 2] = anchor_boxes[:, 0] + anchor_boxes[:, 2] / 2  # xmax
    ab_x1y1_x2y2[:, 3] = anchor_boxes[:, 1] + anchor_boxes[:, 3] / 2  # ymax
    w_ab, h_ab = anchor_boxes[:, 2], anchor_boxes[:, 3]

    # gt_box :
    # We need to expand gt_box(ndarray) to the shape of anchor_boxes(ndarray), in order to compute IoU easily.
    gt_box_expand = np.repeat(gt_box, len(anchor_boxes), axis=0)

    gb_x1y1_x2y2 = np.zeros([len(anchor_boxes), 4])
    gb_x1y1_x2y2[:, 0] = gt_box_expand[:, 0] - gt_box_expand[:, 2] / 2 # xmin
    gb_x1y1_x2y2[:, 1] = gt_box_expand[:, 1] - gt_box_expand[:, 3] / 2 # ymin
    gb_x1y1_x2y2[:, 2] = gt_box_expand[:, 0] + gt_box_expand[:, 2] / 2 # xmax
    gb_x1y1_x2y2[:, 3] = gt_box_expand[:, 1] + gt_box_expand[:, 3] / 2 # ymin
    w_gt, h_gt = gt_box_expand[:, 2], gt_box_expand[:, 3]

    # Then we compute IoU between anchor_box and gt_box
    S_gt = w_gt * h_gt
    S_ab = w_ab * h_ab
    I_w = np.minimum(gb_x1y1_x2y2[:, 2], ab_x1y1_x2y2[:, 2]) - np.maximum(gb_x1y1_x2y2[:, 0], ab_x1y1_x2y2[:, 0])
    I_h = np.minimum(gb_x1y1_x2y2[:, 3], ab_x1y1_x2y2[:, 3]) - np.maximum(gb_x1y1_x2y2[:, 1], ab_x1y1_x2y2[:, 1])
    S_I = I_h * I_w
    U = S_gt + S_ab - S_I
    U[U == 0] = 1e-20
    IoU = S_I / U

    return IoU

def set_anchors(anchor_size):
    """
    Input:
        anchor_size : list -> [[h_1, w_1], [h_2, w_2], ..., [h_n, w_n]].
    Output:
        anchor_boxes : ndarray -> [[0, 0, anchor_w, anchor_h],
                                   [0, 0, anchor_w, anchor_h],
                                   ...
                                   [0, 0, anchor_w, anchor_h]].
    """
    anchor_number = len(anchor_size)
    anchor_boxes = np.zeros([anchor_number, 4])
    for index, size in enumerate(anchor_size):
        anchor_w, anchor_h = size
        anchor_boxes[index] = np.array([0, 0, anchor_w, anchor_h])

    return anchor_boxes

def generate_txtytwth(gt_label, w, h, s, all_anchor_size, ignore_thresh):
    xmin, ymin, xmax, ymax = gt_label[:-1]
    # compute the center, width and height
    c_x = (xmax + xmin) / 2 * w
    c_y = (ymax + ymin) / 2 * h
    box_w = (xmax - xmin) * w
    box_h = (ymax - ymin) * h

    if box_w < 1. or box_h < 1.:
        # print('A dirty data !!!')
        return False

    # map the center, width and height to the feature map size
    c_x_s = c_x / s
    c_y_s = c_y / s
    box_ws = box_w / s
    box_hs = box_h / s

    # the grid cell location
    grid_x = int(c_x_s)
    grid_y = int(c_y_s)
    # generate anchor boxes
    anchor_boxes = set_anchors(all_anchor_size)
    gt_box = np.array([[0, 0, box_ws, box_hs]])
    # compute the IoU
    iou = compute_iou(anchor_boxes, gt_box)
    # We consider those anchor boxes whose IoU is more than ignore thresh,
    iou_mask = (iou > ignore_thresh)

    result = []
    if iou_mask.sum() == 0:
        # We assign the anchor box with highest IoU score.
        index = np.argmax(iou)
        p_w, p_h = all_anchor_size[index]
        tx = c_x_s - grid_x
        ty = c_y_s - grid_y
        tw = np.log(box_ws / p_w)
        th = np.log(box_hs / p_h)
        weight = 2.0 - (box_w / w) * (box_h / h)

        result.append([index, grid_x, grid_y, tx, ty, tw, th, weight, xmin, ymin, xmax, ymax])

        return result

    else:
        # There are more than one anchor boxes whose IoU are higher than ignore thresh.
        # But we only assign only one anchor box whose IoU is the best(objectness target is 1) and ignore other
        # anchor boxes whose(we set their objectness as -1 which means we will ignore them during computing obj loss )
        # iou_ = iou * iou_mask

        # We get the index of the best IoU
        best_index = np.argmax(iou)
        for index, iou_m in enumerate(iou_mask):
            if iou_m:
                if index == best_index:
                    p_w, p_h = all_anchor_size[index]
                    tx = c_x_s - grid_x
                    ty = c_y_s - grid_y
                    tw = np.log(box_ws / p_w)
                    th = np.log(box_hs / p_h)
                    weight = 2.0 - (box_w / w) * (box_h / h)

                    result.append([index, grid_x, grid_y, tx, ty, tw, th, weight, xmin, ymin, xmax, ymax])
                else:
                    # we ignore other anchor boxes even if their iou scores all higher than ignore thresh
                    result.append([index, grid_x, grid_y, 0., 0., 0., 0., -1.0, 0., 0., 0., 0.])

        return result

def gt_creator(input_size, stride, label_lists, anchor_size, ingnore_thresh):
    """
    Input:
        input_size : list -> the size of image in the training stage.
        stride : int or list -> the downSample of the CNN, such as 32, 64 and so on.
        label_list : list -> [[[xmin, ymin, xmax, ymax, cls_ind], ... ], [[xmin, ymin, xmax, ymax, cls_ind], ... ]],
                        and len(label_list) = batch_size;
                            len(label_list[i]) = the number of class instance in a image;
                            (xmin, ymin, xmax, ymax) : the coords of a bbox whose valus is between 0 and 1;
                            cls_ind : the corresponding class label.
    Output:
        gt_tensor : ndarray -> shape = [batch_size, anchor_number, 1+1+4, grid_cell number ]
    """
    assert len(input_size) > 0 and len(label_lists) > 0
    # prepare the all empty gt datas
    batch_size = len(label_lists)
    w = input_size[1]
    h = input_size[0]

    # We  make gt labels by anchor-free method and anchor-based method.
    ws = w // stride
    hs = h // stride
    s = stride

    # We use anchor boxes to build training target.
    all_anchor_size = anchor_size
    anchor_number = len(all_anchor_size)

    gt_tensor = np.zeros([batch_size, hs, ws, anchor_number, 1+1+4+1+4])

    for batch_index in range(batch_size):
        for gt_label in label_lists[batch_index]:
            # get a bbox coords
            gt_class = int(gt_label[-1])
            results = generate_txtytwth(gt_label, w, h, s, all_anchor_size, ingnore_thresh)
            if results:
                for result in results:
                    index, grid_x, grid_y, tx, ty, tw, th, weight, xmin, ymin, xmax, ymax = result
                    if weight > 0.:
                        if grid_y < gt_tensor.shape[1] and grid_x < gt_tensor.shape[2]:
                            gt_tensor[batch_index, grid_y, grid_x, index, 0] = 1.0
                            gt_tensor[batch_index, grid_y, grid_x, index, 1] = gt_class
                            gt_tensor[batch_index, grid_y, grid_x, index, 2:6] = np.array([tx, ty, tw, th])
                            gt_tensor[batch_index, grid_y, grid_x, index, 6] = weight
                            gt_tensor[batch_index, grid_y, grid_x, index, 7:] = np.array([xmin, ymin, xmax, ymax])
                    else:
                        gt_tensor[batch_index, grid_y, grid_x, index, 0] = -1.0
                        gt_tensor[batch_index, grid_y, grid_x, index, 6] = -1.0

    gt_tensor = gt_tensor.reshape(batch_size, hs * ws * anchor_number, 1+1+4+1+4)

    return gt_tensor

def multi_gt_creator(input_size, strides, label_lists, anchor_size, ignore_thresh):
    """creator multi scales gt"""
    # prepare the all empty gt datas
    batch_size = len(label_lists)
    h, w = input_size
    num_scale = len(strides)
    gt_tensor = []

    # generate gt datas
    all_anchor_size = anchor_size
    anchor_number = len(all_anchor_size) // num_scale
    for s in strides:
        gt_tensor.append(np.zeros([batch_size, h//s, w//s, anchor_number, 1+1+4+1+4]))
    for batch_index in range(batch_size):
        for gt_label in label_lists[batch_index]:
            # get a bbox coords
            gt_class = int(gt_label[-1])
            xmin, ymin, xmax, ymax = gt_label[:-1]
            # compute the center, width and height
            c_x = (xmax + xmin) / 2 * w
            c_y = (ymax + ymin) / 2 * h
            box_w = (xmax - xmin) * w
            box_h = (ymax - ymin) * h

            if box_w < 1. or box_h < 1.:
                # print('A dirty data !!!')
                continue

            # compute the IoU
            anchor_boxes = set_anchors(all_anchor_size)
            gt_box = np.array([[0, 0, box_w, box_h]])
            iou = compute_iou(anchor_boxes, gt_box)

            # We only consider those anchor boxes whose IoU is more than ignore thresh,
            iou_mask = (iou > ignore_thresh)

            if iou_mask.sum() == 0:
                # We assign the anchor box with highest IoU score.
                index = np.argmax(iou)
                # s_indx, ab_ind = index // num_scale, index % num_scale
                s_indx = index // anchor_number
                ab_ind = index - s_indx * anchor_number
                # get the corresponding stride
                s = strides[s_indx]
                # get the corresponding anchor box
                p_w, p_h = anchor_boxes[index, 2], anchor_boxes[index, 3]
                # compute the gride cell location
                c_x_s = c_x / s
                c_y_s = c_y / s
                grid_x = int(c_x_s)
                grid_y = int(c_y_s)
                # compute gt labels
                tx = c_x_s - grid_x
                ty = c_y_s - grid_y
                tw = np.log(box_w / p_w)
                th = np.log(box_h / p_h)
                weight = 2.0 - (box_w / w) * (box_h / h)

                if grid_y < gt_tensor[s_indx].shape[1] and grid_x < gt_tensor[s_indx].shape[2]:
                    gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 0] = 1.0
                    gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 1] = gt_class
                    gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 2:6] = np.array([tx, ty, tw, th])
                    gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 6] = weight
                    gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 7:] = np.array([xmin, ymin, xmax, ymax])

            else:
                # There are more than one anchor boxes whose IoU are higher than ignore thresh.
                # But we only assign only one anchor box whose IoU is the best(objectness target is 1) and ignore other
                # anchor boxes whose(we set their objectness as -1 which means we will ignore them during computing obj loss )
                # iou_ = iou * iou_mask

                # We get the index of the best IoU
                best_index = np.argmax(iou)
                for index, iou_m in enumerate(iou_mask):
                    if iou_m:
                        if index == best_index:
                            # s_indx, ab_ind = index // num_scale, index % num_scale
                            s_indx = index // anchor_number
                            ab_ind = index - s_indx * anchor_number
                            # get the corresponding stride
                            s = strides[s_indx]
                            # get the corresponding anchor box
                            p_w, p_h = anchor_boxes[index, 2], anchor_boxes[index, 3]
                            # compute the gride cell location
                            c_x_s = c_x / s
                            c_y_s = c_y / s
                            grid_x = int(c_x_s)
                            grid_y = int(c_y_s)
                            # compute gt labels
                            tx = c_x_s - grid_x
                            ty = c_y_s - grid_y
                            tw = np.log(box_w / p_w)
                            th = np.log(box_h / p_h)
                            weight = 2.0 - (box_w / w) * (box_h / h)

                            if grid_y < gt_tensor[s_indx].shape[1] and grid_x < gt_tensor[s_indx].shape[2]:
                                gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 0] = 1.0
                                gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 1] = gt_class
                                gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 2:6] = np.array([tx, ty, tw, th])
                                gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 6] = weight
                                gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 7:] = np.array([xmin, ymin, xmax, ymax])

                        else:
                            # we ignore other anchor boxes even if their iou scores are higher than ignore thresh
                            # s_indx, ab_ind = index // num_scale, index % num_scale
                            s_indx = index // anchor_number
                            ab_ind = index - s_indx * anchor_number
                            s = strides[s_indx]
                            c_x_s = c_x / s
                            c_y_s = c_y / s
                            grid_x = int(c_x_s)
                            grid_y = int(c_y_s)
                            gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 0] = -1.0
                            gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 6] = -1.0

    gt_tensor = [gt.reshape(batch_size, -1, 1+1+4+1+4) for gt in gt_tensor]
    gt_tensor = np.concatenate(gt_tensor, 1)

    return gt_tensor

def iou_score(bboxes_a, bboxes_b):
    """
        bbox_1 : [B*N, 4] = [x1, y1, x2, y2]
        bbox_2 : [B*N, 4] = [x1, y1, x2, y2]
    """
    tl = torch.max(bboxes_a[:, :2], bboxes_b[:, :2])
    br = torch.min(bboxes_a[:, 2:], bboxes_b[:, 2:])
    area_a = torch.prod(bboxes_a[:, 2:] - bboxes_a[:, :2], 1)
    area_b = torch.prod(bboxes_b[:, 2:] - bboxes_b[:, :2], 1)

    en = (tl < br).type(tl.type()).prod(dim=1)
    area_i = torch.prod(br - tl, 1) * en  # * ((tl < br).all())
    s = area_a + area_b - area_i
    s[s==0] = 1e-12
    cross = area_i
    # a = torch.sum(torch.isinf(s))
    # b = torch.sum(torch.isinf(cross))
    # if a > 0 or b > 0:
    #     print(s, cross)
    #     print("---", a, b)
    return  cross / s

def loss(pred_conf, pred_cls, pred_txtytwth, label, num_classes, obj_loss_f='mse'):
    if obj_loss_f == 'bce':
        # In yolov3, we use bce as conf loss_f
        conf_loss_function = BCELoss(reduction='mean')
        obj = 1.0
        noobj = 1.0
    elif obj_loss_f == 'mse':
        # In yolov2, we use mse as conf loss_f.
        conf_loss_function = MSELoss(reduction='mean')
        obj = 5.0
        noobj = 1.0

    cls_loss_function = nn.CrossEntropyLoss(reduction='none')
    txty_loss_function = nn.BCEWithLogitsLoss(reduction='none')
    twth_loss_function = nn.MSELoss(reduction='none')

    pred_conf = torch.sigmoid(pred_conf[:, :, 0])
    pred_cls = pred_cls.permute(0, 2, 1)
    txty_pred = pred_txtytwth[:, :, :2]
    twth_pred = pred_txtytwth[:, :, 2:]

    gt_conf = label[:, :, 0].float()
    gt_obj = label[:, :, 1].float()
    gt_cls = label[:, :, 2].long()
    gt_txtytwth = label[:, :, 3:-1].float()
    gt_box_scale_weight = label[:, :, -1]
    gt_mask = (gt_box_scale_weight > 0.).float()

    # objectness loss
    pos_loss, neg_loss = conf_loss_function(pred_conf, gt_conf, gt_obj)
    conf_loss = obj * pos_loss + noobj * neg_loss
    if torch.isnan(conf_loss):
        print(pred_conf)
        print(obj)
        print(pos_loss)
        print(noobj)
        print(neg_loss)
        print(pred_conf)
        print(gt_conf)
        print(gt_obj)
        print(torch.sum(torch.isnan(pred_conf)), torch.sum(torch.isnan(gt_conf)), torch.sum(torch.isnan(gt_obj)))
        print("---")

    # class loss
    cls_loss = torch.mean(torch.sum(cls_loss_function(pred_cls, gt_cls) * gt_mask, 1))

    # box loss
    txty_loss = torch.mean(torch.sum(torch.sum(txty_loss_function(txty_pred, gt_txtytwth[:, :, :2]), 2) * gt_box_scale_weight * gt_mask, 1))
    twth_loss = torch.mean(torch.sum(torch.sum(twth_loss_function(twth_pred, gt_txtytwth[:, :, 2:]), 2) * gt_box_scale_weight * gt_mask, 1))

    txtytwth_loss = txty_loss + twth_loss

    total_loss = conf_loss + cls_loss + txtytwth_loss

    return conf_loss, cls_loss, txtytwth_loss, total_loss

if __name__ == "__main__":
    gt_box = np.array([[0.0, 0.0, 10, 10]])
    anchor_boxes = np.array([[0.0, 0.0, 10, 10],
                             [0.0, 0.0, 4, 4],
                             [0.0, 0.0, 8, 8],
                             [0.0, 0.0, 16, 16]
                             ])
    iou = compute_iou(anchor_boxes, gt_box)
    print(iou)