update w_to_s converter

w_s_converter.py

@@ -10,24 +10,10 @@
 
 import os
 import re
-import random
-import math
-import time
-import click
-import legacy
-from typing import List, Optional
-
-import cv2
-import clip
-import dnnlib
+from typing import List
+
 import numpy as np
-import torchvision
-from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
 import torch
-from torch import linalg as LA
-import torch.nn.functional as F
-from PIL import Image
-import matplotlib.pyplot as plt
 
 from torch_utils import misc
 from torch_utils import persistence
@@ -86,56 +72,27 @@ def unravel_index(index, shape):
         out.append(index % dim)
         index = index // dim
     return tuple(reversed(out))
+    
 
-
-def num_range(s: str) -> List[int]:
-    '''Accept either a comma separated list of numbers 'a,b,c' or a range 'a-c' and return as a list of ints.'''
-
-    range_re = re.compile(r'^(\d+)-(\d+)$')
-    m = range_re.match(s)
-    if m:
-        return list(range(int(m.group(1)), int(m.group(2))+1))
-    vals = s.split(',')
-    return [int(x) for x in vals]
-
-
-@click.command()
-@click.pass_context
-@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
-@click.option('--seeds', type=num_range, help='List of random seeds')
-@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=1, show_default=True)
-@click.option('--noise-mode', help='Noise mode', type=click.Choice(['const', 'random', 'none']), default='const', show_default=True)
-@click.option('--projected-w', help='Projection result file', type=str, metavar='FILE')
-@click.option('--projected_s', help='Projection result file', type=str, metavar='FILE')
-@click.option('--outdir', help='Where to save the output images', type=str, required=True, metavar='DIR')
-
-def generate_images(
-    ctx: click.Context,
-    network_pkl: str,
-    seeds: Optional[List[int]],
-    truncation_psi: float,
-    noise_mode: str,
+def w_to_s(
+    G,
     outdir: str,
-    class_idx: Optional[int],
-    projected_w: Optional[str],
-    projected_s: Optional[str]
+    projected_w: str,
+    truncation_psi: float = 0.7,
+    noise_mode: str = "const",
 ):
 
-    print('Loading networks from "%s"...' % network_pkl)
     # Use GPU if available
     if torch.cuda.is_available():
         device = torch.device("cuda")
     else:
         device = torch.device("cpu")
 
-    with dnnlib.util.open_url(network_pkl) as f:
-        G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore
-
     os.makedirs(outdir, exist_ok=True)
 
     # Generate images.
     for i in G.parameters():
-      i.requires_grad = True
+        i.requires_grad = True
 
     ws = np.load(projected_w)['w']
     ws = torch.tensor(ws, device=device)
@@ -145,14 +102,12 @@ def generate_images(
         misc.assert_shape(ws, [None, G.synthesis.num_ws, G.synthesis.w_dim])
         ws = ws.to(torch.float32)
 
-
         w_idx = 0
         for res in G.synthesis.block_resolutions:
             block = getattr(G.synthesis, f'b{res}')
             block_ws.append(ws.narrow(1, w_idx, block.num_conv + block.num_torgb))
             w_idx += block.num_conv
 
-
     styles = torch.zeros(1,26,512, device=device)
     styles_idx = 0
     temp_shapes = []
@@ -160,29 +115,24 @@ def generate_images(
         block = getattr(G.synthesis, f'b{res}')
 
         if res == 4:
-          temp_shape = (block.conv1.affine.weight.shape[0], block.conv1.affine.weight.shape[0], block.torgb.affine.weight.shape[0])
-          styles[0,:1,:] = block.conv1.affine(cur_ws[0,:1,:])
-          styles[0,1:2,:] = block.torgb.affine(cur_ws[0,1:2,:])
+            temp_shape = (block.conv1.affine.weight.shape[0], block.conv1.affine.weight.shape[0], block.torgb.affine.weight.shape[0])
+            styles[0,:1,:] = block.conv1.affine(cur_ws[0,:1,:])
+            styles[0,1:2,:] = block.torgb.affine(cur_ws[0,1:2,:])
 
-          block.conv1.affine = torch.nn.Identity()
-          block.torgb.affine = torch.nn.Identity()
-          styles_idx += 2
+            block.conv1.affine = torch.nn.Identity()
+            block.torgb.affine = torch.nn.Identity()
+            styles_idx += 2
         else:
-          temp_shape = (block.conv0.affine.weight.shape[0], block.conv1.affine.weight.shape[0], block.torgb.affine.weight.shape[0])
-          styles[0,styles_idx:styles_idx+1,:temp_shape[0]] = block.conv0.affine(cur_ws[0,:1,:])
-          styles[0,styles_idx+1:styles_idx+2,:temp_shape[1]] = block.conv1.affine(cur_ws[0,1:2,:])
-          styles[0,styles_idx+2:styles_idx+3,:temp_shape[2]] = block.torgb.affine(cur_ws[0,2:3,:])
-
-          block.conv0.affine = torch.nn.Identity()
-          block.conv1.affine = torch.nn.Identity()
-          block.torgb.affine = torch.nn.Identity()
-          styles_idx += 3
+            temp_shape = (block.conv0.affine.weight.shape[0], block.conv1.affine.weight.shape[0], block.torgb.affine.weight.shape[0])
+            styles[0,styles_idx:styles_idx+1,:temp_shape[0]] = block.conv0.affine(cur_ws[0,:1,:])
+            styles[0,styles_idx+1:styles_idx+2,:temp_shape[1]] = block.conv1.affine(cur_ws[0,1:2,:])
+            styles[0,styles_idx+2:styles_idx+3,:temp_shape[2]] = block.torgb.affine(cur_ws[0,2:3,:])
+
+            block.conv0.affine = torch.nn.Identity()
+            block.conv1.affine = torch.nn.Identity()
+            block.torgb.affine = torch.nn.Identity()
+            styles_idx += 3
         temp_shapes.append(temp_shape)
 
-
     styles = styles.detach()
     np.savez(f'{outdir}/input.npz', s=styles.cpu().numpy())
-
-        
-if __name__ == "__main__":
-    generate_images()