diff options
author | Cédric Deltheil <cedric@moodstocks.com> | 2016-02-03 13:18:53 +0300 |
---|---|---|
committer | Cédric Deltheil <cedric@moodstocks.com> | 2016-02-03 13:18:53 +0300 |
commit | 003f4c71bde741bf4d029385afb8360f535c8f0d (patch) | |
tree | c8421d28536b9f273d9d2ebab6cc16aace31c7b9 | |
parent | 4543acebcd340632fd7a3de309ccab0e79f5df22 (diff) |
load: use magic number to detect image format
-rw-r--r-- | init.lua | 32 |
1 files changed, 31 insertions, 1 deletions
@@ -35,6 +35,13 @@ require 'xlua' require 'dok' require 'libimage' +local startswith = function(str, prefix) + return string.find(str, prefix, 1, true) == 1 +end + +local magicJPG = string.char(0xff, 0xd8, 0xff) +local magicPNG = string.char(0x89, 0x50, 0x4e, 0x47) + ---------------------------------------------------------------------- -- include unit test function -- @@ -331,7 +338,30 @@ local function load(filename, depth, tensortype) {type='string', help='type: byte | float | double'})) dok.error('missing file name', 'image.load') end - local ext = string.match(filename,'%.(%a+)$') + + local ext + + local f, err = io.open(filename, 'rb') + if not f then + error(err) + end + local hdr = f:read(4) or '' + f:close() + + if startswith(hdr, magicJPG) then + ext = 'jpg' + elseif startswith(hdr, magicPNG) then + ext = 'png' + elseif hdr:match('^P[25]') then + ext = 'pgm' + elseif hdr:match('^P[36]') then + ext = 'ppm' + end + + if not ext then + ext = string.match(filename,'%.(%a+)$') + end + local tensor if image.is_supported(ext) then tensor = filetypes[ext].loader(filename, depth, tensortype) |