def main():
parser = argparse.ArgumentParser(description='mean and std computing')
parser.add_argument('--root',
type=str,
default=None,
required=True,
help='path to root folder of the CelebA_Spoof')
parser.add_argument('--img_size',
type=tuple,
default=(128, 128),
required=False,
help='height and width of the image to resize')
args = parser.parse_args()
# transform image
transforms = A.Compose([
A.Resize(*args.img_size, interpolation=cv2.INTER_CUBIC),
A.Normalize(mean=[0, 0, 0], std=[1, 1, 1])
])
root_folder = args.root
train_dataset = CelebASpoofDataset(root_folder,
test_mode=False,
transform=Transform(transforms),
multi_learning=False)
dataloader = DataLoader(train_dataset, batch_size=100, shuffle=True)
mean, std = compute_mean_std(dataloader)
print(f'mean:{mean}, std:{std}')
def main():
# parsing arguments
parser = argparse.ArgumentParser(description='antispoofing training')
parser.add_argument('--draw_graph', default=False, type=bool, required=False,
help='whether or not to draw graphics')
parser.add_argument('--GPU', default=0, type=int, required=False,
help='specify which GPU to use')
parser.add_argument('--config', type=str, default=None, required=True,
help='path to configuration file')
parser.add_argument('--device', type=str, default='cuda',
help='if you want to eval model on cpu, pass "cpu" param')
args = parser.parse_args()
# reading config and manage device
path_to_config = args.config
config = read_py_config(path_to_config)
device = args.device + f':{args.GPU}' if args.device == 'cuda' else 'cpu'
# building model
model = build_model(config, device, strict=True, mode='eval')
model.to(device)
if config.data_parallel.use_parallel:
model = nn.DataParallel(model, **config.data_parallel.parallel_params)
# load snapshot
path_to_experiment = os.path.join(config.checkpoint.experiment_path, config.checkpoint.snapshot_name)
epoch_of_checkpoint = load_checkpoint(path_to_experiment, model, map_location=device, optimizer=None)
# preprocessing, making dataset and loader
normalize = A.Normalize(**config.img_norm_cfg)
test_transform = A.Compose([
A.Resize(**config.resize, interpolation=cv.INTER_CUBIC),
normalize
])
test_transform = Transform(val=test_transform)
test_dataset = make_dataset(config, val_transform=test_transform, mode='eval')
test_loader = DataLoader(dataset=test_dataset, batch_size=100, shuffle=True, num_workers=2)
# computing metrics
auc_, eer, accur, apcer, bpcer, acer, fpr, tpr = evaluate(model, test_loader,
config, device,
compute_accuracy=True)
print((f'eer = {round(eer*100,2)}\n'
+ f'accuracy on test data = {round(np.mean(accur),3)}\n'
+ f'auc = {round(auc_,3)}\n'
+ f'apcer = {round(apcer*100,2)}\n'
+ f'bpcer = {round(bpcer*100,2)}\n'
+ f'acer = {round(acer*100,2)}\n'
+ f'checkpoint made on {epoch_of_checkpoint} epoch'))
# draw graphics if needed
if args.draw_graph:
fnr = 1 - tpr
plot_roc_curve(fpr, tpr, config)
det_curve(fpr, fnr, eer, config)
def eval():
cnt = 0
SAE = 0 # sum of absolute errors
SSE = 0 # sum of square errors
print("Evaluation on {} data".format(args.test_split))
im_ids = data_split[args.test_split]
pbar = tqdm(im_ids)
for im_id in pbar:
anno = annotations[im_id]
bboxes = anno['box_examples_coordinates']
dots = np.array(anno['points'])
rects = list()
for bbox in bboxes:
x1 = bbox[0][0]
y1 = bbox[0][1]
x2 = bbox[2][0]
y2 = bbox[2][1]
rects.append([y1, x1, y2, x2])
image = Image.open('{}/{}'.format(im_dir, im_id))
image.load()
sample = {'image': image, 'lines_boxes': rects}
sample = Transform(sample)
image, boxes = sample['image'].cuda(), sample['boxes'].cuda()
with torch.no_grad():
output = regressor(
extract_features(resnet50_conv, image.unsqueeze(0),
boxes.unsqueeze(0), MAPS, Scales))
gt_cnt = dots.shape[0]
pred_cnt = output.sum().item()
cnt = cnt + 1
err = abs(gt_cnt - pred_cnt)
SAE += err
SSE += err**2
pbar.set_description(
'{:<8}: actual-predicted: {:6d}, {:6.1f}, error: {:6.1f}. Current MAE: {:5.2f}, RMSE: {:5.2f}'
.format(im_id, gt_cnt, pred_cnt, abs(pred_cnt - gt_cnt), SAE / cnt,
(SSE / cnt)**0.5))
print("")
print('On {} data, MAE: {:6.2f}, RMSE: {:6.2f}'.format(
args.test_split, SAE / cnt, (SSE / cnt)**0.5))
return SAE / cnt, (SSE / cnt)**0.5
def forwardVariable(th, linkparam):
"""
Calculate the forward transformation by
angle(theta) and it's link parameters.
"""
T = Transform()
i = 0
for p in linkparam:
if p[2] is None:
T = T * link(p[0], p[1], th[i], p[3])
i += 1
elif p[3] is None:
T = T * link(p[0], p[1], p[2], th[i])
i += 1
else:
T = T * link(p[0], p[1], p[2], p[3])
return T
def sim(self, th, text=""):
# init
trans = [Transform()]
pos = [Translation()]
dirs = [np.identity(4)]
# main
for i in range(len(th)):
L = self.linkparam[i]
T = link(L[0], L[1], th[i], L[3])
trans.append(trans[-1] * T)
pos.append(trans[-1] * pos[0])
dirs.append(trans[-1].mat)
# save it
self.trans.append(trans)
self.poses.append(np.array([p.mat for p in pos]))
self.dirs.append(dirs)
self.text.append(text)
rects1 = list()
for line in lines:
data = line.split()
y1 = int(data[0])
x1 = int(data[1])
y2 = int(data[2])
x2 = int(data[3])
rects1.append([y1, x1, y2, x2])
print("Bounding boxes: ", end="")
print(rects1)
image = Image.open(args.input_image)
image.load()
sample = {'image': image, 'lines_boxes': rects1}
sample = Transform(sample)
image, boxes = sample['image'], sample['boxes']
if use_gpu:
image = image.cuda()
boxes = boxes.cuda()
with torch.no_grad():
features = extract_features(resnet50_conv, image.unsqueeze(0),
boxes.unsqueeze(0), MAPS, Scales)
if not args.adapt:
with torch.no_grad():
output = regressor(features)
else:
features.required_grad = True
def caculate_iou(self,fov):
polygon1 = _sort_vertices_anti_clockwise_and_remove_duplicates(self.points)
polygon2 = _sort_vertices_anti_clockwise_and_remove_duplicates(fov.points)
polygon3 = intersect(polygon1, polygon2)
plot_polygon(polygon1)
plot_polygon(polygon2)
#img = np.zeros([500,500],np.int8)
#cv2.drawContours(img,[np.array(polygon3).reshape(-1,1,2).astype(np.int32)],0,(255,255,255),2)
if len(polygon3) > 0:
plot_polygon(polygon3)
print(cv2.contourArea(np.array(polygon3).reshape(-1,1,2).astype(np.int32))/cv2.contourArea(np.array(polygon1).reshape(-1,1,2).astype(np.int32)))
plt.show()
if __name__ == '__main__':
cam1 = Transform(location=Location(x=-64, y=-126, z=4),rotation=Rotation(pitch=0, yaw=175, roll=0))
cam2 = Transform(location=Location(x=-64, y=-120, z=4),rotation=Rotation(pitch=0, yaw=175, roll=0))
cam3 = Transform(location=Location(x=-64, y=-126.5, z=4),rotation=Rotation(pitch=0, yaw=155, roll=0))
cam4 = Transform(location=Location(x=-64, y=-152, z=4),rotation=Rotation(pitch=0, yaw=110, roll=0))
fov1 = FOV(cam1)
fov2 = FOV(cam2)
fov3 = FOV(cam3)
fov4 = FOV(cam4)
polygon1 = _sort_vertices_anti_clockwise_and_remove_duplicates(fov1.points)
polygon2 = _sort_vertices_anti_clockwise_and_remove_duplicates(fov2.points)
polygon3 = _sort_vertices_anti_clockwise_and_remove_duplicates(fov3.points)
polygon4 = _sort_vertices_anti_clockwise_and_remove_duplicates(fov4.points)
plot_polygon(polygon1)
plot_polygon(polygon2)
polygon12 = intersect(polygon1, polygon2)
plot_polygon(polygon12)
def train(config, device='cuda:0', save_chkpt=True):
''' procedure launching all main functions of training,
validation and testing pipelines'''
# for pipeline testing purposes
save_chkpt = False if config.test_steps else True
# preprocessing data
normalize = A.Normalize(**config.img_norm_cfg)
train_transform_real = A.Compose([
A.Resize(**config.resize, interpolation=cv.INTER_CUBIC),
A.HorizontalFlip(p=0.5),
A.augmentations.transforms.ISONoise(color_shift=(0.15, 0.35),
intensity=(0.2, 0.5),
p=0.2),
A.augmentations.transforms.RandomBrightnessContrast(
brightness_limit=0.2,
contrast_limit=0.2,
brightness_by_max=True,
always_apply=False,
p=0.3),
A.augmentations.transforms.MotionBlur(blur_limit=5, p=0.2), normalize
])
train_transform_spoof = A.Compose([
A.Resize(**config.resize, interpolation=cv.INTER_CUBIC),
A.HorizontalFlip(p=0.5),
A.augmentations.transforms.ISONoise(color_shift=(0.15, 0.35),
intensity=(0.2, 0.5),
p=0.2),
A.augmentations.transforms.RandomBrightnessContrast(
brightness_limit=0.2,
contrast_limit=0.2,
brightness_by_max=True,
always_apply=False,
p=0.3),
A.augmentations.transforms.MotionBlur(blur_limit=5, p=0.2), normalize
])
val_transform = A.Compose(
[A.Resize(**config.resize, interpolation=cv.INTER_CUBIC), normalize])
# load data
sampler = config.data.sampler
if sampler:
num_instances, weights = make_weights(config)
sampler = torch.utils.data.WeightedRandomSampler(weights,
num_instances,
replacement=True)
train_transform = Transform(train_spoof=train_transform_spoof,
train_real=train_transform_real,
val=None)
val_transform = Transform(train_spoof=None,
train_real=None,
val=val_transform)
train_dataset, val_dataset, test_dataset = make_dataset(
config, train_transform, val_transform)
train_loader, val_loader, test_loader = make_loader(train_dataset,
val_dataset,
test_dataset,
config,
sampler=sampler)
# build model and put it to cuda and if it needed then wrap model to data parallel
model = build_model(config, device=device, strict=False, mode='train')
model.to(device)
if config.data_parallel.use_parallel:
model = torch.nn.DataParallel(model,
**config.data_parallel.parallel_params)
# build a criterion
softmax = build_criterion(config, device, task='main').to(device)
cross_entropy = build_criterion(config, device, task='rest').to(device)
bce = nn.BCELoss().to(device)
criterion = (softmax, cross_entropy,
bce) if config.multi_task_learning else softmax
# build optimizer and scheduler for it
optimizer = torch.optim.SGD(model.parameters(), **config.optimizer)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
**config.scheduler)
# create Trainer object and get experiment information
trainer = Trainer(model, criterion, optimizer, device, config,
train_loader, val_loader, test_loader)
trainer.get_exp_info()
# learning epochs
for epoch in range(config.epochs.start_epoch, config.epochs.max_epoch):
if epoch != config.epochs.start_epoch:
scheduler.step()
# train model for one epoch
train_loss, train_accuracy = trainer.train(epoch)
print(
f'epoch: {epoch} train loss: {train_loss} train accuracy: {train_accuracy}'
)
# validate your model
accuracy = trainer.validate()
# eval metrics such as AUC, APCER, BPCER, ACER on val and test dataset according to rule
trainer.eval(epoch, accuracy, save_chkpt=save_chkpt)
# for testing purposes
if config.test_steps:
exit()
# evaluate in the end of training
if config.evaluation:
file_name = 'tests.txt'
trainer.test(file_name=file_name)
def main():
# create model
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# model = SENet(SEResNeXtBottleneck, [3, 4, 6, 3], groups=32, reduction=16).to(device)
if args.net == 0:
model = SEAttentionNet(SEResNeXtBottleneck, [3, 4, 6, 3],
groups=32,
reduction=16,
dropout_p=0.2,
inplanes=64,
input_3x3=False,
downsample_kernel_size=1,
downsample_padding=0)
elif args.net == 1:
model = SENet(SEResNeXtBottleneck, [3, 4, 6, 3],
groups=32,
reduction=16,
dropout_p=0.2,
inplanes=64,
input_3x3=False,
downsample_kernel_size=1,
downsample_padding=0)
elif args.net == 2:
model = SENetHeavyHead(SEResNeXtBottleneck, [3, 4, 6, 3],
groups=32,
reduction=16,
dropout_p=0.2,
inplanes=64,
input_3x3=False,
downsample_kernel_size=1,
downsample_padding=0)
else:
print("Undefined Net Option")
if args.load_model:
model.load_state_dict(torch.load(args.load_model_path))
model = model.to(device)
print("Create Model Done")
# create dataset
if not args.valid_shuffle:
train_images = load_image(args.data_path, args.valid_fold)
valid_images = load_image(args.data_path, args.valid_fold, False)
df = pd.read_csv(args.data_path + '/train.csv')
labels = df[[
'grapheme_root', 'vowel_diacritic', 'consonant_diacritic'
]].values
valid_indices = [i + 50210 * args.valid_fold for i in range(50210)]
train_indices = [i for i in range(50210 * 4) if i not in valid_indices]
train_labels = np.take(labels, train_indices, axis=0)
valid_labels = np.take(labels, valid_indices, axis=0)
else:
images = load_image_shuffle(args.data_path)
train_pd = pd.read_csv(args.data_path + '/train.csv')
labels = train_pd[[
'grapheme_root', 'vowel_diacritic', 'consonant_diacritic'
]].values
num_train = int(images.shape[0] * (1 - args.valid_ratio))
train_indices = random.sample(set(range(images.shape[0])), num_train)
valid_indices = [
i for i in range(images.shape[0]) if i not in train_indices
]
train_images = np.take(images, train_indices, axis=0)
valid_images = np.take(images, valid_indices, axis=0)
train_labels = np.take(labels, train_indices, axis=0)
valid_labels = np.take(labels, valid_indices, axis=0)
num_train = len(train_indices)
train_transform = Transform(
affine=args.affine,
size=(128, 128),
threshold=args.threshold,
sigma=-1.,
blur_ratio=args.blur_ratio,
noise_ratio=args.noise_ratio,
cutout_ratio=args.cutout_ratio,
elastic_distortion_ratio=args.elastic_distortion_ratio,
random_brightness_ratio=args.random_brightness_ratio,
piece_affine_ratio=args.piece_affine_ratio,
ssr_ratio=args.ssr_ratio)
# transform = Transform(size=(image_size, image_size))
train_dataset = BengaliAIDataset(train_images,
train_labels,
transform=train_transform)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
valid_transform = Transform(affine=args.affine,
size=(128, 128),
threshold=args.threshold,
sigma=-1.)
valid_dataset = BengaliAIDataset(valid_images,
valid_labels,
transform=valid_transform)
valid_loader = DataLoader(valid_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
# create train loaders for cutmix
cutmix_transform = Transform(affine=args.affine,
size=(128, 128),
threshold=args.threshold,
sigma=-1.)
cutmix_loaders = []
grapheme_dict = defaultdict(list)
for i in range(train_labels.shape[0]):
label_v, label_c = train_labels[i][1], train_labels[i][2]
grapheme_dict['%2d%2d' % (label_v, label_c)].append(i)
idx, iter_idxes = 0, []
for _, v in grapheme_dict.items():
subimages = np.take(train_images, v, axis=0)
sublabels = np.take(train_labels, v, axis=0)
cutmix_dataset = BengaliAIDataset(subimages,
sublabels,
transform=cutmix_transform)
cutmix_loaders.append(
DataLoader(cutmix_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers))
iter_idxes += [idx] * (len(cutmix_dataset) // args.batch_size)
idx += 1
iter_idxes += [idx] * (len(train_dataset) // args.batch_size)
print("DataLoader Done")
# train code
criterion = nn.CrossEntropyLoss()
if args.optim == 'sgd':
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.0)
else:
optimizer = optim.Adam(model.parameters(), lr=args.lr)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
'min',
factor=0.25,
patience=6,
min_lr=1e-07)
writer = SummaryWriter()
best_score = 0
for epoch in range(args.epochs):
random.shuffle(iter_idxes)
all_iters = [iter(loader) for loader in cutmix_loaders]
all_iters.append(iter(train_loader))
cutmix_losses = [0, 0, 0]
train_losses = [0, 0, 0]
for idx in iter_idxes:
dataiter = all_iters[idx]
inputs, labels = next(dataiter)
inputs, labels = inputs.to(device), labels.to(device)
if idx < len(all_iters) - 1:
train_loss = dotrain(model, optimizer, criterion, inputs,
labels)
for i in range(3):
train_losses[i] += train_loss[i] * inputs.shape[0]
else:
cutmix_loss = docutmixtrain(model, optimizer, criterion,
inputs, labels, args.alpha)
for i in range(3):
cutmix_losses[i] += cutmix_loss[i] * inputs.shape[0]
# # train with grapheme root cutmix
# cutmix_losses = [0, 0, 0]
# for cutmix_loader in cutmix_loaders:
# for inputs, labels in cutmix_loader:
# inputs, labels = inputs.to(device), labels.to(device)
# cutmix_loss = docutmixtrain(model, optimizer, criterion, inputs, labels, args.alpha)
# for i in range(3):
# cutmix_losses[i] += cutmix_loss[i]*inputs.shape[0]
# # train with normal data augmentation
# train_losses = [0, 0, 0]
# for inputs, labels in train_loader:
# inputs, labels = inputs.to(device), labels.to(device)
# train_loss = dotrain(model, optimizer, criterion, inputs, labels)
# for i in range(3):
# train_losses[i] += train_loss[i]*inputs.shape[0]
train_acc, train_scores, train_loss = dovalid(model, train_loader,
device, criterion)
writer.add_scalars('train acc', {
'acc1': train_acc[0],
'acc2': train_acc[1],
'acc3': train_acc[2]
}, epoch)
writer.add_scalars(
'train score', {
'score1': train_scores[0],
'score2': train_scores[1],
'score3': train_scores[2]
}, epoch)
writer.add_scalars('train loss', {
'loss1': train_loss[0],
'loss2': train_loss[1],
'loss3': train_loss[2]
}, epoch)
valid_acc, valid_scores, valid_loss = dovalid(model, valid_loader,
device, criterion)
writer.add_scalars('valid acc', {
'acc1': valid_acc[0],
'acc2': valid_acc[1],
'acc3': valid_acc[2]
}, epoch)
writer.add_scalars(
'valid score', {
'score1': valid_scores[0],
'score2': valid_scores[1],
'score3': valid_scores[2]
}, epoch)
writer.add_scalars('valid loss', {
'loss1': valid_loss[0],
'loss2': valid_loss[1],
'loss3': valid_loss[2]
}, epoch)
print("epoch %d done" % (epoch))
# save model
score = (valid_scores[0] * 2 + valid_scores[1] + valid_scores[2]) / 4
if score > best_score:
torch.save(model.state_dict(), args.save_path + "bengali.pt")
best_score = score
train_losses = [loss / num_train for loss in train_losses]
cutmix_losses = [loss / num_train for loss in cutmix_losses]
if args.verbal:
print("Normal Train Losses: %f, %f, %f" %
(train_losses[0], train_losses[1], train_losses[2]))
print("Cutmix Train Losses: %f, %f, %f" %
(cutmix_losses[0], cutmix_losses[1], cutmix_losses[2]))
print("Train ACC: %f, %f, %f" %
(train_acc[0], train_acc[1], train_acc[2]))
print("Train Scores: %f, %f, %f" %
(train_scores[0], train_scores[1], train_scores[2]))
print("Train Loss: %f, %f, %f" %
(train_loss[0], train_loss[1], train_loss[2]))
print("Valid ACC: %f, %f, %f" %
(valid_acc[0], valid_acc[1], valid_acc[2]))
print("Valid Scores: %f, %f, %f" %
(valid_scores[0], valid_scores[1], valid_scores[2]))
print("Valid Loss: %f, %f, %f" %
(valid_loss[0], valid_loss[1], valid_loss[2]))
print("Best Score: ", best_score)
scheduler_loss = np.average(
train_losses, weights=[2, 1, 1]) + np.average(cutmix_losses,
weights=[2, 1, 1])
scheduler.step(scheduler_loss)
all_iters.clear()
def call(self, source_images, driving_images, tape):
kp_source = self.key_point_detector(source_images)
kp_driving = self.key_point_detector(driving_images)
generated = self.generator(source_images,
kp_source=kp_source,
kp_driving=kp_driving)
# Debug/print
generated.update({'kp_source': kp_source, 'kp_driving': kp_driving})
loss_values = {}
pyramide_real = self.pyramid(driving_images)
pyramide_generated = self.pyramid(generated['prediction'])
# Perceptual loss (Loss for gan generator)
perceptual_loss = 0
for scale in self.scales:
x_vgg = self.vgg(pyramide_generated['prediction_' + str(scale)])
y_vgg = self.vgg(pyramide_real['prediction_' + str(scale)])
for i, weight in enumerate(self.perceptual_weights):
loss = tf.reduce_mean(
tf.abs(x_vgg[i] - tf.stop_gradient(y_vgg[i])))
perceptual_loss += self.perceptual_weights[i] * loss
loss_values['perceptual'] = perceptual_loss
# Gan loss (only one scale used, the original [1])
# We detach the keypoints here so we dont compue its gradients and we use it as input images!!!
discriminator_maps_real, _ = self.discriminator(
driving_images, kp=detach_keypoint(kp_driving))
discriminator_maps_generated, discriminator_pred_map_generated = self.discriminator(
generated['prediction'], kp=detach_keypoint(kp_driving))
# LSGAN G Loss
# Discriminator outputs a pathmap like pix2pix where 1 labels are for real images and 0 labels are for generated images
# Since we want to fool the discriminator we want our generated images to output 1
gan_loss = tf.reduce_mean((discriminator_pred_map_generated - 1)**2)
# same as tf.reduce_mean(tf.keras.losses.mean_squared_error(tf.ones_like(discriminator_pred_map_generated), discriminator_pred_map_generated))
gan_loss += self.loss_weights['generator_gan'] * gan_loss
loss_values['gen_gan'] = gan_loss
# feature_matching loss
feature_matching_loss = tf.reduce_mean(
tf.abs(discriminator_maps_real - discriminator_maps_generated))
feature_matching_loss += self.feature_matching_weights * feature_matching_loss
loss_values['feature_matching'] = feature_matching_loss
# Equivariance Loss
batch_size = driving_images.shape[0]
transform = Transform(batch_size)
transformed_frame = transform.transform_frame(driving_images)
# image Y
# shape batch x height x width x 2
transformed_keypoints = self.key_point_detector(transformed_frame)
# Ty <-R
# Debug/print
generated['transformed_frame'] = transformed_frame
# Debug/print
generated['transformed_kp'] = transformed_keypoints
keypoints_loss = tf.reduce_mean(
tf.abs(kp_driving['value'] -
transform.warp_coordinates(transformed_keypoints['value'])))
loss_values[
'equivariance_value'] = self.equivariance_weights * keypoints_loss
# Here we apply the transformation for a second time and then compute the jacobian
jacobian_transformed = tf.linalg.matmul(
transform.jacobian(transformed_keypoints['value'], tape),
transformed_keypoints['jacobian'])
# Equivariance properties
normed_driving = tf.linalg.inv(
kp_driving['jacobian']) #inverse of Tx <-R
normed_transformed = jacobian_transformed
jacobian_mul = tf.linalg.matmul(normed_driving, normed_transformed)
identity_matrix = tf.cast(tf.reshape(tf.eye(2), [1, 1, 2, 2]),
jacobian_mul.dtype)
jacobian_loss = tf.reduce_mean(tf.abs(identity_matrix - jacobian_mul))
loss_values[
'equivariance_jacobian'] = self.equivariance_weights * jacobian_loss
return loss_values, generated
# DH
linkparam = [[90, 0, 90, 0], [0, 0.02, None, 0], [0, 0.05, None, 0],
[0, 0.08, None, 0], [0, 0.04, 90, 0], [90, 0, 0, 0]]
g = 9.8
M = [0, 0.2, 0.2, 0.2, 0, 0]
I = [0, 0, 0, [[0, 0, 0], [0, 0, 0], [0, 0, 0.01]], 0, 0]
C = [[0, 0, 0], [0.01, 0, 0], [0.025, 0, 0], [0.04, 0, 0], [0.02, 0, 0],
[0, 0, 0]]
# init position
th0 = np.array([0, 0, 0])
init_angle = th0 / 180 * np.pi
# find target position
tan68 = np.arctan(6 / 8) * 180 / np.pi
t_fp = Transform(rot=Rotation(90 - tan68, 1)) * \
Transform(loc=Translation(0.050, 0.035, 0))
t_cf = Transform(rot=Rotation(90, 0), loc=Translation(0, 0, 0.065))
t_bc = Transform(rot=Rotation(90, 0), loc=Translation(0, 0.100, 0.100))
t_bp = t_bc * t_cf * t_fp
print("Target", t_bp)
th1 = gradientVariable(t_bp.mat, linkparam, init_angle)
# set rotation angle that < 180 degrees
th1[(th1 - th0) > 180] -= 360
th1[(th1 - th0) < -180] += 360
vth = [0, *(th1 - th0) / 10, 0, 0]
ath = [0, 0, 0, 0, 0, 0]
print("Angular velcoity", vth)
# init
[ 0, 3, None, 0],
[ 0, 2, 0, 0]]
g = 9.8
th = [ 10, 20, 30, 0]
vth = [ 1, 2, 3, 0]
ath = [0.5, 1, 1.5, 0]
M = [0, 20, 15, 10]
I = [0, [[0, 0, 0], [0, 0, 0], [0, 0, 0.5]], [[0, 0, 0], [0, 0, 0], [0, 0, 0.2]], [[0, 0, 0], [0, 0, 0], [0, 0, 0.1]], 0]
C = [[0, 0, 0], [2, 0, 0], [1.5, 0, 0], [1, 0, 0]]
"""
vth = np.array(vth) / np.pi * 180
ath = np.array(ath) / np.pi * 180
# Result
Ts0 = [Transform()]
Ts = [Transform()]
v = [vTranslate()]
a = [vTranslate(v=Translation(0, g, 0))]
xy = [Translation()]
ac = [Translation()]
F = [Translation()]
N = [Translation()]
f = [Translation()] * (len(linkparam) + 1)
n = [Translation()] * (len(linkparam) + 1)
# joint
for i in range(len(linkparam)):
# set angle of joint
linkparam[i][2] = th[i]
