Mirrors_Framework/FSRCNN-TensorFlow
1
0
Fork 0
mirror of https://github.com/igv/FSRCNN-TensorFlow.git synced 2026-08-18 09:07:31 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

gen.py: generate deconvolution passes

Browse Source 
Usable now, except:
    - No support for shrinking/expanding layers
    - Generated shaders for scaling factors other than 2x are not compatible with mpv, they have to be modified manually
This commit is contained in:
igv
2017-10-07 12:52:05 +03:00
parent 9e421c1d33
commit e786c07a4b
2 changed files with 97 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gen.py
+97 -16
View File
@@ -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
sort.py
-42
View File
@@ -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()