From 9e421c1d33af75315344b2fcc8871918ba568aaf Mon Sep 17 00:00:00 2001
From: igv <igv@users.noreply.github.com>
Date: Fri, 6 Oct 2017 15:44:23 +0300
Subject: [PATCH] Half working script for generating a shader

Doesn't support shrinking/expanding layers and doesn't generate a deconvolution passes yet
---
 gen.py | 104 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 104 insertions(+)
 create mode 100644 gen.py

diff --git a/gen.py b/gen.py
new file mode 100644
index 0000000..c023e33
--- /dev/null
+++ b/gen.py
@@ -0,0 +1,104 @@
+import sys
+from itertools import islice
+
+radius = 1
+
+def get_line_number(phrase, file_name):
+    with open(file_name) as f:
+        for i, line in enumerate(f, 1):
+            if phrase in line:
+                return i
+
+def read_weights(file_name, ln, size=1):
+    content = []
+    with open(file_name) as f:
+        for line in islice(f, ln, ln + size):
+            if line.find('[') != -1:
+                line = line[line.index('[') + 1:]
+            if line.find(']') != -1:
+                line = line[:line.rindex(']')]
+            content.append(line)
+
+    return [x.strip() for x in content]
+
+def header1(file, n, d):
+    file.write('//!HOOK LUMA\n')
+    file.write('//!DESC feature map {}\n'.format((n//4)%(d//4) + 1))
+    file.write('//!BIND LUMA\n')
+    file.write('//!SAVE MODEL{}\n'.format((n//4)%(d//4) + 1))
+    file.write('//!COMPONENTS 4\n')
+
+def header2(file, w, n, s):
+    file.write('//!HOOK LUMA\n')
+    file.write('//!DESC mapping {}_{}\n'.format(w+1, (n//4)%(s//4) + 1))
+    for i in range(s//4):
+        if (w+1) % 2 == 1:
+            file.write('//!BIND MODEL{}\n'.format(i+1))
+        else:
+            file.write('//!BIND MODEL{}{}\n'.format(2, i+1))
+    if (w+1) % 2 == 1:
+        file.write('//!SAVE MODEL{}{}\n'.format(2, (n//4)%(s//4) + 1))
+    else:
+        file.write('//!SAVE MODEL{}\n'.format((n//4)%(s//4) + 1))
+    file.write('//!COMPONENTS 4\n')
+
+def main():
+  if len(sys.argv) == 2:
+    fname=sys.argv[1]
+    d, s, m = [int(i) for i in fname[7:fname.index('.')].split("_")]
+    if s == 0:
+        s = d
+    dst = fname.replace("weights", "FSRCNN_").replace("txt", "glsl")
+    with open(dst, 'w') as file:
+
+        # Feature layer
+        feature_radius = 2
+        ln = get_line_number("w1", fname)
+        weights = read_weights(fname, ln, (feature_radius*2+1)**2)
+        ln = get_line_number("b1", fname)
+        biases = read_weights(fname, ln, 0)
+        ln = get_line_number("alpha1", fname)
+        alphas = read_weights(fname, ln, 0)
+        for n in range(0, d, 4):
+            header1(file, n, d)
+            file.write('vec4 hook()\n')
+            file.write('{\n')
+            file.write('vec4 res = vec4({});\n'.format(",".join(biases[0].strip(",").split(",")[n:n+4])))
+            p = 0
+            for l in range(0, len(weights)):
+                y, x = p%(feature_radius*2+1)-feature_radius, p//(feature_radius*2+1)-feature_radius
+                p += 1
+                file.write('res += vec4({}) * float(LUMA_texOff(vec2({},{})));\n'.format(",".join(weights[l].strip(",").split(",")[n:n+4]), x, y))
+            file.write('res = mix(res, vec4({}) * res, lessThan(res, vec4(0.0)));\n'.format(",".join(alphas[0].strip(",").split(",")[n:n+4])))
+            file.write('return res;\n')
+            file.write('}\n\n')
+
+        # Mapping layers
+        for w in range(m):
+            ln = get_line_number("w{}".format(w + 3), fname)
+            weights = read_weights(fname, ln, s*9)
+            ln = get_line_number("b{}".format(w + 3), fname)
+            biases = read_weights(fname, ln, 0)
+            ln = get_line_number("alpha{}".format(w + 3), fname)
+            alphas = read_weights(fname, ln, 0)
+            for n in range(0, s, 4):
+                header2(file, w, n, s)
+                file.write('vec4 hook()\n')
+                file.write('{\n')
+                file.write('vec4 res = vec4({});\n'.format(",".join(biases[0].strip(",").split(",")[n:n+4])))
+                p = 0
+                for l in range(0, len(weights), 4):
+                    if l % s == 0:
+                        y, x = p%3-1, p//3-1
+                        p += 1
+                    file.write('res += mat4({},{},{},{}) * vec4(MODEL{}_texOff(vec2({},{})));\n'.format(",".join(weights[l].strip(",").split(",")[n:n+4]), ",".join(weights[l+1].strip(",").split(",")[n:n+4]), ",".join(weights[l+2].strip(",").split(",")[n:n+4]), ",".join(weights[l+3].strip(",").split(",")[n:n+4]), (l//4)%(s//4) + 1 + (20 if (w+1) % 2 == 0 else 0), x, y))
+                file.write('res = mix(res, vec4({}) * res, lessThan(res, vec4(0.0)));\n'.format(",".join(alphas[0].strip(",").split(",")[n:n+4])))
+                file.write('return res;\n')
+                file.write('}\n\n')
+
+  else:
+    print("Missing argument: You must specify a file name")
+    return
+
+if __name__ == '__main__':
+  main()