def __init__(self, lm_count, dir_alpha=1):
starting_model_number = 560000 + 310000
weights = torch.load(
f'{Paths.default.models()}/hm2img_{str(starting_model_number).zfill(6)}.pt',
map_location="cpu")
self.enc_dec = StyleGanAutoEncoder().load_state_dict(weights).cuda()
self.dir_alpha = dir_alpha
self.mes_sampler = Uniform2DBarycenterSampler(lm_count, dir_alpha)
self.heatmapper = ToGaussHeatMap(256, 4)
def draw_lm(lm1, lm2):
heatmaper = ToGaussHeatMap(128, 1)
keyptsiki = torch.from_numpy(lm1)[None, ]
tmp1 = heatmaper.forward(keyptsiki).type(torch.float32).sum(1,
keepdim=True)
lm2[:, 0] += 0.1
keyptsiki = torch.from_numpy(lm2)[None, ]
tmp2 = heatmaper.forward(keyptsiki).type(torch.float32).sum(1,
keepdim=True)
tmp = torch.cat([tmp1, tmp2, torch.zeros_like(tmp2)], dim=1)
plt.imshow(tmp[0].permute(1, 2, 0).numpy())
plt.show()
def viz_mes(ms):
heatmaper = ToGaussHeatMap(128, 1)
kapusta = torch.zeros(128, 128)
for m in ms:
keyptsiki = torch.from_numpy(m)[None, ].clamp(0, 1)
tmp = heatmaper.forward(keyptsiki)
kapusta += tmp.sum(axis=(0, 1))
plt.imshow(kapusta)
plt.show()
return kapusta / kapusta.sum()
def sup_loss(pred_mes, target_mes):
heatmapper = ToGaussHeatMap(256, 1)
pred_hm = heatmapper.forward(pred_mes.coord)
pred_hm = pred_hm / (pred_hm.sum(dim=[2, 3], keepdim=True).detach() + 1e-8)
target_hm = heatmapper.forward(target_mes.coord).detach()
target_hm = target_hm / target_hm.sum(dim=[2, 3], keepdim=True).detach()
return Loss(
nn.BCELoss()(pred_hm, target_hm) * 100 +
nn.MSELoss()(pred_mes.coord, target_mes.coord) * 0.5
)
def fill_contour(landmarks: np.ndarray):
heatmapper = ToGaussHeatMap(256, 3)
hm = heatmapper.forward(torch.from_numpy(landmarks)[None, ])[0].numpy()
img_pl = np.zeros((256, 256))
hmi = hm.sum(0)
res = cv.findContours(hmi.astype(np.uint8), 0, cv.CHAIN_APPROX_NONE)
contours = res[-2]
contours = [i.transpose(1, 0, 2) for i in contours]
cv.fillPoly(img_pl, contours, 1)
return img_pl
class ImageBarycenterSampler:
def __init__(self, lm_count, dir_alpha=1):
starting_model_number = 560000 + 310000
weights = torch.load(
f'{Paths.default.models()}/hm2img_{str(starting_model_number).zfill(6)}.pt',
map_location="cpu")
self.enc_dec = StyleGanAutoEncoder().load_state_dict(weights).cuda()
self.dir_alpha = dir_alpha
self.mes_sampler = Uniform2DBarycenterSampler(lm_count, dir_alpha)
self.heatmapper = ToGaussHeatMap(256, 4)
def sample(self, measures: List[np.ndarray],
images: torch.Tensor) -> (np.ndarray, torch.Tensor, np.ndarray):
bc_mes, weights = self.mes_sampler.sample(measures)
with torch.no_grad():
latents = self.enc_dec.encode_latent(images)
bc_latent = reduce(lambda x, y: x + y,
[(latents[i] * weights[i]).type(torch.float32)
for i in range(len(measures))])[None, ]
bc_mes_cuda = torch.from_numpy(bc_mes).cuda().type(
torch.float32)[None, ]
hm = self.heatmapper.forward(bc_mes_cuda).sum(1, keepdim=True)
return bc_mes, self.enc_dec.decode(hm, bc_latent)[0].cpu(), weights
def coord_hm_loss(pred_coord: Tensor, target_hm: Tensor, coef=1.0):
target_hm = target_hm / target_hm.sum(dim=[2, 3], keepdim=True)
target_hm = target_hm.detach()
heatmapper = ToGaussHeatMap(256, 4)
target_coord = UniformMeasure2DFactory.from_heatmap(
target_hm).coord.detach()
# sk = CoordToGaussSkeleton(target_hm.shape[-1], 1)
# pred_sk = sk.forward(pred_coord).sum(dim=1, keepdim=True)
# target_sk = sk.forward(target_coord).sum(dim=1, keepdim=True).detach()
pred_hm = heatmapper.forward(pred_coord).sum(dim=1, keepdim=True)
pred_hm = pred_hm / pred_hm.sum(dim=[2, 3], keepdim=True).detach()
target_hm = heatmapper.forward(target_coord).sum(dim=1,
keepdim=True).detach()
target_hm = target_hm / target_hm.sum(dim=[2, 3], keepdim=True).detach()
return Loss(nn.BCELoss()(pred_hm, target_hm) * coef * 1.5 +
# noviy_hm_loss(pred_sk, target_sk, coef).to_tensor() * 0.5 +
nn.MSELoss()(pred_coord, target_coord) * (0.001 * coef) +
nn.L1Loss()(pred_coord, target_coord) * (0.001 * coef))
def train_content(cont_opt, R_b, R_t, real_img, model, encoder_HG, decoder, generator, style_encoder):
B = real_img.shape[0]
C = 512
heatmapper = ToGaussHeatMap(256, 1)
requires_grad(encoder_HG, True)
coefs = json.load(open("../parameters/content_loss.json"))
encoded = encoder_HG(real_img)
pred_measures: UniformMeasure2D01 = UniformMeasure2D01(encoded["coords"])
heatmap_content = heatmapper.forward(encoded["coords"]).detach()
ll = (
WR.writable("R_b", R_b.__call__)(real_img, pred_measures) * coefs["R_b"] +
WR.writable("R_t", R_t.__call__)(real_img, heatmap_content) * coefs["R_t"]
)
ll.minimize_step(cont_opt)
class GHSparse(nn.Module):
def __init__(self, gh: nn.Module):
super().__init__()
self.gh = gh
self.heatmaper = ToGaussHeatMap(64, 1.0)
def forward(self, image: Tensor):
img_content = self.gh(image)
pred_measures: UniformMeasure2D01 = UniformMeasure2DFactory.from_heatmap(
img_content)
sparce_hm = self.heatmaper.forward(pred_measures.coord * 63)
return pred_measures, sparce_hm
class LandmarksDatasetAugment(Dataset):
def __init__(self, path: str, transform=None):
super().__init__()
self.path = path
self.transform = transform
self.data = make_dataset(path)
self.heatmapper = ToGaussHeatMap(256, 1)
def __getitem__(self, index):
keypts = np.load(self.data[index]) # numpy - [200, 2]
hm = self.heatmapper.forward(torch.tensor(keypts)[None, ])[0]
transformed = self.transform(image=np.zeros_like(
np.array(hm.permute(1, 2, 0))),
mask=np.array(hm.permute(1, 2, 0)))
mask = transformed["mask"]
coord, p = heatmap_to_measure(mask.permute(2, 0, 1)[None])
return coord[0].type(torch.float32) #torch.tensor - [200, 2]
def __len__(self):
return len(self.data)
def content_trainer_with_gan(cont_opt, tuner, encoder_HG, R_b, R_t, model, generator, g_transforms, decoder, style_encoder):
C = 512
heatmapper = ToGaussHeatMap(256, 1)
def do_train(real_img):
B = real_img.shape[0]
requires_grad(encoder_HG, True)
requires_grad(decoder, False)
coefs = json.load(open("../parameters/content_loss.json"))
encoded = encoder_HG(real_img)
pred_measures: UniformMeasure2D01 = UniformMeasure2D01(encoded["coords"])
heatmap_content = heatmapper.forward(encoded["coords"]).detach()
restored = decoder(encoded["skeleton"], style_encoder(real_img))
noise = mixing_noise(B, C, 0.9, real_img.device)
fake, _ = generator(encoded["skeleton"], noise)
fake_content = encoder_HG(fake.detach())["coords"]
ll = (
WR.writable("R_b", R_b.__call__)(real_img, pred_measures) * coefs["R_b"] +
WR.writable("R_t", R_t.__call__)(real_img, heatmap_content) * coefs["R_t"] +
WR.L1("L1 image")(restored, real_img) * coefs["L1 image"] +
WR.writable("fake_content loss", coord_hm_loss)(
fake_content, heatmap_content
) * coefs["fake_content loss"] +
WR.writable("Fake-content D", model.loss.generator_loss)(
real=None,
fake=[fake, encoded["skeleton"].detach()]) * coefs["Fake-content D"]
)
ll.minimize_step(cont_opt)
return do_train
def content_trainer_with_gan(encoder_HG, model, generator, decoder,
style_encoder):
C = 512
heatmapper = ToGaussHeatMap(256, 1)
def do_train(real_img):
B = real_img.shape[0]
coefs = json.load(open("../parameters/content_loss.json"))
requires_grad(encoder_HG, True)
requires_grad(decoder, False)
requires_grad(generator, False)
encoded = encoder_HG(real_img)
pred_measures: UniformMeasure2D01 = encoded["mes"]
heatmap_content = heatmapper.forward(pred_measures.coord).detach()
restored = decoder(encoded["skeleton"], style_encoder(real_img))
noise = mixing_noise(B, C, 0.9, real_img.device)
fake, _ = generator(encoded["skeleton"], noise)
fake_content = encoder_HG(fake.detach())["mes"]
ll = (WR.L1("L1 image")(restored, real_img) * coefs["L1 image"] +
WR.writable("fake_content brule_loss", coord_hm_loss)
(fake_content, heatmap_content) *
coefs["fake_content brule_loss"] +
WR.writable("Fake-content D", model.loss.generator_loss)
(real=None, fake=[fake, encoded["skeleton"].detach()]) *
coefs["Fake-content D"])
ll.minimize_step(model.optimizer.opt_min)
return do_train
starting_model_number = 560000 + args.weights
weights = torch.load(
f'{Paths.default.models()}/hm2img_brule1_{str(starting_model_number).zfill(6)}.pt',
map_location="cpu"
)
enc_dec = StyleGanAutoEncoder().load_state_dict(weights).cuda()
discriminator_img = Discriminator(image_size)
discriminator_img.load_state_dict(weights['di'])
discriminator_img = discriminator_img.cuda()
heatmapper = ToGaussHeatMap(256, 4)
hg = HG_heatmap(heatmapper, num_blocks=1)
hg.load_state_dict(weights['gh'])
hg = hg.cuda()
cont_opt = optim.Adam(hg.parameters(), lr=2e-5, betas=(0, 0.8))
gan_model_tuda = StyleGanModel[HeatmapToImage](enc_dec.generator, StyleGANLoss(discriminator_img), (0.001/4, 0.0015/4))
style_opt = optim.Adam(enc_dec.style_encoder.parameters(), lr=1e-5)
writer = SummaryWriter(f"{Paths.default.board()}/brule1_{int(time.time())}")
WR.writer = writer
batch = next(LazyLoader.w300().loader_train_inf)
test_img = batch["data"].cuda()
test_landmarks = torch.clamp(next(LazyLoader.w300_landmarks(args.data_path).loader_train_inf).cuda(), max=1)
def train(generator, decoder, discriminator, encoder_HG, style_encoder, device, starting_model_number):
latent_size = 512
batch_size = 12
sample_z = torch.randn(8, latent_size, device=device)
MAFL.batch_size = batch_size
MAFL.test_batch_size = 64
Celeba.batch_size = batch_size
test_img = next(LazyLoader.mafl().loader_train_inf)["data"][:8].cuda()
loss_st: StyleGANLoss = StyleGANLoss(discriminator)
model = CondStyleGanModel(generator, loss_st, (0.001, 0.0015))
style_opt = optim.Adam(style_encoder.parameters(), lr=5e-4, betas=(0.9, 0.99))
cont_opt = optim.Adam(encoder_HG.parameters(), lr=2e-5, betas=(0.5, 0.97))
g_transforms: albumentations.DualTransform = albumentations.Compose([
ToNumpy(),
NumpyBatch(albumentations.Compose([
ResizeMask(h=256, w=256),
albumentations.ElasticTransform(p=0.7, alpha=150, alpha_affine=1, sigma=10),
albumentations.ShiftScaleRotate(p=0.7, rotate_limit=15),
ResizeMask(h=64, w=64),
NormalizeMask(dim=(0, 1, 2))
])),
ToTensor(device),
])
R_t = DualTransformRegularizer.__call__(
g_transforms, lambda trans_dict, img:
# rt_loss(encoder_HG(trans_dict['image']), trans_dict['mask'])
stariy_hm_loss(encoder_HG(trans_dict['image']), trans_dict['mask'])
)
R_s = UnoTransformRegularizer.__call__(
g_transforms, lambda trans_dict, img, ltnt:
L1("R_s")(ltnt, style_encoder(trans_dict['image']))
)
barycenter: UniformMeasure2D01 = UniformMeasure2DFactory.load(
f"{Paths.default.models()}/face_barycenter_5").cuda().batch_repeat(batch_size)
R_b = BarycenterRegularizer.__call__(barycenter, 1.0, 2.0, 4.0)
tuner = GoldTuner([0.37, 1.55, 0.9393, 0.1264, 1.7687, 0.8648, 1.8609], device=device, rule_eps=0.01 / 2,
radius=0.1, active=True)
heatmaper = ToGaussHeatMap(64, 1.0)
sparse_bc = heatmaper.forward(barycenter.coord * 63)
sparse_bc = nn.Upsample(scale_factor=4)(sparse_bc).sum(dim=1, keepdim=True).repeat(1, 3, 1, 1) * \
torch.tensor([1.0, 1.0, 0.0], device=device).view(1, 3, 1, 1)
sparse_bc = (sparse_bc - sparse_bc.min()) / sparse_bc.max()
send_images_to_tensorboard(writer, sparse_bc, "BC", 0, normalize=False, range=(0, 1))
trainer_gan = gan_trainer(model, generator, decoder, encoder_HG, style_encoder, R_s, style_opt, heatmaper,
g_transforms)
content_trainer = content_trainer_with_gan(cont_opt, tuner, heatmaper, encoder_HG, R_b, R_t, model, generator,
g_transforms)
supervise_trainer = content_trainer_supervised(cont_opt, encoder_HG, LazyLoader.mafl().loader_train_inf)
for i in range(100000):
counter.update(i)
requires_grad(encoder_HG, False) # REMOVE BEFORE TRAINING
real_img = next(LazyLoader.mafl().loader_train_inf)["data"].to(device) \
if i % 5 == 0 else next(LazyLoader.celeba().loader).to(device)
img_content = encoder_HG(real_img)
pred_measures: UniformMeasure2D01 = UniformMeasure2DFactory.from_heatmap(img_content)
sparse_hm = heatmaper.forward(pred_measures.coord * 63).detach()
trainer_gan(i, real_img, pred_measures.detach(), sparse_hm.detach(), apply_g=False)
supervise_trainer()
if i % 4 == 0:
# real_img = next(LazyLoader.mafl().loader_train_inf)["data"].to(device)
trainer_gan(i, real_img, pred_measures.detach(), sparse_hm.detach(), apply_g=True)
content_trainer(real_img)
if i % 100 == 0:
coefs = json.load(open("../parameters/content_loss.json"))
print(i, coefs)
with torch.no_grad():
# pred_measures_test, sparse_hm_test = encoder_HG(test_img)
content_test = encoder_HG(test_img)
pred_measures_test: UniformMeasure2D01 = UniformMeasure2DFactory.from_heatmap(content_test)
heatmaper_256 = ToGaussHeatMap(256, 2.0)
sparse_hm_test = heatmaper.forward(pred_measures_test.coord * 63)
sparse_hm_test_1 = heatmaper_256.forward(pred_measures_test.coord * 255)
latent_test = style_encoder(test_img)
sparce_mask = sparse_hm_test_1.sum(dim=1, keepdim=True)
sparce_mask[sparce_mask < 0.0003] = 0
iwm = imgs_with_mask(test_img, sparce_mask)
send_images_to_tensorboard(writer, iwm, "REAL", i)
fake_img, _ = generator(sparse_hm_test, [sample_z])
iwm = imgs_with_mask(fake_img, pred_measures_test.toImage(256))
send_images_to_tensorboard(writer, iwm, "FAKE", i)
restored = decoder(sparse_hm_test, latent_test)
iwm = imgs_with_mask(restored, pred_measures_test.toImage(256))
send_images_to_tensorboard(writer, iwm, "RESTORED", i)
content_test_256 = nn.Upsample(scale_factor=4)(sparse_hm_test).sum(dim=1, keepdim=True).repeat(1, 3, 1,
1) * \
torch.tensor([1.0, 1.0, 0.0], device=device).view(1, 3, 1, 1)
content_test_256 = (content_test_256 - content_test_256.min()) / content_test_256.max()
send_images_to_tensorboard(writer, content_test_256, "HM", i, normalize=False, range=(0, 1))
if i % 50 == 0 and i >= 0:
test_loss = liuboff(encoder_HG)
# test_loss = nadbka(encoder_HG)
tuner.update(test_loss)
writer.add_scalar("liuboff", test_loss, i)
if i % 10000 == 0 and i > 0:
torch.save(
{
'g': generator.module.state_dict(),
'd': discriminator.module.state_dict(),
'c': encoder_HG.module.state_dict(),
"s": style_encoder.state_dict()
},
f'{Paths.default.models()}/stylegan2_MAFL_{str(i + starting_model_number).zfill(6)}.pt',
)
from joblib import Parallel, delayed
N = 300
dataset = LazyLoader.cardio().dataset_train
padding = 200
prob = np.ones(padding) / padding
data_ids = np.random.permutation(np.arange(0, 700))[0:N]
def LS(k):
return dataset[k]['keypoints'].numpy()
ls = [LS(k) for k in data_ids]
bc_sampler = Uniform2DBarycenterSampler(padding, dir_alpha=1.0)
bc = bc_sampler.mean(ls)
plt.scatter(bc[:, 0], bc[:, 1])
plt.show()
heatmapper = ToGaussHeatMap(256, 4)
hm = heatmapper.forward(torch.from_numpy(bc)[None, ])[0].sum(0)
np.save(f"{Paths.default.models()}/cardio_200_coord.npy", bc)
plt.imshow(hm.numpy())
plt.show()
def train(generator, decoder, discriminator, encoder_HG, style_encoder, device,
starting_model_number):
latent_size = 512
batch_size = 8
sample_z = torch.randn(8, latent_size, device=device)
Celeba.batch_size = batch_size
W300DatasetLoader.batch_size = batch_size
W300DatasetLoader.test_batch_size = 16
test_img = next(LazyLoader.w300().loader_train_inf)["data"][:8].cuda()
model = CondStyleGanModel(generator, StyleGANLoss(discriminator),
(0.001 / 4, 0.0015 / 4))
style_opt = optim.Adam(style_encoder.parameters(),
lr=5e-4,
betas=(0.9, 0.99))
cont_opt = optim.Adam(encoder_HG.parameters(), lr=3e-5, betas=(0.5, 0.97))
g_transforms: albumentations.DualTransform = albumentations.Compose([
ToNumpy(),
NumpyBatch(
albumentations.Compose([
albumentations.ElasticTransform(p=0.7,
alpha=150,
alpha_affine=1,
sigma=10),
albumentations.ShiftScaleRotate(p=0.9, rotate_limit=15),
])),
ToTensor(device),
])
g_transforms_without_norm: albumentations.DualTransform = albumentations.Compose(
[
ToNumpy(),
NumpyBatch(
albumentations.Compose([
albumentations.ElasticTransform(p=0.3,
alpha=150,
alpha_affine=1,
sigma=10),
albumentations.ShiftScaleRotate(p=0.7, rotate_limit=15),
])),
ToTensor(device),
])
R_t = DualTransformRegularizer.__call__(
g_transforms, lambda trans_dict, img: coord_hm_loss(
encoder_HG(trans_dict['image'])["coords"], trans_dict['mask']))
R_s = UnoTransformRegularizer.__call__(
g_transforms, lambda trans_dict, img, ltnt: WR.L1("R_s")
(ltnt, style_encoder(trans_dict['image'])))
barycenter: UniformMeasure2D01 = UniformMeasure2DFactory.load(
f"{Paths.default.models()}/face_barycenter_68").cuda().batch_repeat(
batch_size)
R_b = BarycenterRegularizer.__call__(barycenter, 1.0, 2.0, 4.0)
tuner = GoldTuner([0.37, 2.78, 0.58, 1.43, 3.23],
device=device,
rule_eps=0.001,
radius=0.3,
active=False)
trainer_gan = gan_trainer(model, generator, decoder, encoder_HG,
style_encoder, R_s, style_opt, g_transforms)
content_trainer = content_trainer_with_gan(cont_opt, tuner, encoder_HG,
R_b, R_t, model, generator,
g_transforms, decoder,
style_encoder)
supervise_trainer = content_trainer_supervised(
cont_opt, encoder_HG,
LazyLoader.w300().loader_train_inf)
for i in range(11000):
WR.counter.update(i)
requires_grad(encoder_HG, False)
real_img = next(LazyLoader.celeba().loader).to(device)
encoded = encoder_HG(real_img)
internal_content = encoded["skeleton"].detach()
trainer_gan(i, real_img, internal_content)
# content_trainer(real_img)
train_content(cont_opt, R_b, R_t, real_img, model, encoder_HG, decoder,
generator, style_encoder)
supervise_trainer()
encoder_ema.accumulate(encoder_HG.module, i, 0.98)
if i % 50 == 0 and i > 0:
encoder_ema.write_to(encoder_HG.module)
if i % 100 == 0:
coefs = json.load(open("../parameters/content_loss.json"))
print(i, coefs)
with torch.no_grad():
# pred_measures_test, sparse_hm_test = encoder_HG(test_img)
encoded_test = encoder_HG(test_img)
pred_measures_test: UniformMeasure2D01 = UniformMeasure2D01(
encoded_test["coords"])
heatmaper_256 = ToGaussHeatMap(256, 1.0)
sparse_hm_test_1 = heatmaper_256.forward(
pred_measures_test.coord)
latent_test = style_encoder(test_img)
sparce_mask = sparse_hm_test_1.sum(dim=1, keepdim=True)
sparce_mask[sparce_mask < 0.0003] = 0
iwm = imgs_with_mask(test_img, sparce_mask)
send_images_to_tensorboard(WR.writer, iwm, "REAL", i)
fake_img, _ = generator(encoded_test["skeleton"], [sample_z])
iwm = imgs_with_mask(fake_img, pred_measures_test.toImage(256))
send_images_to_tensorboard(WR.writer, iwm, "FAKE", i)
restored = decoder(encoded_test["skeleton"], latent_test)
iwm = imgs_with_mask(restored, pred_measures_test.toImage(256))
send_images_to_tensorboard(WR.writer, iwm, "RESTORED", i)
content_test_256 = (encoded_test["skeleton"]).repeat(1, 3, 1, 1) * \
torch.tensor([1.0, 1.0, 0.0], device=device).view(1, 3, 1, 1)
content_test_256 = (content_test_256 - content_test_256.min()
) / content_test_256.max()
send_images_to_tensorboard(WR.writer,
content_test_256,
"HM",
i,
normalize=False,
range=(0, 1))
if i % 50 == 0 and i >= 0:
test_loss = liuboff(encoder_HG)
print("liuboff", test_loss)
# test_loss = nadbka(encoder_HG)
tuner.update(test_loss)
WR.writer.add_scalar("liuboff", test_loss, i)
if i % 10000 == 0 and i > 0:
torch.save(
{
'g': generator.module.state_dict(),
'd': discriminator.module.state_dict(),
'c': encoder_HG.module.state_dict(),
"s": style_encoder.state_dict(),
"e": encoder_ema.storage_model.state_dict()
},
f'{Paths.default.models()}/stylegan2_new_{str(i + starting_model_number).zfill(6)}.pt',
)
ToTensor(device),
])
starting_model_number = 280000
weights = torch.load(
f'{Paths.default.models()}/cardio_part_{str(starting_model_number).zfill(6)}.pt',
map_location="cpu"
)
enc_dec = StyleGanAutoEncoder().load_state_dict(weights, style=False).cuda()
discriminator_img = Discriminator(image_size)
discriminator_img.load_state_dict(weights['di'])
discriminator_img = discriminator_img.cuda()
heatmapper = ToGaussHeatMap(256, 4)
hg = HG_heatmap(heatmapper, num_classes=200)
# hg.load_state_dict(weights['gh'])
hg = hg.cuda()
hm_discriminator = Discriminator(image_size, input_nc=1, channel_multiplier=1)
hm_discriminator.load_state_dict(weights["dh"])
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
])),
ToTensor(device),
])
starting_model_number = 280000
weights = torch.load(
f'{Paths.default.models()}/cardio_brule_unsup_{str(starting_model_number).zfill(6)}.pt',
map_location="cpu")
enc_dec = StyleGanAutoEncoder().load_state_dict(weights, style=False).cuda()
discriminator_img = Discriminator(image_size)
discriminator_img.load_state_dict(weights['di'])
discriminator_img = discriminator_img.cuda()
heatmapper = ToGaussHeatMap(256, 4)
hg = HG_heatmap(heatmapper, num_blocks=1, num_classes=200)
hg.load_state_dict(weights['gh'])
hg = hg.cuda()
cont_opt = optim.Adam(hg.parameters(), lr=2e-5, betas=(0, 0.8))
gan_model_tuda = StyleGanModel[HeatmapToImage](enc_dec.generator,
StyleGANLoss(discriminator_img),
(0.001 / 4, 0.0015 / 4))
style_opt = optim.Adam(enc_dec.style_encoder.parameters(), lr=1e-5)
writer = SummaryWriter(
f"{Paths.default.board()}/brule1_cardio_{int(time.time())}")
WR.writer = writer
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
parents=[DatasetParameters(),
RuntimeParameters(),
ModelParameters()])
args = parser.parse_args()
for k in vars(args):
print(f"{k}: {vars(args)[k]}")
device = torch.device(
f"cuda:{args.cuda}" if torch.cuda.is_available() else "cpu")
torch.cuda.set_device(device)
W300DatasetLoader.batch_size = batch_size
heatmapper = ToGaussHeatMap(256, 4)
hg = HG_softmax2020()
hg = hg.cuda()
hg_opt = optim.Adam(hg.parameters(), lr=4e-5)
writer = SummaryWriter(f"{Paths.default.board()}/hm2img{int(time.time())}")
WR.writer = writer
batch = next(LazyLoader.w300().loader_train_inf)
test_img = batch["data"].cuda()
test_landmarks = torch.clamp(batch["meta"]['keypts_normalized'].cuda(), max=1)
test_hm = heatmapper.forward(test_landmarks).sum(1, keepdim=True).detach()
test_noise = mixing_noise(batch_size, 512, 0.9, device)
# mes_loss = MesBceL2Loss(heatmapper, bce_coef=100000, l2_coef=200)
def train(generator, decoder, discriminator, discriminatorHG, encoder_HG, style_encoder, device, starting_model_number):
latent_size = 512
batch_size = 12
sample_z = torch.randn(8, latent_size, device=device)
Celeba.batch_size = batch_size
W300DatasetLoader.batch_size = batch_size
W300DatasetLoader.test_batch_size = 64
encoder_HG_supervised = HG_softmax2020(num_classes=68, heatmap_size=64)
encoder_HG_supervised.load_state_dict(torch.load(f'{Paths.default.models()}/hg2_e29.pt', map_location="cpu"))
encoder_HG_supervised = encoder_HG_supervised.cuda()
requires_grad(encoder_HG_supervised, False)
style_opt = optim.Adam(style_encoder.parameters(), lr=5e-4, betas=(0.9, 0.99))
test_img = next(LazyLoader.celeba().loader)[:8].cuda()
loss_st: StyleGANLoss = StyleGANLoss(discriminator)
model = CondStyleGanModel(generator, loss_st, (0.001, 0.0015))
loss_st2: StyleGANLoss = StyleGANLoss(discriminatorHG)
model2 = CondStyleGanModel(encoder_HG, loss_st2, (2e-5, 0.0005))
heatmaper = ToGaussHeatMap(64, 1.0)
# wide_heatmaper = HeatMapToGaussHeatMap(64, 5.0)
tuda_trainer = gan_tuda_trainer(model, generator)
obratno_trainer = gan_obratno_trainer(model2, encoder_HG)
tuda_obratno_trainer = gan_tuda_obratno_trainer(generator, encoder_HG, decoder, style_encoder, style_opt)
for i in range(100000):
counter.update(i)
# requires_grad(encoder_HG, False) # REMOVE BEFORE TRAINING
real_img = next(LazyLoader.celeba().loader).to(device)
heatmap = encoder_HG_supervised(real_img).detach()
# heatmap = wide_heatmaper.forward(encoder_HG_supervised(real_img)).detach()
# heatmap = encoder_HG_supervised(next(LazyLoader.celeba().loader).to(device)).detach()
# tuda_trainer(real_img, heatmap)
obratno_trainer(real_img, heatmap)
# tuda_obratno_trainer(real_img, heatmap)
# pred_hm = encoder_HG(real_img)
# stariy_hm_loss(pred_hm, heatmap).minimize_step(model2.optimizer.opt_min)
if i % 100 == 0:
with torch.no_grad():
content_test = encoder_HG(test_img)
pred_measures_test: UniformMeasure2D01 = UniformMeasure2DFactory.from_heatmap(content_test)
heatmaper_256 = ToGaussHeatMap(256, 2.0)
sparse_hm_test = heatmaper.forward(pred_measures_test.coord * 63)
sparse_hm_test_1 = heatmaper_256.forward(pred_measures_test.coord * 255)
latent_test = style_encoder(test_img)
sparce_mask = sparse_hm_test_1.sum(dim=1, keepdim=True)
sparce_mask[sparce_mask < 0.0003] = 0
iwm = imgs_with_mask(test_img, sparce_mask)
send_images_to_tensorboard(writer, iwm, "REAL", i)
fake_img, _ = generator(sparse_hm_test, [sample_z])
iwm = imgs_with_mask(fake_img, pred_measures_test.toImage(256))
send_images_to_tensorboard(writer, iwm, "FAKE", i)
restored = decoder(sparse_hm_test, latent_test)
iwm = imgs_with_mask(restored, pred_measures_test.toImage(256))
send_images_to_tensorboard(writer, iwm, "RESTORED", i)
content_test_256 = nn.Upsample(scale_factor=4)(sparse_hm_test).sum(dim=1, keepdim=True).repeat(1, 3, 1, 1) * \
torch.tensor([1.0, 1.0, 0.0], device=device).view(1, 3, 1, 1)
content_test_256 = (content_test_256 - content_test_256.min()) / content_test_256.max()
send_images_to_tensorboard(writer, content_test_256, "HM", i, normalize=False, range=(0, 1))
if i % 50 == 0 and i >= 0:
test_loss = liuboff(encoder_HG)
print("liuboff", test_loss)
# test_loss = nadbka(encoder_HG)
# tuner.update(test_loss)
writer.add_scalar("liuboff", test_loss, i)
if i % 10000 == 0 and i > 0:
torch.save(
{
'g': generator.module.state_dict(),
'd': discriminator.module.state_dict(),
'c': encoder_HG.module.state_dict(),
"s": style_encoder.state_dict(),
"d2": discriminatorHG.state_dict()
},
f'{Paths.default.models()}/cyclegan_{str(i + starting_model_number).zfill(6)}.pt',
)
def __init__(self, gh: nn.Module):
super().__init__()
self.gh = gh
self.heatmaper = ToGaussHeatMap(64, 1.0)
# style_encoder.load_state_dict(weights['s'])
encoder_HG.load_state_dict(weights['c'])
generator = generator.cuda()
# discriminator = discriminator.to(device)
encoder_HG = encoder_HG.cuda()
# style_encoder = style_encoder.cuda()
decoder = CondGenDecode(generator)
# generator = nn.DataParallel(generator, [0, 1, 3])
# discriminator = nn.DataParallel(discriminator, [0, 1, 3])
# encoder_HG = nn.DataParallel(encoder_HG, [0, 1, 3])
# decoder = nn.DataParallel(decoder, [0, 1, 3])
test_img = next(LazyLoader.celeba().loader)[:4].cuda()
heatmaper = ToGaussHeatMap(64, 1.5)
sample_z = torch.randn(4, 512, device=device)
noise = mixing_noise(4, 512, 0.9, device)
LazyLoader.w300_save = None
W300DatasetLoader.test_batch_size = 4
w300_test = next(iter(LazyLoader.w300().test_loader))
w300_test_image = w300_test['data'].to(device)[:4]
w300_test_mask = ProbabilityMeasureFabric(256).from_coord_tensor(
w300_test["meta"]["keypts_normalized"][:4].to(device))
sparse_hm = heatmaper.forward(w300_test_mask.coord * 63).detach()
g_transforms: albumentations.DualTransform = albumentations.Compose([
ToNumpy(),
NumpyBatch(
albumentations.Compose([
def LS(k):
return dataset[k]["meta"]['keypts_normalized'].numpy()
ls = []
images = []
for k in range(N):
dk = dataset[k]
ls.append(dk["meta"]['keypts_normalized'].numpy())
images.append(dk["data"])
# bc_sampler = Uniform2DBarycenterSampler(padding, dir_alpha=1.0)
bc_sampler = ImageBarycenterSampler(padding, dir_alpha=4.0)
heatmapper = ToGaussHeatMap(256, 4)
# [51, 52, 85, 86]
# [1, 44, 34, 94, 10, 72]
# [20, 68, 96, 88]
# [7, 17, 65, 42, 84]
# [71, 72, 73]
# [0, 75, 54, 48, 37]
# [13, 77, 32, 38, 99]
# [9, 40, 36, 66, 67, 78]
# [3, 35, 18, 55, 63, 25, 15]
data_img = []
data_lm = []
def plot_img_with_lm(img: torch.Tensor, lm: torch.Tensor, nrows=4, ncols=4):
def __init__(self, path: str, transform=None):
super().__init__()
self.path = path
self.transform = transform
self.data = make_dataset(path)
self.heatmapper = ToGaussHeatMap(256, 1)
def train(generator, decoder, discriminator, encoder_HG, style_encoder, device,
starting_model_number):
latent_size = 512
batch_size = 24
sample_z = torch.randn(8, latent_size, device=device)
Celeba.batch_size = batch_size
test_img = next(LazyLoader.celeba().loader)[:8].cuda()
loss_st: StyleGANLoss = StyleGANLoss(discriminator)
model = CondStyleGanModel(generator, loss_st, (0.001, 0.0015))
style_opt = optim.Adam(style_encoder.parameters(),
lr=5e-4,
betas=(0.9, 0.99))
cont_opt = optim.Adam(encoder_HG.parameters(), lr=4e-5, betas=(0.5, 0.97))
g_transforms: albumentations.DualTransform = albumentations.Compose([
ToNumpy(),
NumpyBatch(
albumentations.Compose([
ResizeMask(h=256, w=256),
albumentations.ElasticTransform(p=0.7,
alpha=150,
alpha_affine=1,
sigma=10),
albumentations.ShiftScaleRotate(p=0.7, rotate_limit=15),
ResizeMask(h=64, w=64),
NormalizeMask(dim=(0, 1, 2))
])),
ToTensor(device),
])
R_t = DualTransformRegularizer.__call__(
g_transforms, lambda trans_dict, img: hm_svoego_roda_loss(
encoder_HG(trans_dict['image']), trans_dict['mask'], 1000, 0.3))
R_s = UnoTransformRegularizer.__call__(
g_transforms, lambda trans_dict, img, ltnt: L1("R_s")
(ltnt, style_encoder(trans_dict['image'])))
barycenter: UniformMeasure2D01 = UniformMeasure2DFactory.load(
f"{Paths.default.models()}/face_barycenter_68").cuda().batch_repeat(
batch_size)
# plt.imshow(barycenter.toImage(256)[0][0].detach().cpu().numpy())
# plt.show()
R_b = BarycenterRegularizer.__call__(barycenter, 1.0, 2.0, 3.0)
# 4.5, 1.2, 1.12, 1.4, 0.07, 2.2
# 1.27, 3.55, 5.88, 3.83, 2.17, 0.22, 1.72
tuner = GoldTuner([2.2112, 2.3467, 3.8438, 3.2202, 2.0494, 0.0260, 5.8378],
device=device,
rule_eps=0.03,
radius=1,
active=True)
# tuner_verka = GoldTuner([3.0, 1.2, 2.0], device=device, rule_eps=0.05, radius=1, active=True)
best_igor = 100
heatmaper = ToGaussHeatMap(64, 1.5)
trainer_gan = gan_trainer(model, generator, decoder, encoder_HG,
style_encoder, R_s, style_opt, heatmaper,
g_transforms)
content_trainer = content_trainer_with_gan(cont_opt, tuner, heatmaper,
encoder_HG, R_b, R_t, model,
generator)
for i in range(100000):
counter.update(i)
requires_grad(encoder_HG, False) # REMOVE BEFORE TRAINING
real_img = next(LazyLoader.celeba().loader).to(device)
img_content = encoder_HG(real_img).detach()
pred_measures: UniformMeasure2D01 = UniformMeasure2DFactory.from_heatmap(
img_content)
sparce_hm = heatmaper.forward(pred_measures.coord * 63).detach()
trainer_gan(i, real_img, img_content, sparce_hm)
if i % 3 == 0:
real_img = next(LazyLoader.celeba().loader).to(device)
content_trainer(real_img)
if i % 100 == 0:
coefs = json.load(open("../parameters/content_loss_sup.json"))
print(i, coefs)
with torch.no_grad():
content_test = encoder_HG(test_img)
latent_test = style_encoder(test_img)
pred_measures = UniformMeasure2DFactory.from_heatmap(
content_test)
sparce_hm = heatmaper.forward(pred_measures.coord *
63).detach()
iwm = imgs_with_mask(test_img, pred_measures.toImage(256))
send_images_to_tensorboard(writer, iwm, "REAL", i)
fake_img, _ = generator(sparce_hm, [sample_z])
iwm = imgs_with_mask(fake_img, pred_measures.toImage(256))
send_images_to_tensorboard(writer, iwm, "FAKE", i)
restored = decoder(sparce_hm, latent_test)
iwm = imgs_with_mask(restored, pred_measures.toImage(256))
send_images_to_tensorboard(writer, iwm, "RESTORED", i)
content_test_256 = nn.Upsample(
scale_factor=4)(content_test).sum(dim=1, keepdim=True)
content_test_256 = content_test_256 / content_test_256.max()
send_images_to_tensorboard(writer, content_test_256, "HM", i)
if i % 50 == 0 and i > 0:
test_loss = verka(encoder_HG)
tuner.update(test_loss)
writer.add_scalar("verka", test_loss, i)
if i % 10000 == 0 and i > 0:
torch.save(
{
'g': generator.module.state_dict(),
'd': discriminator.module.state_dict(),
'c': encoder_HG.module.state_dict(),
"s": style_encoder.state_dict()
},
f'{Paths.default.models()}/stylegan2_new_{str(i + starting_model_number).zfill(6)}.pt',
)
torch.cuda.set_device(device)
HumanLoader.batch_size = batch_size
starting_model_number = 90000 + 210000
weights = torch.load(
f'{Paths.default.models()}/human_{str(starting_model_number).zfill(6)}.pt',
map_location="cpu")
enc_dec = StyleGanAutoEncoder(
hm_nc=measure_size, image_size=image_size).load_state_dict(weights).cuda()
discriminator_img = Discriminator(image_size)
discriminator_img.load_state_dict(weights['di'])
discriminator_img = discriminator_img.cuda()
heatmapper = ToGaussHeatMap(image_size, 2)
hg = HG_heatmap(heatmapper,
num_classes=measure_size,
image_size=image_size,
num_blocks=4)
hg.load_state_dict(weights['gh'])
hg = hg.cuda()
hm_discriminator = Discriminator(image_size,
input_nc=measure_size,
channel_multiplier=1)
hm_discriminator.load_state_dict(weights["dh"])
hm_discriminator = hm_discriminator.cuda()
gan_model_tuda = StyleGanModel[HeatmapToImage](enc_dec.generator,
StyleGANLoss(discriminator_img),
(0.001 / 4, 0.0015 / 4))
# f'{Paths.default.nn()}/stylegan2_w300_{str(starting_model_number).zfill(6)}.pt',
map_location="cpu")
generator.load_state_dict(weights['g'])
style_encoder.load_state_dict(weights['s'])
encoder_HG.load_state_dict(weights['e'])
encoder_ema.storage_model.load_state_dict(weights['e'])
generator = generator.cuda()
encoder_HG = encoder_HG.cuda()
style_encoder = style_encoder.cuda()
decoder = ConditionalDecoder(generator)
test_img = next(LazyLoader.w300().loader_train_inf)["data"][:8].cuda()
with torch.no_grad():
# pred_measures_test, sparse_hm_test = encoder_HG(test_img)
encoded_test = encoder_HG(test_img)
pred_measures_test: UniformMeasure2D01 = UniformMeasure2D01(
encoded_test["coords"])
heatmaper_256 = ToGaussHeatMap(256, 1.0)
sparse_hm_test_1 = heatmaper_256.forward(pred_measures_test.coord)
latent_test = style_encoder(test_img)
sparce_mask = sparse_hm_test_1.sum(dim=1, keepdim=True)
sparce_mask[sparce_mask < 0.0003] = 0
iwm = imgs_with_mask(test_img, sparce_mask)
send_images_to_tensorboard(WR.writer, iwm, "REAL", i)
args = parser.parse_args()
for k in vars(args):
print(f"{k}: {vars(args)[k]}")
device = torch.device(f"cuda:{args.cuda}" if torch.cuda.is_available() else "cpu")
torch.cuda.set_device(device)
Cardio.batch_size = batch_size
starting_model_number = 264000
weights = torch.load(
f'{Paths.default.models()}/cardio_brule_sup_{str(starting_model_number).zfill(6)}.pt',
map_location="cpu"
)
heatmapper = ToGaussHeatMap(256, 4)
hg = HG_heatmap(heatmapper, num_blocks=1, num_classes=200)
hg.load_state_dict(weights['gh'])
hg = hg.cuda()
cont_opt = optim.Adam(hg.parameters(), lr=2e-5, betas=(0, 0.8))
# gan_model_tuda = StyleGanModel[HeatmapToImage](enc_dec.generator, StyleGANLoss(discriminator_img), (0.001/4, 0.0015/4))
# style_opt = optim.Adam(enc_dec.style_encoder.parameters(), lr=1e-5)
writer = SummaryWriter(f"{Paths.default.board()}/brule1_cardio_{int(time.time())}")
WR.writer = writer
batch = next(iter(LazyLoader.cardio().test_loader))
test_img = batch["image"].cuda()
