diff options
Diffstat (limited to 'src/img/png.nim')
| -rw-r--r-- | src/img/png.nim | 413 |
1 files changed, 413 insertions, 0 deletions
diff --git a/src/img/png.nim b/src/img/png.nim new file mode 100644 index 00000000..9bcbd348 --- /dev/null +++ b/src/img/png.nim @@ -0,0 +1,413 @@ +import math + +import bindings/zlib +import img/bitmap +import types/color + +import lib/endians2 + +type PNGWriter = object + buf: pointer + i: int + outlen: int + +func pngInt(i: uint32): auto = + doAssert i < uint32(2) ^ 31 + return i.toBytesBE() + +func oq(writer: PNGWriter): ptr UncheckedArray[uint8] = + cast[ptr UncheckedArray[uint8]](writer.buf) + +proc writeStr[T](writer: var PNGWriter, s: T) = + if writer.outlen < writer.i + s.len: + writer.outlen = writer.i + s.len + writer.buf = realloc(writer.buf, writer.outlen) + copyMem(addr writer.oq[writer.i], unsafeAddr s[0], s.len) + writer.i += s.len + +proc writeInt(writer: var PNGWriter, i: uint32) = + writer.writeStr(i.toBytesBE()) + +proc writePngInt(writer: var PNGWriter, i: uint32) = + doAssert i < uint32(2) ^ 31 + writer.writeInt(i) + +proc writeChunk[T](writer: var PNGWriter, t: string, data: T) = + var crc = uint32(crc32(0, cast[ptr uint8](unsafeAddr t[0]), cuint(t.len))) + if data.len > 0: + crc = uint32(crc32(crc, cast[ptr uint8](unsafeAddr data[0]), + cuint(data.len))) + writer.writePngInt(uint32(data.len)) + writer.writeStr(t) + if data.len > 0: + writer.writeStr(data) + writer.writeInt(uint32(crc)) + +type PNGColorType {.size: sizeof(uint8).} = enum + GRAYSCALE = 0 + TRUECOLOR = 2 + INDEXED_COLOR = 3 + GRAYSCALE_WITH_ALPHA = 4 + TRUECOLOR_WITH_ALPHA = 6 + +func u8toc(x: openArray[uint8]): string = + #TODO ew + var s = newString(x.len) + copyMem(addr s[0], unsafeAddr x[0], x.len) + return s + +const PNGSignature = "\x89PNG\r\n\x1A\n" +proc writeIHDR(writer: var PNGWriter, width, height: uint32, + bitDepth: uint8, colorType: PNGColorType, + compressionMethod, filterMethod, interlaceMethod: uint8) = + writer.writeStr(PNGSignature) + let ihdr = u8toc(pngInt(width)) & + u8toc(pngInt(height)) & + char(bitDepth) & + char(uint8(colorType)) & + char(compressionMethod) & + char(filterMethod) & + char(interlaceMethod) + writer.writeChunk("IHDR", ihdr) + +proc writeIDAT(writer: var PNGWriter, bmp: Bitmap) = + #TODO smaller idat chunks + # +1 height for filter + var idat = newSeq[uint8]((bmp.width + 1) * bmp.height * 4) + var j = 0 # idat pointer + for k in 0 ..< bmp.px.len: + if k mod int(bmp.width) == 0: + # begin row + # For now, filter is always 0. TODO implement other filters + inc j + let p = bmp.px[k] + idat[j] = uint8(p.r) + idat[j + 1] = uint8(p.g) + idat[j + 2] = uint8(p.b) + idat[j + 3] = uint8(p.a) + j += 4 + var hlen = compressBound(culong(idat.len)) + var oidat = newSeq[uint8](int(hlen)) + let res = compress(addr oidat[0], addr hlen, addr idat[0], culong(idat.len)) + doAssert res == Z_OK #TODO error handling... + oidat.setLen(int(hlen)) + writer.writeChunk("IDAT", oidat) + +proc toPNG*(bmp: Bitmap, outlen: var int): pointer = + var writer = PNGWriter( + buf: alloc(PNGSignature.len), + outlen: PNGSignature.len + ) + writer.writeIHDR(uint32(bmp.width), uint32(bmp.height), 8, + TRUECOLOR_WITH_ALPHA, 0, 0, 0) + writer.writeIDAT(bmp) + writer.writeChunk("IEND", "") + outlen = writer.outlen + return writer.buf + +type PNGReader = object + bmp: Bitmap + iq: ptr UncheckedArray[uint8] + limit: int + i: int + bitDepth: uint8 + colorType: PNGColorType + background: RGBColor + isend: bool + idatBuf: seq[uint8] + uprow: seq[uint8] + idatAt: int + hasstrm: bool + strm: z_stream + strmend: bool + atline: int + +func width(reader: PNGReader): int {.inline.} = int(reader.bmp.width) + +func height(reader: PNGReader): int {.inline.} = int(reader.bmp.height) + +func spp(reader: PNGReader): int = + case reader.colorType + of TRUECOLOR: return 3 + of GRAYSCALE: return 1 + of INDEXED_COLOR: return 1 + of GRAYSCALE_WITH_ALPHA: return 2 + of TRUECOLOR_WITH_ALPHA: return 4 + +func scanlen(reader: PNGReader): int {.inline.} = + let w = reader.width + 1 + return (w * reader.spp * int(reader.bitDepth) + 7) div 8 + +proc handleError(reader: var PNGReader, msg: string) = + reader.bmp = nil + if reader.hasstrm: + discard inflateEnd(addr reader.strm) + +template err(reader: var PNGReader, msg: string) = + reader.handleError(msg) + return + +template readStr(reader: var PNGReader, L: int): string = + if reader.i + L > reader.limit: + reader.err "too short" + var s = newString(L) + copyMem(addr s[0], addr reader.iq[reader.i], L) + reader.i += L + s + +template readU8(reader: var PNGReader): uint8 = + if reader.i > reader.limit: + reader.err "too short" + let x = reader.iq[reader.i] + inc reader.i + x + +template readU32(reader: var PNGReader): uint32 = + if reader.i + 4 > reader.limit: + reader.err "too short" + let x = fromBytesBE(uint32, toOpenArray(reader.iq, reader.i, reader.i + 3)) + reader.i += 4 + x + +template readPNGInt(reader: var PNGReader): uint32 = + let x = reader.readU32() + if x >= uint32(2) ^ 31: + reader.err "int too large" + x + +template readColorType(reader: var PNGReader): PNGColorType = + case reader.readU8() + of 0u8: GRAYSCALE + of 2u8: TRUECOLOR + of 3u8: INDEXED_COLOR + of 4u8: GRAYSCALE_WITH_ALPHA + of 6u8: TRUECOLOR_WITH_ALPHA + else: reader.err "unknown color type" + +func bitDepthValid(colorType: PNGColorType, bitDepth: uint8): bool = + case colorType + of GRAYSCALE: + return int(bitDepth) in [1, 2, 4, 8, 16] + of INDEXED_COLOR: + return int(bitDepth) in [1, 2, 4, 8] + of TRUECOLOR, GRAYSCALE_WITH_ALPHA, TRUECOLOR_WITH_ALPHA: + return int(bitDepth) in [8, 16] + +proc readIHDR(reader: var PNGReader) = + if reader.readStr(PNGSignature.len) != PNGSignature: + reader.err "wrong signature" + if reader.readPNGInt() != 13: + reader.err "invalid header length" + if reader.readStr(4) != "IHDR": + reader.err "invalid header chunk" + let width = reader.readPNGInt() + let height = reader.readPNGInt() + reader.bitDepth = reader.readU8() + reader.colorType = reader.readColorType() + if not bitDepthValid(reader.colorType, reader.bitDepth): + reader.err "invalid bit depth" + let compressionMethod = reader.readU8() + if compressionMethod != 0: + reader.err "unknown compression method" + let filterMethod = reader.readU8() + if filterMethod != 0: + reader.err "unknown filter method" + let interlaceMethod = reader.readU8() + if interlaceMethod != 0: + reader.err "unknown interlace method" + let crc = crc32(0, addr reader.iq[reader.i - 17], 17) + if uint32(crc) != reader.readU32(): reader.err "wrong crc" + reader.bmp = newBitmap(width, height) + +proc readbKGD(reader: var PNGReader) = + case reader.colorType + of GRAYSCALE, GRAYSCALE_WITH_ALPHA: + discard reader.readU8() #TODO bit depth > 8 + reader.background = gray(reader.readU8()) + of TRUECOLOR, TRUECOLOR_WITH_ALPHA: + discard reader.readU8() #TODO bit depth > 8 + let r = reader.readU8() + discard reader.readU8() + let g = reader.readU8() + discard reader.readU8() + let b = reader.readU8() + reader.background = rgb(r, g, b) + of INDEXED_COLOR: + discard #TODO + +proc unfilter(reader: var PNGReader, irow: openArray[uint8], bpp: int) = + # none, sub, up -> replace uprow directly + # average, paeth -> copy to temp array, then replace uprow + let fil = irow[0] + let w = reader.width + case fil + of 0u8: # none + copyMem(addr reader.uprow[0], unsafeAddr irow[1], w) + of 1u8: # sub + for i in 1 ..< irow.len: + let j = i - 1 # skip filter byte + reader.uprow[j] = irow[i] + if j - bpp >= 0: + reader.uprow[j] += irow[j - bpp] + of 2u8: # up + for i in 1 ..< irow.len: + let j = i - 1 # skip filter byte + reader.uprow[j] += irow[i] + of 3u8: # average + reader.err "average not implemented yet" + of 4u8: # paeth + reader.err "paeth not implemented yet" + else: + eprint fil + reader.err "got invalid filter" + +proc writepxs(reader: var PNGReader, crow: var openArray[RGBAColor]) = + case reader.colorType + of GRAYSCALE: + var i = 0 + var j = 0 + for x in 0 ..< crow.len: + let u = reader.uprow[i] + let n = case reader.bitDepth + of 1: ((u shr (7 - j)) and 1) * 255 + of 2: ((u shr (6 - j)) and 3) * 85 + of 4: ((u shr (6 - j)) and 15) * 17 + of 8: u + of 16: u # we ignore the lower 8 bits. + else: 0 + j += int(reader.bitDepth) + i += j div 8 + j = j mod 8 + crow[x] = rgba(n, n, n, 255u8) + of TRUECOLOR: + let step = int(reader.bitdepth) div 8 + var i = 0 + for x in 0 ..< crow.len: + let r = reader.uprow[i] + i += step + let g = reader.uprow[i] + i += step + let b = reader.uprow[i] + i += step + crow[x] = rgba(r, g, b, 255u8) + of INDEXED_COLOR: discard #TODO + of GRAYSCALE_WITH_ALPHA: + let step = int(reader.bitDepth) div 8 + var i = 0 + for x in 0 ..< crow.len: + let n = reader.uprow[i] + i += step + let a = reader.uprow[i] + i += step + crow[x] = rgba(n, n, n, a) + of TRUECOLOR_WITH_ALPHA: + let step = int(reader.bitDepth) div 8 + for x in 0 ..< crow.len: + let r = reader.uprow[x * step] + let g = reader.uprow[(x + 1) * step] + let b = reader.uprow[(x + 2) * step] + let a = reader.uprow[(x + 3) * step] + crow[x] = rgba(r, g, b, a) + +proc readIDAT(reader: var PNGReader) = + if reader.idatAt == reader.idatBuf.len: + reader.err "idat buffer already filled" + if reader.strmend: + reader.err "stream already ended" + reader.strm.avail_in = cuint(reader.limit - reader.i) + reader.strm.next_in = addr reader.iq[reader.i] + let olen = reader.idatBuf.len - reader.idatAt + reader.strm.avail_out = cuint(olen) + reader.strm.next_out = addr reader.idatBuf[reader.idatAt] + let res = inflate(addr reader.strm, Z_NO_FLUSH) + doAssert res != Z_STREAM_ERROR + case res + of Z_NEED_DICT, Z_DATA_ERROR, Z_MEM_ERROR, Z_BUF_ERROR: + # Z_BUF_ERROR is fatal here, as outlen is at least as large as idat. + reader.err "error decompressing idat stream" + of Z_STREAM_END: + reader.strmend = true + of Z_OK: + if reader.strm.avail_out == 0: + reader.err "not enough space for output; is width or height wrong?" + else: doAssert false + reader.idatAt = int(reader.strm.total_out) + reader.i = reader.limit + let maxline = reader.idatAt div int(reader.scanlen) + let bmp = reader.bmp + let bps = if reader.bitDepth <= 8: 1 else: 2 # else 16 bit + let bpp = bps * reader.spp + let sl = int(reader.scanlen) + for y in reader.atline ..< maxline: + let yi = y * sl + assert yi + sl - 1 < reader.idatAt + reader.unfilter(toOpenArray(reader.idatBuf, yi, yi + sl - 1), bpp) + if unlikely(reader.bmp == nil): return + let yj = y * reader.width + reader.writepxs(toOpenArray(bmp.px, yj, yj + reader.width - 1)) + +proc readIEND(reader: var PNGReader) = + if reader.i < reader.limit: + reader.err "IEND too long" + reader.isend = true + +proc readUnknown(reader: var PNGReader, s: string) = + if (int(s[0]) and 0x20) == 0: + reader.err "unrecognized critical chunk " & s + #else: eprint "warning: unknown chunk " & s #debug + reader.i = reader.limit + +proc zlibAlloc(opaque: pointer, items: cuint, size: cuint): pointer {.cdecl.} = + return alloc(items * size) + +proc zlibFree(opaque: pointer, address: pointer) {.cdecl.} = + dealloc(address) + +proc initZStream(reader: var PNGReader) = + let bps = max(int(reader.bitDepth) div 8, 1) + reader.idatBuf = newSeq[uint8](reader.scanlen * reader.height * bps) + reader.uprow = newSeq[uint8](reader.width * bps) + reader.strm = z_stream( + zalloc: zlibAlloc, + zfree: zlibFree + ) + let ret = inflateInit(addr reader.strm) + if ret != Z_OK: + reader.err "failed to init inflate: " & $ret + reader.hasstrm = true + +proc fromPNG*(iq: openArray[uint8]): Bitmap = + if iq.len == 0: return + var reader = PNGReader( + iq: cast[ptr UncheckedArray[uint8]](unsafeAddr iq[0]), + limit: iq.len + ) + reader.readIHDR() + if reader.bmp == nil: return + if reader.width == 0 or reader.height == 0: + reader.err "invalid zero sized png" + if reader.colorType == INDEXED_COLOR: + reader.err "indexed color not implemented yet" + reader.initZStream() + while reader.i < iq.len and not reader.isend: + let len = int(reader.readPNGInt()) + if reader.i + len > iq.len: + reader.err "chunk too long" + let j = reader.i + let t = reader.readStr(4) + reader.limit = reader.i + len + case t + of "IHDR": reader.err "IHDR expected to be first chunk" + of "IDAT": reader.readIDAT() + of "IEND": reader.readIEND() + of "bKGD": reader.readbKGD() + else: reader.readUnknown(t) + if reader.bmp == nil: return + let crc = crc32(0, unsafeAddr iq[j], cuint(len + 4)) + reader.limit = iq.len + let y = reader.readU32() + if uint32(crc) != y: + reader.err "wrong crc" + if not reader.isend: + reader.err "IEND not found" + return reader.bmp |