diff --git a/gen.py b/gen.py index c023e33..abcd1b1 100644 --- a/gen.py +++ b/gen.py @@ -1,8 +1,11 @@ import sys from itertools import islice +scale = 2 radius = 1 +dsize = radius * scale * 2 + 1 + def get_line_number(phrase, file_name): with open(file_name) as f: for i, line in enumerate(f, 1): @@ -21,34 +24,48 @@ def read_weights(file_name, ln, size=1): return [x.strip() for x in content] -def header1(file, n, d): +def base_header(file): file.write('//!HOOK LUMA\n') + file.write('//!WHEN OUTPUT.w LUMA.w / {0}.400 > OUTPUT.h LUMA.h / {0}.400 > *\n'.format(scale - 1)) + +def header1(file, n, d): + base_header(file) 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') + base_header(file) 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('//!BIND MODEL{}\n'.format(i+1 + (0 if w % 2 == 0 else 20))) + file.write('//!SAVE MODEL{}\n'.format((n//4)%(s//4) + 1 + (20 if w % 2 == 0 else 0))) file.write('//!COMPONENTS 4\n') +def header3(file, m, n, d): + base_header(file) + file.write('//!DESC sub-pixel convolution {}\n'.format((n//4)%(d//4) + 1)) + file.write('//!BIND MODEL{}\n'.format((n//4)%(d//4) + 1 + (20 if m % 2 == 1 else 0))) + file.write('//!SAVE MODEL{}\n'.format((n//4)%(d//4) + 1 + (20 if m % 2 == 1 else 0))) + file.write('//!COMPONENTS 4\n') + +def header4(file, m, d): + base_header(file) + file.write('//!WIDTH LUMA.w {} *\n'.format(scale)) + file.write('//!HEIGHT LUMA.h {} *\n'.format(scale)) + file.write('//!DESC aggregation\n') + for i in range(d//4): + file.write('//!BIND MODEL{}\n'.format(i+1 + (20 if m % 2 == 1 else 0))) + file.write('//!OFFSET -{}.0 -{}.0\n'.format(scale//2, scale//2)) + 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") + dst = fname.replace("_", "-").replace("weights", "FSRCNNX_x{}_".format(scale)).replace("txt", "glsl") with open(dst, 'w') as file: # Feature layer @@ -56,9 +73,9 @@ def main(): 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) + biases = read_weights(fname, ln) ln = get_line_number("alpha1", fname) - alphas = read_weights(fname, ln, 0) + alphas = read_weights(fname, ln) for n in range(0, d, 4): header1(file, n, d) file.write('vec4 hook()\n') @@ -78,9 +95,9 @@ def main(): 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) + biases = read_weights(fname, ln) ln = get_line_number("alpha{}".format(w + 3), fname) - alphas = read_weights(fname, ln, 0) + alphas = read_weights(fname, ln) for n in range(0, s, 4): header2(file, w, n, s) file.write('vec4 hook()\n') @@ -91,11 +108,75 @@ def main(): 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 += 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 % 2 == 1 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') + # Sub-pixel convolution + ln = get_line_number("w{}".format(m + 4), fname) + weights = read_weights(fname, ln, dsize**2) + + x=list(reversed(range(scale))) + if dsize % 2 == 1: + x=x[-1:]+x[:-1] + xy = [] + for i in x: + for j in x: + xy.append([j, i]) + + id = [] + for i in range(0, len(xy)): + xi, yi = xy[i] + for y in range(yi, dsize, scale): + for x in range(xi, dsize, scale): + id.append(y + x * dsize) + + weights = list(reversed(weights)) + sort = [weights[id[l]].strip(",") for l in range(0, len(id))] + for n in range(0, d, 4): + header3(file, m, n, d) + file.write('vec4 hook()\n') + file.write('{\n') + file.write('vec4 res = vec4(0);\n') + total = 0 + for i in range(scale): + for j in range(scale): + file.write('res[{}] +=\n'.format(i * scale + j)) + s2 = radius*2+1 if i == 0 and dsize % 2 == 1 else radius*2 + for yi, y in enumerate(range(-radius + (0 if i == 0 and dsize % 2 == 1 else 1), radius + 1)): + s1 = radius*2+1 if j == 0 and dsize % 2 == 1 else radius*2 + for xi, x in enumerate(range(-radius + (0 if j == 0 and dsize % 2 == 1 else 1), radius + 1)): + l = yi * s1 + xi + file.write('dot(vec4({}), MODEL{}_texOff(vec2({},{}))){}\n'.format(",".join(sort[l+total].strip(",").split(",")[n:n+4]), + (n//4)%(d//4) + 1 + (20 if m % 2 == 1 else 0), x, y, ';' if l == s1 * s2 - 1 else '+')) + total = total + l + 1 + file.write('return res;\n') + file.write('}\n\n') + + # Aggregation + ln = get_line_number("b{}".format(m + 4), fname) + biases = read_weights(fname, ln) + header4(file, m, d) + file.write('vec4 hook()\n') + file.write('{\n') + file.write('float res = {};\n'.format(biases[0])) + v = 1 + (20 if m % 2 == 1 else 0) + file.write('vec2 fcoord = fract(MODEL{}_pos * MODEL{}_size);\n'.format(v, v)) + file.write('vec2 base = MODEL{}_pos + (vec2(0.5) - fcoord) * MODEL{}_pt;\n'.format(v, v)) + file.write('ivec2 index = ivec2(fcoord * vec2({}));\n'.format(scale)) + file.write('res += (MODEL{}_tex(base)'.format(v)) + for i in range(d//4-1): + file.write('+MODEL{}_tex(base)'.format(i + 1 + v)) + file.write(')[index.y * {} + index.x];\n'.format(scale)) + file.write('return vec4(res, 0, 0, 1);\n') + file.write('}\n') + else: print("Missing argument: You must specify a file name") return diff --git a/sort.py b/sort.py deleted file mode 100644 index 329666d..0000000 --- a/sort.py +++ /dev/null @@ -1,42 +0,0 @@ -import sys - -def main(): - scale = 2 - radius = 2 - size = radius * scale * 2 + 1 - d = 64 #size of the feature layer - - if len(sys.argv) == 2: - fname=sys.argv[1] - with open(fname) as f: - content = f.readlines() - content = [x.strip() for x in content] - - x=list(reversed(range(scale))) - x=x[-1:]+x[:-1] - xy = [] - for i in x: - for j in x: - xy.append([j, i]) - - m = [] - for i in range(0, len(xy)): - xi, yi = xy[i] - for y in range(yi, size, scale): - for x in range(xi, size, scale): - m.append(y + x * size) - #print(m) - content = list(reversed(content)) - sort = [content[m[l]].strip(",") for l in range(0, len(m))] - with open('sorted.txt', 'w') as file: - for i in range(0, d, 4): - for l in range(0, len(sort)): - m=sort[l].strip(",").split(",") - file.write(",".join(m[i:i+4])+"\n") - file.write("\n") - else: - print("Missing argument: You must specify a file name that contains the deconvolution weights") - return - -if __name__ == '__main__': - main()