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