import pickle
import numpy as np
from tqdm import tqdm

# Linear
print('-' * 20 + 'Linear Eval' + '-' * 20)

joint_path = '/mnt/netdisk/linlilang/NoisyDINO/work_dir/master/ntu60_cs/aimclr_joint/linear_eval/'
bone_path = '/mnt/netdisk/linlilang/NoisyDINO/work_dir/master/ntu60_cs/aimclr_bone/linear_eval/'
motion_path = '/mnt/netdisk/linlilang/NoisyDINO/work_dir/master/ntu60_cs/aimclr_motion/linear_eval/'

label = open('/mnt/netdisk/linlilang/CrosSCLR/data/NTU-RGB-D/xsub/val_label.pkl', 'rb')
label = np.array(pickle.load(label))
r1 = open(joint_path + 'test_result.pkl', 'rb')
r1 = list(pickle.load(r1).items())
r2 = open(bone_path + 'test_result.pkl', 'rb')
r2 = list(pickle.load(r2).items())
r3 = open(motion_path + 'test_result.pkl', 'rb')
r3 = list(pickle.load(r3).items())

alpha = [0.4, 0.4, 0.4] # joint, bone, motion

right_num = total_num = right_num_5 = 0
for i in tqdm(range(len(label[0]))):
    _, l = label[:, i]
    _, r11 = r1[i]
    _, r22 = r2[i]
    _, r33 = r3[i]
    r = r11 * alpha[0] + r22 * alpha[1] + r33 * alpha[2]
    rank_5 = r.argsort()[-5:]
    right_num_5 += int(int(l) in rank_5)
    r = np.argmax(r)
    right_num += int(r == int(l))
    total_num += 1
acc = right_num / total_num
acc5 = right_num_5 / total_num
print('top1: ', acc)
print('top5: ', acc5)
ensemble_ntu_cs.py
LanglandsLin/ActCLR
import pickle
import numpy as np
from tqdm import tqdm

# Linear
print('-' * 20 + 'Linear Eval' + '-' * 20)

joint_path = '/mnt/netdisk/linlilang/NoisyDINO/work_dir/master/ntu60_cv/aimclr_joint/adaptive_finetune_eval/'
bone_path = '/mnt/netdisk/linlilang/NoisyDINO/work_dir/master/ntu60_cv/aimclr_bone/adaptive_finetune_eval/'
motion_path = '/mnt/netdisk/linlilang/NoisyDINO/work_dir/master/ntu60_cv/aimclr_motion/adaptive_finetune_eval/'

label = open('/mnt/netdisk/linlilang/CrosSCLR/data/NTU-RGB-D/xview/val_label.pkl', 'rb')
label = np.array(pickle.load(label))
r1 = open(joint_path + 'test_result.pkl', 'rb')
r1 = list(pickle.load(r1).items())
r2 = open(bone_path + 'test_result.pkl', 'rb')
r2 = list(pickle.load(r2).items())
r3 = open(motion_path + 'test_result.pkl', 'rb')
r3 = list(pickle.load(r3).items())

alpha = [0.6, 0.6, 0.4] # joint, bone, motion

right_num = total_num = right_num_5 = 0
for i in tqdm(range(len(label[0]))):
    _, l = label[:, i]
    _, r11 = r1[i]
    _, r22 = r2[i]
    _, r33 = r3[i]
    r = r11 * alpha[0] + r22 * alpha[1] + r33 * alpha[2]
    rank_5 = r.argsort()[-5:]
    right_num_5 += int(int(l) in rank_5)
    r = np.argmax(r)
    right_num += int(r == int(l))
    total_num += 1
acc = right_num / total_num
acc5 = right_num_5 / total_num
print('top1: ', acc)
print('top5: ', acc5)
ensemble_ntu_cv.py
LanglandsLin/ActCLR
import numpy as np
import random
import pickle
import cv2

# torch
import torch


def _y_transmat(thetas):
    tms = np.zeros((0, 3, 3))
    thetas = thetas * np.pi / 180
    for theta in thetas:
        tm = np.zeros((3, 3))
        tm[0, 0] = np.cos(theta)
        tm[0, 2] = -np.sin(theta)
        tm[1, 1] = 1
        tm[2, 0] = np.sin(theta)
        tm[2, 2] = np.cos(theta)
        tm = tm[np.newaxis, :, :]
        tms = np.concatenate((tms, tm), axis=0)
    return tms


def parallel_skeleton(ins_data):
    right_shoulder = ins_data[:, :, 8]  # 9th joint
    left_shoulder = ins_data[:, :, 4]  # 5tf joint
    vec = right_shoulder - left_shoulder
    vec[1, :] = 0
    # print(vec.shape)
    l2_norm = np.sqrt(np.sum(np.square(vec), axis=0))
    theta = vec[0, :] / (l2_norm + 0.0001)
    # print(l2_norm)
    thetas = np.arccos(theta) * (180 / np.pi)
    isv = np.sum(vec[2, :])
    if isv >= 0:
        thetas = -thetas
    # print (thetas)
    y_tms = _y_transmat(thetas)
    # print(y_tms)
    new_skel = np.zeros(shape=(0, 25, 3))
    # print(new_skel.shape)
    ins_data = ins_data.transpose(1, 2, 0)
    # print(ins_data.shape, new_skel.shape)
    for ind, each_s in enumerate(ins_data):
        # print(each_s.shape)
        r = np.reshape(each_s, newshape=(25, 3))
        r = np.transpose(r)
        r = np.dot(y_tms[ind], r)
        r_t = np.transpose(r)
        r_t = np.reshape(r_t, newshape=(1, -1, 3))
        # print(new_skel.shape, r_t.shape)
        new_skel = np.concatenate((new_skel, r_t), axis=0)
    return new_skel, ins_data


class SimpleLoader(torch.utils.data.Dataset):
    """ Loader for skeleton-based action recognition
    Arguments:
        data_path: the path to '.npy' data, the shape of data should be (N, C, T, V, M)
        label_path: the path to label
    """

    def __init__(self,
                 data_path,
                 label_path,
                 debug=False,
                 mmap=True,
                 data_type='relative',
                 displacement=False,
                 t_length=200,
                 y_rotation=True,
                 sampling='resize'):
        self.debug = debug
        self.data_path = data_path
        self.label_path = label_path
        self.data_type = data_type
        if displacement is not False:
            assert isinstance(displacement, int)
        self.displacement = displacement
        self.max_length = 300
        self.t_length = t_length if t_length < self.max_length else self.max_length
        self.sampling = sampling
        self.y_rotation = y_rotation
        self.mmap = mmap
        self.load_data(mmap)

    def load_data(self, mmap):
        # data: N C V T M
        # load data
        position_path = self.data_path.replace('_data.', '_position.')
        motion_path = self.data_path.replace('_data.', '_motion.')
        print(position_path)
        print(motion_path)
        print(self.label_path)

        if mmap:
            self.data = np.load(position_path, mmap_mode='r')
            self.motion = np.load(motion_path, mmap_mode='r') if self.displacement > 0 else None
            self.label = np.load(self.label_path).reshape(-1)
        else:
            self.data = np.load(position_path)
            self.motion = np.load(motion_path) if self.displacement > 0 else None
            self.label = np.load(self.label_path).reshape(-1)
        self.N, self.C, self.T, self.V, self.M = self.data.shape

    def __len__(self):
        return len(self.label)

    def __getitem__(self, index):
        # get data
        position = np.array(self.data[index])
        motion = np.array(self.motion[index]) if self.displacement > 0 else None
        label = np.array(self.label[index])

        # return motion_data, label
        if motion is not None:
            return position, motion, label
        else:
            return position, label


class NTUMotionProcessor(torch.utils.data.Dataset):
    """ Feeder for skeleton-based action recognition
    Arguments:
        data_path: the path to '.npy' data, the shape of data should be (N, C, T, V, M)
        label_path: the path to label
    """

    def __init__(self,
                 data_path,
                 label_path,
                 mmap=True,
                 data_type='normal',
                 t_length=300,
                 y_rotation=False,
                 displacement=False,
                 sampling='force_crop'):
        self.data_path = data_path
        self.label_path = label_path
        self.data_type = data_type
        self.max_length = 300
        self.displacement = displacement
        self.sampling = sampling
        self.y_rotation = y_rotation
        self.t_length = t_length if t_length < self.max_length else self.max_length
        neighbor_1base = [(21, 2), (21, 3), (21, 5), (21, 9), (3, 4),
                          (5, 6), (6, 7), (7, 8), (8, 23), (23, 22),
                          (9, 10), (10, 11), (11, 12), (12, 25), (25, 24),
                          (2, 1), (1, 13), (1, 17), (13, 14), (14, 15),
                          (15, 16), (17, 18), (18, 19), (19, 20)]
        self.neighbor_1base = [(i - 1, j - 1) for (i, j) in neighbor_1base]
        self.load_data(mmap)

    def load_data(self, mmap):
        # data: N C V T M
        # load label
        with open(self.label_path, 'rb') as f:
            print(self.label_path)
            self.sample_name, self.label = pickle.load(f)
            self.label = np.array(self.label)
            self.label = self.label - self.label.min()

        # load data
        print(self.data_path)
        if mmap:
            self.data = np.load(self.data_path, mmap_mode='r')
        else:
            self.data = np.load(self.data_path)

        self.N, self.C, self.T, self.V, self.M = self.data.shape

    def __len__(self):
        return len(self.label)

    def __getitem__(self, index):
        # get data
        data_numpy = np.array(self.data[index])
        label = self.label[index]
        
        # get motion (displacement)
        if self.displacement > 0:
            motion_data = np.zeros_like(data_numpy)
            motion_data[:, :-self.displacement, :, :] = \
                data_numpy[:, self.displacement:, :, :] - data_numpy[:, :-self.displacement, :, :]
        else:
            motion_data = None

        # get skeleton sequence length, person number
        length = self.get_length(data_numpy)
        num = self.get_person_num(data_numpy)

        # relative coordinate
        if self.data_type == 'relative':
            # center = 21
            C, T, V, M = data_numpy.shape
            center = data_numpy[:, :, 20, :].reshape([C, T, 1, M])
            data_numpy = data_numpy - center
        elif self.data_type == 'normal':
            pass
        else:
            raise TypeError('Invalid data type: %s' % self.data_type)

        # crop length
        if self.sampling == 'force_crop':
            if self.displacement > 0:
                motion_data = motion_data[:, :self.t_length, :, :]
            data_numpy = data_numpy[:, :self.t_length, :, :]
        elif self.sampling == 'resize':
            data_numpy = self.real_resize(data_numpy, length, self.t_length)
            motion_data = self.real_resize(motion_data, length, self.t_length)
        else:
            raise TypeError('Invalid sampling type: ' % self.sampling)

        # y rotation
        if self.y_rotation:
            for n in range(2):
                tmp_new, tmp_old = parallel_skeleton(data_numpy[:, :, :, n])
                data_numpy[:, :, :, n] = tmp_new.transpose(2, 0, 1)

        # get actual length
        length = self.t_length if length > self.t_length else length

        # return motion_data, label
        if motion_data is not None:
            return data_numpy, motion_data, label
        else:
            return data_numpy, label

    @staticmethod
    def get_length(data):
        length = (abs(data[:, :, 0, 0]).sum(axis=0) != 0).sum()
        return length

    @staticmethod
    def get_person_num(data):
        num = (abs(data[:, :, 0, :]).sum(axis=0).sum(axis=0) != 0) * 1.0
        return num

    @staticmethod
    def real_resize(data_numpy, length, crop_size):
        C, T, V, M = data_numpy.shape
        new_data = np.zeros([C, crop_size, V, M])
        for i in range(M):
            tmp = cv2.resize(data_numpy[:, :length, :, i].transpose(
                [1, 2, 0]), (V, crop_size), interpolation=cv2.INTER_LINEAR)
            tmp = tmp.transpose([2, 0, 1])
            new_data[:, :, :, i] = tmp
        return new_data.astype(np.float32)

    def get_relative_data(self, data_numpy):
        if self.data_type == 'relative':
            new_data = data_numpy
        elif self.data_type == 'multi_relative':
            anchor = [4, 8, 12, 16]  # 5,9,13,17
            C, T, V, M = data_numpy.shape
            new_data = np.zeros([C * 4, T, V, M])
            for i, anc in enumerate(anchor):
                center = data_numpy[:, :, anc, :].reshape([C, T, 1, M])
                new_data[i * 3:(i + 1) * 3, :, :, :] = data_numpy - center
        else:
            raise TypeError('Invalid data type: %s' % self.data_type)
        return new_data.astype(np.float32)

feeder/NTUDatasets.py
LanglandsLin/ActCLR
import random
import numpy as np
import pickle, torch
from . import tools


class Feeder_single(torch.utils.data.Dataset):
    """ Feeder for single inputs """

    def __init__(self, data_path, label_path, shear_amplitude=0.5, temperal_padding_ratio=6, mmap=True):
        self.data_path = data_path
        self.label_path = label_path

        self.shear_amplitude = shear_amplitude
        self.temperal_padding_ratio = temperal_padding_ratio

        self.load_data(mmap)

    def load_data(self, mmap):
        # load label
        with open(self.label_path, 'rb') as f:
            self.sample_name, self.label = pickle.load(f)

        # load data
        if mmap:
            self.data = np.load(self.data_path, mmap_mode='r')
        else:
            self.data = np.load(self.data_path)

    def __len__(self):
        return len(self.label)

    def __getitem__(self, index):
        # get data
        data_numpy = np.array(self.data[index])
        label = self.label[index]

        # processing
        data = self._aug(data_numpy)
        return data, label

    def _aug(self, data_numpy):
        if self.temperal_padding_ratio > 0:
            data_numpy = tools.temperal_crop(data_numpy, self.temperal_padding_ratio)

        if self.shear_amplitude > 0:
            data_numpy = tools.shear(data_numpy, self.shear_amplitude)

        return data_numpy


class Feeder_triple(torch.utils.data.Dataset):
    """ Feeder for triple inputs """

    def __init__(self, data_path, label_path, shear_amplitude=0.5, temperal_padding_ratio=6, mmap=True,
                 aug_method='12345', strong_aug_method='1234589'):
        self.data_path = data_path
        self.label_path = label_path
        self.aug_method = aug_method
        self.strong_aug_method = strong_aug_method

        self.shear_amplitude = shear_amplitude
        self.temperal_padding_ratio = temperal_padding_ratio

        self.load_data(mmap)

    def load_data(self, mmap):
        # load label
        with open(self.label_path, 'rb') as f:
            self.sample_name, self.label = pickle.load(f)

        # load data
        if mmap:
            self.data = np.load(self.data_path, mmap_mode='r')
        else:
            self.data = np.load(self.data_path)

    def __len__(self):
        return len(self.label)

    def __getitem__(self, index):
        # get data
        data_numpy = np.array(self.data[index])

        # processing
        data1 = self._strong_aug(data_numpy)
        data2 = self._aug(data_numpy)
        data3 = self._weak_aug(data_numpy)
        data4 = self._weak_aug(data_numpy)
        return data1, data2, data3, data4
    
    def _weak_aug(self, data_numpy):
        if self.temperal_padding_ratio > 0:
            data_numpy = tools.temperal_crop(data_numpy, self.temperal_padding_ratio)

        if self.shear_amplitude > 0:
            data_numpy = tools.shear(data_numpy, self.shear_amplitude)
        return data_numpy

    def _aug(self, data_numpy):
        if self.temperal_padding_ratio > 0:
            data_numpy = tools.temperal_crop(data_numpy, self.temperal_padding_ratio)
        if self.shear_amplitude > 0:
            data_numpy = tools.shear(data_numpy, self.shear_amplitude)
        if '1' in self.aug_method:
            data_numpy = tools.random_spatial_flip(data_numpy)
        if '2' in self.aug_method:
            data_numpy = tools.random_rotate(data_numpy)
        if '3' in self.aug_method:
            data_numpy = tools.gaus_noise(data_numpy)
        if '4' in self.aug_method:
            data_numpy = tools.gaus_filter(data_numpy)
        if '5' in self.aug_method:
            data_numpy = tools.axis_mask(data_numpy)
        if '6' in self.aug_method:
            data_numpy = tools.random_time_flip(data_numpy)
        if '7' in self.aug_method:
            data_numpy = tools.dropout(data_numpy)
        if '8' in self.aug_method:
            data_numpy = tools.interpolate(data_numpy, 1)
        if '9' in self.aug_method:
            data_numpy = tools.rescale(data_numpy, 0.5)
            
        return data_numpy
    
    # you can choose different combinations
    def _strong_aug(self, data_numpy):
        if self.temperal_padding_ratio > 0:
            data_numpy = tools.temperal_crop(data_numpy, self.temperal_padding_ratio)
        if self.shear_amplitude > 0:
            data_numpy = tools.shear(data_numpy, self.shear_amplitude)
        if '1' in self.strong_aug_method:
            data_numpy = tools.random_spatial_flip(data_numpy)
        if '2' in self.strong_aug_method:
            data_numpy = tools.random_rotate(data_numpy)
        if '3' in self.strong_aug_method:
            data_numpy = tools.gaus_noise(data_numpy, std=0.05)
        if '4' in self.strong_aug_method:
            data_numpy = tools.gaus_filter(data_numpy)
        if '5' in self.strong_aug_method:
            data_numpy = tools.axis_mask(data_numpy)
        if '6' in self.strong_aug_method:
            data_numpy = tools.random_time_flip(data_numpy)
        if '7' in self.strong_aug_method:
            data_numpy = tools.dropout(data_numpy)
        if '8' in self.strong_aug_method:
            data_numpy = tools.interpolate(data_numpy, 1)
        if '9' in self.strong_aug_method:
            data_numpy = tools.rescale(data_numpy, 0.5)

        return data_numpy
