benchmark.lua 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807
  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. local abab = (num_a + num_b - a_and_b)
  214. if abab > 0 then
  215. return (a_and_b / abab)
  216. else
  217. return 1
  218. end
  219. end}
  220. else
  221. error("unknown metric: " .. metric)
  222. end
  223. else
  224. return nil
  225. end
  226. end
  227. local function benchmark(opt, x, model1, model2)
  228. local mse1, mse2, am1, am2
  229. local won = {0, 0}
  230. local model1_mse = 0
  231. local model2_mse = 0
  232. local baseline_mse = 0
  233. local model1_psnr = 0
  234. local model2_psnr = 0
  235. local baseline_psnr = 0
  236. local model1_time = 0
  237. local model2_time = 0
  238. local scale_f = reconstruct.scale
  239. local image_f = reconstruct.image
  240. local detail_fp = nil
  241. local am = nil
  242. local model1_am = 0
  243. local model2_am = 0
  244. if opt.method == "user" or opt.method == "diff" then
  245. am = create_metric(opt.metric)
  246. end
  247. if opt.save_info then
  248. detail_fp = io.open(path.join(opt.output_dir, "benchmark_details.txt"), "w")
  249. end
  250. if opt.tta then
  251. scale_f = function(model, scale, x, block_size, batch_size)
  252. return reconstruct.scale_tta(model, opt.tta_level,
  253. scale, x, block_size, batch_size)
  254. end
  255. image_f = function(model, x, block_size, batch_size)
  256. return reconstruct.image_tta(model, opt.tta_level,
  257. x, block_size, batch_size)
  258. end
  259. end
  260. for i = 1, #x do
  261. if i % 10 == 0 then
  262. collectgarbage()
  263. end
  264. local basename = x[i].basename
  265. local input, model1_output, model2_output, baseline_output, ground_truth
  266. if opt.method == "scale" then
  267. input = transform_scale(iproc.byte2float(x[i].y), opt)
  268. ground_truth = iproc.byte2float(x[i].y)
  269. if opt.force_cudnn and i == 1 then -- run cuDNN benchmark first
  270. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  271. if model2 then
  272. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  273. end
  274. end
  275. t = sys.clock()
  276. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  277. model1_time = model1_time + (sys.clock() - t)
  278. if model2 then
  279. t = sys.clock()
  280. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  281. model2_time = model2_time + (sys.clock() - t)
  282. end
  283. baseline_output = baseline_scale(input, opt.baseline_filter)
  284. elseif opt.method == "scale4" then
  285. input = transform_scale4(iproc.byte2float(x[i].y), opt)
  286. ground_truth = iproc.byte2float(x[i].y)
  287. if opt.force_cudnn and i == 1 then -- run cuDNN benchmark first
  288. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  289. if model2 then
  290. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  291. end
  292. end
  293. t = sys.clock()
  294. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  295. model1_output = scale_f(model1, 2.0, model1_output, opt.crop_size, opt.batch_size)
  296. model1_time = model1_time + (sys.clock() - t)
  297. if model2 then
  298. t = sys.clock()
  299. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  300. model2_output = scale_f(model2, 2.0, model2_output, opt.crop_size, opt.batch_size)
  301. model2_time = model2_time + (sys.clock() - t)
  302. end
  303. baseline_output = baseline_scale4(input, opt.baseline_filter)
  304. elseif opt.method == "noise" then
  305. input = transform_jpeg(iproc.byte2float(x[i].y), opt)
  306. ground_truth = iproc.byte2float(x[i].y)
  307. if opt.force_cudnn and i == 1 then
  308. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  309. if model2 then
  310. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  311. end
  312. end
  313. t = sys.clock()
  314. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  315. model1_time = model1_time + (sys.clock() - t)
  316. if model2 then
  317. t = sys.clock()
  318. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  319. model2_time = model2_time + (sys.clock() - t)
  320. end
  321. baseline_output = input
  322. elseif opt.method == "noise_scale" then
  323. input = transform_scale_jpeg(iproc.byte2float(x[i].y), opt)
  324. ground_truth = iproc.byte2float(x[i].y)
  325. if opt.force_cudnn and i == 1 then
  326. if model1.noise_scale_model then
  327. model1_output = scale_f(model1.noise_scale_model, 2.0,
  328. input, opt.crop_size, opt.batch_size)
  329. else
  330. if model1.noise_model then
  331. model1_output = image_f(model1.noise_model, input, opt.crop_size, opt.batch_size)
  332. else
  333. model1_output = input
  334. end
  335. model1_output = scale_f(model1.scale_model, 2.0, model1_output,
  336. opt.crop_size, opt.batch_size)
  337. end
  338. if model2 then
  339. if model2.noise_scale_model then
  340. model2_output = scale_f(model2.noise_scale_model, 2.0,
  341. input, opt.crop_size, opt.batch_size)
  342. else
  343. if model2.noise_model then
  344. model2_output = image_f(model2.noise_model, input,
  345. opt.crop_size, opt.batch_size)
  346. else
  347. model2_output = input
  348. end
  349. model2_output = scale_f(model2.scale_model, 2.0, model2_output,
  350. opt.crop_size, opt.batch_size)
  351. end
  352. end
  353. end
  354. t = sys.clock()
  355. if model1.noise_scale_model then
  356. model1_output = scale_f(model1.noise_scale_model, 2.0,
  357. input, opt.crop_size, opt.batch_size)
  358. else
  359. if model1.noise_model then
  360. model1_output = image_f(model1.noise_model, input, opt.crop_size, opt.batch_size)
  361. else
  362. model1_output = input
  363. end
  364. model1_output = scale_f(model1.scale_model, 2.0, model1_output,
  365. opt.crop_size, opt.batch_size)
  366. end
  367. model1_time = model1_time + (sys.clock() - t)
  368. if model2 then
  369. t = sys.clock()
  370. if model2.noise_scale_model then
  371. model2_output = scale_f(model2.noise_scale_model, 2.0,
  372. input, opt.crop_size, opt.batch_size)
  373. else
  374. if model2.noise_model then
  375. model2_output = image_f(model2.noise_model, input,
  376. opt.crop_size, opt.batch_size)
  377. else
  378. model2_output = input
  379. end
  380. model2_output = scale_f(model2.scale_model, 2.0, model2_output,
  381. opt.crop_size, opt.batch_size)
  382. end
  383. model2_time = model2_time + (sys.clock() - t)
  384. end
  385. baseline_output = baseline_scale(input, opt.baseline_filter)
  386. elseif opt.method == "user" then
  387. input = iproc.byte2float(x[i].x)
  388. ground_truth = iproc.byte2float(x[i].y)
  389. local y_scale = ground_truth:size(2) / input:size(2)
  390. if y_scale > 1 then
  391. if opt.force_cudnn and i == 1 then
  392. model1_output = scale_f(model1, y_scale, input, opt.crop_size, opt.batch_size)
  393. if model2 then
  394. model2_output = scale_f(model2, y_scale, input, opt.crop_size, opt.batch_size)
  395. end
  396. end
  397. t = sys.clock()
  398. model1_output = scale_f(model1, y_scale, input, opt.crop_size, opt.batch_size)
  399. model1_time = model1_time + (sys.clock() - t)
  400. if model2 then
  401. t = sys.clock()
  402. model2_output = scale_f(model2, y_scale, input, opt.crop_size, opt.batch_size)
  403. model2_time = model2_time + (sys.clock() - t)
  404. end
  405. else
  406. if opt.force_cudnn and i == 1 then
  407. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  408. if model2 then
  409. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  410. end
  411. end
  412. t = sys.clock()
  413. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  414. model1_time = model1_time + (sys.clock() - t)
  415. if model2 then
  416. t = sys.clock()
  417. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  418. model2_time = model2_time + (sys.clock() - t)
  419. end
  420. end
  421. elseif opt.method == "diff" then
  422. input = iproc.byte2float(x[i].x)
  423. ground_truth = iproc.byte2float(x[i].y)
  424. model1_output = input
  425. end
  426. if opt.border > 0 then
  427. ground_truth = remove_border(ground_truth, opt.border)
  428. model1_output = remove_border(model1_output, opt.border)
  429. end
  430. if am then
  431. am1 = am.func(ground_truth, model1_output)
  432. model1_am = model1_am + am1
  433. else
  434. mse1 = MSE(ground_truth, model1_output, opt.color)
  435. model1_mse = model1_mse + mse1
  436. model1_psnr = model1_psnr + MSE2PSNR(mse1)
  437. end
  438. local won_model = 1
  439. if model2 then
  440. if opt.border > 0 then
  441. model2_output = remove_border(model2_output, opt.border)
  442. end
  443. if am then
  444. am2 = am.func(ground_truth, model2_output)
  445. model2_am = model2_am + am2
  446. else
  447. mse2 = MSE(ground_truth, model2_output, opt.color)
  448. model2_mse = model2_mse + mse2
  449. model2_psnr = model2_psnr + MSE2PSNR(mse2)
  450. end
  451. if am then
  452. if am1 < am2 then
  453. won[1] = won[1] + 1
  454. elseif am1 > am2 then
  455. won[2] = won[2] + 1
  456. won_model = 2
  457. end
  458. else
  459. if mse1 < mse2 then
  460. won[1] = won[1] + 1
  461. elseif mse1 > mse2 then
  462. won[2] = won[2] + 1
  463. won_model = 2
  464. end
  465. end
  466. if detail_fp then
  467. if am then
  468. detail_fp:write(string.format("%s,%f,%d\n", x[i].basename, am1, am2, won_model))
  469. else
  470. detail_fp:write(string.format("%s,%f,%f,%d\n", x[i].basename,
  471. MSE2PSNR(mse1), MSE2PSNR(mse2), won_model))
  472. end
  473. end
  474. else
  475. if detail_fp then
  476. if am then
  477. detail_fp:write(string.format("%s,%f\n", x[i].basename, am1))
  478. else
  479. detail_fp:write(string.format("%s,%f\n", x[i].basename, MSE2PSNR(mse1)))
  480. end
  481. end
  482. end
  483. if baseline_output then
  484. baseline_output = remove_border(baseline_output, opt.border)
  485. mse = MSE(ground_truth, baseline_output, opt.color)
  486. baseline_mse = baseline_mse + mse
  487. baseline_psnr = baseline_psnr + MSE2PSNR(mse)
  488. end
  489. if opt.save_image then
  490. if opt.save_baseline_image and baseline_output then
  491. image.save(path.join(opt.output_dir, string.format("%s_baseline.png", basename)),
  492. baseline_output)
  493. end
  494. if model1_output then
  495. image.save(path.join(opt.output_dir, string.format("%s_model1.png", basename)),
  496. model1_output)
  497. end
  498. if model2_output then
  499. image.save(path.join(opt.output_dir, string.format("%s_model2.png", basename)),
  500. model2_output)
  501. end
  502. end
  503. if opt.show_progress or i == #x then
  504. if am then
  505. if model2 then
  506. io.stdout:write(
  507. string.format("%d/%d; model1_time=%.2f, model2_time=%.2f, model1_%s=%.3f, model2_%s=%.3f \r",
  508. i, #x,
  509. model1_time,
  510. model2_time,
  511. am.name, model1_am / i, am.name, model2_am / i
  512. ))
  513. else
  514. io.stdout:write(
  515. string.format("%d/%d; model1_time=%.2f, model1_%s=%.3f \r",
  516. i, #x,
  517. model1_time,
  518. am.name, model1_am / i
  519. ))
  520. end
  521. else
  522. if model2 then
  523. if baseline_output then
  524. io.stdout:write(
  525. 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",
  526. i, #x,
  527. model1_time,
  528. model2_time,
  529. math.sqrt(baseline_mse / i),
  530. math.sqrt(model1_mse / i), math.sqrt(model2_mse / i),
  531. baseline_psnr / i,
  532. model1_psnr / i, model2_psnr / i,
  533. won[1], won[2]
  534. ))
  535. else
  536. io.stdout:write(
  537. 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",
  538. i, #x,
  539. model1_time,
  540. model2_time,
  541. math.sqrt(model1_mse / i), math.sqrt(model2_mse / i),
  542. model1_psnr / i, model2_psnr / i,
  543. won[1], won[2]
  544. ))
  545. end
  546. else
  547. if baseline_output then
  548. io.stdout:write(
  549. string.format("%d/%d; model1_time=%.2f, baseline_rmse=%.3f, model1_rmse=%.3f, baseline_psnr=%.3f, model1_psnr=%.3f \r",
  550. i, #x,
  551. model1_time,
  552. math.sqrt(baseline_mse / i), math.sqrt(model1_mse / i),
  553. baseline_psnr / i, model1_psnr / i
  554. ))
  555. else
  556. io.stdout:write(
  557. string.format("%d/%d; model1_time=%.2f, model1_rmse=%.3f, model1_psnr=%.3f \r",
  558. i, #x,
  559. model1_time,
  560. math.sqrt(model1_mse / i), model1_psnr / i
  561. ))
  562. end
  563. end
  564. end
  565. io.stdout:flush()
  566. end
  567. end
  568. if opt.save_info then
  569. local fp = io.open(path.join(opt.output_dir, "benchmark.txt"), "w")
  570. fp:write("options : " .. cjson.encode(opt) .. "\n")
  571. if baseline_psnr > 0 then
  572. fp:write(string.format("baseline: RMSE = %.3f, PSNR = %.3f\n",
  573. math.sqrt(baseline_mse / #x), baseline_psnr / #x))
  574. end
  575. if model1_psnr > 0 then
  576. fp:write(string.format("model1 : RMSE = %.3f, PSNR = %.3f, evaluation time = %.3f\n",
  577. math.sqrt(model1_mse / #x), model1_psnr / #x, model1_time))
  578. end
  579. if model2_psnr > 0 then
  580. fp:write(string.format("model2 : RMSE = %.3f, PSNR = %.3f, evaluation time = %.3f\n",
  581. math.sqrt(model2_mse / #x), model2_psnr / #x, model2_time))
  582. end
  583. if model1_am > 0 then
  584. fp:write(string.format("model1 : %s = %.3f, evaluation time = %.3f\n",
  585. math.sqrt(model1_am / #x), model1_time))
  586. end
  587. if model2_am > 0 then
  588. fp:write(string.format("model2 : %s = %.3f, evaluation time = %.3f\n",
  589. math.sqrt(model2_am / #x), model2_time))
  590. end
  591. fp:close()
  592. if detail_fp then
  593. detail_fp:close()
  594. end
  595. end
  596. io.stdout:write("\n")
  597. end
  598. local function load_data_from_dir(test_dir)
  599. local test_x = {}
  600. local files = dir.getfiles(test_dir, "*.*")
  601. for i = 1, #files do
  602. local name = path.basename(files[i])
  603. local e = path.extension(name)
  604. local base = name:sub(0, name:len() - e:len())
  605. local img = image_loader.load_byte(files[i])
  606. if img then
  607. table.insert(test_x, {y = iproc.crop_mod4(img),
  608. basename = base})
  609. end
  610. if i % 10 == 0 then
  611. if opt.show_progress then
  612. xlua.progress(i, #files)
  613. end
  614. collectgarbage()
  615. end
  616. end
  617. return test_x
  618. end
  619. local function load_data_from_file(test_file)
  620. local test_x = {}
  621. local files = utils.split(file.read(test_file), "\n")
  622. for i = 1, #files do
  623. local name = path.basename(files[i])
  624. local e = path.extension(name)
  625. local base = name:sub(0, name:len() - e:len())
  626. local img = image_loader.load_byte(files[i])
  627. if img then
  628. table.insert(test_x, {y = iproc.crop_mod4(img),
  629. basename = base})
  630. end
  631. if i % 10 == 0 then
  632. if opt.show_progress then
  633. xlua.progress(i, #files)
  634. end
  635. collectgarbage()
  636. end
  637. end
  638. return test_x
  639. end
  640. local function get_basename(f)
  641. local name = path.basename(f)
  642. local e = path.extension(name)
  643. local base = name:sub(0, name:len() - e:len())
  644. return base
  645. end
  646. local function load_user_data(y_dir, y_file, x_dir, x_file)
  647. local test = {}
  648. local y_files
  649. local x_files
  650. if y_file:len() > 0 then
  651. y_files = utils.split(file.read(y_file), "\n")
  652. else
  653. y_files = dir.getfiles(y_dir, "*.*")
  654. end
  655. if x_file:len() > 0 then
  656. x_files = utils.split(file.read(x_file), "\n")
  657. else
  658. x_files = dir.getfiles(x_dir, "*.*")
  659. end
  660. local basename_db = {}
  661. for i = 1, #y_files do
  662. basename_db[get_basename(y_files[i])] = {y = y_files[i]}
  663. end
  664. for i = 1, #x_files do
  665. local key = get_basename(x_files[i])
  666. if basename_db[key] then
  667. basename_db[key].x = x_files[i]
  668. else
  669. error(string.format("%s is not found in %s", key, y_dir))
  670. end
  671. end
  672. for i = 1, #y_files do
  673. local key = get_basename(y_files[i])
  674. local d = basename_db[key]
  675. if not (d.x and d.y) then
  676. error(string.format("%s is not found in %s", key, x_dir))
  677. end
  678. end
  679. for i = 1, #y_files do
  680. local key = get_basename(y_files[i])
  681. local x = image_loader.load_byte(basename_db[key].x)
  682. local y = image_loader.load_byte(basename_db[key].y)
  683. if x and y then
  684. table.insert(test, {y = y,
  685. x = x,
  686. basename = key})
  687. end
  688. if i % 10 == 0 then
  689. if opt.show_progress then
  690. xlua.progress(i, #y_files)
  691. end
  692. collectgarbage()
  693. end
  694. end
  695. return test
  696. end
  697. function load_noise_scale_model(model_dir, noise_level, force_cudnn)
  698. local f = path.join(model_dir, string.format("noise%d_scale2.0x_model.t7", opt.noise_level))
  699. local s1, noise_scale = pcall(w2nn.load_model, f, force_cudnn)
  700. local model = {}
  701. if not s1 then
  702. f = path.join(model_dir, string.format("noise%d_model.t7", opt.noise_level))
  703. local noise
  704. s1, noise = pcall(w2nn.load_model, f, force_cudnn)
  705. if not s1 then
  706. model.noise_model = nil
  707. print(model_dir .. "'s noise model is not found. benchmark will use only scale model.")
  708. else
  709. model.noise_model = noise
  710. end
  711. f = path.join(model_dir, "scale2.0x_model.t7")
  712. local scale
  713. s1, scale = pcall(w2nn.load_model, f, force_cudnn)
  714. if not s1 then
  715. error(model_dir .. ": load error")
  716. return nil
  717. end
  718. model.scale_model = scale
  719. else
  720. model.noise_scale_model = noise_scale
  721. end
  722. return model
  723. end
  724. if opt.show_progress then
  725. print(opt)
  726. end
  727. if opt.method == "scale" or opt.method == "scale4" then
  728. local f1 = path.join(opt.model1_dir, "scale2.0x_model.t7")
  729. local f2 = path.join(opt.model2_dir, "scale2.0x_model.t7")
  730. local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
  731. local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
  732. if not s1 then
  733. error("Load error: " .. f1)
  734. end
  735. if not s2 then
  736. model2 = nil
  737. end
  738. local test_x
  739. if opt.file:len() > 0 then
  740. test_x = load_data_from_file(opt.file)
  741. else
  742. test_x = load_data_from_dir(opt.dir)
  743. end
  744. benchmark(opt, test_x, model1, model2)
  745. elseif opt.method == "noise" then
  746. local f1 = path.join(opt.model1_dir, string.format("noise%d_model.t7", opt.noise_level))
  747. local f2 = path.join(opt.model2_dir, string.format("noise%d_model.t7", opt.noise_level))
  748. local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
  749. local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
  750. if not s1 then
  751. error("Load error: " .. f1)
  752. end
  753. if not s2 then
  754. model2 = nil
  755. end
  756. local test_x
  757. if opt.file:len() > 0 then
  758. test_x = load_data_from_file(opt.file)
  759. else
  760. test_x = load_data_from_dir(opt.dir)
  761. end
  762. benchmark(opt, test_x, model1, model2)
  763. elseif opt.method == "noise_scale" then
  764. local model2 = nil
  765. local model1 = load_noise_scale_model(opt.model1_dir, opt.noise_level, opt.force_cudnn)
  766. if opt.model2_dir:len() > 0 then
  767. model2 = load_noise_scale_model(opt.model2_dir, opt.noise_level, opt.force_cudnn)
  768. end
  769. local test_x
  770. if opt.file:len() > 0 then
  771. test_x = load_data_from_file(opt.file)
  772. else
  773. test_x = load_data_from_dir(opt.dir)
  774. end
  775. benchmark(opt, test_x, model1, model2)
  776. elseif opt.method == "user" then
  777. local f1 = path.join(opt.model1_dir, string.format("%s_model.t7", opt.name))
  778. local f2 = path.join(opt.model2_dir, string.format("%s_model.t7", opt.name))
  779. local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
  780. local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
  781. if not s1 then
  782. error("Load error: " .. f1)
  783. end
  784. if not s2 then
  785. model2 = nil
  786. end
  787. local test = load_user_data(opt.y_dir, opt.y_file, opt.x_dir, opt.x_file)
  788. benchmark(opt, test, model1, model2)
  789. elseif opt.method == "diff" then
  790. local test = load_user_data(opt.y_dir, opt.y_file, opt.x_dir, opt.x_file)
  791. benchmark(opt, test, nil, nil)
  792. end