implemented generating 16 images instead of 1

image_generator.py

@@ -18,6 +18,7 @@ import PIL.Image
 import torch
 from networks_fastgan import MyGenerator
 import random
+import cv2
 #----------------------------------------------------------------------------
 
 def parse_range(s: Union[str, List]) -> List[int]:
@@ -74,21 +75,8 @@ def generate_images(
     outdir = "out",
     translate = "0,0",
     rotate = 0,
+    number_of_images = 16
 ):
-    """Generate images using pretrained network pickle.
-
-    Examples:
-
-    \b
-    # Generate an image using pre-trained AFHQv2 model ("Ours" in Figure 1, left).
-    python gen_images.py --outdir=out --trunc=1 --seeds=2 \\
-        --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl
-
-    \b
-    # Generate uncurated images with truncation using the MetFaces-U dataset
-    python gen_images.py --outdir=out --trunc=0.7 --seeds=600-605 \\
-        --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-metfacesu-1024x1024.pkl
-    """
     model_owner = "huggan"
     #inputs = gr.inputs.Radio(["Abstract Expressionism", "Impressionism", "Cubism", "Minimalism", "Pop Art", "Color Field", "Hana Hanak houses"])
     model_path_dict = {
@@ -106,13 +94,13 @@ def generate_images(
     
     model_path = model_owner + "/" + model_path_dict[model_path]
     print(model_path)
-    seeds = parse_range(seeds)
     print(seeds)
-    seeds=[random.randint(1,200)]
+    seeds=random.randint(1,230)
+    seeds =f"{seeds}-{seeds+number_of_images-1}"
+    seeds = parse_range(seeds)
     device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
     G = MyGenerator.from_pretrained(model_path)
     os.makedirs(outdir, exist_ok=True)
-
     # Labels.
     label = torch.zeros([1, G.c_dim], device=device)
     """
@@ -125,7 +113,11 @@ def generate_images(
             print ('warn: --class=lbl ignored when running on an unconditional network')
     """
 
+    print(f"z dimenzija mi je: {G.z_dim}")
     # Generate images.
+
+    #imgs_row = np.array()
+    #imgs_complete = np.array()
     for seed_idx, seed in enumerate(seeds):
         print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx, len(seeds)))
         z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device).float()
@@ -136,12 +128,22 @@ def generate_images(
             m = make_transform(translate, rotate)
             m = np.linalg.inv(m)
             G.synthesis.input.transform.copy_(torch.from_numpy(m))
-
+        
         img = G(z, label, truncation_psi=truncation_psi, noise_mode=noise_mode)
         img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
-        img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
-        #PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(f'{outdir}/seed{seed:04d}.png')
-    return img
+        print(seed_idx)
+        #first image
+        if seed_idx == 0:
+            imgs_row = img[0].cpu().numpy()
+        else:
+            imgs_row = np.hstack((imgs_row, img[0].cpu().numpy()))
+        #img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
+        PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(f'{outdir}/seed{seed:04d}.png')
+    #napravi vsplit i toe to ka
+    imgs_complete = np.vstack(np.hsplit(imgs_row, 4))
+    #cv2.imshow("lalaxd", imgs_complete)
+    #cv2.waitKey()
+    return PIL.Image.fromarray(imgs_complete, 'RGB')
 
 
 #----------------------------------------------------------------------------