Source file src/image/png/writer.go

     1	// Copyright 2009 The Go Authors. All rights reserved.
     2	// Use of this source code is governed by a BSD-style
     3	// license that can be found in the LICENSE file.
     4	
     5	package png
     6	
     7	import (
     8		"bufio"
     9		"compress/zlib"
    10		"encoding/binary"
    11		"hash/crc32"
    12		"image"
    13		"image/color"
    14		"io"
    15		"strconv"
    16	)
    17	
    18	// Encoder configures encoding PNG images.
    19	type Encoder struct {
    20		CompressionLevel CompressionLevel
    21	
    22		// BufferPool optionally specifies a buffer pool to get temporary
    23		// EncoderBuffers when encoding an image.
    24		BufferPool EncoderBufferPool
    25	}
    26	
    27	// EncoderBufferPool is an interface for getting and returning temporary
    28	// instances of the EncoderBuffer struct. This can be used to reuse buffers
    29	// when encoding multiple images.
    30	type EncoderBufferPool interface {
    31		Get() *EncoderBuffer
    32		Put(*EncoderBuffer)
    33	}
    34	
    35	// EncoderBuffer holds the buffers used for encoding PNG images.
    36	type EncoderBuffer encoder
    37	
    38	type encoder struct {
    39		enc     *Encoder
    40		w       io.Writer
    41		m       image.Image
    42		cb      int
    43		err     error
    44		header  [8]byte
    45		footer  [4]byte
    46		tmp     [4 * 256]byte
    47		cr      [nFilter][]uint8
    48		pr      []uint8
    49		zw      *zlib.Writer
    50		zwLevel int
    51		bw      *bufio.Writer
    52	}
    53	
    54	type CompressionLevel int
    55	
    56	const (
    57		DefaultCompression CompressionLevel = 0
    58		NoCompression      CompressionLevel = -1
    59		BestSpeed          CompressionLevel = -2
    60		BestCompression    CompressionLevel = -3
    61	
    62		// Positive CompressionLevel values are reserved to mean a numeric zlib
    63		// compression level, although that is not implemented yet.
    64	)
    65	
    66	type opaquer interface {
    67		Opaque() bool
    68	}
    69	
    70	// Returns whether or not the image is fully opaque.
    71	func opaque(m image.Image) bool {
    72		if o, ok := m.(opaquer); ok {
    73			return o.Opaque()
    74		}
    75		b := m.Bounds()
    76		for y := b.Min.Y; y < b.Max.Y; y++ {
    77			for x := b.Min.X; x < b.Max.X; x++ {
    78				_, _, _, a := m.At(x, y).RGBA()
    79				if a != 0xffff {
    80					return false
    81				}
    82			}
    83		}
    84		return true
    85	}
    86	
    87	// The absolute value of a byte interpreted as a signed int8.
    88	func abs8(d uint8) int {
    89		if d < 128 {
    90			return int(d)
    91		}
    92		return 256 - int(d)
    93	}
    94	
    95	func (e *encoder) writeChunk(b []byte, name string) {
    96		if e.err != nil {
    97			return
    98		}
    99		n := uint32(len(b))
   100		if int(n) != len(b) {
   101			e.err = UnsupportedError(name + " chunk is too large: " + strconv.Itoa(len(b)))
   102			return
   103		}
   104		binary.BigEndian.PutUint32(e.header[:4], n)
   105		e.header[4] = name[0]
   106		e.header[5] = name[1]
   107		e.header[6] = name[2]
   108		e.header[7] = name[3]
   109		crc := crc32.NewIEEE()
   110		crc.Write(e.header[4:8])
   111		crc.Write(b)
   112		binary.BigEndian.PutUint32(e.footer[:4], crc.Sum32())
   113	
   114		_, e.err = e.w.Write(e.header[:8])
   115		if e.err != nil {
   116			return
   117		}
   118		_, e.err = e.w.Write(b)
   119		if e.err != nil {
   120			return
   121		}
   122		_, e.err = e.w.Write(e.footer[:4])
   123	}
   124	
   125	func (e *encoder) writeIHDR() {
   126		b := e.m.Bounds()
   127		binary.BigEndian.PutUint32(e.tmp[0:4], uint32(b.Dx()))
   128		binary.BigEndian.PutUint32(e.tmp[4:8], uint32(b.Dy()))
   129		// Set bit depth and color type.
   130		switch e.cb {
   131		case cbG8:
   132			e.tmp[8] = 8
   133			e.tmp[9] = ctGrayscale
   134		case cbTC8:
   135			e.tmp[8] = 8
   136			e.tmp[9] = ctTrueColor
   137		case cbP8:
   138			e.tmp[8] = 8
   139			e.tmp[9] = ctPaletted
   140		case cbP4:
   141			e.tmp[8] = 4
   142			e.tmp[9] = ctPaletted
   143		case cbP2:
   144			e.tmp[8] = 2
   145			e.tmp[9] = ctPaletted
   146		case cbP1:
   147			e.tmp[8] = 1
   148			e.tmp[9] = ctPaletted
   149		case cbTCA8:
   150			e.tmp[8] = 8
   151			e.tmp[9] = ctTrueColorAlpha
   152		case cbG16:
   153			e.tmp[8] = 16
   154			e.tmp[9] = ctGrayscale
   155		case cbTC16:
   156			e.tmp[8] = 16
   157			e.tmp[9] = ctTrueColor
   158		case cbTCA16:
   159			e.tmp[8] = 16
   160			e.tmp[9] = ctTrueColorAlpha
   161		}
   162		e.tmp[10] = 0 // default compression method
   163		e.tmp[11] = 0 // default filter method
   164		e.tmp[12] = 0 // non-interlaced
   165		e.writeChunk(e.tmp[:13], "IHDR")
   166	}
   167	
   168	func (e *encoder) writePLTEAndTRNS(p color.Palette) {
   169		if len(p) < 1 || len(p) > 256 {
   170			e.err = FormatError("bad palette length: " + strconv.Itoa(len(p)))
   171			return
   172		}
   173		last := -1
   174		for i, c := range p {
   175			c1 := color.NRGBAModel.Convert(c).(color.NRGBA)
   176			e.tmp[3*i+0] = c1.R
   177			e.tmp[3*i+1] = c1.G
   178			e.tmp[3*i+2] = c1.B
   179			if c1.A != 0xff {
   180				last = i
   181			}
   182			e.tmp[3*256+i] = c1.A
   183		}
   184		e.writeChunk(e.tmp[:3*len(p)], "PLTE")
   185		if last != -1 {
   186			e.writeChunk(e.tmp[3*256:3*256+1+last], "tRNS")
   187		}
   188	}
   189	
   190	// An encoder is an io.Writer that satisfies writes by writing PNG IDAT chunks,
   191	// including an 8-byte header and 4-byte CRC checksum per Write call. Such calls
   192	// should be relatively infrequent, since writeIDATs uses a bufio.Writer.
   193	//
   194	// This method should only be called from writeIDATs (via writeImage).
   195	// No other code should treat an encoder as an io.Writer.
   196	func (e *encoder) Write(b []byte) (int, error) {
   197		e.writeChunk(b, "IDAT")
   198		if e.err != nil {
   199			return 0, e.err
   200		}
   201		return len(b), nil
   202	}
   203	
   204	// Chooses the filter to use for encoding the current row, and applies it.
   205	// The return value is the index of the filter and also of the row in cr that has had it applied.
   206	func filter(cr *[nFilter][]byte, pr []byte, bpp int) int {
   207		// We try all five filter types, and pick the one that minimizes the sum of absolute differences.
   208		// This is the same heuristic that libpng uses, although the filters are attempted in order of
   209		// estimated most likely to be minimal (ftUp, ftPaeth, ftNone, ftSub, ftAverage), rather than
   210		// in their enumeration order (ftNone, ftSub, ftUp, ftAverage, ftPaeth).
   211		cdat0 := cr[0][1:]
   212		cdat1 := cr[1][1:]
   213		cdat2 := cr[2][1:]
   214		cdat3 := cr[3][1:]
   215		cdat4 := cr[4][1:]
   216		pdat := pr[1:]
   217		n := len(cdat0)
   218	
   219		// The up filter.
   220		sum := 0
   221		for i := 0; i < n; i++ {
   222			cdat2[i] = cdat0[i] - pdat[i]
   223			sum += abs8(cdat2[i])
   224		}
   225		best := sum
   226		filter := ftUp
   227	
   228		// The Paeth filter.
   229		sum = 0
   230		for i := 0; i < bpp; i++ {
   231			cdat4[i] = cdat0[i] - pdat[i]
   232			sum += abs8(cdat4[i])
   233		}
   234		for i := bpp; i < n; i++ {
   235			cdat4[i] = cdat0[i] - paeth(cdat0[i-bpp], pdat[i], pdat[i-bpp])
   236			sum += abs8(cdat4[i])
   237			if sum >= best {
   238				break
   239			}
   240		}
   241		if sum < best {
   242			best = sum
   243			filter = ftPaeth
   244		}
   245	
   246		// The none filter.
   247		sum = 0
   248		for i := 0; i < n; i++ {
   249			sum += abs8(cdat0[i])
   250			if sum >= best {
   251				break
   252			}
   253		}
   254		if sum < best {
   255			best = sum
   256			filter = ftNone
   257		}
   258	
   259		// The sub filter.
   260		sum = 0
   261		for i := 0; i < bpp; i++ {
   262			cdat1[i] = cdat0[i]
   263			sum += abs8(cdat1[i])
   264		}
   265		for i := bpp; i < n; i++ {
   266			cdat1[i] = cdat0[i] - cdat0[i-bpp]
   267			sum += abs8(cdat1[i])
   268			if sum >= best {
   269				break
   270			}
   271		}
   272		if sum < best {
   273			best = sum
   274			filter = ftSub
   275		}
   276	
   277		// The average filter.
   278		sum = 0
   279		for i := 0; i < bpp; i++ {
   280			cdat3[i] = cdat0[i] - pdat[i]/2
   281			sum += abs8(cdat3[i])
   282		}
   283		for i := bpp; i < n; i++ {
   284			cdat3[i] = cdat0[i] - uint8((int(cdat0[i-bpp])+int(pdat[i]))/2)
   285			sum += abs8(cdat3[i])
   286			if sum >= best {
   287				break
   288			}
   289		}
   290		if sum < best {
   291			filter = ftAverage
   292		}
   293	
   294		return filter
   295	}
   296	
   297	func zeroMemory(v []uint8) {
   298		for i := range v {
   299			v[i] = 0
   300		}
   301	}
   302	
   303	func (e *encoder) writeImage(w io.Writer, m image.Image, cb int, level int) error {
   304		if e.zw == nil || e.zwLevel != level {
   305			zw, err := zlib.NewWriterLevel(w, level)
   306			if err != nil {
   307				return err
   308			}
   309			e.zw = zw
   310			e.zwLevel = level
   311		} else {
   312			e.zw.Reset(w)
   313		}
   314		defer e.zw.Close()
   315	
   316		bitsPerPixel := 0
   317	
   318		switch cb {
   319		case cbG8:
   320			bitsPerPixel = 8
   321		case cbTC8:
   322			bitsPerPixel = 24
   323		case cbP8:
   324			bitsPerPixel = 8
   325		case cbP4:
   326			bitsPerPixel = 4
   327		case cbP2:
   328			bitsPerPixel = 2
   329		case cbP1:
   330			bitsPerPixel = 1
   331		case cbTCA8:
   332			bitsPerPixel = 32
   333		case cbTC16:
   334			bitsPerPixel = 48
   335		case cbTCA16:
   336			bitsPerPixel = 64
   337		case cbG16:
   338			bitsPerPixel = 16
   339		}
   340	
   341		// cr[*] and pr are the bytes for the current and previous row.
   342		// cr[0] is unfiltered (or equivalently, filtered with the ftNone filter).
   343		// cr[ft], for non-zero filter types ft, are buffers for transforming cr[0] under the
   344		// other PNG filter types. These buffers are allocated once and re-used for each row.
   345		// The +1 is for the per-row filter type, which is at cr[*][0].
   346		b := m.Bounds()
   347		sz := 1 + (bitsPerPixel*b.Dx()+7)/8
   348		for i := range e.cr {
   349			if cap(e.cr[i]) < sz {
   350				e.cr[i] = make([]uint8, sz)
   351			} else {
   352				e.cr[i] = e.cr[i][:sz]
   353			}
   354			e.cr[i][0] = uint8(i)
   355		}
   356		cr := e.cr
   357		if cap(e.pr) < sz {
   358			e.pr = make([]uint8, sz)
   359		} else {
   360			e.pr = e.pr[:sz]
   361			zeroMemory(e.pr)
   362		}
   363		pr := e.pr
   364	
   365		gray, _ := m.(*image.Gray)
   366		rgba, _ := m.(*image.RGBA)
   367		paletted, _ := m.(*image.Paletted)
   368		nrgba, _ := m.(*image.NRGBA)
   369	
   370		for y := b.Min.Y; y < b.Max.Y; y++ {
   371			// Convert from colors to bytes.
   372			i := 1
   373			switch cb {
   374			case cbG8:
   375				if gray != nil {
   376					offset := (y - b.Min.Y) * gray.Stride
   377					copy(cr[0][1:], gray.Pix[offset:offset+b.Dx()])
   378				} else {
   379					for x := b.Min.X; x < b.Max.X; x++ {
   380						c := color.GrayModel.Convert(m.At(x, y)).(color.Gray)
   381						cr[0][i] = c.Y
   382						i++
   383					}
   384				}
   385			case cbTC8:
   386				// We have previously verified that the alpha value is fully opaque.
   387				cr0 := cr[0]
   388				stride, pix := 0, []byte(nil)
   389				if rgba != nil {
   390					stride, pix = rgba.Stride, rgba.Pix
   391				} else if nrgba != nil {
   392					stride, pix = nrgba.Stride, nrgba.Pix
   393				}
   394				if stride != 0 {
   395					j0 := (y - b.Min.Y) * stride
   396					j1 := j0 + b.Dx()*4
   397					for j := j0; j < j1; j += 4 {
   398						cr0[i+0] = pix[j+0]
   399						cr0[i+1] = pix[j+1]
   400						cr0[i+2] = pix[j+2]
   401						i += 3
   402					}
   403				} else {
   404					for x := b.Min.X; x < b.Max.X; x++ {
   405						r, g, b, _ := m.At(x, y).RGBA()
   406						cr0[i+0] = uint8(r >> 8)
   407						cr0[i+1] = uint8(g >> 8)
   408						cr0[i+2] = uint8(b >> 8)
   409						i += 3
   410					}
   411				}
   412			case cbP8:
   413				if paletted != nil {
   414					offset := (y - b.Min.Y) * paletted.Stride
   415					copy(cr[0][1:], paletted.Pix[offset:offset+b.Dx()])
   416				} else {
   417					pi := m.(image.PalettedImage)
   418					for x := b.Min.X; x < b.Max.X; x++ {
   419						cr[0][i] = pi.ColorIndexAt(x, y)
   420						i += 1
   421					}
   422				}
   423	
   424			case cbP4, cbP2, cbP1:
   425				pi := m.(image.PalettedImage)
   426	
   427				var a uint8
   428				var c int
   429				for x := b.Min.X; x < b.Max.X; x++ {
   430					a = a<<uint(bitsPerPixel) | pi.ColorIndexAt(x, y)
   431					c++
   432					if c == 8/bitsPerPixel {
   433						cr[0][i] = a
   434						i += 1
   435						a = 0
   436						c = 0
   437					}
   438				}
   439				if c != 0 {
   440					for c != 8/bitsPerPixel {
   441						a = a << uint(bitsPerPixel)
   442						c++
   443					}
   444					cr[0][i] = a
   445				}
   446	
   447			case cbTCA8:
   448				if nrgba != nil {
   449					offset := (y - b.Min.Y) * nrgba.Stride
   450					copy(cr[0][1:], nrgba.Pix[offset:offset+b.Dx()*4])
   451				} else {
   452					// Convert from image.Image (which is alpha-premultiplied) to PNG's non-alpha-premultiplied.
   453					for x := b.Min.X; x < b.Max.X; x++ {
   454						c := color.NRGBAModel.Convert(m.At(x, y)).(color.NRGBA)
   455						cr[0][i+0] = c.R
   456						cr[0][i+1] = c.G
   457						cr[0][i+2] = c.B
   458						cr[0][i+3] = c.A
   459						i += 4
   460					}
   461				}
   462			case cbG16:
   463				for x := b.Min.X; x < b.Max.X; x++ {
   464					c := color.Gray16Model.Convert(m.At(x, y)).(color.Gray16)
   465					cr[0][i+0] = uint8(c.Y >> 8)
   466					cr[0][i+1] = uint8(c.Y)
   467					i += 2
   468				}
   469			case cbTC16:
   470				// We have previously verified that the alpha value is fully opaque.
   471				for x := b.Min.X; x < b.Max.X; x++ {
   472					r, g, b, _ := m.At(x, y).RGBA()
   473					cr[0][i+0] = uint8(r >> 8)
   474					cr[0][i+1] = uint8(r)
   475					cr[0][i+2] = uint8(g >> 8)
   476					cr[0][i+3] = uint8(g)
   477					cr[0][i+4] = uint8(b >> 8)
   478					cr[0][i+5] = uint8(b)
   479					i += 6
   480				}
   481			case cbTCA16:
   482				// Convert from image.Image (which is alpha-premultiplied) to PNG's non-alpha-premultiplied.
   483				for x := b.Min.X; x < b.Max.X; x++ {
   484					c := color.NRGBA64Model.Convert(m.At(x, y)).(color.NRGBA64)
   485					cr[0][i+0] = uint8(c.R >> 8)
   486					cr[0][i+1] = uint8(c.R)
   487					cr[0][i+2] = uint8(c.G >> 8)
   488					cr[0][i+3] = uint8(c.G)
   489					cr[0][i+4] = uint8(c.B >> 8)
   490					cr[0][i+5] = uint8(c.B)
   491					cr[0][i+6] = uint8(c.A >> 8)
   492					cr[0][i+7] = uint8(c.A)
   493					i += 8
   494				}
   495			}
   496	
   497			// Apply the filter.
   498			// Skip filter for NoCompression and paletted images (cbP8) as
   499			// "filters are rarely useful on palette images" and will result
   500			// in larger files (see http://www.libpng.org/pub/png/book/chapter09.html).
   501			f := ftNone
   502			if level != zlib.NoCompression && cb != cbP8 && cb != cbP4 && cb != cbP2 && cb != cbP1 {
   503				// Since we skip paletted images we don't have to worry about
   504				// bitsPerPixel not being a multiple of 8
   505				bpp := bitsPerPixel / 8
   506				f = filter(&cr, pr, bpp)
   507			}
   508	
   509			// Write the compressed bytes.
   510			if _, err := e.zw.Write(cr[f]); err != nil {
   511				return err
   512			}
   513	
   514			// The current row for y is the previous row for y+1.
   515			pr, cr[0] = cr[0], pr
   516		}
   517		return nil
   518	}
   519	
   520	// Write the actual image data to one or more IDAT chunks.
   521	func (e *encoder) writeIDATs() {
   522		if e.err != nil {
   523			return
   524		}
   525		if e.bw == nil {
   526			e.bw = bufio.NewWriterSize(e, 1<<15)
   527		} else {
   528			e.bw.Reset(e)
   529		}
   530		e.err = e.writeImage(e.bw, e.m, e.cb, levelToZlib(e.enc.CompressionLevel))
   531		if e.err != nil {
   532			return
   533		}
   534		e.err = e.bw.Flush()
   535	}
   536	
   537	// This function is required because we want the zero value of
   538	// Encoder.CompressionLevel to map to zlib.DefaultCompression.
   539	func levelToZlib(l CompressionLevel) int {
   540		switch l {
   541		case DefaultCompression:
   542			return zlib.DefaultCompression
   543		case NoCompression:
   544			return zlib.NoCompression
   545		case BestSpeed:
   546			return zlib.BestSpeed
   547		case BestCompression:
   548			return zlib.BestCompression
   549		default:
   550			return zlib.DefaultCompression
   551		}
   552	}
   553	
   554	func (e *encoder) writeIEND() { e.writeChunk(nil, "IEND") }
   555	
   556	// Encode writes the Image m to w in PNG format. Any Image may be
   557	// encoded, but images that are not image.NRGBA might be encoded lossily.
   558	func Encode(w io.Writer, m image.Image) error {
   559		var e Encoder
   560		return e.Encode(w, m)
   561	}
   562	
   563	// Encode writes the Image m to w in PNG format.
   564	func (enc *Encoder) Encode(w io.Writer, m image.Image) error {
   565		// Obviously, negative widths and heights are invalid. Furthermore, the PNG
   566		// spec section 11.2.2 says that zero is invalid. Excessively large images are
   567		// also rejected.
   568		mw, mh := int64(m.Bounds().Dx()), int64(m.Bounds().Dy())
   569		if mw <= 0 || mh <= 0 || mw >= 1<<32 || mh >= 1<<32 {
   570			return FormatError("invalid image size: " + strconv.FormatInt(mw, 10) + "x" + strconv.FormatInt(mh, 10))
   571		}
   572	
   573		var e *encoder
   574		if enc.BufferPool != nil {
   575			buffer := enc.BufferPool.Get()
   576			e = (*encoder)(buffer)
   577	
   578		}
   579		if e == nil {
   580			e = &encoder{}
   581		}
   582		if enc.BufferPool != nil {
   583			defer enc.BufferPool.Put((*EncoderBuffer)(e))
   584		}
   585	
   586		e.enc = enc
   587		e.w = w
   588		e.m = m
   589	
   590		var pal color.Palette
   591		// cbP8 encoding needs PalettedImage's ColorIndexAt method.
   592		if _, ok := m.(image.PalettedImage); ok {
   593			pal, _ = m.ColorModel().(color.Palette)
   594		}
   595		if pal != nil {
   596			if len(pal) <= 2 {
   597				e.cb = cbP1
   598			} else if len(pal) <= 4 {
   599				e.cb = cbP2
   600			} else if len(pal) <= 16 {
   601				e.cb = cbP4
   602			} else {
   603				e.cb = cbP8
   604			}
   605		} else {
   606			switch m.ColorModel() {
   607			case color.GrayModel:
   608				e.cb = cbG8
   609			case color.Gray16Model:
   610				e.cb = cbG16
   611			case color.RGBAModel, color.NRGBAModel, color.AlphaModel:
   612				if opaque(m) {
   613					e.cb = cbTC8
   614				} else {
   615					e.cb = cbTCA8
   616				}
   617			default:
   618				if opaque(m) {
   619					e.cb = cbTC16
   620				} else {
   621					e.cb = cbTCA16
   622				}
   623			}
   624		}
   625	
   626		_, e.err = io.WriteString(w, pngHeader)
   627		e.writeIHDR()
   628		if pal != nil {
   629			e.writePLTEAndTRNS(pal)
   630		}
   631		e.writeIDATs()
   632		e.writeIEND()
   633		return e.err
   634	}
   635	

View as plain text