mirror of
https://github.com/igv/FSRCNN-TensorFlow.git
synced 2026-08-18 17:17:29 +08:00
Half working script for generating a shader
Doesn't support shrinking/expanding layers and doesn't generate a deconvolution passes yet
This commit is contained in:
@@ -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()
|
||||
Reference in New Issue
Block a user