benchmark.lua 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802
  1. require 'pl'
  2. local __FILE__ = (function() return string.gsub(debug.getinfo(2, 'S').source, "^@", "") end)()
  3. package.path = path.join(path.dirname(__FILE__), "..", "lib", "?.lua;") .. package.path
  4. require 'xlua'
  5. require 'w2nn'
  6. local iproc = require 'iproc'
  7. local reconstruct = require 'reconstruct'
  8. local image_loader = require 'image_loader'
  9. local gm = require 'graphicsmagick'
  10. local cjson = require 'cjson'
  11. local cmd = torch.CmdLine()
  12. cmd:text()
  13. cmd:text("waifu2x-benchmark")
  14. cmd:text("Options:")
  15. cmd:option("-dir", "./data/test", 'test image directory')
  16. cmd:option("-file", "", 'test image file list')
  17. cmd:option("-model1_dir", "./models/anime_style_art_rgb", 'model1 directory')
  18. cmd:option("-model2_dir", "", 'model2 directory (optional)')
  19. cmd:option("-method", "scale", '(scale|noise|noise_scale|user|diff|scale4)')
  20. cmd:option("-filter", "Catrom", "downscaling filter (Box|Lanczos|Catrom(Bicubic))")
  21. cmd:option("-resize_blur", 1.0, 'blur parameter for resize')
  22. cmd:option("-color", "y", '(rgb|y|r|g|b)')
  23. cmd:option("-noise_level", 1, 'model noise level')
  24. cmd:option("-jpeg_quality", 75, 'jpeg quality')
  25. cmd:option("-jpeg_times", 1, 'jpeg compression times')
  26. cmd:option("-jpeg_quality_down", 5, 'value of jpeg quality to decrease each times')
  27. cmd:option("-range_bug", 0, 'Reproducing the dynamic range bug that is caused by MATLAB\'s rgb2ycbcr(1|0)')
  28. cmd:option("-save_image", 0, 'save converted images')
  29. cmd:option("-save_baseline_image", 0, 'save baseline images')
  30. cmd:option("-output_dir", "./", 'output directroy')
  31. cmd:option("-show_progress", 1, 'show progressbar')
  32. cmd:option("-baseline_filter", "Catrom", 'baseline interpolation (Box|Lanczos|Catrom(Bicubic))')
  33. cmd:option("-save_info", 0, 'save score and parameters to benchmark.txt')
  34. cmd:option("-save_all", 0, 'group -save_info, -save_image and -save_baseline_image option')
  35. cmd:option("-thread", -1, 'number of CPU threads')
  36. cmd:option("-tta", 0, 'use tta')
  37. cmd:option("-tta_level", 8, 'tta level')
  38. cmd:option("-crop_size", 128, 'patch size per process')
  39. cmd:option("-batch_size", 1, 'batch_size')
  40. cmd:option("-force_cudnn", 0, 'use cuDNN backend')
  41. cmd:option("-yuv420", 0, 'use yuv420 jpeg')
  42. cmd:option("-name", "", 'model name for user method')
  43. cmd:option("-x_dir", "", 'input image for user method')
  44. cmd:option("-y_dir", "", 'groundtruth image for user method. filename must be the same as x_dir')
  45. cmd:option("-x_file", "", 'input image for user method')
  46. cmd:option("-y_file", "", 'groundtruth image for user method. filename must be the same as x_file')
  47. cmd:option("-border", 0, 'border px that will removed')
  48. cmd:option("-metric", "", '(jaccard)')
  49. local function to_bool(settings, name)
  50. if settings[name] == 1 then
  51. settings[name] = true
  52. else
  53. settings[name] = false
  54. end
  55. end
  56. local opt = cmd:parse(arg)
  57. torch.setdefaulttensortype('torch.FloatTensor')
  58. if cudnn then
  59. cudnn.fastest = true
  60. cudnn.benchmark = true
  61. end
  62. to_bool(opt, "force_cudnn")
  63. to_bool(opt, "yuv420")
  64. to_bool(opt, "save_all")
  65. to_bool(opt, "tta")
  66. if opt.save_all then
  67. opt.save_image = true
  68. opt.save_info = true
  69. opt.save_baseline_image = true
  70. else
  71. to_bool(opt, "save_image")
  72. to_bool(opt, "save_info")
  73. to_bool(opt, "save_baseline_image")
  74. end
  75. to_bool(opt, "show_progress")
  76. if opt.thread > 0 then
  77. torch.setnumthreads(tonumber(opt.thread))
  78. end
  79. if opt.output_dir:len() > 0 then
  80. dir.makepath(opt.output_dir)
  81. end
  82. local function rgb2y_matlab(x)
  83. local y = torch.Tensor(1, x:size(2), x:size(3)):zero()
  84. x = iproc.byte2float(x)
  85. y:add(x[1] * 65.481)
  86. y:add(x[2] * 128.553)
  87. y:add(x[3] * 24.966)
  88. y:add(16.0)
  89. return y:byte():float()
  90. end
  91. local function RGBMSE(x1, x2)
  92. x1 = iproc.float2byte(x1):float()
  93. x2 = iproc.float2byte(x2):float()
  94. return (x1 - x2):pow(2):mean()
  95. end
  96. local function CHMSE(x1, x2, ch)
  97. x1 = iproc.float2byte(x1):float()
  98. x2 = iproc.float2byte(x2):float()
  99. return (x1[ch] - x2[ch]):pow(2):mean()
  100. end
  101. local function YMSE(x1, x2)
  102. if opt.range_bug == 1 then
  103. local x1_2 = rgb2y_matlab(x1)
  104. local x2_2 = rgb2y_matlab(x2)
  105. return (x1_2 - x2_2):pow(2):mean()
  106. else
  107. local x1_2 = image.rgb2y(x1):mul(255.0)
  108. local x2_2 = image.rgb2y(x2):mul(255.0)
  109. return (x1_2 - x2_2):pow(2):mean()
  110. end
  111. end
  112. local function MSE(x1, x2, color)
  113. if color == "y" then
  114. return YMSE(x1, x2)
  115. elseif color == "r" then
  116. return CHMSE(x1, x2, 1)
  117. elseif color == "g" then
  118. return CHMSE(x1, x2, 2)
  119. elseif color == "b" then
  120. return CHMSE(x1, x2, 3)
  121. else
  122. return RGBMSE(x1, x2)
  123. end
  124. end
  125. local function PSNR(x1, x2, color)
  126. local mse = math.max(MSE(x1, x2, color), 1)
  127. return 10 * math.log10((255.0 * 255.0) / mse)
  128. end
  129. local function MSE2PSNR(mse)
  130. return 10 * math.log10((255.0 * 255.0) / math.max(mse, 1))
  131. end
  132. local function transform_jpeg(x, opt)
  133. for i = 1, opt.jpeg_times do
  134. jpeg = gm.Image(x, "RGB", "DHW")
  135. jpeg:format("jpeg")
  136. if opt.yuv420 then
  137. jpeg:samplingFactors({2.0, 1.0, 1.0})
  138. else
  139. jpeg:samplingFactors({1.0, 1.0, 1.0})
  140. end
  141. blob, len = jpeg:toBlob(opt.jpeg_quality - (i - 1) * opt.jpeg_quality_down)
  142. jpeg:fromBlob(blob, len)
  143. x = jpeg:toTensor("byte", "RGB", "DHW")
  144. end
  145. return iproc.byte2float(x)
  146. end
  147. local function baseline_scale(x, filter)
  148. return iproc.scale(x,
  149. x:size(3) * 2.0,
  150. x:size(2) * 2.0,
  151. filter)
  152. end
  153. local function baseline_scale4(x, filter)
  154. return iproc.scale(x,
  155. x:size(3) * 4.0,
  156. x:size(2) * 4.0,
  157. filter)
  158. end
  159. local function transform_scale(x, opt)
  160. return iproc.scale(x,
  161. x:size(3) * 0.5,
  162. x:size(2) * 0.5,
  163. opt.filter, opt.resize_blur)
  164. end
  165. local function transform_scale4(x, opt)
  166. return iproc.scale(x,
  167. x:size(3) * 0.25,
  168. x:size(2) * 0.25,
  169. opt.filter, opt.resize_blur)
  170. end
  171. local function transform_scale_jpeg(x, opt)
  172. x = iproc.scale(x,
  173. x:size(3) * 0.5,
  174. x:size(2) * 0.5,
  175. opt.filter, opt.resize_blur)
  176. for i = 1, opt.jpeg_times do
  177. jpeg = gm.Image(x, "RGB", "DHW")
  178. jpeg:format("jpeg")
  179. if opt.yuv420 then
  180. jpeg:samplingFactors({2.0, 1.0, 1.0})
  181. else
  182. jpeg:samplingFactors({1.0, 1.0, 1.0})
  183. end
  184. blob, len = jpeg:toBlob(opt.jpeg_quality - (i - 1) * opt.jpeg_quality_down)
  185. jpeg:fromBlob(blob, len)
  186. x = jpeg:toTensor("byte", "RGB", "DHW")
  187. end
  188. return iproc.byte2float(x)
  189. end
  190. local function remove_border(x, border)
  191. return iproc.crop(x,
  192. border, border,
  193. x:size(3) - border,
  194. x:size(2) - border)
  195. end
  196. local function create_metric(metric)
  197. if metric and metric:len() > 0 then
  198. if metric == "jaccard" then
  199. return {
  200. name = "jaccard",
  201. func = function (a, b)
  202. local ga = iproc.rgb2y(a)
  203. local gb = iproc.rgb2y(b)
  204. local ba = torch.Tensor():resizeAs(ga)
  205. local bb = torch.Tensor():resizeAs(gb)
  206. ba:zero()
  207. bb:zero()
  208. ba[torch.gt(ga, 0.5)] = 1.0
  209. bb[torch.gt(gb, 0.5)] = 1.0
  210. local num_a = ba:sum()
  211. local num_b = bb:sum()
  212. local a_and_b = ba:cmul(bb):sum()
  213. return (a_and_b / (num_a + num_b - a_and_b))
  214. end}
  215. else
  216. error("unknown metric: " .. metric)
  217. end
  218. else
  219. return nil
  220. end
  221. end
  222. local function benchmark(opt, x, model1, model2)
  223. local mse1, mse2, am1, am2
  224. local won = {0, 0}
  225. local model1_mse = 0
  226. local model2_mse = 0
  227. local baseline_mse = 0
  228. local model1_psnr = 0
  229. local model2_psnr = 0
  230. local baseline_psnr = 0
  231. local model1_time = 0
  232. local model2_time = 0
  233. local scale_f = reconstruct.scale
  234. local image_f = reconstruct.image
  235. local detail_fp = nil
  236. local am = nil
  237. local model1_am = 0
  238. local model2_am = 0
  239. if opt.method == "user" or opt.method == "diff" then
  240. am = create_metric(opt.metric)
  241. end
  242. if opt.save_info then
  243. detail_fp = io.open(path.join(opt.output_dir, "benchmark_details.txt"), "w")
  244. end
  245. if opt.tta then
  246. scale_f = function(model, scale, x, block_size, batch_size)
  247. return reconstruct.scale_tta(model, opt.tta_level,
  248. scale, x, block_size, batch_size)
  249. end
  250. image_f = function(model, x, block_size, batch_size)
  251. return reconstruct.image_tta(model, opt.tta_level,
  252. x, block_size, batch_size)
  253. end
  254. end
  255. for i = 1, #x do
  256. if i % 10 == 0 then
  257. collectgarbage()
  258. end
  259. local basename = x[i].basename
  260. local input, model1_output, model2_output, baseline_output, ground_truth
  261. if opt.method == "scale" then
  262. input = transform_scale(iproc.byte2float(x[i].y), opt)
  263. ground_truth = iproc.byte2float(x[i].y)
  264. if opt.force_cudnn and i == 1 then -- run cuDNN benchmark first
  265. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  266. if model2 then
  267. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  268. end
  269. end
  270. t = sys.clock()
  271. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  272. model1_time = model1_time + (sys.clock() - t)
  273. if model2 then
  274. t = sys.clock()
  275. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  276. model2_time = model2_time + (sys.clock() - t)
  277. end
  278. baseline_output = baseline_scale(input, opt.baseline_filter)
  279. elseif opt.method == "scale4" then
  280. input = transform_scale4(iproc.byte2float(x[i].y), opt)
  281. ground_truth = iproc.byte2float(x[i].y)
  282. if opt.force_cudnn and i == 1 then -- run cuDNN benchmark first
  283. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  284. if model2 then
  285. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  286. end
  287. end
  288. t = sys.clock()
  289. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  290. model1_output = scale_f(model1, 2.0, model1_output, opt.crop_size, opt.batch_size)
  291. model1_time = model1_time + (sys.clock() - t)
  292. if model2 then
  293. t = sys.clock()
  294. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  295. model2_output = scale_f(model2, 2.0, model2_output, opt.crop_size, opt.batch_size)
  296. model2_time = model2_time + (sys.clock() - t)
  297. end
  298. baseline_output = baseline_scale4(input, opt.baseline_filter)
  299. elseif opt.method == "noise" then
  300. input = transform_jpeg(iproc.byte2float(x[i].y), opt)
  301. ground_truth = iproc.byte2float(x[i].y)
  302. if opt.force_cudnn and i == 1 then
  303. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  304. if model2 then
  305. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  306. end
  307. end
  308. t = sys.clock()
  309. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  310. model1_time = model1_time + (sys.clock() - t)
  311. if model2 then
  312. t = sys.clock()
  313. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  314. model2_time = model2_time + (sys.clock() - t)
  315. end
  316. baseline_output = input
  317. elseif opt.method == "noise_scale" then
  318. input = transform_scale_jpeg(iproc.byte2float(x[i].y), opt)
  319. ground_truth = iproc.byte2float(x[i].y)
  320. if opt.force_cudnn and i == 1 then
  321. if model1.noise_scale_model then
  322. model1_output = scale_f(model1.noise_scale_model, 2.0,
  323. input, opt.crop_size, opt.batch_size)
  324. else
  325. if model1.noise_model then
  326. model1_output = image_f(model1.noise_model, input, opt.crop_size, opt.batch_size)
  327. else
  328. model1_output = input
  329. end
  330. model1_output = scale_f(model1.scale_model, 2.0, model1_output,
  331. opt.crop_size, opt.batch_size)
  332. end
  333. if model2 then
  334. if model2.noise_scale_model then
  335. model2_output = scale_f(model2.noise_scale_model, 2.0,
  336. input, opt.crop_size, opt.batch_size)
  337. else
  338. if model2.noise_model then
  339. model2_output = image_f(model2.noise_model, input,
  340. opt.crop_size, opt.batch_size)
  341. else
  342. model2_output = input
  343. end
  344. model2_output = scale_f(model2.scale_model, 2.0, model2_output,
  345. opt.crop_size, opt.batch_size)
  346. end
  347. end
  348. end
  349. t = sys.clock()
  350. if model1.noise_scale_model then
  351. model1_output = scale_f(model1.noise_scale_model, 2.0,
  352. input, opt.crop_size, opt.batch_size)
  353. else
  354. if model1.noise_model then
  355. model1_output = image_f(model1.noise_model, input, opt.crop_size, opt.batch_size)
  356. else
  357. model1_output = input
  358. end
  359. model1_output = scale_f(model1.scale_model, 2.0, model1_output,
  360. opt.crop_size, opt.batch_size)
  361. end
  362. model1_time = model1_time + (sys.clock() - t)
  363. if model2 then
  364. t = sys.clock()
  365. if model2.noise_scale_model then
  366. model2_output = scale_f(model2.noise_scale_model, 2.0,
  367. input, opt.crop_size, opt.batch_size)
  368. else
  369. if model2.noise_model then
  370. model2_output = image_f(model2.noise_model, input,
  371. opt.crop_size, opt.batch_size)
  372. else
  373. model2_output = input
  374. end
  375. model2_output = scale_f(model2.scale_model, 2.0, model2_output,
  376. opt.crop_size, opt.batch_size)
  377. end
  378. model2_time = model2_time + (sys.clock() - t)
  379. end
  380. baseline_output = baseline_scale(input, opt.baseline_filter)
  381. elseif opt.method == "user" then
  382. input = iproc.byte2float(x[i].x)
  383. ground_truth = iproc.byte2float(x[i].y)
  384. local y_scale = ground_truth:size(2) / input:size(2)
  385. if y_scale > 1 then
  386. if opt.force_cudnn and i == 1 then
  387. model1_output = scale_f(model1, y_scale, input, opt.crop_size, opt.batch_size)
  388. if model2 then
  389. model2_output = scale_f(model2, y_scale, input, opt.crop_size, opt.batch_size)
  390. end
  391. end
  392. t = sys.clock()
  393. model1_output = scale_f(model1, y_scale, input, opt.crop_size, opt.batch_size)
  394. model1_time = model1_time + (sys.clock() - t)
  395. if model2 then
  396. t = sys.clock()
  397. model2_output = scale_f(model2, y_scale, input, opt.crop_size, opt.batch_size)
  398. model2_time = model2_time + (sys.clock() - t)
  399. end
  400. else
  401. if opt.force_cudnn and i == 1 then
  402. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  403. if model2 then
  404. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  405. end
  406. end
  407. t = sys.clock()
  408. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  409. model1_time = model1_time + (sys.clock() - t)
  410. if model2 then
  411. t = sys.clock()
  412. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  413. model2_time = model2_time + (sys.clock() - t)
  414. end
  415. end
  416. elseif opt.method == "diff" then
  417. input = iproc.byte2float(x[i].x)
  418. ground_truth = iproc.byte2float(x[i].y)
  419. model1_output = input
  420. end
  421. if opt.border > 0 then
  422. ground_truth = remove_border(ground_truth, opt.border)
  423. model1_output = remove_border(model1_output, opt.border)
  424. end
  425. if am then
  426. am1 = am.func(ground_truth, model1_output)
  427. model1_am = model1_am + am1
  428. else
  429. mse1 = MSE(ground_truth, model1_output, opt.color)
  430. model1_mse = model1_mse + mse1
  431. model1_psnr = model1_psnr + MSE2PSNR(mse1)
  432. end
  433. local won_model = 1
  434. if model2 then
  435. if opt.border > 0 then
  436. model2_output = remove_border(model2_output, opt.border)
  437. end
  438. if am then
  439. am2 = am.func(ground_truth, model2_output)
  440. model2_am = model2_am + am2
  441. else
  442. mse2 = MSE(ground_truth, model2_output, opt.color)
  443. model2_mse = model2_mse + mse2
  444. model2_psnr = model2_psnr + MSE2PSNR(mse2)
  445. end
  446. if am then
  447. if am1 < am2 then
  448. won[1] = won[1] + 1
  449. elseif am1 > am2 then
  450. won[2] = won[2] + 1
  451. won_model = 2
  452. end
  453. else
  454. if mse1 < mse2 then
  455. won[1] = won[1] + 1
  456. elseif mse1 > mse2 then
  457. won[2] = won[2] + 1
  458. won_model = 2
  459. end
  460. end
  461. if detail_fp then
  462. if am then
  463. detail_fp:write(string.format("%s,%f,%d\n", x[i].basename, am1, am2, won_model))
  464. else
  465. detail_fp:write(string.format("%s,%f,%f,%d\n", x[i].basename,
  466. MSE2PSNR(mse1), MSE2PSNR(mse2), won_model))
  467. end
  468. end
  469. else
  470. if detail_fp then
  471. if am then
  472. detail_fp:write(string.format("%s,%f\n", x[i].basename, am1))
  473. else
  474. detail_fp:write(string.format("%s,%f\n", x[i].basename, MSE2PSNR(mse1)))
  475. end
  476. end
  477. end
  478. if baseline_output then
  479. baseline_output = remove_border(baseline_output, opt.border)
  480. mse = MSE(ground_truth, baseline_output, opt.color)
  481. baseline_mse = baseline_mse + mse
  482. baseline_psnr = baseline_psnr + MSE2PSNR(mse)
  483. end
  484. if opt.save_image then
  485. if opt.save_baseline_image and baseline_output then
  486. image.save(path.join(opt.output_dir, string.format("%s_baseline.png", basename)),
  487. baseline_output)
  488. end
  489. if model1_output then
  490. image.save(path.join(opt.output_dir, string.format("%s_model1.png", basename)),
  491. model1_output)
  492. end
  493. if model2_output then
  494. image.save(path.join(opt.output_dir, string.format("%s_model2.png", basename)),
  495. model2_output)
  496. end
  497. end
  498. if opt.show_progress or i == #x then
  499. if am then
  500. if model2 then
  501. io.stdout:write(
  502. string.format("%d/%d; model1_time=%.2f, model2_time=%.2f, model1_%s=%.3f, model2_%s=%.3f \r",
  503. i, #x,
  504. model1_time,
  505. model2_time,
  506. am.name, model1_am / i, am.name, model2_am / i
  507. ))
  508. else
  509. io.stdout:write(
  510. string.format("%d/%d; model1_time=%.2f, model1_%s=%.3f \r",
  511. i, #x,
  512. model1_time,
  513. am.name, model1_am / i
  514. ))
  515. end
  516. else
  517. if model2 then
  518. if baseline_output then
  519. io.stdout:write(
  520. string.format("%d/%d; model1_time=%.2f, model2_time=%.2f, baseline_rmse=%.3f, model1_rmse=%.3f, model2_rmse=%.3f, baseline_psnr=%.3f, model1_psnr=%.3f, model2_psnr=%.3f, model1_won=%d, model2_won=%d \r",
  521. i, #x,
  522. model1_time,
  523. model2_time,
  524. math.sqrt(baseline_mse / i),
  525. math.sqrt(model1_mse / i), math.sqrt(model2_mse / i),
  526. baseline_psnr / i,
  527. model1_psnr / i, model2_psnr / i,
  528. won[1], won[2]
  529. ))
  530. else
  531. io.stdout:write(
  532. string.format("%d/%d; model1_time=%.2f, model2_time=%.2f, model1_rmse=%.3f, model2_rmse=%.3f, model1_psnr=%.3f, model2_psnr=%.3f, model1_own=%d, model2_won=%d \r",
  533. i, #x,
  534. model1_time,
  535. model2_time,
  536. math.sqrt(model1_mse / i), math.sqrt(model2_mse / i),
  537. model1_psnr / i, model2_psnr / i,
  538. won[1], won[2]
  539. ))
  540. end
  541. else
  542. if baseline_output then
  543. io.stdout:write(
  544. string.format("%d/%d; model1_time=%.2f, baseline_rmse=%.3f, model1_rmse=%.3f, baseline_psnr=%.3f, model1_psnr=%.3f \r",
  545. i, #x,
  546. model1_time,
  547. math.sqrt(baseline_mse / i), math.sqrt(model1_mse / i),
  548. baseline_psnr / i, model1_psnr / i
  549. ))
  550. else
  551. io.stdout:write(
  552. string.format("%d/%d; model1_time=%.2f, model1_rmse=%.3f, model1_psnr=%.3f \r",
  553. i, #x,
  554. model1_time,
  555. math.sqrt(model1_mse / i), model1_psnr / i
  556. ))
  557. end
  558. end
  559. end
  560. io.stdout:flush()
  561. end
  562. end
  563. if opt.save_info then
  564. local fp = io.open(path.join(opt.output_dir, "benchmark.txt"), "w")
  565. fp:write("options : " .. cjson.encode(opt) .. "\n")
  566. if baseline_psnr > 0 then
  567. fp:write(string.format("baseline: RMSE = %.3f, PSNR = %.3f\n",
  568. math.sqrt(baseline_mse / #x), baseline_psnr / #x))
  569. end
  570. if model1_psnr > 0 then
  571. fp:write(string.format("model1 : RMSE = %.3f, PSNR = %.3f, evaluation time = %.3f\n",
  572. math.sqrt(model1_mse / #x), model1_psnr / #x, model1_time))
  573. end
  574. if model2_psnr > 0 then
  575. fp:write(string.format("model2 : RMSE = %.3f, PSNR = %.3f, evaluation time = %.3f\n",
  576. math.sqrt(model2_mse / #x), model2_psnr / #x, model2_time))
  577. end
  578. if model1_am > 0 then
  579. fp:write(string.format("model1 : %s = %.3f, evaluation time = %.3f\n",
  580. math.sqrt(model1_am / #x), model1_time))
  581. end
  582. if model2_am > 0 then
  583. fp:write(string.format("model2 : %s = %.3f, evaluation time = %.3f\n",
  584. math.sqrt(model2_am / #x), model2_time))
  585. end
  586. fp:close()
  587. if detail_fp then
  588. detail_fp:close()
  589. end
  590. end
  591. io.stdout:write("\n")
  592. end
  593. local function load_data_from_dir(test_dir)
  594. local test_x = {}
  595. local files = dir.getfiles(test_dir, "*.*")
  596. for i = 1, #files do
  597. local name = path.basename(files[i])
  598. local e = path.extension(name)
  599. local base = name:sub(0, name:len() - e:len())
  600. local img = image_loader.load_byte(files[i])
  601. if img then
  602. table.insert(test_x, {y = iproc.crop_mod4(img),
  603. basename = base})
  604. end
  605. if opt.show_progress then
  606. xlua.progress(i, #files)
  607. end
  608. if i % 10 == 0 then
  609. collectgarbage()
  610. end
  611. end
  612. return test_x
  613. end
  614. local function load_data_from_file(test_file)
  615. local test_x = {}
  616. local files = utils.split(file.read(test_file), "\n")
  617. for i = 1, #files do
  618. local name = path.basename(files[i])
  619. local e = path.extension(name)
  620. local base = name:sub(0, name:len() - e:len())
  621. local img = image_loader.load_byte(files[i])
  622. if img then
  623. table.insert(test_x, {y = iproc.crop_mod4(img),
  624. basename = base})
  625. end
  626. if opt.show_progress then
  627. xlua.progress(i, #files)
  628. end
  629. if i % 10 == 0 then
  630. collectgarbage()
  631. end
  632. end
  633. return test_x
  634. end
  635. local function get_basename(f)
  636. local name = path.basename(f)
  637. local e = path.extension(name)
  638. local base = name:sub(0, name:len() - e:len())
  639. return base
  640. end
  641. local function load_user_data(y_dir, y_file, x_dir, x_file)
  642. local test = {}
  643. local y_files
  644. local x_files
  645. if y_file:len() > 0 then
  646. y_files = utils.split(file.read(y_file), "\n")
  647. else
  648. y_files = dir.getfiles(y_dir, "*.*")
  649. end
  650. if x_file:len() > 0 then
  651. x_files = utils.split(file.read(x_file), "\n")
  652. else
  653. x_files = dir.getfiles(x_dir, "*.*")
  654. end
  655. local basename_db = {}
  656. for i = 1, #y_files do
  657. basename_db[get_basename(y_files[i])] = {y = y_files[i]}
  658. end
  659. for i = 1, #x_files do
  660. local key = get_basename(x_files[i])
  661. if basename_db[key] then
  662. basename_db[key].x = x_files[i]
  663. else
  664. error(string.format("%s is not found in %s", key, y_dir))
  665. end
  666. end
  667. for i = 1, #y_files do
  668. local key = get_basename(y_files[i])
  669. local d = basename_db[key]
  670. if not (d.x and d.y) then
  671. error(string.format("%s is not found in %s", key, x_dir))
  672. end
  673. end
  674. for i = 1, #y_files do
  675. local key = get_basename(y_files[i])
  676. local x = image_loader.load_byte(basename_db[key].x)
  677. local y = image_loader.load_byte(basename_db[key].y)
  678. if x and y then
  679. table.insert(test, {y = y,
  680. x = x,
  681. basename = key})
  682. end
  683. if opt.show_progress then
  684. xlua.progress(i, #y_files)
  685. end
  686. if i % 10 == 0 then
  687. collectgarbage()
  688. end
  689. end
  690. return test
  691. end
  692. function load_noise_scale_model(model_dir, noise_level, force_cudnn)
  693. local f = path.join(model_dir, string.format("noise%d_scale2.0x_model.t7", opt.noise_level))
  694. local s1, noise_scale = pcall(w2nn.load_model, f, force_cudnn)
  695. local model = {}
  696. if not s1 then
  697. f = path.join(model_dir, string.format("noise%d_model.t7", opt.noise_level))
  698. local noise
  699. s1, noise = pcall(w2nn.load_model, f, force_cudnn)
  700. if not s1 then
  701. model.noise_model = nil
  702. print(model_dir .. "'s noise model is not found. benchmark will use only scale model.")
  703. else
  704. model.noise_model = noise
  705. end
  706. f = path.join(model_dir, "scale2.0x_model.t7")
  707. local scale
  708. s1, scale = pcall(w2nn.load_model, f, force_cudnn)
  709. if not s1 then
  710. error(model_dir .. ": load error")
  711. return nil
  712. end
  713. model.scale_model = scale
  714. else
  715. model.noise_scale_model = noise_scale
  716. end
  717. return model
  718. end
  719. if opt.show_progress then
  720. print(opt)
  721. end
  722. if opt.method == "scale" or opt.method == "scale4" then
  723. local f1 = path.join(opt.model1_dir, "scale2.0x_model.t7")
  724. local f2 = path.join(opt.model2_dir, "scale2.0x_model.t7")
  725. local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
  726. local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
  727. if not s1 then
  728. error("Load error: " .. f1)
  729. end
  730. if not s2 then
  731. model2 = nil
  732. end
  733. local test_x
  734. if opt.file:len() > 0 then
  735. test_x = load_data_from_file(opt.file)
  736. else
  737. test_x = load_data_from_dir(opt.dir)
  738. end
  739. benchmark(opt, test_x, model1, model2)
  740. elseif opt.method == "noise" then
  741. local f1 = path.join(opt.model1_dir, string.format("noise%d_model.t7", opt.noise_level))
  742. local f2 = path.join(opt.model2_dir, string.format("noise%d_model.t7", opt.noise_level))
  743. local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
  744. local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
  745. if not s1 then
  746. error("Load error: " .. f1)
  747. end
  748. if not s2 then
  749. model2 = nil
  750. end
  751. local test_x
  752. if opt.file:len() > 0 then
  753. test_x = load_data_from_file(opt.file)
  754. else
  755. test_x = load_data_from_dir(opt.dir)
  756. end
  757. benchmark(opt, test_x, model1, model2)
  758. elseif opt.method == "noise_scale" then
  759. local model2 = nil
  760. local model1 = load_noise_scale_model(opt.model1_dir, opt.noise_level, opt.force_cudnn)
  761. if opt.model2_dir:len() > 0 then
  762. model2 = load_noise_scale_model(opt.model2_dir, opt.noise_level, opt.force_cudnn)
  763. end
  764. local test_x
  765. if opt.file:len() > 0 then
  766. test_x = load_data_from_file(opt.file)
  767. else
  768. test_x = load_data_from_dir(opt.dir)
  769. end
  770. benchmark(opt, test_x, model1, model2)
  771. elseif opt.method == "user" then
  772. local f1 = path.join(opt.model1_dir, string.format("%s_model.t7", opt.name))
  773. local f2 = path.join(opt.model2_dir, string.format("%s_model.t7", opt.name))
  774. local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
  775. local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
  776. if not s1 then
  777. error("Load error: " .. f1)
  778. end
  779. if not s2 then
  780. model2 = nil
  781. end
  782. local test = load_user_data(opt.y_dir, opt.y_file, opt.x_dir, opt.x_file)
  783. benchmark(opt, test, model1, model2)
  784. elseif opt.method == "diff" then
  785. local test = load_user_data(opt.y_dir, opt.y_file, opt.x_dir, opt.x_file)
  786. benchmark(opt, test, nil, nil)
  787. end