def main():
# create spout object
spout = Spout(silent=True, n_rec=3, n_send=3)
# create receiver
spout.createReceiver('input1', id=0)
spout.createReceiver('input2', id=1)
spout.createReceiver('input3', id=2)
# create sender
spout.createSender('output1', id=0)
spout.createSender('output2', id=1)
spout.createSender('output3', id=2)
while True:
# check on close window
spout.check()
# receive data
data = spout.receive(id=0)
data1 = spout.receive(id=1)
data2 = spout.receive(id=2)
# send data
if random.random() > .9:
spout.send(data1)
spout.send(data, 1)
else:
spout.send(data)
spout.send(data1, 1)
spout.send(data2, id=2)
def main():
# create spout object
spout = Spout(silent=False)
# create receiver
spout.createReceiver('input')
# create sender
spout.createSender('output')
while True:
# check on close window
spout.check()
# receive data
data = spout.receive()
# send data
spout.send(data)
def generate_images(network_pkl):
global z
tflib.init_tf()
print('Loading networks from "%s"...' % network_pkl)
with dnnlib.util.open_url(network_pkl) as fp:
_G, _D, Gs = pickle.load(fp)
Gs_kwargs = {
'output_transform': dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True),
'randomize_noise': False
}
noise_vars = [var for name, var in Gs.components.synthesis.vars.items(
) if name.startswith('noise')]
label = np.zeros([1] + Gs.input_shapes[1][1:])
# create spout object
spout = Spout(silent=False, width=512, height=512)
# create sender
spout.createSender('output')
while True:
# check on close window
spout.check()
server.handle_request()
Gs_kwargs['truncation_psi'] = psi
# GENERATION:
seed = 74 # random.randrange(9999)
rnd = np.random.RandomState(seed)
# z = rnd.randn(1, *Gs.input_shape[1:]) # [minibatch, component]
noise_rnd = np.random.RandomState(1) # fix noise
tflib.set_vars({var: noise_rnd.randn(*var.shape.as_list())
for var in noise_vars}) # [height, width]
image = Gs.run(z, label, **Gs_kwargs)
# send data
spout.send(image[0])
import time
import dlib # requires cmake
import numpy as np
from align_face import align_face
from Library.Spout import Spout
spout = Spout(silent = False, width = 1044, height = 1088)
spout.createReceiver('input')
spout.createSender('output')
# python projector.py --target=out/seed0002.png --project-in-wplus --save-video --num-steps=1000 --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl
while True :
# check on close window
spout.check()
# receive data
data = spout.receive()
print(data.shape)
data = align_face(data)
print(data.shape)
spout.send(data)
def run_projection(
ctx: click.Context,
network_pkl: str,
initial_learning_rate: float,
w_avg_samples: int = 10000,
initial_noise_factor: float = 0.05,
regularize_noise_weight: float = 1e5,
):
"""
Project given image to the latent space of pretrained network pickle.
Adapted from stylegan3-fun/projector.py
"""
torch.manual_seed(42)
# Load networks.
print('Loading networks from "%s"...' % network_pkl)
device = torch.device('cuda')
with dnnlib.util.open_url(network_pkl) as fp:
G = legacy.load_network_pkl(fp)['G_ema'].requires_grad_(False).to(
device)
# Open Spout stream for target images
spout = Spout(silent=False, width=1044,
height=1088) # TODO set W and H to 720p ?
spout.createReceiver('input')
spout.createSender('output')
# Stabilize the latent space to make things easier (for StyleGAN3's config t and r models)
gen_utils.anchor_latent_space(G)
# == Adapted from project() in stylegan3-fun/projector.py == #
G = copy.deepcopy(G).eval().requires_grad_(False).to(device)
# Compute w stats.
z_samples = np.random.RandomState(123).randn(w_avg_samples, G.z_dim)
w_samples = G.mapping(torch.from_numpy(z_samples).to(device),
None) # [N, L, C]
print('Projecting in W+ latent space...')
w_avg = torch.mean(w_samples, dim=0, keepdim=True) # [1, L, C]
w_std = (torch.sum((w_samples - w_avg)**2) / w_avg_samples)**0.5
# Setup noise inputs (only for StyleGAN2 models)
noise_buffs = {
name: buf
for (name, buf) in G.synthesis.named_buffers() if 'noise_const' in name
}
w_noise_scale = w_std * initial_noise_factor # noise scale is constant
lr = initial_learning_rate # learning rate is constant
# Load the VGG16 feature detector.
url = 'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/metrics/vgg16.pkl'
vgg16 = metric_utils.get_feature_detector(url, device=device)
w_opt = w_avg.clone().detach().requires_grad_(True)
optimizer = torch.optim.Adam([w_opt] + list(noise_buffs.values()),
betas=(0.9, 0.999),
lr=initial_learning_rate)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# Init noise.
for buf in noise_buffs.values():
buf[:] = torch.randn_like(buf)
buf.requires_grad = True
# == End setup from project() == #
# Project continuously
while True:
# check on close window
spout.check()
# receive data
data = spout.receive()
# data = align_face(data, G.img_resolution) ## TOO SLOW !!!
#print(data.shape)
# Features for target image. Reshape to 256x256 if it's larger to use with VGG16
# target = np.array(target, dtype=np.uint8)
target = torch.tensor(data.transpose([2, 0, 1]), device=device)
target = target.unsqueeze(0).to(device).to(torch.float32)
if target.shape[2] > 256:
target = F.interpolate(target, size=(256, 256), mode='area')
target_features = vgg16(target, resize_images=False, return_lpips=True)
# Synth images from opt_w.
w_noise = torch.randn_like(w_opt) * w_noise_scale
ws = w_opt + w_noise
synth_images = G.synthesis(ws, noise_mode='const')
# Downsample image to 256x256 if it's larger than that. VGG was built for 224x224 images.
synth_images = (synth_images + 1) * (255 / 2)
if synth_images.shape[2] > 256:
synth_images = F.interpolate(synth_images,
size=(256, 256),
mode='area')
# Features for synth images.
synth_features = vgg16(synth_images,
resize_images=False,
return_lpips=True)
dist = (target_features - synth_features).square().sum()
# Noise regularization.
reg_loss = 0.0
for v in noise_buffs.values():
noise = v[None, None, :, :] # must be [1,1,H,W] for F.avg_pool2d()
while True:
reg_loss += (noise *
torch.roll(noise, shifts=1, dims=3)).mean()**2
reg_loss += (noise *
torch.roll(noise, shifts=1, dims=2)).mean()**2
if noise.shape[2] <= 8:
break
noise = F.avg_pool2d(noise, kernel_size=2)
loss = dist + reg_loss * regularize_noise_weight
# Print in the same line (avoid cluttering the commandline)
message = f'dist {dist:.7e} | loss {loss.item():.7e}'
print(message, end='\r')
# Step
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
# Normalize noise.
with torch.no_grad():
for buf in noise_buffs.values():
buf -= buf.mean()
buf *= buf.square().mean().rsqrt()
# Produce image
data = gen_utils.w_to_img(G,
dlatents=w_opt.detach()[0],
noise_mode='const')[0]
spout.send(data)
def main():
# create spout object
spout = Spout(silent=False)
# create receiver
spout.createReceiver('vvvvideo')
# create sender
spout.createSender('output')
opt = TestOptions().parse() # get test options
# hard-code some parameters for test
opt.num_threads = 0 # test code only supports num_threads = 0
opt.batch_size = 1 # test code only supports batch_size = 1
opt.serial_batches = True # disable data shuffling; comment this line if results on randomly chosen images are needed.
opt.no_flip = True # no flip; comment this line if results on flipped images are needed.
opt.display_id = -1 # no visdom display; the test code saves the results to a HTML file.
model = create_model(opt) # create a model given opt.model and other options
model.setup(opt) # regular setup: load and print networks; create schedulers
dataset = create_dataset(opt) # create a dataset given opt.dataset_mode and other options
pix2pix_image = np.zeros((256, 256, 3), np.uint8)
start_time = time.time()
seconds_to_wait = 1
frame_counter = 0
while True:
# check on close window
spout.check()
# receive data
frame = spout.receive()
if len(frame) > 10:
#pil_image = Image.frombytes('RGBA', (model_image_size, model_image_size), frame, 'raw').convert('RGB')
pil_image = Image.fromarray(frame)
# learntest
# model.update_learning_rate()
for i, data in enumerate(dataset):
if i >= opt.num_test: # only apply our model to opt.num_test images.
break
model.set_input(data, pil_image) # unpack data from data loader
# learntest
# model.optimize_parameters()
model.test() # run inference
visuals = model.get_current_visuals() # get image results
for label, im_data in visuals.items():
im = util.tensor2im(im_data)
open_cv_image = numpy.array(im)
#open_cv_image = open_cv_image[:, :, ::-1].copy()
pix2pix_image = open_cv_image
# dim = (256,256)
# blank_image = cv2.resize(blank_image, dim, interpolation = cv2.INTER_AREA)
# send data
spout.send(pix2pix_image)
#cv2.imshow("image", pix2pix_image)
#cv2.waitKey(1)
counter_and_time = print_fps(frame_counter, start_time, seconds_to_wait)
frame_counter = counter_and_time[0]
start_time = counter_and_time[1]