feeder/ntu_feeder.py
LanglandsLin/ActCLR
import os
import sys
import pickle

import torch
import argparse
import numpy as np
from numpy.lib.format import open_memmap
from NTUDatasets import NTUMotionProcessor


max_body = 2
num_joint = 25
max_frame = 50
batch_size = 64

toolbar_width = 30


def print_toolbar(rate, annotation=''):
    # setup toolbar
    sys.stdout.write("{}[".format(annotation))
    for i in range(toolbar_width):
        if i * 1.0 / toolbar_width > rate:
            sys.stdout.write(' ')
        else:
            sys.stdout.write('-')
        sys.stdout.flush()
    sys.stdout.write(']\r')


def end_toolbar():
    sys.stdout.write("\n")


def gendata(dataset_path, out_path, benchmark, part='eval'):
    dataset = NTUMotionProcessor(
        '{}/{}_data.npy'.format(os.path.join(dataset_path, benchmark), part),
        '{}/{}_label.pkl'.format(os.path.join(dataset_path, benchmark), part),
        data_type='relative',
        t_length=max_frame,
        y_rotation=True,
        sampling='resize',
        displacement=1,
        mmap=True)

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

    f_position = open_memmap(
        '{}/{}_position.npy'.format(out_path, part),
        dtype='float32',
        mode='w+',
        shape=(dataset.N, 3, max_frame, num_joint, max_body))

    f_motion = open_memmap(
        '{}/{}_motion.npy'.format(out_path, part),
        dtype='float32',
        mode='w+',
        shape=(dataset.N, 3, max_frame, num_joint, max_body))

    f_label = open_memmap(
        '{}/{}_label.npy'.format(out_path, part),
        dtype='int64',
        mode='w+',
        shape=(dataset.N, 1))

    index = 0
    for i, (data, motion, label) in enumerate(data_loader):
        print_toolbar(i * 1.0 / len(data_loader),
                      '({:>5}/{:<5}) Processing {:>5}-{:<5} data: '.format(
                          i + 1, len(data_loader), benchmark, part))
        length = label.shape[0]
        if i * batch_size != index:
            print(i, index)
        f_position[index:(index+length), :, :, :, :] = data.numpy()
        f_motion[index:(index+length), :, :, :, :] = motion.numpy()
        f_label[index:(index+length), :] = label.numpy().reshape(-1, 1)
        index += length
    end_toolbar()


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='NTU-RGB-D Data Converter.')
    parser.add_argument('--dataset_path', default='/data/gty/action_dataset/ntu60')
    parser.add_argument('--out_folder', default='/data/gty/action_dataset/ntu60_frame50')

    benchmark = ['xsub', 'xview']
    part = ['train', 'val']
    arg = parser.parse_args()

    for b in benchmark:
        for p in part:
            out_path = os.path.join(arg.out_folder, b)
            if not os.path.exists(out_path):
                os.makedirs(out_path)
            gendata(arg.dataset_path, out_path, benchmark=b, part=p)
feeder/preprocess_ntu.py
LanglandsLin/ActCLR
import os
import sys
import pickle

import torch
import argparse
import numpy as np
from numpy.lib.format import open_memmap
from NTUDatasets import NTUMotionProcessor


max_body = 2
num_joint = 25
max_frame = 50
batch_size = 64

toolbar_width = 30


def print_toolbar(rate, annotation=''):
    # setup toolbar
    sys.stdout.write("{}[".format(annotation))
    for i in range(toolbar_width):
        if i * 1.0 / toolbar_width > rate:
            sys.stdout.write(' ')
        else:
            sys.stdout.write('-')
        sys.stdout.flush()
    sys.stdout.write(']\r')


def end_toolbar():
    sys.stdout.write("\n")


def gendata(dataset_path, out_path, benchmark, part='eval'):
    dataset = NTUMotionProcessor(
        '{}/{}_data.npy'.format(os.path.join(dataset_path, benchmark), part),
        '{}/{}_label.pkl'.format(os.path.join(dataset_path, benchmark), part),
        data_type='relative',
        t_length=max_frame,
        y_rotation=True,
        sampling='resize',
        displacement=1,
        mmap=True)

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

    f_position = open_memmap(
        '{}/{}_position.npy'.format(out_path, part),
        dtype='float32',
        mode='w+',
        shape=(dataset.N, 3, max_frame, num_joint, max_body))

    f_motion = open_memmap(
        '{}/{}_motion.npy'.format(out_path, part),
        dtype='float32',
        mode='w+',
        shape=(dataset.N, 3, max_frame, num_joint, max_body))

    f_label = open_memmap(
        '{}/{}_label.npy'.format(out_path, part),
        dtype='int64',
        mode='w+',
        shape=(dataset.N, 1))

    index = 0
    for i, (data, motion, label) in enumerate(data_loader):
        print_toolbar(i * 1.0 / len(data_loader),
                      '({:>5}/{:<5}) Processing {:>5}-{:<5} data: '.format(
                          i + 1, len(data_loader), benchmark, part))
        length = label.shape[0]
        if i * batch_size != index:
            print(i, index)
        f_position[index:(index+length), :, :, :, :] = data.numpy()
        f_motion[index:(index+length), :, :, :, :] = motion.numpy()
        f_label[index:(index+length), :] = label.numpy().reshape(-1, 1)
        index += length
    end_toolbar()


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='NTU-RGB-D Data Converter.')
    parser.add_argument('--dataset_path', default='E:/2021/datasets/crossclr_pkummd/pku_part1')
    parser.add_argument('--out_folder', default='E:/2021/datasets/crossclr_pkummd/pku_part1_frame50')

    benchmark = ['xsub', 'xview']
    part = ['train', 'val']
    arg = parser.parse_args()

    for b in benchmark:
        for p in part:
            out_path = os.path.join(arg.out_folder, b)
            if not os.path.exists(out_path):
                os.makedirs(out_path)
            gendata(arg.dataset_path, out_path, benchmark=b, part=p)
feeder/preprocess_pku.py
LanglandsLin/ActCLR
import math
import torch
import random
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from math import sin, cos
import cv2


transform_order = {
    'ntu': [0, 1, 2, 3, 8, 9, 10, 11, 4, 5, 6, 7, 16, 17, 18, 19, 12, 13, 14, 15, 20, 23, 24, 21, 22]
}

trunk_ori_index = [4, 3, 21, 2, 1]
left_hand_ori_index = [9, 10, 11, 12, 24, 25]
right_hand_ori_index = [5, 6, 7, 8, 22, 23]
left_leg_ori_index = [17, 18, 19, 20]
right_leg_ori_index = [13, 14, 15, 16]

trunk = [i - 1 for i in trunk_ori_index]
left_hand = [i - 1 for i in left_hand_ori_index]
right_hand = [i - 1 for i in right_hand_ori_index]
left_leg = [i - 1 for i in left_leg_ori_index]
right_leg = [i - 1 for i in right_leg_ori_index]
body_parts = [trunk, left_hand, right_hand, left_leg, right_leg]

# def shear(data_numpy, r=0.5):
#     s1_list = [random.uniform(-r, r), random.uniform(-r, r), random.uniform(-r, r)]
#     s2_list = [random.uniform(-r, r), random.uniform(-r, r), random.uniform(-r, r)]

#     R = np.array([[1,          s1_list[0], s2_list[0]],
#                   [s1_list[1], 1,          s2_list[1]],
#                   [s1_list[2], s2_list[2], 1        ]])

#     R = R.transpose()
#     data_numpy = np.dot(data_numpy.transpose([1, 2, 3, 0]), R)
#     data_numpy = data_numpy.transpose(3, 0, 1, 2)
#     return data_numpy

def shear(data_numpy, r=0.5):

    center_trans = data_numpy[:,:, 20, :].copy()
    data_numpy[:,:,20,:] = 0.0

    s1_list = [random.uniform(-r, r), random.uniform(-r, r), random.uniform(-r, r)]
    s2_list = [random.uniform(-r, r), random.uniform(-r, r), random.uniform(-r, r)]

    R = np.array([[1,          s1_list[0], s2_list[0]],
                  [s1_list[1], 1,          s2_list[1]],
                  [s1_list[2], s2_list[2], 1        ]])

    R = R.transpose()
    data_numpy = np.dot(data_numpy.transpose([1, 2, 3, 0]), R)
    data_numpy = data_numpy.transpose(3, 0, 1, 2)

    data_numpy[:,:,20,:] = center_trans
    
    return data_numpy


def interpolate(data_numpy, interpolate_ratio=0.0):
    C, T, V, M = data_numpy.shape

    if np.random.rand() > 0.5:
        return data_numpy
    rand_ratio = np.random.rand() * interpolate_ratio + 1 / (interpolate_ratio + 1)
    interpolate_size = int(T * rand_ratio)
    interpolate_size = max(1, interpolate_size)
    new_data = np.zeros((C, T, V, M))

    for i in range(M):
        tmp = cv2.resize(data_numpy[:, :, :, i].transpose(
            [1, 2, 0]), (V, interpolate_size), interpolation=cv2.INTER_LINEAR)

        tmp = tmp.transpose([2, 0, 1])

        new_data[:, :interpolate_size, :, i] = tmp[:, :T]

    return new_data


def temperal_crop(data_numpy, temperal_padding_ratio=6):
    C, T, V, M = data_numpy.shape
    padding_len = T // temperal_padding_ratio
    frame_start = np.random.randint(0, padding_len * 2 + 1)
    data_numpy = np.concatenate((data_numpy[:, :padding_len][:, ::-1],
                                 data_numpy,
                                 data_numpy[:, -padding_len:][:, ::-1]),
                                axis=1)
    data_numpy = data_numpy[:, frame_start:frame_start + T]
    return data_numpy

# def rescale(data_numpy, scale_rate=0.0):
#     C, T, V, M = data_numpy.shape
#     if np.random.rand() > 0.5:
#         return data_numpy
#     rand_ratio = np.random.rand(*data_numpy.shape) * scale_rate + 1 / (scale_rate + 1)
#     mean = np.mean(data_numpy, axis=(2,3), keepdims=True)
#     std = np.std(data_numpy, axis=(2,3), keepdims=True)
#     data_numpy = (data_numpy - mean) / (std + 1e-3) * (std * rand_ratio + 1e-3) + (mean * rand_ratio)
#     return data_numpy

def rescale(data_numpy, scale_rate=0.0):
    Bone = [(21, 21), (2, 21), (3, 21), (9, 21), (5, 21), (1, 2), (4, 3), (6, 5), (7, 6), (8, 7),
            (10, 9), (11, 10), (12, 11), (13, 1), (14, 13), (15, 14), (16, 15), (17, 1),
            (18, 17), (19, 18), (20, 19), (23, 8), (22, 23), (25, 12), (24, 25)]

    bone = np.zeros_like(data_numpy)

    for v1, v2 in Bone:
        bone[:, :, v1 - 1, :] = data_numpy[:, :, v1 - 1, :] - data_numpy[:, :, v2 - 1, :]

    for v1, v2 in Bone:
        data_numpy[:, :, v1 - 1, :] = data_numpy[:, :, v2 - 1, :] + bone[:, :, v1 - 1, :] * scale_rate * random.random()


    # C, T, V, M = data_numpy.shape
    # mean = np.mean(data_numpy, axis=(3), keepdims=True)
    # std = np.std(data_numpy, axis=(2), keepdims=True)
    # data_numpy = (data_numpy - mean) / (std + 1e-3) * (scale_rate * random.random() + 1e-3) + (scale_rate * random.random())
    return data_numpy


def random_spatial_flip(seq, p=0.5):
    if random.random() < p:
        # Do the left-right transform C,T,V,M
        index = transform_order['ntu']
        trans_seq = seq[:, :, index, :]
        return trans_seq
    else:
        return seq


def random_time_flip(seq, p=0.5):
    T = seq.shape[1]
    if random.random() < p:
        time_range_order = [i for i in range(T)]
        time_range_reverse = list(reversed(time_range_order))
        return seq[:, time_range_reverse, :, :]
    else:
        return seq
    
