File size: 801 Bytes
c1c640d
 
 
 
 
329d335
c1c640d
 
 
329d335
fc626a1
caf14e5
329d335
c1c640d
 
329d335
07be559
c1c640d
329d335
c1c640d
 
 
 
 
 
 
 
 
 
8f4be28
c1c640d
caf14e5
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import os
import math
import numpy as np
import onnxruntime as ort
from PIL import Image

LATENT_FEATURES = 512
RESOLUTION = 256
LAST_INDEX = math.log2(RESOLUTION) - 2

MODEL_PATH = os.path.join("model", "batik_stylegan.onnx")
model = ort.InferenceSession(MODEL_PATH)

alpha = np.array([1.0], dtype=np.float32)
steps = np.array([LAST_INDEX], dtype=np.int64)

def generate_stylegan():
    z = np.random.randn(1, LATENT_FEATURES).astype(np.float32)

    output = model.run(None, {
        'z': z,
        'alpha': alpha,
        'steps': steps
    })[0]

    image = output.squeeze(0)
    image = (image * 0.5 + 0.5) * 255
    image = image.astype(np.uint8)
    image = np.transpose(image, (1, 2, 0))
    pil_img = Image.fromarray(image, 'RGB')
    
    return pil_img.resize((512, 512), Image.LANCZOS)