import sys, os
from dataset.d300w import ThreeHundredW
import ot
from dataset.lazy_loader import Cardio, LazyLoader
from parameters.path import Paths
import numpy as np
from torch.utils.data import Subset
image_size = 256
padding = 200
N = 300
prob = np.ones(padding) / padding
dataset_train = LazyLoader.cardio_landmarks("cardio_300/lm").dataset_train
def load_landmarks(k):
return dataset_train[k].numpy()
landmarks = [load_landmarks(k) for k in range(N)]
def compute_w2(i, j):
M_ij = ot.dist(landmarks[i], landmarks[j])
D_ij = ot.emd2(prob, prob, M_ij)
return D_ij
hm_discriminator = hm_discriminator.cuda()
gan_model_tuda = StyleGanModel[HeatmapToImage](enc_dec.generator, StyleGANLoss(discriminator_img), (0.001/4, 0.0015/4))
gan_model_obratno = StyleGanModel[HG_skeleton](hg, StyleGANLoss(hm_discriminator, r1=3), (2e-5, 0.0015/4))
style_opt = optim.Adam(enc_dec.style_encoder.parameters(), lr=1e-5)
print(f"board path: {Paths.default.board()}/cardio{int(time.time())}")
writer = SummaryWriter(f"{Paths.default.board()}/cardio{int(time.time())}")
WR.writer = writer
#%%
test_batch = next(LazyLoader.cardio().loader_train_inf)
test_img = test_batch["image"].cuda()
test_landmarks = next(LazyLoader.cardio_landmarks(args.data_path).loader_train_inf).cuda()
test_measure = UniformMeasure2D01(torch.clamp(test_landmarks, max=1))
test_hm = heatmapper.forward(test_measure.coord).sum(1, keepdim=True).detach()
test_noise = mixing_noise(batch_size, 512, 0.9, device)
psp_loss = PSPLoss(id_lambda=0).cuda()
mes_loss = MesBceWasLoss(heatmapper, bce_coef=100000, was_coef=2000)
image_accumulator = Accumulator(enc_dec.generator, decay=0.99, write_every=100)
hm_accumulator = Accumulator(hg, decay=0.99, write_every=100)
for i in range(100000):
WR.counter.update(i)
try:
from dataset.lazy_loader import LazyLoader, W300DatasetLoader, CelebaWithKeyPoints, Celeba
from sklearn.neighbors import NearestNeighbors
import numpy as np
import matplotlib.pyplot as plt
from dataset.toheatmap import ToGaussHeatMap
from dataset.probmeasure import UniformMeasure2D01
import pandas as pd
import networkx as nx
import ot
from barycenters.sampler import Uniform2DBarycenterSampler, Uniform2DAverageSampler
from parameters.path import Paths
from joblib import Parallel, delayed
N = 701
dataset = LazyLoader.cardio_landmarks(f"cardio_{N}/lm").dataset_train
D = np.load(f"{Paths.default.models()}/cardio_graph{N}.npy")
padding = 200
prob = np.ones(padding) / padding
NS = 1000
def LS(k):
return dataset[k].numpy()
ls = np.asarray([LS(k) for k in range(N)])
def viz_mes(ms):
heatmaper = ToGaussHeatMap(128, 1)
gan_model_tuda = StyleGanModel[HeatmapToImage](enc_dec.generator,
StyleGANLoss(discriminator_img),
(0.001 / 4, 0.0015 / 4))
gan_model_obratno = StyleGanModel[HG_skeleton](hg,
StyleGANLoss(hm_discriminator),
(2e-5, 0.0015 / 4))
style_opt = optim.Adam(enc_dec.style_encoder.parameters(), lr=1e-5)
writer = SummaryWriter(f"{Paths.default.board()}/cardio{int(time.time())}")
WR.writer = writer
test_batch = next(LazyLoader.cardio().loader_train_inf)
test_img = test_batch["image"].cuda()
test_landmarks = next(LazyLoader.cardio_landmarks().loader_train_inf).cuda()
test_hm = heatmapper.forward(test_landmarks).sum(1, keepdim=True).detach()
test_noise = mixing_noise(batch_size, 512, 0.9, device)
psp_loss = PSPLoss(id_lambda=0).cuda()
mes_loss = MesBceWasLoss(heatmapper, bce_coef=1000000, was_coef=4000)
image_accumulator = Accumulator(enc_dec.generator, decay=0.99, write_every=100)
hm_accumulator = Accumulator(hg, decay=0.99, write_every=100)
for i in range(100000):
WR.counter.update(i)
real_img = next(LazyLoader.cardio().loader_train_inf)["image"].cuda()
landmarks = next(LazyLoader.cardio_landmarks().loader_train_inf).cuda()