def ske_swap(data1, data2):
    '''
    swap a batch skeleton
    T   64 --> 32 --> 16    # 8n
    S   25 --> 25 --> 25 (5 parts)
    '''
    spa_l, spa_u, tem_l, tem_u = 2, 3, 15, 35
    
    T = data1.shape[1]

    # ------ Spatial ------ #
    Cs = random.randint(spa_l, spa_u)
    # sample the parts index
    parts_idx = random.sample(body_parts, Cs)
    # generate spa_idx
    spa_idx = []
    for part_idx in parts_idx:
        spa_idx += part_idx
    spa_idx.sort()

    # ------ Temporal ------ #
    Ct = random.randint(tem_l, tem_u)
    tem_idx = random.randint(0, T - Ct)
    rt = Ct

    p = random.random()

    if p > 0.5:
        # begin swap
        data1[:, tem_idx: tem_idx + rt, spa_idx, :] = data2[:, tem_idx: tem_idx + rt, spa_idx, :]
    elif p <= 0.5 and p > 0.25:
        data2 = data2[:, ::(T // rt), spa_idx, :]
        t = data2.shape[1]
        rt = min(rt, t)
        data1[:, tem_idx: tem_idx + rt, spa_idx, :] = data2[:, :rt, :, :]
    elif p <= 0.25:
        lamb = random.random()
        data1 = data1 * (1 - lamb) + data2 * lamb

    return data1


def random_rotate(seq):
    def rotate(seq, axis, angle):
        # x
        if axis == 0:
            R = np.array([[1, 0, 0],
                              [0, cos(angle), sin(angle)],
                              [0, -sin(angle), cos(angle)]])
        # y
        if axis == 1:
            R = np.array([[cos(angle), 0, -sin(angle)],
                              [0, 1, 0],
                              [sin(angle), 0, cos(angle)]])

        # z
        if axis == 2:
            R = np.array([[cos(angle), sin(angle), 0],
                              [-sin(angle), cos(angle), 0],
                              [0, 0, 1]])
        R = R.T
        temp = np.matmul(seq, R)
        return temp

    new_seq = seq.copy()
    # C, T, V, M -> T, V, M, C
    new_seq = np.transpose(new_seq, (1, 2, 3, 0))
    total_axis = [0, 1, 2]
    main_axis = random.randint(0, 2)
    for axis in total_axis:
        if axis == main_axis:
            rotate_angle = random.uniform(0, 30)
            rotate_angle = math.radians(rotate_angle)
            new_seq = rotate(new_seq, axis, rotate_angle)
        else:
            rotate_angle = random.uniform(0, 1)
            rotate_angle = math.radians(rotate_angle)
            new_seq = rotate(new_seq, axis, rotate_angle)

    new_seq = np.transpose(new_seq, (3, 0, 1, 2))

    return new_seq


def gaus_noise(data_numpy, mean= 0, std=0.01, p=0.5):
    if random.random() < p:
        temp = data_numpy.copy()
        C, T, V, M = data_numpy.shape
        noise = np.random.normal(mean, std, size=(C, T, V, M))
        return temp + noise
    else:
        return data_numpy


def gaus_filter(data_numpy):
    g = GaussianBlurConv(3)
    return g(data_numpy)

def dropout(data_numpy, p=0.5):
    # if random.random() < p:
    mask = np.random.uniform(size=data_numpy.shape) > p
    return mask * data_numpy
    # else:
    #     return data_numpy


class GaussianBlurConv(nn.Module):
    def __init__(self, channels=3, kernel = 15, sigma = [0.1, 2]):
        super(GaussianBlurConv, self).__init__()
        self.channels = channels
        self.kernel = kernel
        self.min_max_sigma = sigma
        radius = int(kernel / 2)
        self.kernel_index = np.arange(-radius, radius + 1)

    def __call__(self, x):
        sigma = random.uniform(self.min_max_sigma[0], self.min_max_sigma[1])
        blur_flter = np.exp(-np.power(self.kernel_index, 2.0) / (2.0 * np.power(sigma, 2.0)))
        kernel = torch.from_numpy(blur_flter).unsqueeze(0).unsqueeze(0)
        # kernel =  kernel.float()
        kernel = kernel.double()
        kernel = kernel.repeat(self.channels, 1, 1, 1) # (3,1,1,5)
        self.weight = nn.Parameter(data=kernel, requires_grad=False)

        prob = np.random.random_sample()
        x = torch.from_numpy(x)
        if prob < 0.5:
            x = x.permute(3,0,2,1) # M,C,V,T
            x = F.conv2d(x, self.weight, padding=(0, int((self.kernel - 1) / 2 )),   groups=self.channels)
            x = x.permute(1,-1,-2, 0) #C,T,V,M

        return x.numpy()

class Zero_out_axis(object):
    def __init__(self, axis = None):
        self.first_axis = axis


    def __call__(self, data_numpy):
        if self.first_axis != None:
            axis_next = self.first_axis
        else:
            axis_next = random.randint(0,2)

        temp = data_numpy.copy()
        C, T, V, M = data_numpy.shape
        x_new = np.zeros((T, V, M))
        temp[axis_next] = x_new
        return temp

def axis_mask(data_numpy, p=0.5):
    am = Zero_out_axis()
    if random.random() < p:
        return am(data_numpy)
    else:
        return data_numpy

if __name__ == '__main__':
    data_seq = np.ones((3, 50, 25, 2))
    data_seq = axis_mask(data_seq)
    print(data_seq.shape)

feeder/tools.py
LanglandsLin/ActCLR
import argparse
import sys
import os
import shutil
import zipfile
import time

# torchlight
import torchlight
from torchlight import import_class

from processor.processor import init_seed
init_seed(0)

def save_src(target_path):
    code_root = os.getcwd()
    srczip = zipfile.ZipFile('./src.zip', 'w')
    for root, dirnames, filenames in os.walk(code_root):
            for filename in filenames:
                if filename.split('\n')[0].split('.')[-1] == 'py':
                    srczip.write(os.path.join(root, filename).replace(code_root, '.'))
                if filename.split('\n')[0].split('.')[-1] == 'yaml':
                    srczip.write(os.path.join(root, filename).replace(code_root, '.'))
                if filename.split('\n')[0].split('.')[-1] == 'ipynb':
                    srczip.write(os.path.join(root, filename).replace(code_root, '.'))
    srczip.close()
    save_path = os.path.join(target_path, 'src_%s.zip' % time.strftime("%Y-%m-%d_%H_%M_%S", time.localtime()))
    shutil.copy('./src.zip', save_path)


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Processor collection')
    processors = dict()

    processors['pretrain_actclr'] = import_class('processor.pretrain_actclr.ActCLR_Processor')
    processors['linear_evaluation'] = import_class('processor.linear_evaluation.LE_Processor')
    processors['finetune_evaluation'] = import_class('processor.finetune_evaluation.FT_Processor')

    # add sub-parser
    subparsers = parser.add_subparsers(dest='processor')
    for k, p in processors.items():
        subparsers.add_parser(k, parents=[p.get_parser()])

    parser.add_argument("--local_rank", default=0, type=int)
    
    # read arguments
    arg = parser.parse_args()

    # start
    Processor = processors[arg.processor]
    for i in range(len(sys.argv)):
        if sys.argv[i] == "--config":
            sys_argv = sys.argv[i:]
            break

    p = Processor(sys_argv)

    p.init_environment()
    
    if p.arg.phase == 'train':
        # save src
        save_src(p.arg.work_dir)
            
    print(int(os.environ["LOCAL_RANK"]))

    p.start(int(os.environ["LOCAL_RANK"]))
main.py
LanglandsLin/ActCLR
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchlight import import_class
import copy
import random

trunk_ori_index = [4, 3, 21, 2, 1]
left_hand_ori_index = [9, 10, 11, 12, 24, 25]
right_hand_ori_index = [5, 6, 7, 8, 22, 23]
left_leg_ori_index = [17, 18, 19, 20]
right_leg_ori_index = [13, 14, 15, 16]

trunk = [i - 1 for i in trunk_ori_index]
left_hand = [i - 1 for i in left_hand_ori_index]
right_hand = [i - 1 for i in right_hand_ori_index]
left_leg = [i - 1 for i in left_leg_ori_index]
right_leg = [i - 1 for i in right_leg_ori_index]
body_parts = [trunk, left_hand, right_hand, left_leg, right_leg]

@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor)
        for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output

def kld(inputs, targets):
        inputs = F.log_softmax(inputs / 0.1, dim=1)
        targets = F.softmax(targets / 0.04, dim=1)
        return F.kl_div(inputs, targets, reduction='batchmean')

class actclr(nn.Module):
    def __init__(self, base_encoder=None, pretrain=True, feature_dim=128, queue_size=32768,
                 momentum=0.999, Temperature=0.07, mlp=True, in_channels=3, hidden_channels=64,
                 hidden_dim=256, num_class=60, dropout=0.5,
                 graph_args={'layout': 'ntu-rgb+d', 'strategy': 'spatial'},
                 edge_importance_weighting=True, **kwargs):
        """
        K: queue size; number of negative keys (default: 32768)
        m: momentum of updating key encoder (default: 0.999)
        T: softmax temperature (default: 0.07)
        """
        super().__init__()
        base_encoder = import_class(base_encoder)
        self.pretrain = pretrain

        if not self.pretrain:
            self.encoder_q = base_encoder(in_channels=in_channels, hidden_channels=hidden_channels,
                                          hidden_dim=hidden_dim, num_class=num_class,
                                          dropout=dropout, graph_args=graph_args,
                                          edge_importance_weighting=edge_importance_weighting,
                                          **kwargs)
        else:
            self.K = queue_size
            self.m = momentum
            self.T = Temperature

            self.encoder_q = base_encoder(in_channels=in_channels, hidden_channels=hidden_channels,
                                          hidden_dim=hidden_dim, num_class=feature_dim,
                                          dropout=dropout, graph_args=graph_args,
                                          edge_importance_weighting=edge_importance_weighting,
                                          **kwargs)
            self.encoder_k = base_encoder(in_channels=in_channels, hidden_channels=hidden_channels,
                                          hidden_dim=hidden_dim, num_class=feature_dim,
                                          dropout=dropout, graph_args=graph_args,
                                          edge_importance_weighting=edge_importance_weighting,
                                          **kwargs)

            if mlp:  # hack: brute-force replacement
                dim_mlp = self.encoder_q.fc.weight.shape[1]
                self.encoder_q.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp),
                                                  nn.ReLU(),
                                                  self.encoder_q.fc)
                self.encoder_k.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp),
                                                  nn.ReLU(),
                                                  self.encoder_k.fc)

            for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
                param_k.data.copy_(param_q.data)  # initialize
                param_k.requires_grad = False  # not update by gradient

            # create the queue
            self.register_buffer("queue", torch.randn(feature_dim, queue_size))
            self.queue = F.normalize(self.queue, dim=0)
            self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
            
    @torch.no_grad()
    def _momentum_update_key_encoder(self):
        """
        Momentum update of the key encoder
        """
        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys):
        keys = concat_all_gather(keys)
        batch_size = keys.shape[0]
        ptr = int(self.queue_ptr)
        self.queue[:, ptr:ptr + batch_size] = keys.T
        ptr = (ptr + batch_size) % self.K  # move pointer
        self.queue_ptr[0] = ptr

    @torch.no_grad()
    def update_ptr(self, batch_size):
        assert self.K % batch_size == 0  # for simplicity
        self.queue_ptr[0] = (self.queue_ptr[0] + batch_size) % self.K


    @torch.no_grad()
    def ske_swap(self, x, epoch=None):
        '''
        swap a batch skeleton
        T   64 --> 32 --> 16    # 8n
        S   25 --> 25 --> 25 (5 parts)
        '''
        if epoch is None:
            spa_l, spa_u, tem_l, tem_u = 2, 3, 3, 5
        else:
            if epoch < 200: 
                spa_l, spa_u, tem_l, tem_u = 2, 3, 3, 5
            else:
                spa_l, spa_u, tem_l, tem_u = 1, 2, 1, 3
        
        N, C, T, V, M = x.size()
        tem_downsample_ratio = 4

        # generate swap swap idx
        idx = torch.arange(N)
        n = torch.randint(1, N - 1, (1,))
        randidx = (idx + n) % N

        # ------ Spatial ------ #
        Cs = random.randint(spa_l, spa_u)
        # sample the parts index
        parts_idx = random.sample(body_parts, Cs)
        # generate spa_idx
        spa_idx = []
        for part_idx in parts_idx:
            spa_idx += part_idx
        spa_idx.sort()

        # ------ Temporal ------ #
        Ct = random.randint(tem_l, tem_u)
        tem_idx = random.randint(0, T // tem_downsample_ratio - Ct)
        rt = Ct * tem_downsample_ratio

        xst = x.clone()

        p = random.random()

        if p > 0.5:
            # begin swap
            xst[:, :, tem_idx * tem_downsample_ratio: tem_idx * tem_downsample_ratio + rt, spa_idx, :] = \
                xst[randidx][:, :, tem_idx * tem_downsample_ratio: tem_idx * tem_downsample_ratio + rt, spa_idx, :]
            # generate mask
            mask = torch.zeros((T + tem_downsample_ratio - 1) // tem_downsample_ratio, V).cuda()
            mask[tem_idx:tem_idx + Ct, spa_idx] = 1
        elif p <= 0.5 and p > 0.25:
            N, C, T, V, M = xst.size()
            xst_temp = xst[:, :, : 2 * rt]
            xst_temp = xst_temp.permute(0, 4, 3, 1, 2).contiguous()
            xst_temp = xst_temp.view(N * M, V * C, -1)
            xst_temp = torch.nn.functional.interpolate(xst_temp, size=rt)
            xst_temp = xst_temp.view(N, M, V, C, rt)
            xst_temp = xst_temp.permute(0, 3, 4, 2, 1).contiguous()
            xst[:, :, tem_idx * tem_downsample_ratio: tem_idx * tem_downsample_ratio + rt, spa_idx, :] = \
                    xst_temp[randidx][:, :, :, spa_idx, :]
            mask_c = torch.zeros((T + tem_downsample_ratio - 1) // tem_downsample_ratio, V).cuda()
            mask_c[tem_idx:tem_idx + Ct, spa_idx] = 1
            mask_p = torch.zeros((T + tem_downsample_ratio - 1) // tem_downsample_ratio, V).cuda()
            mask_p[:, spa_idx] = 1
            mask = (mask_p, mask_c)
        elif p <= 0.25:
            lamb = random.random()
            xst = xst * (1 - lamb) + xst[randidx] * lamb
            mask = torch.zeros((T + tem_downsample_ratio - 1) // tem_downsample_ratio, V).cuda() + lamb

        return randidx, xst, mask
    
    @torch.no_grad()
    def actionlet_swap(self, x, y, mask):
        '''
        swap a batch skeleton
        T   64 --> 32 --> 16    # 8n
        S   25 --> 25 --> 25 (5 parts)
        '''
        N, C, T, V, M = x.size()
        
        mask = nn.functional.interpolate(mask, (T, V), mode='bilinear', align_corners=True)

        x = x.permute(0, 4, 1, 2, 3).contiguous()
        x = x.view(N * M, C, T, V)
        
        y = y.permute(0, 4, 1, 2, 3).contiguous()
        y = y.view(N * M, C, T, V)
        
        z = x * mask + y * (1 - mask)
        
        z = z.view(N, M, C, T, V)
        z = z.permute(0, 2, 3, 4, 1).contiguous()
    
        return z

    @torch.no_grad()
    def ske_trans(self, x, epoch=None):
        N, C, T, V, M = x.size()

        # generate swap swap idx
        idx = torch.arange(N)
        n = torch.randint(1, N - 1, (1,))
        randidx = (idx + n) % N

        xst = x.clone()
        # # begin swap
        # xst = (xst - xst.mean(dim=(3), keepdim=True)) / (xst.std(dim=(2), keepdim=True) + 1e-3) * xst[randidx].std(dim=(2), keepdim=True) + xst[randidx].mean(dim=(3), keepdim=True)
        # return xst
        Bone = [(21, 21), (2, 21), (3, 21), (9, 21), (5, 21), (1, 2), (4, 3), (6, 5), (7, 6), (8, 7),
            (10, 9), (11, 10), (12, 11), (13, 1), (14, 13), (15, 14), (16, 15), (17, 1),
            (18, 17), (19, 18), (20, 19), (23, 8), (22, 23), (25, 12), (24, 25)]


        for v1, v2 in Bone:
            xst[:, :, :, v1 - 1, :] = x[:, :, :, v1 - 1, :] - x[:, :, :, v2 - 1, :]

        for v1, v2 in Bone:
            x[:, :, :, v1 - 1, :] = x[:, :, :, v2 - 1, :] + xst[:, :, :, v1 - 1, :] * (torch.sqrt((xst[randidx, :, :, v1 - 1, :] ** 2).sum(dim=(-1), keepdim=True) + 1e-5) / torch.sqrt((xst[:, :, :, v1 - 1, :] ** 2).sum(dim=(-1), keepdim=True) + 1e-5)) ** (100 / (epoch + 400))


        return x

    def select_topk(self, q, k):
        k[k > 0.95] = -1
        k_topk, idx_topk = torch.topk(k, 8192, dim=-1)
        q_topk = torch.gather(q, -1, idx_topk)
        return q_topk, k_topk

    @torch.no_grad()
    def _batch_shuffle_ddp(self, x):
        """
        Batch shuffle, for making use of BatchNorm.
        *** Only support DistributedDataParallel (DDP) model. ***
        """
        # gather from all gpus
        batch_size_this = x.shape[0]
        x_gather = concat_all_gather(x)
        batch_size_all = x_gather.shape[0]

        num_gpus = batch_size_all // batch_size_this

        # random shuffle index
        idx_shuffle = torch.randperm(batch_size_all).cuda()

        # broadcast to all gpus
        torch.distributed.broadcast(idx_shuffle, src=0)

        # index for restoring
        idx_unshuffle = torch.argsort(idx_shuffle)

        # shuffled index for this gpu
        gpu_idx = torch.distributed.get_rank()
        idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]

        return x_gather[idx_this], idx_unshuffle

    @torch.no_grad()
    def _batch_unshuffle_ddp(self, x, idx_unshuffle):
        """
        Undo batch shuffle.
        *** Only support DistributedDataParallel (DDP) model. ***
        """
        # gather from all gpus
        batch_size_this = x.shape[0]
        x_gather = concat_all_gather(x)
        batch_size_all = x_gather.shape[0]

        num_gpus = batch_size_all // batch_size_this

        # restored index for this gpu
        gpu_idx = torch.distributed.get_rank()
        idx_this = idx_unshuffle.view(num_gpus, -1)[gpu_idx]

        return x_gather[idx_this]
    
    def _mask_pool(self, x, mask, randint):
        '''
        :param x:       N M C T V
        :param mask:    T V (bool) # paste True
        :return:        p [N C]   c [N C]
        '''
        if len(mask) == 2:
            p = (x[randint] * mask[0]).sum(-1).sum(-1) / (x[randint] + 1e-5).sum(-1).sum(-1)
            p = p.mean(-1).mean(-1)
            c = (x * (1 - mask[1])).sum(-1).sum(-1) / (x.sum(-1).sum(-1) + 1e-5)
            c = c.mean(-1).mean(-1)
            p, c = p /(p + c), c/(p + c)
            return p, c
            
        else:
            p = (x[randint] * mask).sum(-1).sum(-1) / (x[randint] + 1e-5).sum(-1).sum(-1)
            p = p.mean(-1).mean(-1)
            c = (x * (1 - mask)).sum(-1).sum(-1) / (x.sum(-1).sum(-1) + 1e-5)
            c = c.mean(-1).mean(-1)
            p, c = p /(p + c), c/(p + c)
            
            if random.random() < 0.001:
                print(f"p, c:{p,c}")
                print(f"lamb:{mask.mean()}")

            return p, c

    
    def forward(self, im_q_extreme, im_q, im_p=None, im_k=None, epoch=None, iter=None):
        """
        Input:
            im_q: a batch of query sequences
            im_k: a batch of key sequences
            im_q_extreme: a batch of extremely augmented query sequences
        """

        if not self.pretrain:
            return self.encoder_q(im_q)
        

        # self._weight_noisy()
        # Compute key features
        with torch.no_grad():  # no gradient to keys
            self._momentum_update_key_encoder()  # update the key encoder
            im_k, idx_k_unshuffle = self._batch_shuffle_ddp(im_k)
            
            im_k_mean = im_k.mean(dim=0, keepdim=True)
            [k_mean_feat] = self.encoder_k(im_k_mean)
            
            k, ks, _ = self.encoder_k(im_k, key=True, mean_feat=k_mean_feat)  # keys: NxC
            k = k * 0.5 + ks * 0.5
            
            k = F.normalize(k, dim=1)
            k = self._batch_unshuffle_ddp(k, idx_k_unshuffle)
            
            kn = k @ self.queue.clone().detach()
            
            _, _, qm = self.encoder_q(im_q, key=True, mean_feat=k_mean_feat)
            _, _, pm = self.encoder_q(im_p, key=True, mean_feat=k_mean_feat)
            
        N, C, T, V, M = im_q.shape
        [q] = self.encoder_q(self.actionlet_swap(im_q, im_q_extreme, qm))  # NxC
        q = F.normalize(q, dim=1)
        qn = (q @ self.queue.clone().detach()) ** 2
        qn_topk, kn_topk = self.select_topk(qn, kn)
        
        im_pt = self.ske_trans(im_p, epoch)
        randidx, im_ps, mask = self.ske_swap(im_pt, epoch)
        [qt] = self.encoder_q(self.actionlet_swap(im_pt, im_ps, pm))
        qt = F.normalize(qt, dim=1)
        qtn = (qt @ self.queue.clone().detach()) ** 2
        qtn_topk, kn_topk = self.select_topk(qtn, kn)
        
        randidx, im_pc, mask = self.ske_swap(im_p)
        p, c = self._mask_pool((pm.view(N, M, -1, (T + 3) // 4, V)/0.1).exp(), mask, randidx)
        kc = k * c + k[randidx] * p
        kc = F.normalize(kc, dim=1)
        [pc] = self.encoder_q(im_pc)
        pc = F.normalize(pc, dim=1)
        pcn = (pc @ self.queue.clone().detach()) ** 2
        kcn = kc @ self.queue.clone().detach()
        pcn_topk, kcn_topk = self.select_topk(pcn, kcn)

        self._dequeue_and_enqueue(k)
        global_loss = kld(qn_topk, kn_topk) + kld(qtn_topk, kn_topk) + kld(pcn_topk, kcn_topk)
        return global_loss
net/actclr.py
LanglandsLin/ActCLR
import torch
import torch.nn as nn
import torch.nn.functional as F

from net.utils.tgcn import ConvTemporalGraphical
from net.utils.graph import Graph

import random

trunk_ori_index = [4, 3, 21, 2, 1]
left_hand_ori_index = [9, 10, 11, 12, 24, 25]
right_hand_ori_index = [5, 6, 7, 8, 22, 23]
left_leg_ori_index = [17, 18, 19, 20]
right_leg_ori_index = [13, 14, 15, 16]

trunk = [i - 1 for i in trunk_ori_index]
left_hand = [i - 1 for i in left_hand_ori_index]
right_hand = [i - 1 for i in right_hand_ori_index]
left_leg = [i - 1 for i in left_leg_ori_index]
right_leg = [i - 1 for i in right_leg_ori_index]
body_parts = [trunk, left_hand, right_hand, left_leg, right_leg]

class TemporaryGrad(object):
    def __enter__(self):
        self.prev = torch.is_grad_enabled()
        torch.set_grad_enabled(True)

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        torch.set_grad_enabled(self.prev)

class Model(nn.Module):
    r"""Spatial temporal graph convolutional networks."""

    def __init__(self, in_channels, hidden_channels, hidden_dim, num_class, graph_args,
                 edge_importance_weighting, **kwargs):
        super().__init__()

        # load graph
        self.graph = Graph(**graph_args)
        A = torch.tensor(self.graph.A, dtype=torch.float32, requires_grad=False)
        self.register_buffer('A', A)

        # build networks
        spatial_kernel_size = A.size(0)
        temporal_kernel_size = 9
        kernel_size = (temporal_kernel_size, spatial_kernel_size)
        self.data_bn = nn.BatchNorm1d(in_channels * A.size(1))
        kwargs0 = {k: v for k, v in kwargs.items() if k != 'dropout'}
        self.st_gcn_networks = nn.ModuleList((
            st_gcn(in_channels, hidden_channels, kernel_size, 1, residual=False, **kwargs0),
            st_gcn(hidden_channels, hidden_channels, kernel_size, 1, **kwargs),
            st_gcn(hidden_channels, hidden_channels, kernel_size, 1, **kwargs),
            st_gcn(hidden_channels, hidden_channels, kernel_size, 1, **kwargs),
            st_gcn(hidden_channels, hidden_channels * 2, kernel_size, 2, **kwargs),
            st_gcn(hidden_channels * 2, hidden_channels * 2, kernel_size, 1, **kwargs),
            st_gcn(hidden_channels * 2, hidden_channels * 2, kernel_size, 1, **kwargs),
            st_gcn(hidden_channels * 2, hidden_channels * 4, kernel_size, 2, **kwargs),
            st_gcn(hidden_channels * 4, hidden_channels * 4, kernel_size, 1, **kwargs),
            st_gcn(hidden_channels * 4, hidden_dim, kernel_size, 1, **kwargs),
        ))
        self.fc = nn.Linear(hidden_dim, num_class)

        # initialize parameters for edge importance weighting
        if edge_importance_weighting:
            self.edge_importance = nn.ParameterList([
                nn.Parameter(torch.ones(self.A.size()))
                for i in self.st_gcn_networks
            ])
        else:
            self.edge_importance = [1] * len(self.st_gcn_networks)

    def _mask_pool(self, x, mask):
        '''
        :param x:       N C T V
        :param mask:    T V (bool) # paste True
        :return:        p [N C]   c [N C]
        '''
        p = (x * mask).sum(-1).sum(-1) / (mask.sum(-1).sum(-1) + 1e-5)
        c = (x * (1 - mask)).sum(-1).sum(-1) / ((1 - mask).sum(-1).sum(-1) + 1e-5)

        return p, c
    
    def _generate_mask(self, x, N, M, mean_feat):
        y = x.detach() # NM C T V
        y.requires_grad = True
        with TemporaryGrad():
            z = F.avg_pool2d(y, y.size()[2:])
            z = z.view(N, M, -1).mean(dim=1)
            z = self.fc(z)
            s = (F.normalize(z,dim=-1) * F.normalize(mean_feat,dim=-1)).sum(-1).mean()
            (1-s).backward()
            dz_dy = y.grad
        dz_dy_mean = dz_dy.mean(dim=(2,3))
        cam = dz_dy_mean[:,:,None,None] * y
        cam = F.relu(cam.mean(dim=1, keepdim=True)) / (1e-8)
        for body_part in body_parts:
            cam[:, :, :, body_part] = cam[:, :, :, body_part].mean(dim=-1, keepdim=True)
        cam = cam.clip(0, 1)
        return (cam).float().detach()
        

    def forward(self, x, key=False, mean_feat=None):

        # data normalization
        N, C, T, V, M = x.size()
        x = x.permute(0, 4, 3, 1, 2).contiguous()
        x = x.view(N * M, V * C, T)
        x = self.data_bn(x)
        x = x.view(N, M, V, C, T)
        x = x.permute(0, 1, 3, 4, 2).contiguous()
        x = x.view(N * M, C, T, V)

        # forward
        for gcn, importance in zip(self.st_gcn_networks, self.edge_importance):
            x, _ = gcn(x, self.A * importance)
            
        output = []
        
        if key:
            m = self._generate_mask(x, N, M, mean_feat)
            pc, _ = self._mask_pool(x, m)
            pc = pc.view(N, M, -1).mean(dim=1)
            pc = self.fc(pc)
            pc = pc.view(pc.size(0), -1)
            
            x = F.avg_pool2d(x, x.size()[2:])
            x = x.view(N, M, -1).mean(dim=1)

            # prediction
            x = self.fc(x)
            x = x.view(x.size(0), -1)
            
            output.extend([x, pc, m])
            
        else:
                
            x = F.avg_pool2d(x, x.size()[2:])
            x = x.view(N, M, -1).mean(dim=1)

            # prediction
            x = self.fc(x)
            x = x.view(x.size(0), -1)
            
            output.append(x)

        return output


class st_gcn(nn.Module):
    r"""Applies a spatial temporal graph convolution over an input graph sequence.

    Args:
        in_channels (int): Number of channels in the input sequence data
        out_channels (int): Number of channels produced by the convolution
        kernel_size (tuple): Size of the temporal convolving kernel and graph convolving kernel
        stride (int, optional): Stride of the temporal convolution. Default: 1
        dropout (int, optional): Dropout rate of the final output. Default: 0
        residual (bool, optional): If ``True``, applies a residual mechanism. Default: ``True``

    Shape:
        - Input[0]: Input graph sequence in :math:`(N, in_channels, T_{in}, V)` format
        - Input[1]: Input graph adjacency matrix in :math:`(K, V, V)` format
        - Output[0]: Outpu graph sequence in :math:`(N, out_channels, T_{out}, V)` format
        - Output[1]: Graph adjacency matrix for output data in :math:`(K, V, V)` format

        where
            :math:`N` is a batch size,
            :math:`K` is the spatial kernel size, as :math:`K == kernel_size[1]`,
            :math:`T_{in}/T_{out}` is a length of input/output sequence,
            :math:`V` is the number of graph nodes.

    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride=1,
                 dropout=0,
                 residual=True):
        super().__init__()

        assert len(kernel_size) == 2
        assert kernel_size[0] % 2 == 1
        padding = ((kernel_size[0] - 1) // 2, 0)

        self.gcn = ConvTemporalGraphical(in_channels, out_channels,
                                         kernel_size[1])

        self.tcn = nn.Sequential(
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(
                out_channels,
                out_channels,
                (kernel_size[0], 1),
                (stride, 1),
                padding,
            ),
            nn.BatchNorm2d(out_channels),
            nn.Dropout(dropout, inplace=True),
        )

        if not residual:
            self.residual = lambda x: 0

        elif (in_channels == out_channels) and (stride == 1):
            self.residual = lambda x: x

        else:
            self.residual = nn.Sequential(
                nn.Conv2d(
                    in_channels,
                    out_channels,
                    kernel_size=1,
                    stride=(stride, 1)),
                nn.BatchNorm2d(out_channels),
            )

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

    def forward(self, x, A):

        res = self.residual(x)
        x, A = self.gcn(x, A)
        x = self.tcn(x) + res

        return self.relu(x), A
net/st_gcn.py
LanglandsLin/ActCLR
import numpy as np

class Graph():
    """ The Graph to model the skeletons extracted by the openpose

    Args:
        strategy (string): must be one of the follow candidates
        - uniform: Uniform Labeling
        - distance: Distance Partitioning
        - spatial: Spatial Configuration
        For more information, please refer to the section 'Partition Strategies'
            in our paper (https://arxiv.org/abs/1801.07455).

        layout (string): must be one of the follow candidates
        - openpose: Is consists of 18 joints. For more information, please
            refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose#output
        - ntu-rgb+d: Is consists of 25 joints. For more information, please
            refer to https://github.com/shahroudy/NTURGB-D

        max_hop (int): the maximal distance between two connected nodes
        dilation (int): controls the spacing between the kernel points

    """

    def __init__(self,
                 layout='openpose',
                 strategy='uniform',
                 max_hop=1,
                 dilation=1):
        self.max_hop = max_hop
        self.dilation = dilation

        self.get_edge(layout)
        self.hop_dis = get_hop_distance(
            self.num_node, self.edge, max_hop=max_hop)
        self.get_adjacency(strategy)

    def __str__(self):
        return self.A

    def get_edge(self, layout):
        if layout == 'openpose':
            self.num_node = 18
            self_link = [(i, i) for i in range(self.num_node)]
            neighbor_link = [(4, 3), (3, 2), (7, 6), (6, 5), (13, 12), (12, 11),
                             (10, 9), (9, 8), (11, 5), (8, 2), (5, 1), (2, 1),
                             (0, 1), (15, 0), (14, 0), (17, 15), (16, 14)]
            self.edge = self_link + neighbor_link
            self.center = 1
        elif layout == 'ntu-rgb+d':
            self.num_node = 25
            self_link = [(i, i) for i in range(self.num_node)]
            neighbor_1base = [(1, 2), (2, 21), (3, 21), (4, 3), (5, 21),
                              (6, 5), (7, 6), (8, 7), (9, 21), (10, 9),
                              (11, 10), (12, 11), (13, 1), (14, 13), (15, 14),
                              (16, 15), (17, 1), (18, 17), (19, 18), (20, 19),
                              (22, 23), (23, 8), (24, 25), (25, 12)]
            neighbor_link = [(i - 1, j - 1) for (i, j) in neighbor_1base]
            self.edge = self_link + neighbor_link
            self.center = 21 - 1
        elif layout == 'nw-ucla':
            self.num_node = 20
            self_link = [(i, i) for i in range(self.num_node)]
            neighbor_1base = [(1, 2), (2, 3), (4, 3), (5, 3), (6, 5), (7, 6),
                              (8, 7), (9, 3), (10, 9), (11, 10), (12, 11), (13, 1),
                              (14, 13), (15, 14), (16, 15), (17, 1), (18, 17), (19, 18),
                              (20, 19)]
            neighbor_link = [(i - 1, j - 1) for (i, j) in neighbor_1base]
            self.edge = self_link + neighbor_link
            self.center = 3 - 1
        else:
            raise ValueError("Do Not Exist This Layout.")

    def get_adjacency(self, strategy):
        valid_hop = range(0, self.max_hop + 1, self.dilation)
        adjacency = np.zeros((self.num_node, self.num_node))
        for hop in valid_hop:
            adjacency[self.hop_dis == hop] = 1
        normalize_adjacency = normalize_digraph(adjacency)

        if strategy == 'uniform':
            A = np.zeros((1, self.num_node, self.num_node))
            A[0] = normalize_adjacency
            self.A = A
        elif strategy == 'distance':
            A = np.zeros((len(valid_hop), self.num_node, self.num_node))
            for i, hop in enumerate(valid_hop):
                A[i][self.hop_dis == hop] = normalize_adjacency[self.hop_dis ==
                                                                hop]
            self.A = A
        elif strategy == 'spatial':
            A = []
            for hop in valid_hop:
                a_root = np.zeros((self.num_node, self.num_node))
                a_close = np.zeros((self.num_node, self.num_node))
                a_further = np.zeros((self.num_node, self.num_node))
                for i in range(self.num_node):
                    for j in range(self.num_node):
                        if self.hop_dis[j, i] == hop:
                            if self.hop_dis[j, self.center] == self.hop_dis[
                                    i, self.center]:
                                a_root[j, i] = normalize_adjacency[j, i]
                            elif self.hop_dis[j, self.
                                              center] > self.hop_dis[i, self.
                                                                     center]:
                                a_close[j, i] = normalize_adjacency[j, i]
                            else:
                                a_further[j, i] = normalize_adjacency[j, i]
                if hop == 0:
                    A.append(a_root)
                else:
                    A.append(a_root + a_close)
                    A.append(a_further)
            A = np.stack(A)
            self.A = A
        else:
            raise ValueError("Do Not Exist This Strategy")


def get_hop_distance(num_node, edge, max_hop=1):
    A = np.zeros((num_node, num_node))
    for i, j in edge:
        A[j, i] = 1
        A[i, j] = 1

    # compute hop steps
    hop_dis = np.zeros((num_node, num_node)) + np.inf
    transfer_mat = [np.linalg.matrix_power(A, d) for d in range(max_hop + 1)]
    arrive_mat = (np.stack(transfer_mat) > 0)
    for d in range(max_hop, -1, -1):
        hop_dis[arrive_mat[d]] = d
    return hop_dis


def normalize_digraph(A):
    Dl = np.sum(A, 0)
    num_node = A.shape[0]
    Dn = np.zeros((num_node, num_node))
    for i in range(num_node):
        if Dl[i] > 0:
            Dn[i, i] = Dl[i]**(-1)
    AD = np.dot(A, Dn)
    return AD


def normalize_undigraph(A):
    Dl = np.sum(A, 0)
    num_node = A.shape[0]
    Dn = np.zeros((num_node, num_node))
    for i in range(num_node):
        if Dl[i] > 0:
            Dn[i, i] = Dl[i]**(-0.5)
    DAD = np.dot(np.dot(Dn, A), Dn)
    return DAD

net/utils/graph.py
LanglandsLin/ActCLR
import sys
import argparse
import yaml
import math
import numpy as np

# torch
import torch
import torch.nn as nn
import torch.optim as optim

# torchlight
import torchlight
from torchlight import str2bool
from torchlight import DictAction
from torchlight import import_class

from .processor import Processor

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv1d') != -1:
        m.weight.data.normal_(0.0, 0.02)
        if m.bias is not None:
            m.bias.data.fill_(0)
    elif classname.find('Conv2d') != -1:
        m.weight.data.normal_(0.0, 0.02)
        if m.bias is not None:
            m.bias.data.fill_(0)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

class FT_Processor(Processor):
    """
        Processor for Finetune Evaluation.
    """

    def load_model(self):
        self.model = self.io.load_model(self.arg.model,
                                        **(self.arg.model_args))
        self.model.apply(weights_init)

        self.loss = nn.CrossEntropyLoss()
        
    def load_optimizer(self):
        if self.arg.optimizer == 'SGD':
            self.optimizer = optim.SGD(
                self.model.parameters(),
                lr=self.arg.base_lr,
                momentum=0.9,
                nesterov=self.arg.nesterov,
                weight_decay=self.arg.weight_decay)
        elif self.arg.optimizer == 'Adam':
            self.optimizer = optim.Adam(
                self.model.parameters(),
                lr=self.arg.base_lr,
                weight_decay=self.arg.weight_decay)
        else:
            raise ValueError()

    def adjust_lr(self):
        if self.arg.optimizer == 'SGD' and self.arg.step:
            lr = self.arg.base_lr * (
                0.1**np.sum(self.meta_info['epoch']>= np.array(self.arg.step)))
            for param_group in self.optimizer.param_groups:
                param_group['lr'] = lr
            self.lr = lr
        else:
            self.lr = self.arg.base_lr

    def show_topk(self, k):
        rank = self.result.argsort()
        hit_top_k = [l in rank[i, -k:] for i, l in enumerate(self.label)]
        accuracy = sum(hit_top_k) * 1.0 / len(hit_top_k)
        self.io.print_log('\tTop{}: {:.2f}%'.format(k, 100 * accuracy))

    def show_best(self, k):
        rank = self.result.argsort()
        hit_top_k = [l in rank[i, -k:] for i, l in enumerate(self.label)]
        accuracy = 100 * sum(hit_top_k) * 1.0 / len(hit_top_k)
        accuracy = round(accuracy, 5)
        self.current_result = accuracy
        if self.best_result <= accuracy:
            self.best_result = accuracy
        self.io.print_log('\tBest Top{}: {:.2f}%'.format(k, self.best_result))

    def train(self, epoch):
        self.model.train()
        self.adjust_lr()
        loader = self.data_loader['train']
        loss_value = []
        result_frag = []
        label_frag = []

        for data, label in loader:
            self.global_step += 1
            # get data
            data = data.float().to(self.dev, non_blocking=True)
            label = label.long().to(self.dev, non_blocking=True)
            
            if self.arg.stream == 'joint':
                pass
            elif self.arg.stream == 'motion':
                motion = torch.zeros_like(data)

                motion[:, :, :-1, :, :] = data[:, :, 1:, :, :] - data[:, :, :-1, :, :]

                data = motion
            elif self.arg.stream == 'bone':
                Bone = [(1, 2), (2, 21), (3, 21), (4, 3), (5, 21), (6, 5), (7, 6), (8, 7), (9, 21),
                        (10, 9), (11, 10), (12, 11), (13, 1), (14, 13), (15, 14), (16, 15), (17, 1),
                        (18, 17), (19, 18), (20, 19), (21, 21), (22, 23), (23, 8), (24, 25), (25, 12)]

                bone = torch.zeros_like(data)

                for v1, v2 in Bone:
                    bone[:, :, :, v1 - 1, :] = data[:, :, :, v1 - 1, :] - data[:, :, :, v2 - 1, :]
                    
                data = bone
            else:
                raise ValueError

            # forward
            output = self.model(None, data)
            loss = self.loss(output, label)

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            result_frag.append(output.data.cpu().numpy())
            label_frag.append(label.data.cpu().numpy())

            # statistics
            self.iter_info['loss'] = loss.data.item()
            self.iter_info['lr'] = '{:.6f}'.format(self.lr)
            loss_value.append(self.iter_info['loss'])
            self.show_iter_info()
            self.meta_info['iter'] += 1
            self.train_log_writer(epoch)

        self.result = np.concatenate(result_frag)
        self.label = np.concatenate(label_frag)

        for k in self.arg.show_topk:
            self.show_topk(k)

        self.epoch_info['train_mean_loss']= np.mean(loss_value)
        self.train_writer.add_scalar('loss', self.epoch_info['train_mean_loss'], epoch)
        self.show_epoch_info()

    def test(self, epoch):
        self.model.eval()
        loader = self.data_loader['test']
        loss_value = []
        result_frag = []
        label_frag = []

        for data, label in loader:
            # get data
            data = data.float().to(self.dev, non_blocking=True)
            label = label.long().to(self.dev, non_blocking=True)
            
            if self.arg.stream == 'joint':
                pass
            elif self.arg.stream == 'motion':
                motion = torch.zeros_like(data)

                motion[:, :, :-1, :, :] = data[:, :, 1:, :, :] - data[:, :, :-1, :, :]

                data = motion
            elif self.arg.stream == 'bone':
                Bone = [(1, 2), (2, 21), (3, 21), (4, 3), (5, 21), (6, 5), (7, 6), (8, 7), (9, 21),
                        (10, 9), (11, 10), (12, 11), (13, 1), (14, 13), (15, 14), (16, 15), (17, 1),
                        (18, 17), (19, 18), (20, 19), (21, 21), (22, 23), (23, 8), (24, 25), (25, 12)]

                bone = torch.zeros_like(data)

                for v1, v2 in Bone:
                    bone[:, :, :, v1 - 1, :] = data[:, :, :, v1 - 1, :] - data[:, :, :, v2 - 1, :]
                    
                data = bone
            else:
                raise ValueError

            # inference
            with torch.no_grad():
                output = self.model(None, data)
            result_frag.append(output.data.cpu().numpy())

            # get loss
            loss = self.loss(output, label)
            loss_value.append(loss.item())
            label_frag.append(label.data.cpu().numpy())

        self.result = np.concatenate(result_frag)
        self.label = np.concatenate(label_frag)

        self.eval_info['eval_mean_loss']= np.mean(loss_value)
        self.show_eval_info()

        # show top-k accuracy 
        for k in self.arg.show_topk:
            self.show_topk(k)
        self.show_best(1)

        self.eval_log_writer(epoch)

    @staticmethod
    def get_parser(add_help=False):
        # parameter priority: command line > config > default
        parent_parser = Processor.get_parser(add_help=False)
        parser = argparse.ArgumentParser(
            add_help=add_help,
            parents=[parent_parser],
            description='Spatial Temporal Graph Convolution Network')

        # evaluation
        parser.add_argument('--show_topk', type=int, default=[1, 5], nargs='+', help='which Top K accuracy will be shown')
        # optim
        parser.add_argument('--base_lr', type=float, default=0.01, help='initial learning rate')
        parser.add_argument('--step', type=int, default=[], nargs='+', help='the epoch where optimizer reduce the learning rate')
        parser.add_argument('--optimizer', default='SGD', help='type of optimizer')
        parser.add_argument('--nesterov', type=str2bool, default=True, help='use nesterov or not')
        parser.add_argument('--weight_decay', type=float, default=0.0001, help='weight decay for optimizer')
        parser.add_argument('--stream', type=str, default='joint', help='the stream of input')
        parser.add_argument('--mining_epoch', type=int, default=1e6,
                            help='the starting epoch of nearest neighbor mining')
        parser.add_argument('--topk', type=int, default=1, help='topk samples in nearest neighbor mining')

        return parser

processor/finetune_evaluation.py
LanglandsLin/ActCLR
import sys
import argparse
import yaml
import numpy as np

# torch
import torch
import torch.nn as nn

# torchlight
import torchlight
from torchlight import str2bool
from torchlight import DictAction
from torchlight import import_class

class IO():
    """
        IO Processor
    """

    def __init__(self, argv=None):

        self.load_arg(argv)
        self.init_environment()
        self.load_model()
        self.load_weights()
        self.gpu()

    def load_arg(self, argv=None):
        parser = self.get_parser()

        # load arg form config file
        p = parser.parse_args(argv)
        if p.config is not None:
            # load config file
            with open(p.config, 'r') as f:
                default_arg = yaml.load(f, Loader=yaml.FullLoader)

            # update parser from config file
            key = vars(p).keys()
            for k in default_arg.keys():
                if k not in key:
                    print('Unknown Arguments: {}'.format(k))
                    assert k in key

            parser.set_defaults(**default_arg)

        self.arg = parser.parse_args(argv)

    def init_environment(self):
        self.io = torchlight.IO(
            self.arg.work_dir,
            save_log=self.arg.save_log,
            print_log=self.arg.print_log)
        self.io.save_arg(self.arg)

        # gpu
        # if self.arg.use_gpu:
        #     gpus = torchlight.visible_gpu(self.arg.device)
        #     torchlight.occupy_gpu(gpus)
        #     self.gpus = gpus
        #     self.dev = "cuda:0"
        # else:
        #     self.dev = "cpu"

    def load_model(self):
        self.model = self.io.load_model(self.arg.model,
                                        **(self.arg.model_args))

    def load_weights(self):
        if self.arg.weights:
            self.model = self.io.load_weights(self.model, self.arg.weights,
                                              self.arg.ignore_weights, self.arg.rename_weights)

    def gpu(self):
        # move modules to gpu
        self.model = self.model.to(self.dev)
        for name, value in vars(self).items():
            cls_name = str(value.__class__)
            if cls_name.find('torch.nn.modules') != -1:
                setattr(self, name, value.to(self.dev))

        # model parallel
        if self.arg.use_gpu and len(self.gpus) > 1:
            self.model = nn.DataParallel(self.model, device_ids=self.gpus)

    def start(self):
        self.io.print_log('Parameters:\n{}\n'.format(str(vars(self.arg))))

    @staticmethod
    def get_parser(add_help=False):
        # parameter priority: command line > config > default
        parser = argparse.ArgumentParser( add_help=add_help, description='IO Processor')

        parser.add_argument('-w', '--work_dir', default='./work_dir/levi', help='the work folder for storing results')
        parser.add_argument('-c', '--config', default=None, help='path to the configuration file')

        # processor
        parser.add_argument('--use_gpu', type=str2bool, default=True, help='use GPUs or not')
        parser.add_argument('--device', type=int, default=0, nargs='+', help='the indexes of GPUs for training or testing')

        # visualize and debug
        parser.add_argument('--print_log', type=str2bool, default=True, help='print logging or not')
        parser.add_argument('--save_log', type=str2bool, default=True, help='save logging or not')

        # model
        parser.add_argument('--model', default=None, help='the model will be used')
        parser.add_argument('--model_args', action=DictAction, default=dict(), help='the arguments of model')
        parser.add_argument('--weights', default=None, help='the weights for network initialization')
        parser.add_argument('--ignore_weights', type=str, default=[], nargs='+', help='the name of weights which will be ignored in the initialization')
        parser.add_argument('--rename_weights', type=str, default=[], nargs='+', help='the name of weights which will be ignored in the initialization')
        
        return parser

processor/io.py
LanglandsLin/ActCLR
import torch

# knn monitor as in InstDisc https://arxiv.org/abs/1805.01978
# implementation follows http://github.com/zhirongw/lemniscate.pytorch and https://github.com/leftthomas/SimCLR
def knn_predict(feature, feature_bank, feature_labels, classes, knn_k, knn_t):
    # compute cos similarity between each feature vector and feature bank ---> [B, N]
    sim_matrix = torch.mm(feature, feature_bank)
    # [B, K]
    sim_weight, sim_indices = sim_matrix.topk(k=knn_k, dim=-1)
    # [B, K]
    sim_labels = torch.gather(feature_labels.expand(feature.size(0), -1), dim=-1, index=sim_indices)
    sim_weight = (sim_weight / knn_t).exp()

    # counts for each class
    one_hot_label = torch.zeros(feature.size(0) * knn_k, classes, device=sim_labels.device)
    # [B*K, C]
    one_hot_label = one_hot_label.scatter(dim=-1, index=sim_labels.view(-1, 1), value=1.0)
    # weighted score ---> [B, C]
    pred_scores = torch.sum(one_hot_label.view(feature.size(0), -1, classes) * sim_weight.unsqueeze(dim=-1), dim=1)

    pred_labels = pred_scores.argsort(dim=-1, descending=True)
    return pred_labels
processor/knn_monitor.py
LanglandsLin/ActCLR
import sys
import argparse
import yaml
import math
import numpy as np

# torch
import torch
import torch.nn as nn
import torch.optim as optim

# torchlight
import torchlight
from torchlight import str2bool
from torchlight import DictAction
from torchlight import import_class

from .processor import Processor

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv1d') != -1 or classname.find('Conv2d') != -1 or classname.find('Linear') != -1:
        m.weight.data.normal_(0.0, 0.02)
        if m.bias is not None:
            m.bias.data.fill_(0)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

class LE_Processor(Processor):
    """
        Processor for Linear Evaluation.
    """

    def load_model(self):
        self.model = self.io.load_model(self.arg.model,
                                        **(self.arg.model_args))
        self.model.apply(weights_init)

        for name, param in self.model.encoder_q.named_parameters():
            if name not in ['fc.weight', 'fc.bias']:
                param.requires_grad = False
        self.num_grad_layers = 2
        if hasattr(self.model, 'encoder_q_motion'):
            for name, param in self.model.encoder_q_motion.named_parameters():
                if name not in ['fc.weight', 'fc.bias']:
                    param.requires_grad = False
            self.num_grad_layers += 2
        if hasattr(self.model, 'encoder_q_bone'):
            for name, param in self.model.encoder_q_bone.named_parameters():
                if name not in ['fc.weight', 'fc.bias']:
                    param.requires_grad = False
            self.num_grad_layers += 2

        self.loss = nn.CrossEntropyLoss()
        
    def load_optimizer(self):
        parameters = list(filter(lambda p: p.requires_grad, self.model.parameters()))
        assert len(parameters) == self.num_grad_layers
        if self.arg.optimizer == 'SGD':
            self.optimizer = optim.SGD(
                parameters,
                lr=self.arg.base_lr,
                momentum=0.9,
                nesterov=self.arg.nesterov,
                weight_decay=self.arg.weight_decay)
        elif self.arg.optimizer == 'Adam':
            self.optimizer = optim.Adam(
                parameters,
                lr=self.arg.base_lr,
                weight_decay=self.arg.weight_decay)
        else:
            raise ValueError()

    def adjust_lr(self):
        if self.arg.optimizer == 'SGD' and self.arg.step:
            lr = self.arg.base_lr * (
                0.1**np.sum(self.meta_info['epoch'] > np.array(self.arg.step)))
            for param_group in self.optimizer.param_groups:
                param_group['lr'] = lr
            self.lr = lr
        else:
            self.lr = self.arg.base_lr

    def show_topk(self, k):
        rank = self.result.argsort()
        hit_top_k = [l in rank[i, -k:] for i, l in enumerate(self.label)]
        accuracy = sum(hit_top_k) * 1.0 / len(hit_top_k)
        self.io.print_log('\tTop{}: {:.2f}%'.format(k, 100 * accuracy))

    def show_best(self, k):
        rank = self.result.argsort()
        hit_top_k = [l in rank[i, -k:] for i, l in enumerate(self.label)]
        accuracy = 100 * sum(hit_top_k) * 1.0 / len(hit_top_k)
        accuracy = round(accuracy, 5)
        self.current_result = accuracy
        if self.best_result <= accuracy:
            self.best_result = accuracy
        self.io.print_log('\tBest Top{}: {:.2f}%'.format(k, self.best_result))

    def confusion_matrix(self):
        predict = np.argmax(self.result, axis=1)
        matrix = np.zeros((60, 60))
        for i, j in zip(predict, self.label):
            matrix[i, j] += 1
        matrix = matrix / matrix.sum(axis=1)
        for i in range(60):
            self.io.print_log('\t{}: {:.2f}%'.format(i, matrix[i, i]))

    def train(self, epoch):
        self.model.eval()
        self.adjust_lr()
        loader = self.data_loader['train']
        loss_value = []

        for data, label in loader:
            self.global_step += 1
            # get data
            data = data.float().to(self.dev, non_blocking=True)
            label = label.long().to(self.dev, non_blocking=True)
            
            if self.arg.stream == 'joint':
                pass
            elif self.arg.stream == 'motion':
                motion = torch.zeros_like(data)

                motion[:, :, :-1, :, :] = data[:, :, 1:, :, :] - data[:, :, :-1, :, :]

                data = motion
            elif self.arg.stream == 'bone':
                Bone = [(1, 2), (2, 21), (3, 21), (4, 3), (5, 21), (6, 5), (7, 6), (8, 7), (9, 21),
                        (10, 9), (11, 10), (12, 11), (13, 1), (14, 13), (15, 14), (16, 15), (17, 1),
                        (18, 17), (19, 18), (20, 19), (21, 21), (22, 23), (23, 8), (24, 25), (25, 12)]

                bone = torch.zeros_like(data)

                for v1, v2 in Bone:
                    bone[:, :, :, v1 - 1, :] = data[:, :, :, v1 - 1, :] - data[:, :, :, v2 - 1, :]
                    
                data = bone
            else:
                raise ValueError

            # forward
            output = self.model(None, data)
            loss = self.loss(output, label)

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            # statistics
            self.iter_info['loss'] = loss.data.item()
            self.iter_info['lr'] = '{:.6f}'.format(self.lr)
            loss_value.append(self.iter_info['loss'])
            self.show_iter_info()
            self.meta_info['iter'] += 1
            self.train_log_writer(epoch)

        self.epoch_info['train_mean_loss']= np.mean(loss_value)
        if self.local_rank == 0:
            self.train_writer.add_scalar('loss', self.epoch_info['train_mean_loss'], epoch)
        self.show_epoch_info()

    def test(self, epoch):
        self.model.eval()
        loader = self.data_loader['test']
        loss_value = []
        result_frag = []
        label_frag = []

        for data, label in loader:
            # get data
            data = data.float().to(self.dev, non_blocking=True)
            label = label.long().to(self.dev, non_blocking=True)
            
            if self.arg.stream == 'joint':
                pass
            elif self.arg.stream == 'motion':
                motion = torch.zeros_like(data)

                motion[:, :, :-1, :, :] = data[:, :, 1:, :, :] - data[:, :, :-1, :, :]

                data = motion
            elif self.arg.stream == 'bone':
                Bone = [(1, 2), (2, 21), (3, 21), (4, 3), (5, 21), (6, 5), (7, 6), (8, 7), (9, 21),
                        (10, 9), (11, 10), (12, 11), (13, 1), (14, 13), (15, 14), (16, 15), (17, 1),
                        (18, 17), (19, 18), (20, 19), (21, 21), (22, 23), (23, 8), (24, 25), (25, 12)]

                bone = torch.zeros_like(data)

                for v1, v2 in Bone:
                    bone[:, :, :, v1 - 1, :] = data[:, :, :, v1 - 1, :] - data[:, :, :, v2 - 1, :]
                    
                data = bone
            else:
                raise ValueError

            # inference
            with torch.no_grad():
                output = self.model(None, data)
            result_frag.append(output.data.cpu().numpy())

            # get loss
            loss = self.loss(output, label)
            loss_value.append(loss.item())
            label_frag.append(label.data.cpu().numpy())

        self.result = np.concatenate(result_frag)
        self.label = np.concatenate(label_frag)

        self.eval_info['eval_mean_loss']= np.mean(loss_value)
        self.show_eval_info()

        # show top-k accuracy 
        for k in self.arg.show_topk:
            self.show_topk(k)
        self.show_best(1)

        self.confusion_matrix()

        if self.local_rank == 0:
            self.eval_log_writer(epoch)

    @staticmethod
    def get_parser(add_help=False):
        # parameter priority: command line > config > default
        parent_parser = Processor.get_parser(add_help=False)
        parser = argparse.ArgumentParser(
            add_help=add_help,
            parents=[parent_parser],
            description='Spatial Temporal Graph Convolution Network')

        # evaluation
        parser.add_argument('--show_topk', type=int, default=[1, 5], nargs='+', help='which Top K accuracy will be shown')
        # optim
        parser.add_argument('--base_lr', type=float, default=0.01, help='initial learning rate')
        parser.add_argument('--step', type=int, default=[], nargs='+', help='the epoch where optimizer reduce the learning rate')
        parser.add_argument('--optimizer', default='SGD', help='type of optimizer')
        parser.add_argument('--nesterov', type=str2bool, default=True, help='use nesterov or not')
        parser.add_argument('--weight_decay', type=float, default=0.0001, help='weight decay for optimizer')
        parser.add_argument('--stream', type=str, default='joint', help='the stream of input')
        parser.add_argument('--mining_epoch', type=int, default=1e6,
                            help='the starting epoch of nearest neighbor mining')
        parser.add_argument('--topk', type=int, default=1, help='topk samples in nearest neighbor mining')

        return parser

processor/linear_evaluation.py
LanglandsLin/ActCLR
import sys
import argparse
import yaml
import math
import numpy as np

# torch
import torch
import torch.nn as nn
import torch.optim as optim

# torchlight
import torchlight
from torchlight import str2bool
from torchlight import DictAction
from torchlight import import_class

from .processor import Processor

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv1d') != -1 or classname.find('Conv2d') != -1 or classname.find('Linear') != -1:
        m.weight.data.normal_(0.0, 0.02)
        if m.bias is not None:
            m.bias.data.fill_(0)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

class PT_Processor(Processor):
    """
        Processor for Pretraining.
    """

    def load_model(self):
        self.model = self.io.load_model(self.arg.model,
                                        **(self.arg.model_args))
        self.model.apply(weights_init)
        self.loss = nn.CrossEntropyLoss()
        self.re_criterion = torch.nn.L1Loss(reduction='none')
        
    def load_optimizer(self):
        if self.arg.optimizer == 'SGD':
            self.optimizer = optim.SGD(
                self.model.parameters(),
                lr=self.arg.base_lr,
                momentum=0.9,
                nesterov=self.arg.nesterov,
                weight_decay=self.arg.weight_decay)
        elif self.arg.optimizer == 'Adam':
            self.optimizer = optim.Adam(
                self.model.parameters(),
                lr=self.arg.base_lr,
                weight_decay=self.arg.weight_decay)
        else:
            raise ValueError()

    def adjust_lr(self):
        if self.arg.optimizer == 'SGD' and self.arg.step:
            lr = self.arg.base_lr * (
                0.1**np.sum(self.meta_info['epoch'] > np.array(self.arg.step)))
            for param_group in self.optimizer.param_groups:
                param_group['lr'] = lr
            self.lr = lr
        else:
            self.lr = self.arg.base_lr

    def train(self, epoch):
        self.model.train()
        self.adjust_lr()
        loader = self.data_loader['train']
        loss_value = []

        for [data1, data2, data3], label in loader:
            self.global_step += 1
            # get data
            data1 = data1.float().to(self.dev, non_blocking=True)
            data2 = data2.float().to(self.dev, non_blocking=True)
            data3 = data3.float().to(self.dev, non_blocking=True)
            label = label.long().to(self.dev, non_blocking=True)

            # forward
            output, target = self.model(data1, data2, data3)
            loss = self.loss(output, target)

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            # statistics
            self.iter_info['loss'] = loss.data.item()
            self.iter_info['lr'] = '{:.6f}'.format(self.lr)
            loss_value.append(self.iter_info['loss'])
            self.show_iter_info()
            self.meta_info['iter'] += 1
            self.train_log_writer(epoch)

        self.epoch_info['train_mean_loss']= np.mean(loss_value)
        self.train_writer.add_scalar('loss', self.epoch_info['train_mean_loss'], epoch)
        self.show_epoch_info()

    @staticmethod
    def get_parser(add_help=False):
        # parameter priority: command line > config > default
        parent_parser = Processor.get_parser(add_help=False)
        parser = argparse.ArgumentParser(
            add_help=add_help,
            parents=[parent_parser],
            description='Spatial Temporal Graph Convolution Network')

        parser.add_argument('--base_lr', type=float, default=0.01, help='initial learning rate')
        parser.add_argument('--step', type=int, default=[], nargs='+', help='the epoch where optimizer reduce the learning rate')
        parser.add_argument('--optimizer', default='SGD', help='type of optimizer')
        parser.add_argument('--nesterov', type=str2bool, default=True, help='use nesterov or not')
        parser.add_argument('--weight_decay', type=float, default=0.0001, help='weight decay for optimizer')

        return parser

processor/pretrain.py
LanglandsLin/ActCLR
import sys
import argparse
import yaml
import math
import random
import numpy as np

# torch
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.distributed as dist

# torchlight
import torchlight
from torchlight import str2bool
from torchlight import DictAction
from torchlight import import_class

from .processor import Processor
from .pretrain import PT_Processor

from .knn_monitor import knn_predict
from tqdm import tqdm

@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor)
        for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output


class ActCLR_Processor(PT_Processor):
    """
        Processor for AimCLR Pre-training.
    """

    def train(self, epoch):
        self.model.train()
        self.adjust_lr()
        loader = self.data_loader['train']
        
        if epoch <= 100:
            loader.dataset.aug_method = random.sample(['1', '2', '3', '4', '5'], 5)
            loader.dataset.strong_aug_method = random.sample(['1', '2', '3', '4', '5', '8', '9'], 7)
        elif epoch > 100 and epoch <= 200:
            loader.dataset.aug_method = random.sample(['1', '2', '3', '4', '5'], 3)
            loader.dataset.strong_aug_method = random.sample(['1', '2', '3', '4', '5', '8', '9'], 5)
        elif epoch > 200 and epoch <= 300:
            loader.dataset.aug_method = random.sample(['1', '2', '3', '4', '5'], 1)
            loader.dataset.strong_aug_method = random.sample(['1', '2', '3', '4', '5', '8', '9'], 3)
        elif epoch > 300:
            loader.dataset.temperal_padding_ratio = -1
            loader.dataset.shear_amplitude = -1
            loader.dataset.aug_method = ''
            loader.dataset.strong_aug_method = ''
        
        loss_value = []

        for data1, data2, data3, data4 in tqdm(loader):
            self.global_step += 1
            # get data
            data1 = data1.float().to(self.dev, non_blocking=True)
            data2 = data2.float().to(self.dev, non_blocking=True)
            data3 = data3.float().to(self.dev, non_blocking=True)
            data4 = data4.float().to(self.dev, non_blocking=True)    


            data1 = self.view_gen(data1)
            data2 = self.view_gen(data2)
            data3 = self.view_gen(data3)
            data4 = self.view_gen(data4)

            # forward
            loss = self.model(data1, data2, data3, data4, epoch=epoch)

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            # statistics
            self.iter_info['loss'] = loss.data.item()
            self.iter_info['lr'] = '{:.6f}'.format(self.lr)
            loss_value.append(self.iter_info['loss'])
            self.show_iter_info()
            self.meta_info['iter'] += 1
            if self.local_rank == 0:
                self.train_log_writer(epoch)

        self.epoch_info['train_mean_loss'] = np.mean(loss_value)
        if self.local_rank == 0:
            self.train_writer.add_scalar('loss', self.epoch_info['train_mean_loss'], epoch)
        self.show_epoch_info()

    def view_gen(self, data):
        if self.arg.stream == 'joint':
            pass
        elif self.arg.stream == 'motion':
            motion = torch.zeros_like(data)

            motion[:, :, :-1, :, :] = data[:, :, 1:, :, :] - data[:, :, :-1, :, :]

            data = motion
        elif self.arg.stream == 'bone':
            Bone = [(1, 2), (2, 21), (3, 21), (4, 3), (5, 21), (6, 5), (7, 6), (8, 7), (9, 21),
                    (10, 9), (11, 10), (12, 11), (13, 1), (14, 13), (15, 14), (16, 15), (17, 1),
                    (18, 17), (19, 18), (20, 19), (21, 21), (22, 23), (23, 8), (24, 25), (25, 12)]

            bone = torch.zeros_like(data)

            for v1, v2 in Bone:
                bone[:, :, :, v1 - 1, :] = data[:, :, :, v1 - 1, :] - data[:, :, :, v2 - 1, :]
                bone[:, :, :, v1 - 1, :] = data[:, :, :, v1 - 1, :] - data[:, :, :, v2 - 1, :]

            data = bone
        else:
            raise ValueError

        return data
    
    @torch.no_grad()
    def knn_monitor(self, epoch):
        self.model.module.encoder_q.eval()
        feature_bank, label_bank = [], []
        with torch.no_grad():
            # generate feature bank
            for data, label in tqdm(self.data_loader['mem_train'], desc='Feature extracting'):
                data = data.float().to(self.dev, non_blocking=True)
                label = label.long().to(self.dev, non_blocking=True)

                data = self.view_gen(data)
                [feature] = self.model.module.encoder_q(data)
                feature = F.normalize(feature, dim=1)
                feature_bank.append(feature)
                label_bank.append(label)
            # [D, N]
            feature_bank = concat_all_gather(torch.cat(feature_bank, dim=0)).t().contiguous()
            # [N]
            feature_labels = concat_all_gather(torch.cat(label_bank)).to(feature_bank.device)
            # loop test data to predict the label by weighted knn search
            for i in self.arg.knn_k:
                total_top1, total_top5, total_num = 0, 0, 0
                test_bar = tqdm(self.data_loader['mem_test'], desc='kNN-{}'.format(i))
                for data, label in test_bar:
                    data = data.float().to(self.dev, non_blocking=True)
                    label = label.float().to(self.dev, non_blocking=True)

                    data = self.view_gen(data)

                    [feature] = self.model.module.encoder_q(data)
                    feature = concat_all_gather(F.normalize(feature, dim=1))

                    pred_labels = knn_predict(feature, feature_bank, feature_labels, self.arg.knn_classes, i,
                                              self.arg.knn_t)

                    total_num += feature.size(0)
                    total_top1 += (pred_labels[:, 0] == concat_all_gather(label)).float().sum().item()
                    test_bar.set_postfix({'k': i, 'Accuracy': total_top1 / total_num * 100})
                acc = total_top1 / total_num * 100

                self.knn_results[i][epoch] = acc
                if self.local_rank == 0:
                    self.train_writer.add_scalar('KNN-{}'.format(i), acc, epoch)

    @staticmethod
    def get_parser(add_help=False):
        # parameter priority: command line > config > default
        parent_parser = Processor.get_parser(add_help=False)
        parser = argparse.ArgumentParser(
            add_help=add_help,
            parents=[parent_parser],
            description='Spatial Temporal Graph Convolution Network')

        parser.add_argument('--base_lr', type=float, default=0.01, help='initial learning rate')
        parser.add_argument('--step', type=int, default=[], nargs='+', help='the epoch where optimizer reduce the learning rate')
        parser.add_argument('--optimizer', default='SGD', help='type of optimizer')
        parser.add_argument('--nesterov', type=str2bool, default=True, help='use nesterov or not')
        parser.add_argument('--weight_decay', type=float, default=0.0001, help='weight decay for optimizer')
        parser.add_argument('--stream', type=str, default='joint', help='the stream of input')
        parser.add_argument('--mining_epoch', type=int, default=1e6,
                            help='the starting epoch of nearest neighbor mining')
        parser.add_argument('--topk', type=int, default=1, help='topk samples in nearest neighbor mining')

        return parser

processor/pretrain_actclr.py
LanglandsLin/ActCLR
import sys
import argparse
import yaml
import os
import shutil
import numpy as np
import random
import math

# torch
import torch
import torch.nn as nn
import torch.optim as optim

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# torchlight
import torchlight
from torchlight import str2bool
from torchlight import DictAction
from torchlight import import_class

from .io import IO
from tensorboardX import SummaryWriter

import subprocess

def init_seed(seed=1):
    torch.cuda.manual_seed_all(seed)
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    # torch.backends.cudnn.enabled = False
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


class Processor(IO):
    """
        Base Processor
    """

    def __init__(self, argv=None):

        self.load_arg(argv)

        self.global_step = 0

    def train_log_writer(self, epoch):
        self.train_writer.add_scalar('batch_loss', self.iter_info['loss'], self.global_step)
        self.train_writer.add_scalar('lr', self.lr, self.global_step)
        self.train_writer.add_scalar('epoch', epoch, self.global_step)

    def eval_log_writer(self, epoch):
        self.val_writer.add_scalar('eval_loss', self.eval_info['eval_mean_loss'], epoch)
        self.val_writer.add_scalar('current_result', self.current_result, epoch)
        self.val_writer.add_scalar('best_result', self.best_result, epoch)

    def init_environment(self):
        super().init_environment()
        self.result = dict()
        self.iter_info = dict()
        self.epoch_info = dict()
        self.eval_info = dict()
        self.meta_info = dict(epoch=0, iter=0)
        dist.init_process_group(backend='nccl')

    def load_optimizer(self):
        pass

    def load_data(self):
        self.data_loader = dict()

        if self.arg.train_feeder_args:
            train_feeder = import_class(self.arg.train_feeder)
            train_dataset = train_feeder(**self.arg.train_feeder_args)
            self.train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
            self.data_loader['train'] = torch.utils.data.DataLoader(
                dataset=train_dataset,
                batch_size=self.arg.batch_size,
                num_workers=self.arg.num_worker * torchlight.ngpu(
                    self.arg.device),
                drop_last=True,
                sampler=self.train_sampler)

        if self.arg.test_feeder_args:
            test_feeder = import_class(self.arg.test_feeder)
            test_dataset = test_feeder(**self.arg.test_feeder_args)
            self.test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset)
            self.data_loader['test'] = torch.utils.data.DataLoader(
                dataset=test_dataset,
                batch_size=self.arg.test_batch_size,
                num_workers=self.arg.num_worker * torchlight.ngpu(
                    self.arg.device),
                drop_last=True,
                sampler=self.test_sampler)

        if self.arg.mem_train_feeder_args:
            mem_train_feeder = import_class(self.arg.mem_train_feeder)
            mem_train_dataset = mem_train_feeder(**self.arg.mem_train_feeder_args)
            self.mem_train_sampler = torch.utils.data.distributed.DistributedSampler(mem_train_dataset)
            self.data_loader['mem_train'] = torch.utils.data.DataLoader(
                dataset=mem_train_dataset,
                batch_size=self.arg.batch_size,
                num_workers=self.arg.num_worker * torchlight.ngpu(
                    self.arg.device),
                drop_last=False,
                sampler=self.mem_train_sampler)

        if self.arg.mem_test_feeder_args:
            mem_test_feeder = import_class(self.arg.mem_test_feeder)
            mem_test_dataset = mem_test_feeder(**self.arg.mem_test_feeder_args)
            self.mem_test_sampler = torch.utils.data.distributed.DistributedSampler(mem_test_dataset)
            self.data_loader['mem_test'] = torch.utils.data.DataLoader(
                dataset=mem_test_dataset,
                batch_size=self.arg.test_batch_size,
                num_workers=self.arg.num_worker * torchlight.ngpu(
                    self.arg.device),
                drop_last=False,
                sampler=self.mem_test_sampler)

    def show_epoch_info(self):
        for k, v in self.epoch_info.items():
            self.io.print_log('\t{}: {}'.format(k, v))
        if self.arg.pavi_log:
            self.io.log('train', self.meta_info['iter'], self.epoch_info)

    def show_eval_info(self):
        for k, v in self.eval_info.items():
            self.io.print_log('\t{}: {}'.format(k, v))
        if self.arg.pavi_log:
            self.io.log('eval', self.meta_info['iter'], self.eval_info)

    def show_iter_info(self):
        if self.meta_info['iter'] % self.arg.log_interval == 0:
            info ='\tIter {} Done.'.format(self.meta_info['iter'])
            for k, v in self.iter_info.items():
                if isinstance(v, float):
                    info = info + ' | {}: {:.4f}'.format(k, v)
                else:
                    info = info + ' | {}: {}'.format(k, v)

            self.io.print_log(info)

            if self.arg.pavi_log:
                self.io.log('train', self.meta_info['iter'], self.iter_info)

    def train(self):
        for _ in range(100):
            self.iter_info['loss'] = 0
            self.show_iter_info()
            self.meta_info['iter'] += 1
        self.epoch_info['train_mean_loss'] = 0
        self.show_epoch_info()

    def test(self):
        for _ in range(100):
            self.iter_info['loss'] = 1
            self.show_iter_info()
        self.eval_info['test_mean_loss'] = 1
        self.show_eval_info()

    def print_networks(self, net, print_flag=False):
        self.io.print_log('---------- Networks initialized -------------')
        num_params = 0
        for param in net.parameters():
            num_params += param.numel()
        if print_flag:
            self.io.print_log(net)
        self.io.print_log('[Network] Total number of parameters : %.3f M' % (num_params / 1e6))
        self.io.print_log('-----------------------------------------------')

    def start(self, local_rank):
        if local_rank == 0:
            # get the output of `git diff`
            diff_output = subprocess.check_output(['git', 'diff'])
            # print the output
            self.io.print_log(diff_output.decode('utf-8'))
        
        
        self.local_rank = local_rank
        torch.cuda.set_device(local_rank)
        self.device = torch.cuda.device(local_rank)
        self.dev = local_rank
        self.gpus = list(range(torch.distributed.get_world_size()))
        
        if self.arg.phase == 'train' and local_rank == 0:
            if os.path.isdir(self.arg.work_dir + '/train'):
                print('log_dir: ', self.arg.work_dir, 'already exist')
                shutil.rmtree(self.arg.work_dir + '/train')
                shutil.rmtree(self.arg.work_dir + '/val')
                print('Dir removed: ', self.arg.work_dir + '/train')
                print('Dir removed: ', self.arg.work_dir + '/val')
            self.train_writer = SummaryWriter(os.path.join(self.arg.work_dir, 'train'), 'train')
            self.val_writer = SummaryWriter(os.path.join(self.arg.work_dir, 'val'), 'val')
        # else:
        #     self.train_writer = self.val_writer = SummaryWriter(os.path.join(self.arg.work_dir, 'test'), 'test')
            
        self.load_model()
        self.load_weights()
        # self.gpu()
        self.load_data()
        self.load_optimizer()
            
        if local_rank == 0:    
            self.io.print_log('Parameters:\n{}\n'.format(str(vars(self.arg))))
            self.print_networks(self.model)

        self.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.model)
        
        self.model = self.model.cuda()
        
        self.model = DDP(
                self.model,
                device_ids=[local_rank],
                output_device=local_rank,
                find_unused_parameters=False
            )

        # training phase
        if self.arg.phase == 'train':
            self.global_step = self.arg.start_epoch * len(self.data_loader['train'])
            self.meta_info['iter'] = self.global_step
            self.best_result = 0.0
            self.knn_results = dict()
            self.KNN_epoch_results = dict()
            for k in self.arg.knn_k:
                self.knn_results[k] = dict()
            
            for epoch in range(self.arg.start_epoch, self.arg.num_epoch):
                self.train_sampler.set_epoch(epoch)
                self.meta_info['epoch'] = epoch + 1

                # training
                if local_rank == 0:   
                    self.io.print_log('Training epoch: {}'.format(epoch + 1))
                self.train(epoch + 1)

                # save model
                if self.arg.save_interval == -1:
                    pass
                elif (((epoch + 1) % self.arg.save_interval == 0) or (
                        epoch + 1 == self.arg.num_epoch)) and dist.get_rank() == 0:
                    filename = 'epoch{}_model.pt'.format(epoch + 1)
                    self.io.save_model(self.model, filename)

                # evaluation
                if self.arg.eval_interval == -1:
                    pass
                elif ((epoch + 1) % self.arg.eval_interval == 0) or (
                        epoch + 1 == self.arg.num_epoch):
                    if local_rank == 0:   
                        self.io.print_log('Eval epoch: {}'.format(epoch + 1))
                    self.test(epoch + 1)
                    self.io.print_log("current %.2f%%, best %.2f%%" % 
                                            (self.current_result, self.best_result))

                    # save best model
                    filename = 'epoch%.3d_acc%.2f_model.pt' % (epoch + 1, self.current_result)
                    self.io.save_model(self.model, filename)
                    if self.current_result >= self.best_result:
                        filename = 'best_model.pt'
                        self.io.save_model(self.model, filename)
                        # save the output of model
                        if self.arg.save_result:
                            result_dict = dict(
                                zip(self.data_loader['test'].dataset.sample_name,
                                    self.result))
                            self.io.save_pkl(result_dict, 'test_result.pkl')

                if self.arg.knn_interval == -1:
                    pass
                elif ((epoch + 1) % self.arg.knn_interval == 0) or (epoch + 1 == self.arg.num_epoch):
                    if local_rank == 0:   
                        self.io.print_log('KNN eval epoch {}'.format(epoch + 1))
                    self.knn_monitor(epoch + 1)

                    for k in self.arg.knn_k:
                        self.io.print_log("KNN - {} current: {:.2f}%, best: {:.2f}%".format(
                            k, self.knn_results[k][epoch + 1], max(self.knn_results[k].values())
                            ))
                        if epoch + 1 in self.arg.KNN_show:
                            if epoch + 1 not in self.KNN_epoch_results:
                                self.KNN_epoch_results[epoch + 1] = dict()
                            self.KNN_epoch_results[epoch + 1][k] = [self.knn_results[k][epoch + 1], max(self.knn_results[k].values())]

                    if (epoch + 1 == self.arg.num_epoch):
                        self.io.print_log('*' * 10 + ' KNN Result ' + '*' * 10)
                        for show_epoch in self.arg.KNN_show:
                            if show_epoch in self.KNN_epoch_results:
                                for k in self.arg.knn_k:
                                    self.io.print_log('\t{}-{}: cur-{:.2f}%, best-{:.2f}%'.format(
                                        show_epoch, k, self.KNN_epoch_results[show_epoch][k][0],
                                        self.KNN_epoch_results[show_epoch][k][1]))

        # test phase
        elif self.arg.phase == 'test':
            
            # the path of weights must be appointed
            if self.arg.weights is None:
                raise ValueError('Please appoint --weights.')
            self.io.print_log('Model:   {}.'.format(self.arg.model))
            self.io.print_log('Weights: {}.'.format(self.arg.weights))

            self.best_result = 0.0

            # evaluation
            self.io.print_log('Evaluation Start:')
            self.test(1)
            self.io.print_log('Done.\n')

            # save the output of model
            if self.arg.save_result:
                result_dict = dict(
                    zip(self.data_loader['test'].dataset.sample_name,
                        self.result))
                self.io.save_pkl(result_dict, 'test_result_%.3d.pkl'% (epoch + 1))

    @staticmethod
    def get_parser(add_help=False):

        # region arguments yapf: disable
        # parameter priority: command line > config > default
        parser = argparse.ArgumentParser(add_help=add_help, description='Base Processor')

        parser.add_argument('-w', '--work_dir', default='./work_dir/tmp', help='the work folder for storing results')
        parser.add_argument('-c', '--config', default=None, help='path to the configuration file')

        # processor
        parser.add_argument('--phase', default='train', help='must be train or test')
        parser.add_argument('--save_result', type=str2bool, default=True,
                            help='if ture, the output of the model will be stored')
        parser.add_argument('--start_epoch', type=int, default=0, help='start training from which epoch')
        parser.add_argument('--num_epoch', type=int, default=80, help='stop training in which epoch')
        parser.add_argument('--use_gpu', type=str2bool, default=True, help='use GPUs or not')
        parser.add_argument('--device', type=int, default=0, nargs='+',
                            help='the indexes of GPUs for training or testing')

        # visulize and debug
        parser.add_argument('--log_interval', type=int, default=100,
                            help='the interval for printing messages (#iteration)')
        parser.add_argument('--save_interval', type=int, default=10, help='the interval for storing models (#epoch)')
        parser.add_argument('--eval_interval', type=int, default=5, help='the interval for evaluating models (#epoch)')
        parser.add_argument('--save_log', type=str2bool, default=True, help='save logging or not')
        parser.add_argument('--print_log', type=str2bool, default=True, help='print logging or not')
        parser.add_argument('--pavi_log', type=str2bool, default=False, help='logging on pavi or not')

        parser.add_argument('--knn_interval', type=int, default=5, help='the interval for knn models (#epoch)')

        # feeder
        parser.add_argument('--train_feeder', default='feeder.feeder', help='train data loader will be used')
        parser.add_argument('--test_feeder', default='feeder.feeder', help='test data loader will be used')
        parser.add_argument('--num_worker', type=int, default=4, help='the number of worker per gpu for data loader')
        parser.add_argument('--train_feeder_args', action=DictAction, default=dict(),
                            help='the arguments of data loader for training')
        parser.add_argument('--test_feeder_args', action=DictAction, default=dict(),
                            help='the arguments of data loader for test')
        parser.add_argument('--mem_train_feeder', default='feeder.feeder',
                            help='memory train data loader will be used for knn')
        parser.add_argument('--mem_test_feeder', default='feeder.feeder',
                            help='memory test data loader will be used for knn')
        parser.add_argument('--mem_train_feeder_args', action=DictAction, default=dict(),
                            help='the arguments of data loader for training')
        parser.add_argument('--mem_test_feeder_args', action=DictAction, default=dict(),
                            help='the arguments of data loader for test')

        parser.add_argument('--batch_size', type=int, default=256, help='training batch size')
        parser.add_argument('--test_batch_size', type=int, default=256, help='test batch size')
        parser.add_argument('--debug', action="store_true", help='less data, faster loading')

        # model
        parser.add_argument('--model', default=None, help='the model will be used')
        parser.add_argument('--model_args', action=DictAction, default=dict(), help='the arguments of model')
        parser.add_argument('--weights', default=None, help='the weights for network initialization')
        parser.add_argument('--ignore_weights', type=str, default=[], nargs='+', help='the name of weights which will be ignored in the initialization')
        parser.add_argument('--rename_weights', type=str, default=[], nargs='+', help='the name of weights which will be ignored in the initialization')
                            
        parser.add_argument('--knn_k', type=int, default=[], nargs='+', help='KNN-K')
        parser.add_argument('--knn_classes', type=int, default=60, help='use cosine lr schedule')
        parser.add_argument('--knn_t', type=float, default=0.1, help='use cosine lr schedule')
        parser.add_argument('--KNN_show', type=int, default=[], nargs='+',
                            help='the epoch to show the best KNN result')
        # endregion yapf: enable

        return parser

processor/processor.py
LanglandsLin/ActCLR
import os
import sys
import pickle

import argparse
import numpy as np
from numpy.lib.format import open_memmap

from utils.ntu_read_skeleton import read_xyz

training_subjects = [
    1, 2, 4, 5, 8, 9, 13, 14, 15, 16, 17, 18, 19, 25, 27, 28, 31, 34, 35, 38
]
training_cameras = [2, 3]
max_body = 2
num_joint = 25
max_frame = 300
toolbar_width = 30

def print_toolbar(rate, annotation=''):
    # setup toolbar
    sys.stdout.write("{}[".format(annotation))
    for i in range(toolbar_width):
        if i * 1.0 / toolbar_width > rate:
            sys.stdout.write(' ')
        else:
            sys.stdout.write('-')
        sys.stdout.flush()
    sys.stdout.write(']\r')


def end_toolbar():
    sys.stdout.write("\n")


def gendata(data_path,
            out_path,
            ignored_sample_path=None,
            benchmark='xview',
            part='eval'):
    if ignored_sample_path != None:
        with open(ignored_sample_path, 'r') as f:
            ignored_samples = [
                line.strip() + '.skeleton' for line in f.readlines()
            ]
    else:
        ignored_samples = []
    sample_name = []
    sample_label = []
    dir_list = os.listdir(data_path)
    dir_list.sort()
    for filename in dir_list:
        if filename in ignored_samples:
            continue
        action_class = int(
            filename[filename.find('A') + 1:filename.find('A') + 4])
        subject_id = int(
            filename[filename.find('P') + 1:filename.find('P') + 4])
        camera_id = int(
            filename[filename.find('C') + 1:filename.find('C') + 4])

        if benchmark == 'xview':
            istraining = (camera_id in training_cameras)
        elif benchmark == 'xsub':
            istraining = (subject_id in training_subjects)
        else:
            raise ValueError()

        if part == 'train':
            issample = istraining
        elif part == 'val':
            issample = not (istraining)
        else:
            raise ValueError()

        if issample:
            sample_name.append(filename)
            sample_label.append(action_class - 1)

    with open('{}/{}_label.pkl'.format(out_path, part), 'wb') as f:
        pickle.dump((sample_name, list(sample_label)), f)
    # np.save('{}/{}_label.npy'.format(out_path, part), sample_label)

    fp = open_memmap(
        '{}/{}_data.npy'.format(out_path, part),
        dtype='float32',
        mode='w+',
        shape=(len(sample_label), 3, max_frame, num_joint, max_body))

    for i, s in enumerate(sample_name):
        print_toolbar(i * 1.0 / len(sample_label),
                      '({:>5}/{:<5}) Processing {:>5}-{:<5} data: '.format(
                          i + 1, len(sample_name), benchmark, part))
        data = read_xyz(
            os.path.join(data_path, s), max_body=max_body, num_joint=num_joint)
        fp[i, :, 0:data.shape[1], :, :] = data
    end_toolbar()


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='NTU-RGB-D Data Converter.')
    parser.add_argument(
        '--data_path', default='/data/gty/action_dataset/original/nturgb+d_skeletons/')
    parser.add_argument(
        '--ignored_sample_path',
        default='/data/gty/action_dataset/original/NTU_RGBD_samples_with_missing_skeletons.txt')
    parser.add_argument('--out_folder', default='/data/gty/action_dataset/ntu60')

    benchmark = ['xsub', 'xview']
    part = ['train', 'val']
    arg = parser.parse_args()

    for b in benchmark:
        for p in part:
            out_path = os.path.join(arg.out_folder, b)
            if not os.path.exists(out_path):
                os.makedirs(out_path)
            gendata(
                arg.data_path,
                out_path,
                arg.ignored_sample_path,
                benchmark=b,
                part=p)

tools/ntu_gendata.py
LanglandsLin/ActCLR
import os
import sys
import pickle

import argparse
import numpy as np
from numpy.lib.format import open_memmap

from utils.pku_read_skeleton import read_xyz

view_dic = {1:"L", 2: "M", 3:"R"}

max_body = 2
num_joint = 25
max_frame = 300
toolbar_width = 30

def print_toolbar(rate, annotation=''):
    # setup toolbar
    sys.stdout.write("{}[".format(annotation))
    for i in range(toolbar_width):
        if i * 1.0 / toolbar_width > rate:
            sys.stdout.write(' ')
        else:
            sys.stdout.write('-')
        sys.stdout.flush()
    sys.stdout.write(']\r')


def end_toolbar():
    sys.stdout.write("\n")


def gendata(data_path,
            out_path,
            ignored_sample_path=None,
            benchmark='xview',
            part='eval'):
    training_split_file_folder = 'E:/2021/datasets/original_data/pkummd/v1/'
    if benchmark=='xview':
        training_split_file = training_split_file_folder + "cross-view.txt"
    else:
        training_split_file = training_split_file_folder + "cross-subject.txt"
    f = open(training_split_file, "r")
    f.readline()
    training_split = f.readline()
    f.readline()
    testing_split = f.readline()

    training_set = training_split.split(',')
    training_set = [s.strip() for s in training_set]
    training_set = [s for s in training_set if s]

    testing_set = testing_split.split(',')
    testing_set = [s.strip() for s in testing_set]
    testing_set = [s for s in testing_set if s]

    if ignored_sample_path != None:
        with open(ignored_sample_path, 'r') as f:
            ignored_samples = [
                line.strip() + '.skeleton' for line in f.readlines()
            ]
    else:
        ignored_samples = []
    sample_name = []
    sample_label = []
    for filename in os.listdir(data_path):
        if filename in ignored_samples:
            continue
        file_id = int(
            filename[filename.find('F') + 1:filename.find('F') + 4])
        view_id = int(
            filename[filename.find('V') + 1:filename.find('V') + 4])
        action_class = int(
            filename[filename.find('C') + 1:filename.find('C') + 4])

        split_train_test_name = '0'+filename[filename.find('F') + 1:filename.find('F') + 4] + '-' + view_dic[view_id]

        if (split_train_test_name in training_set) and (split_train_test_name in testing_set):
            raise ValueError()
        istraining = (split_train_test_name in training_set)

        if part == 'train':
            issample = istraining
        elif part == 'val':
            issample = not (istraining)
        else:
            raise ValueError()

        if issample:
            sample_name.append(filename)
            sample_label.append(action_class - 1)

    with open('{}/{}_label.pkl'.format(out_path, part), 'wb') as f:
        pickle.dump((sample_name, list(sample_label)), f)
    # np.save('{}/{}_label.npy'.format(out_path, part), sample_label)

    fp = open_memmap(
        '{}/{}_data.npy'.format(out_path, part),
        dtype='float32',
        mode='w+',
        shape=(len(sample_label), 3, max_frame, num_joint, max_body))

    for i, s in enumerate(sample_name):
        print_toolbar(i * 1.0 / len(sample_label),
                      '({:>5}/{:<5}) Processing {:>5}-{:<5} data: '.format(
                          i + 1, len(sample_name), benchmark, part))
        data = read_xyz(
            os.path.join(data_path, s), max_body=max_body, num_joint=num_joint)
        fp[i, :, 0:data.shape[1], :, :] = data
    end_toolbar()


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='PKU-MMD Part1 Data Converter.')
    parser.add_argument(
        '--data_path', default='E:/2021/datasets/original_data/pkummd/skeleton_pku_v1')
    parser.add_argument('--out_folder', default='E:/2021/datasets/crossclr_pkummd/pku_part1')

    parser.add_argument('--ignored_sample_path',
                        default='E:/2021/datasets/original_data/pkummd/v1/samples_with_missing_skeletons.txt')

    benchmark = ['xsub', 'xview']
    part = ['train', 'val']
    arg = parser.parse_args()

    for b in benchmark:
        for p in part:
            out_path = os.path.join(arg.out_folder, b)
            if not os.path.exists(out_path):
                os.makedirs(out_path)
            gendata(
                arg.data_path,
                out_path,
                arg.ignored_sample_path,
                benchmark=b,
                part=p)
tools/pku_part1_gendata.py
LanglandsLin/ActCLR
import os
import numpy as np

skeleton_folder = 'E:/2021/datasets/original_data/pkummd/v1/PKU_Skeleton_Renew/'
label_folder = 'E:/2021/datasets/original_data/pkummd/v1/Train_Label_PKU_final/'
save_path = 'E:/2021/datasets/original_data/pkummd/skeleton_pku_v1/'

view_dic = {'L': '001', 'M': '002', 'R': '003'}

total_file_num = 0


def array_to_skl_lines(skeleton_frame):
    skeleton_lines = ''
    for i in range(0, 75, 3):
        x = skeleton_frame[i:i + 3]
        x = map(lambda x: str(x), x)
        line = " ".join(a for a in x)
        skeleton_lines = skeleton_lines + line + '\n'
    return skeleton_lines


def transfer_and_save(file_id, file_view):
    global total_file_num
    if file_id / 10 < 1:
        file_id_ = '00' + str(file_id)
    elif file_id / 100 < 1:
        file_id_ = '0' + str(file_id)
    else:
        file_id_ = str(file_id)
    sample_file = '0' + file_id_ + '-' + file_view + '.txt'

    if not os.path.exists(label_folder + sample_file):
        print(label_folder + sample_file + ' not exist!')
        return
    else:
        print(label_folder + sample_file + ' do!')
        total_file_num += 1
    skl_list = np.loadtxt(label_folder + sample_file, delimiter=",")
    skl_file = skeleton_folder + sample_file

    skeleton_frames = np.loadtxt(skl_file)

    for i in range(0, skl_list.shape[0]):
        sample_length = int(skl_list[i][2] - skl_list[i][1]) + 1
        sample_lines = str(sample_length) + '\n'  # first lines, length of the action sample
        skeleton_sample = skeleton_frames[int(skl_list[i][1]) - 1:int(skl_list[i][2]), :150]

        for skl_id in range(0, skeleton_sample.shape[0]):
            skeleton_frame = skeleton_sample[skl_id]
            if np.sum(skeleton_frame[75:150]) == 0:
                sample_lines = sample_lines + '1\n'
                sample_lines = sample_lines + '6 1 0 0 1 1 0 -0.437266 -0.117168 2'
                sample_lines = sample_lines + '\n25\n'

                skeleton_lines = array_to_skl_lines(skeleton_frame[0:75])
                sample_lines = sample_lines + skeleton_lines
            else:
                sample_lines = sample_lines + '2\n'
                sample_lines = sample_lines + '6 1 0 0 1 1 0 -0.437266 -0.117168 2'
                sample_lines = sample_lines + '\n25\n'

                skeleton_lines_1 = array_to_skl_lines(skeleton_frame[0:75])
                skeleton_lines_2 = array_to_skl_lines(skeleton_frame[75:150])
                sample_lines = sample_lines + skeleton_lines_1 + skeleton_lines_2

        class_id = int(skl_list[i][0])
        if class_id / 10 < 1:
            class_id_ = '00' + str(class_id)
        elif class_id / 100 < 1:
            class_id_ = '0' + str(class_id)
        else:
            class_id_ = str(class_id)
        if (i + 1) / 10 < 1:
            index = '00' + str(i + 1)
        elif (i + 1) / 100 < 1:
            index = '0' + str(i + 1)
        else:
            index = str(i + 1)
        view_id = view_dic[file_view]
        save_name = 'F' + file_id_ + 'V' + view_id + 'C' + class_id_ + 'L' + index
        with open(save_path + save_name + ".skeleton", "w") as text_file:
            text_file.write(sample_lines)


for file_id in range(2, 365):
    for file_view in ['L', 'M', 'R']:
        transfer_and_save(file_id, file_view)

print('The total preprocess file:', total_file_num)

tools/pku_part1_skeleton.py
LanglandsLin/ActCLR
import os
import sys
import pickle

import argparse
import numpy as np
from numpy.lib.format import open_memmap

from utils.pku_read_skeleton import read_xyz

view_dic = {1:"L", 2: "M", 3:"R"}

max_body = 2
num_joint = 25
max_frame = 300
toolbar_width = 30

def print_toolbar(rate, annotation=''):
    # setup toolbar
    sys.stdout.write("{}[".format(annotation))
    for i in range(toolbar_width):
        if i * 1.0 / toolbar_width > rate:
            sys.stdout.write(' ')
        else:
            sys.stdout.write('-')
        sys.stdout.flush()
    sys.stdout.write(']\r')


def end_toolbar():
    sys.stdout.write("\n")


def gendata(data_path,
            out_path,
            ignored_sample_path=None,
            benchmark='xview',
            part='eval'):
    training_split_file_folder = 'E:/2021/datasets/original_data/pkummd/v2/'

    if benchmark=='xview':
        training_split_file = training_split_file_folder + "cross_view_v2.txt"
    else:
        training_split_file = training_split_file_folder + "cross_subject_v2.txt"

    f = open(training_split_file, "r")
    f.readline()
    training_split = f.readline()
    f.readline()
    testing_split = f.readline()

    training_set = training_split.split(',')
    training_set = [s.strip() for s in training_set]
    training_set = [s for s in training_set if s]

    testing_set = testing_split.split(',')
    testing_set = [s.strip() for s in testing_set]
    testing_set = [s for s in testing_set if s]

    if ignored_sample_path != None:
        with open(ignored_sample_path, 'r') as f:
            ignored_samples = [
                line.strip() + '.skeleton' for line in f.readlines()
            ]
    else:
        ignored_samples = []

    sample_name = []
    sample_label = []
    for filename in os.listdir(data_path):
        if filename in ignored_samples:
            continue
        a_id = int(
            filename[filename.find('A') + 1:filename.find('A') + 4])
        n_id = int(
            filename[filename.find('N') + 1:filename.find('N') + 4])
        view_id = int(
            filename[filename.find('V') + 1:filename.find('V') + 4])
        action_class = int(
            filename[filename.find('C') + 1:filename.find('C') + 4])

        split_train_test_name = 'A' + filename[filename.find('A') + 2:filename.find('A') + 4] \
                                + 'N' + filename[filename.find('N') + 2:filename.find('N') + 4] \
                                + '-' + view_dic[view_id]

        if (split_train_test_name in training_set) and (split_train_test_name in testing_set):
            raise ValueError()
        istraining = (split_train_test_name in training_set)

        if part == 'train':
            issample = istraining
        elif part == 'val':
            issample = not (istraining)
        else:
            raise ValueError()

        if issample:
            sample_name.append(filename)
            sample_label.append(action_class - 1)

    with open('{}/{}_label.pkl'.format(out_path, part), 'wb') as f:
        pickle.dump((sample_name, list(sample_label)), f)
    # np.save('{}/{}_label.npy'.format(out_path, part), sample_label)

    fp = open_memmap(
        '{}/{}_data.npy'.format(out_path, part),
        dtype='float32',
        mode='w+',
        shape=(len(sample_label), 3, max_frame, num_joint, max_body))

    for i, s in enumerate(sample_name):
        print_toolbar(i * 1.0 / len(sample_label),
                      '({:>5}/{:<5}) Processing {:>5}-{:<5} data: '.format(
                          i + 1, len(sample_name), benchmark, part))
        data = read_xyz(
            os.path.join(data_path, s), max_body=max_body, num_joint=num_joint)
        fp[i, :, 0:data.shape[1], :, :] = data
    end_toolbar()


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='PKU-MMD Part1 Data Converter.')
    parser.add_argument(
        '--data_path', default='E:/2021/datasets/original_data/pkummd/skeleton_pku_v2')
    parser.add_argument('--out_folder', default='E:/2021/datasets/crossclr_pkummd/pku_part2')

    parser.add_argument('--ignored_sample_path',
                        default='E:/2021/datasets/original_data/pkummd/v2/samples_with_missing_skeletons_v2.txt')

    benchmark = ['xsub', 'xview']
    part = ['train', 'val']
    arg = parser.parse_args()

    for b in benchmark:
        for p in part:
            out_path = os.path.join(arg.out_folder, b)
            if not os.path.exists(out_path):
                os.makedirs(out_path)
            gendata(
                arg.data_path,
                out_path,
                arg.ignored_sample_path,
                benchmark=b,
                part=p)
tools/pku_part2_gendata.py
LanglandsLin/ActCLR
import os
import numpy as np

skeleton_folder = 'E:/2021/datasets/original_data/pkummd/v2/skeleton/'
label_folder = 'E:/2021/datasets/original_data/pkummd/v2/label/'
save_path = 'E:/2021/datasets/original_data/pkummd/skeleton_pku_v2/'

view_dic = {'L': '001', 'M': '002', 'R': '003'}
total_file_num = 0


def array_to_skl_lines(skeleton_frame):
    skeleton_lines = ''
    for i in range(0, 75, 3):
        x = skeleton_frame[i:i + 3]
        x = map(lambda x: str(x), x)
        line = " ".join(a for a in x)
        skeleton_lines = skeleton_lines + line + '\n'
    return skeleton_lines


def transfer_and_save(a_id, n_id, file_view):
    global total_file_num
    if a_id < 10:
        a_id = '0' + str(a_id)
    else:
        a_id = str(a_id)
    if n_id < 10:
        n_id = '0' + str(n_id)
    else:
        n_id = str(n_id)
    sample_file = 'A' + a_id + 'N' + n_id + '-' + file_view + '.txt'

    if not os.path.exists(label_folder + sample_file):
        print(label_folder + sample_file + ' not exist!')
        return
    else:
        print(label_folder + sample_file + ' do!')
        total_file_num += 1
    skl_list = np.loadtxt(label_folder + sample_file, delimiter=",")
    skl_file = skeleton_folder + sample_file

    skeleton_frames = np.loadtxt(skl_file)

    for i in range(0, skl_list.shape[0]):
        sample_length = int(skl_list[i][2] - skl_list[i][1]) + 1
        sample_lines = str(sample_length) + '\n'  # first lines, length of the action sample
        skeleton_sample = skeleton_frames[int(skl_list[i][1]) - 1:int(skl_list[i][2]), :150]

        for skl_id in range(0, skeleton_sample.shape[0]):
            skeleton_frame = skeleton_sample[skl_id]
            if np.sum(skeleton_frame[75:150]) == 0:
                sample_lines = sample_lines + '1\n'
                sample_lines = sample_lines + '6 1 0 0 1 1 0 -0.437266 -0.117168 2'
                sample_lines = sample_lines + '\n25\n'

                skeleton_lines = array_to_skl_lines(skeleton_frame[0:75])
                sample_lines = sample_lines + skeleton_lines
            else:
                sample_lines = sample_lines + '2\n'
                sample_lines = sample_lines + '6 1 0 0 1 1 0 -0.437266 -0.117168 2'
                sample_lines = sample_lines + '\n25\n'

                skeleton_lines_1 = array_to_skl_lines(skeleton_frame[0:75])
                skeleton_lines_2 = array_to_skl_lines(skeleton_frame[75:150])
                sample_lines = sample_lines + skeleton_lines_1 + skeleton_lines_2

        class_id = int(skl_list[i][0])
        if class_id / 10 < 1:
            class_id_ = '00' + str(class_id)
        elif class_id / 100 < 1:
            class_id_ = '0' + str(class_id)
        else:
            class_id_ = str(class_id)
        if (i + 1) / 10 < 1:
            index = '00' + str(i + 1)
        elif (i + 1) / 100 < 1:
            index = '0' + str(i + 1)
        else:
            index = str(i + 1)
        view_id = view_dic[file_view]
        save_name = 'A0' + a_id + 'N0' + n_id + 'V' + view_id + 'C' + class_id_ + 'L' + index
        with open(save_path + save_name + ".skeleton", "w") as text_file:
            text_file.write(sample_lines)


for a_id in range(1, 27):
    for n_id in range(1, 14):
        for file_view in ['L', 'M', 'R']:
            transfer_and_save(a_id, n_id, file_view)

print('The total preprocess file:', total_file_num)

tools/pku_part2_skeleton.py
LanglandsLin/ActCLR
from . import ntu_read_skeleton
from . import pku_read_skeleton
tools/utils/__init__.py
LanglandsLin/ActCLR
#!/usr/bin/env python
import argparse
import os
import sys
import traceback
import time
import warnings
import pickle
from collections import OrderedDict
import yaml
import numpy as np
# torch
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import copy

with warnings.catch_warnings():
    warnings.filterwarnings("ignore",category=FutureWarning)
    import h5py

class IO():
    def __init__(self, work_dir, save_log=True, print_log=True):
        self.work_dir = work_dir
        self.save_log = save_log
        self.print_to_screen = print_log
        self.cur_time = time.time()
        self.split_timer = {}
        self.pavi_logger = None
        self.session_file = None
        self.model_text = ''
        
    # PaviLogger is removed in this version
    def log(self, *args, **kwargs):
        pass
    #     try:
    #         if self.pavi_logger is None:
    #             from torchpack.runner.hooks import PaviLogger
    #             url = 'http://pavi.parrotsdnn.org/log'
    #             with open(self.session_file, 'r') as f:
    #                 info = dict(
    #                     session_file=self.session_file,
    #                     session_text=f.read(),
    #                     model_text=self.model_text)
    #             self.pavi_logger = PaviLogger(url)
    #             self.pavi_logger.connect(self.work_dir, info=info)
    #         self.pavi_logger.log(*args, **kwargs)
    #     except:  #pylint: disable=W0702
    #         pass

    def load_model(self, model, **model_args):
        Model = import_class(model)
        model = Model(**model_args)
        self.model_text += '\n\n' + str(model)
        return model

    def load_weights(self, model, weights_path, ignore_weights=None, rename_weights=None, output=False):
        if ignore_weights is None:
            ignore_weights = []
        if isinstance(ignore_weights, str):
            ignore_weights = [ignore_weights]

        if rename_weights is None:
            rename_weights = []
        if isinstance(rename_weights, str):
            rename_weights = [rename_weights]

        self.print_log('Load weights from {}.'.format(weights_path))
        weights = torch.load(weights_path)
        if 'state_dict' in weights.keys():
            weights = weights['state_dict']
        weights = OrderedDict([[k.split('module.')[-1],
                                v.cpu()] for k, v in weights.items()])

        for r in rename_weights:
            i, j = r.split(".")
            rename_name = list()
            if i == '_':
                for w in list(weights.keys()):
                    weights[j + '.' + w] = weights.pop(w)
                    if output:
                        self.print_log('Add [{}] to [{}].'.format(j, w))
            elif j == '_':
                for w in list(weights.keys()):
                    weights[w.split(i + '.')[-1]] = weights.pop(w)
                    if output:
                        self.print_log('Del [{}] from [{}].'.format(i, w))
            else:
                for w in weights.keys():
                    if w.find(i) == 0:
                        rename_name.append(w)
                for n in rename_name:
                    weights[n.replace(i, j)] = copy.deepcopy(weights[n])
                    if output:
                        self.print_log('Repalce [{}] weights with [{}].'.format(n.replace(i, j), n))

        # filter weights
        for i in ignore_weights:
            ignore_name = list()
            for w in weights:
                if w.find(i) == 0:
                    ignore_name.append(w)
            for n in ignore_name:
                weights.pop(n)
                if output:
                    self.print_log('Filter [{}] remove weights [{}].'.format(i,n))

        if output:
            for w in weights:
                self.print_log('Load weights [{}].'.format(w))

        try:
            model.load_state_dict(weights)
        except (KeyError, RuntimeError):
            state = model.state_dict()
            diff = list(set(state.keys()).difference(set(weights.keys())))
            if output:
                for d in diff:
                    self.print_log('Can not find weights [{}].'.format(d))
            state.update(weights)
            model.load_state_dict(state)
        return model

    def save_pkl(self, result, filename):
        with open('{}/{}'.format(self.work_dir, filename), 'wb') as f:
            pickle.dump(result, f)

    def save_h5(self, result, filename):
        with h5py.File('{}/{}'.format(self.work_dir, filename), 'w') as f:
            for k in result.keys():
                f[k] = result[k]

    def save_model(self, model, name):
        model_path = '{}/{}'.format(self.work_dir, name)
        state_dict = model.state_dict()
        weights = OrderedDict([[''.join(k.split('module.')),
                                v.cpu()] for k, v in state_dict.items()])
        torch.save(weights, model_path)
        self.print_log('The model has been saved as {}.'.format(model_path))

    def save_arg(self, arg):

        self.session_file = '{}/config.yaml'.format(self.work_dir)

        # save arg
        arg_dict = vars(arg)
        if not os.path.exists(self.work_dir):
            os.makedirs(self.work_dir)
        with open(self.session_file, 'w') as f:
            f.write('# command line: {}\n\n'.format(' '.join(sys.argv)))
            yaml.dump(arg_dict, f, default_flow_style=False, indent=4)

    def print_log(self, str, print_time=True):
        if print_time:
            # localtime = time.asctime(time.localtime(time.time()))
            str = time.strftime("[%m.%d.%y|%X] ", time.localtime()) + str

        if self.print_to_screen:
            print(str)
        if self.save_log:
            with open('{}/log.txt'.format(self.work_dir), 'a') as f:
                print(str, file=f)

    def init_timer(self, *name):
        self.record_time()
        self.split_timer = {k: 0.0000001 for k in name}

    def check_time(self, name):
        self.split_timer[name] += self.split_time()

    def record_time(self):
        self.cur_time = time.time()
        return self.cur_time

    def split_time(self):
        split_time = time.time() - self.cur_time
        self.record_time()
        return split_time

    def print_timer(self):
        proportion = {
            k: '{:02d}%'.format(int(round(v * 100 / sum(self.split_timer.values()))))
            for k, v in self.split_timer.items()
        }
        self.print_log('Time consumption:')
        for k in proportion:
            self.print_log(
                '\t[{}][{}]: {:.4f}'.format(k, proportion[k],self.split_timer[k])
                )


def str2bool(v):
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')


def str2dict(v):
    return eval('dict({})'.format(v))  #pylint: disable=W0123


def _import_class_0(name):
    components = name.split('.')
    mod = __import__(components[0])
    for comp in components[1:]:
        mod = getattr(mod, comp)
    return mod


def import_class(import_str):
    mod_str, _sep, class_str = import_str.rpartition('.')
    __import__(mod_str)
    try:
        return getattr(sys.modules[mod_str], class_str)
    except AttributeError:
        raise ImportError('Class %s cannot be found (%s)' %
                          (class_str,
                           traceback.format_exception(*sys.exc_info())))


class DictAction(argparse.Action):
    def __init__(self, option_strings, dest, nargs=None, **kwargs):
        if nargs is not None:
            raise ValueError("nargs not allowed")
        super(DictAction, self).__init__(option_strings, dest, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        input_dict = eval('dict({})'.format(values))  #pylint: disable=W0123
        output_dict = getattr(namespace, self.dest)
        for k in input_dict:
            output_dict[k] = input_dict[k]
        setattr(namespace, self.dest, output_dict)

torchlight/io.py
LanglandsLin/ActCLR