Main

data.sh

@@ -0,0 +1,62 @@
+#----------------------------------SETTINGS------------------------------------#
+
+PROJECT_ID='for-nikita' #ID проекта
+BUCKET_NAME='for-nikita' #Имя Bucket
+LOAD='DATASETS' #DATASETS (загрузка датасета из GS) \
+                #или IMAGES (создание датасета и его загрузка на GS).
+PATH_TO_DATASETS='style-GAN2/datasets/256-gray' #Путь до датасета на GS
+DATASET_FOLDER="img-gray-256" # Каталог для сохранения датасета на GS
+RESOLUTION=256 #Разрешение img 
+MODE="gray" #gray or rgb
+COUNT=5000 #Кол-во изображений для тестового датасета. 
+           #Не передавать аргумент для использования всех изображений
+
+#------------------------------------------------------------------------------#
+
+gcloud auth login
+gcloud config set project $PROJECT_ID
+
+cd ./stylegan2/
+
+if [ $LOAD = "IMAGES" ]; then
+	gsutil cp gs://$BUCKET_NAME/images.zip ./images.zip;
+  unzip -q ./images.zip -d ./;
+fi
+
+echo $0
+
+full_path=$(realpath $0)
+dir_path=$(dirname $full_path)
+echo $dir_path
+
+PATH_TO_IMG="${dir_path}/images/"
+
+if [ $LOAD = "IMAGES" ]; then
+  echo Обработка изображений
+  mkdir -p ./images/custom/;
+  python3 preprocessing.py --resolution $RESOLUTION\
+                           --mode $MODE\
+                           --count $COUNT\
+                           --path $PATH_TO_IMG
+fi
+
+if [ $LOAD = "IMAGES" ]; then
+  echo Создание датасета
+  python3 dataset_tool.py create_from_images ./datasets/custom ./images/custom
+  rm -rf ./images
+  echo Копирование датасета на GS 
+  gsutil -m cp -r ./datasets/custom/*.tfrecords gs://for-nikita/style-GAN2/datasets/$DATASET_FOLDER/
+fi
+
+if [ $LOAD = "DATASETS" ]; then
+  echo Копирование датасета из GS 
+	mkdir -p ./datasets/custom/;
+  gsutil -m cp -r gs://${BUCKET_NAME}/${PATH_TO_DATASETS}/*.tfrecords ./datasets/custom;
+fi
+
+
+
+
+
+
+

preprocessing.py

@@ -0,0 +1,64 @@
+import cv2
+from PIL import Image
+from tqdm import tqdm
+import numpy as np
+import os
+import random
+import argparse
+import matplotlib.pyplot as plt
+
+parser = argparse.ArgumentParser(description='Train model')
+parser.add_argument("--resolution", type=int)
+parser.add_argument("--mode", type=str)
+parser.add_argument("--count", type=int)
+parser.add_argument("--path", type=str)
+args = parser.parse_args()
+
+
+def rgba2rgb(rgba, background=(255, 255, 255)):
+    row, col, ch = rgba.shape
+    if ch == 3:
+        return rgba
+    assert ch == 4, 'RGBA image has 4 channels.'
+
+    rgb = np.zeros((row, col, 3), dtype='float32')
+    r, g, b, a = rgba[:, :, 0], rgba[:, :, 1], rgba[:, :, 2], rgba[:, :, 3]
+
+    R, G, B = background
+
+    rgb[:, :, 0] = r * a + (1.0 - a) * R
+    rgb[:, :, 1] = g * a + (1.0 - a) * G
+    rgb[:, :, 2] = b * a + (1.0 - a) * B
+
+    return np.asarray(rgb, dtype='uint8')
+
+
+def rgb2gray(rgb):
+    return np.dot(rgb[..., :3], [0.2989, 0.5870, 0.1140])
+
+
+def preprocessing(resolution=256, mode="gray", count=None, path=None):
+    files = os.listdir(path)
+    if not count:
+        count = len(files)
+    for i in tqdm(range(0, count)):
+        if files[i][-3:] != 'png':
+            continue
+        try:
+            img = plt.imread(os.path.join(path, files[i]))
+            img = rgba2rgb(img)
+            if mode == "gray":
+                img = rgb2gray(img)
+            img = cv2.resize(img, (resolution, resolution))
+            cv2.imwrite("{}custom/{}.jpeg".format(path,files[i][:-4]), img)
+        except:
+            pass
+
+
+if __name__ == "__main__":
+    preprocessing(
+        resolution=args.resolution,
+        mode=args.mode,
+        count=args.count,
+        path=args.path
+    )

run_generator.py

@@ -33,7 +33,12 @@ def generate_images(network_pkl, seeds, truncation_psi):
         z = rnd.randn(1, *Gs.input_shape[1:]) # [minibatch, component]
         tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]
         images = Gs.run(z, None, **Gs_kwargs) # [minibatch, height, width, channel]
-        PIL.Image.fromarray(images[0], 'RGB').save(dnnlib.make_run_dir_path('seed%04d.png' % seed))
+        if images[0].shape[2] == 1:
+          H, W = images[0].shape[0], images[0].shape[1]
+          img = images[0].reshape(H, W)
+        else:
+          img = images[0]
+        PIL.Image.fromarray(img).save(dnnlib.make_run_dir_path('seed%04d.png' % seed))
 
 #----------------------------------------------------------------------------
 
@@ -68,7 +73,12 @@ def style_mixing_example(network_pkl, row_seeds, col_seeds, truncation_psi, col_
 
     print('Saving images...')
     for (row_seed, col_seed), image in image_dict.items():
-        PIL.Image.fromarray(image, 'RGB').save(dnnlib.make_run_dir_path('%d-%d.png' % (row_seed, col_seed)))
+        if image.shape == 1:
+          H, W = image.shape[0], image.shape[1]
+          img = image.reshape(H, W)
+        else:
+          img = image
+        PIL.Image.fromarray(image).save(dnnlib.make_run_dir_path('%d-%d.png' % (row_seed, col_seed)))
 
     print('Saving image grid...')
     _N, _C, H, W = Gs.output_shape