writer.go 10.8 KB
Newer Older
Nigel Tao's avatar
Nigel Tao committed
1 2 3 4 5 6 7
// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package png

import (
8 9 10 11
	"bufio"
	"compress/zlib"
	"hash/crc32"
	"image"
12
	"image/color"
13 14
	"io"
	"strconv"
Nigel Tao's avatar
Nigel Tao committed
15 16 17
)

type encoder struct {
Nigel Tao's avatar
Nigel Tao committed
18 19 20
	w      io.Writer
	m      image.Image
	cb     int
21
	err    error
Nigel Tao's avatar
Nigel Tao committed
22 23
	header [8]byte
	footer [4]byte
24
	tmp    [4 * 256]byte
Nigel Tao's avatar
Nigel Tao committed
25 26 27 28
}

// Big-endian.
func writeUint32(b []uint8, u uint32) {
29 30 31 32
	b[0] = uint8(u >> 24)
	b[1] = uint8(u >> 16)
	b[2] = uint8(u >> 8)
	b[3] = uint8(u >> 0)
Nigel Tao's avatar
Nigel Tao committed
33 34
}

35 36 37 38
type opaquer interface {
	Opaque() bool
}

Nigel Tao's avatar
Nigel Tao committed
39 40
// Returns whether or not the image is fully opaque.
func opaque(m image.Image) bool {
41 42 43
	if o, ok := m.(opaquer); ok {
		return o.Opaque()
	}
44 45 46
	b := m.Bounds()
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
47
			_, _, _, a := m.At(x, y).RGBA()
48
			if a != 0xffff {
49
				return false
Nigel Tao's avatar
Nigel Tao committed
50 51 52
			}
		}
	}
53
	return true
Nigel Tao's avatar
Nigel Tao committed
54 55
}

Nigel Tao's avatar
Nigel Tao committed
56 57 58
// The absolute value of a byte interpreted as a signed int8.
func abs8(d uint8) int {
	if d < 128 {
59
		return int(d)
Nigel Tao's avatar
Nigel Tao committed
60
	}
61
	return 256 - int(d)
Nigel Tao's avatar
Nigel Tao committed
62 63
}

Nigel Tao's avatar
Nigel Tao committed
64 65
func (e *encoder) writeChunk(b []byte, name string) {
	if e.err != nil {
66
		return
Nigel Tao's avatar
Nigel Tao committed
67
	}
68
	n := uint32(len(b))
Nigel Tao's avatar
Nigel Tao committed
69
	if int(n) != len(b) {
70 71
		e.err = UnsupportedError(name + " chunk is too large: " + strconv.Itoa(len(b)))
		return
Nigel Tao's avatar
Nigel Tao committed
72
	}
73
	writeUint32(e.header[:4], n)
74 75 76 77 78 79 80
	e.header[4] = name[0]
	e.header[5] = name[1]
	e.header[6] = name[2]
	e.header[7] = name[3]
	crc := crc32.NewIEEE()
	crc.Write(e.header[4:8])
	crc.Write(b)
81
	writeUint32(e.footer[:4], crc.Sum32())
Nigel Tao's avatar
Nigel Tao committed
82

83
	_, e.err = e.w.Write(e.header[:8])
Nigel Tao's avatar
Nigel Tao committed
84
	if e.err != nil {
85
		return
Nigel Tao's avatar
Nigel Tao committed
86
	}
87
	_, e.err = e.w.Write(b)
Nigel Tao's avatar
Nigel Tao committed
88
	if e.err != nil {
89
		return
Nigel Tao's avatar
Nigel Tao committed
90
	}
91
	_, e.err = e.w.Write(e.footer[:4])
Nigel Tao's avatar
Nigel Tao committed
92 93 94
}

func (e *encoder) writeIHDR() {
95 96 97
	b := e.m.Bounds()
	writeUint32(e.tmp[0:4], uint32(b.Dx()))
	writeUint32(e.tmp[4:8], uint32(b.Dy()))
Nigel Tao's avatar
Nigel Tao committed
98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121
	// Set bit depth and color type.
	switch e.cb {
	case cbG8:
		e.tmp[8] = 8
		e.tmp[9] = ctGrayscale
	case cbTC8:
		e.tmp[8] = 8
		e.tmp[9] = ctTrueColor
	case cbP8:
		e.tmp[8] = 8
		e.tmp[9] = ctPaletted
	case cbTCA8:
		e.tmp[8] = 8
		e.tmp[9] = ctTrueColorAlpha
	case cbG16:
		e.tmp[8] = 16
		e.tmp[9] = ctGrayscale
	case cbTC16:
		e.tmp[8] = 16
		e.tmp[9] = ctTrueColor
	case cbTCA16:
		e.tmp[8] = 16
		e.tmp[9] = ctTrueColorAlpha
	}
122 123 124
	e.tmp[10] = 0 // default compression method
	e.tmp[11] = 0 // default filter method
	e.tmp[12] = 0 // non-interlaced
125
	e.writeChunk(e.tmp[:13], "IHDR")
Nigel Tao's avatar
Nigel Tao committed
126 127
}

128
func (e *encoder) writePLTEAndTRNS(p color.Palette) {
Nigel Tao's avatar
Nigel Tao committed
129
	if len(p) < 1 || len(p) > 256 {
130 131
		e.err = FormatError("bad palette length: " + strconv.Itoa(len(p)))
		return
Nigel Tao's avatar
Nigel Tao committed
132
	}
133 134
	last := -1
	for i, c := range p {
135 136 137 138 139
		c1 := color.NRGBAModel.Convert(c).(color.NRGBA)
		e.tmp[3*i+0] = c1.R
		e.tmp[3*i+1] = c1.G
		e.tmp[3*i+2] = c1.B
		if c1.A != 0xff {
140 141
			last = i
		}
142
		e.tmp[3*256+i] = c1.A
143
	}
144 145 146
	e.writeChunk(e.tmp[:3*len(p)], "PLTE")
	if last != -1 {
		e.writeChunk(e.tmp[3*256:3*256+1+last], "tRNS")
147 148 149
	}
}

Nigel Tao's avatar
Nigel Tao committed
150 151 152 153 154 155
// An encoder is an io.Writer that satisfies writes by writing PNG IDAT chunks,
// including an 8-byte header and 4-byte CRC checksum per Write call. Such calls
// should be relatively infrequent, since writeIDATs uses a bufio.Writer.
//
// This method should only be called from writeIDATs (via writeImage).
// No other code should treat an encoder as an io.Writer.
156
func (e *encoder) Write(b []byte) (int, error) {
157
	e.writeChunk(b, "IDAT")
Nigel Tao's avatar
Nigel Tao committed
158
	if e.err != nil {
159
		return 0, e.err
Nigel Tao's avatar
Nigel Tao committed
160
	}
161
	return len(b), nil
Nigel Tao's avatar
Nigel Tao committed
162 163 164
}

// Chooses the filter to use for encoding the current row, and applies it.
Nigel Tao's avatar
Nigel Tao committed
165
// The return value is the index of the filter and also of the row in cr that has had it applied.
166
func filter(cr *[nFilter][]byte, pr []byte, bpp int) int {
Nigel Tao's avatar
Nigel Tao committed
167 168 169 170
	// We try all five filter types, and pick the one that minimizes the sum of absolute differences.
	// This is the same heuristic that libpng uses, although the filters are attempted in order of
	// estimated most likely to be minimal (ftUp, ftPaeth, ftNone, ftSub, ftAverage), rather than
	// in their enumeration order (ftNone, ftSub, ftUp, ftAverage, ftPaeth).
171 172 173 174 175 176 177
	cdat0 := cr[0][1:]
	cdat1 := cr[1][1:]
	cdat2 := cr[2][1:]
	cdat3 := cr[3][1:]
	cdat4 := cr[4][1:]
	pdat := pr[1:]
	n := len(cdat0)
Nigel Tao's avatar
Nigel Tao committed
178 179

	// The up filter.
180
	sum := 0
Nigel Tao's avatar
Nigel Tao committed
181
	for i := 0; i < n; i++ {
182 183
		cdat2[i] = cdat0[i] - pdat[i]
		sum += abs8(cdat2[i])
Nigel Tao's avatar
Nigel Tao committed
184
	}
185 186
	best := sum
	filter := ftUp
Nigel Tao's avatar
Nigel Tao committed
187 188

	// The Paeth filter.
189
	sum = 0
Nigel Tao's avatar
Nigel Tao committed
190
	for i := 0; i < bpp; i++ {
191 192
		cdat4[i] = cdat0[i] - paeth(0, pdat[i], 0)
		sum += abs8(cdat4[i])
Nigel Tao's avatar
Nigel Tao committed
193 194
	}
	for i := bpp; i < n; i++ {
195 196
		cdat4[i] = cdat0[i] - paeth(cdat0[i-bpp], pdat[i], pdat[i-bpp])
		sum += abs8(cdat4[i])
Nigel Tao's avatar
Nigel Tao committed
197
		if sum >= best {
198
			break
Nigel Tao's avatar
Nigel Tao committed
199 200 201
		}
	}
	if sum < best {
202 203
		best = sum
		filter = ftPaeth
Nigel Tao's avatar
Nigel Tao committed
204 205 206
	}

	// The none filter.
207
	sum = 0
Nigel Tao's avatar
Nigel Tao committed
208
	for i := 0; i < n; i++ {
209
		sum += abs8(cdat0[i])
Nigel Tao's avatar
Nigel Tao committed
210
		if sum >= best {
211
			break
Nigel Tao's avatar
Nigel Tao committed
212 213 214
		}
	}
	if sum < best {
215 216
		best = sum
		filter = ftNone
Nigel Tao's avatar
Nigel Tao committed
217 218 219
	}

	// The sub filter.
220
	sum = 0
Nigel Tao's avatar
Nigel Tao committed
221
	for i := 0; i < bpp; i++ {
222 223
		cdat1[i] = cdat0[i]
		sum += abs8(cdat1[i])
Nigel Tao's avatar
Nigel Tao committed
224 225
	}
	for i := bpp; i < n; i++ {
226 227
		cdat1[i] = cdat0[i] - cdat0[i-bpp]
		sum += abs8(cdat1[i])
Nigel Tao's avatar
Nigel Tao committed
228
		if sum >= best {
229
			break
Nigel Tao's avatar
Nigel Tao committed
230 231 232
		}
	}
	if sum < best {
233 234
		best = sum
		filter = ftSub
Nigel Tao's avatar
Nigel Tao committed
235 236 237
	}

	// The average filter.
238
	sum = 0
Nigel Tao's avatar
Nigel Tao committed
239
	for i := 0; i < bpp; i++ {
240 241
		cdat3[i] = cdat0[i] - pdat[i]/2
		sum += abs8(cdat3[i])
Nigel Tao's avatar
Nigel Tao committed
242 243
	}
	for i := bpp; i < n; i++ {
244 245
		cdat3[i] = cdat0[i] - uint8((int(cdat0[i-bpp])+int(pdat[i]))/2)
		sum += abs8(cdat3[i])
Nigel Tao's avatar
Nigel Tao committed
246
		if sum >= best {
247
			break
Nigel Tao's avatar
Nigel Tao committed
248 249 250
		}
	}
	if sum < best {
251 252
		best = sum
		filter = ftAverage
Nigel Tao's avatar
Nigel Tao committed
253 254
	}

255
	return filter
Nigel Tao's avatar
Nigel Tao committed
256 257
}

258
func writeImage(w io.Writer, m image.Image, cb int) error {
259
	zw := zlib.NewWriter(w)
260
	defer zw.Close()
Nigel Tao's avatar
Nigel Tao committed
261

262
	bpp := 0 // Bytes per pixel.
263

Nigel Tao's avatar
Nigel Tao committed
264 265
	switch cb {
	case cbG8:
Mathieu Lonjaret's avatar
Mathieu Lonjaret committed
266
		bpp = 1
Nigel Tao's avatar
Nigel Tao committed
267
	case cbTC8:
268
		bpp = 3
Nigel Tao's avatar
Nigel Tao committed
269
	case cbP8:
270
		bpp = 1
Nigel Tao's avatar
Nigel Tao committed
271
	case cbTCA8:
272
		bpp = 4
Nigel Tao's avatar
Nigel Tao committed
273 274 275 276 277 278
	case cbTC16:
		bpp = 6
	case cbTCA16:
		bpp = 8
	case cbG16:
		bpp = 2
Nigel Tao's avatar
Nigel Tao committed
279
	}
Nigel Tao's avatar
Nigel Tao committed
280 281 282 283 284
	// cr[*] and pr are the bytes for the current and previous row.
	// cr[0] is unfiltered (or equivalently, filtered with the ftNone filter).
	// cr[ft], for non-zero filter types ft, are buffers for transforming cr[0] under the
	// other PNG filter types. These buffers are allocated once and re-used for each row.
	// The +1 is for the per-row filter type, which is at cr[*][0].
285
	b := m.Bounds()
286
	var cr [nFilter][]uint8
287
	for i := range cr {
288
		cr[i] = make([]uint8, 1+bpp*b.Dx())
289
		cr[i][0] = uint8(i)
Nigel Tao's avatar
Nigel Tao committed
290
	}
291
	pr := make([]uint8, 1+bpp*b.Dx())
Nigel Tao's avatar
Nigel Tao committed
292

293 294 295 296 297
	gray, _ := m.(*image.Gray)
	rgba, _ := m.(*image.RGBA)
	paletted, _ := m.(*image.Paletted)
	nrgba, _ := m.(*image.NRGBA)

298
	for y := b.Min.Y; y < b.Max.Y; y++ {
Nigel Tao's avatar
Nigel Tao committed
299
		// Convert from colors to bytes.
300
		i := 1
Nigel Tao's avatar
Nigel Tao committed
301 302
		switch cb {
		case cbG8:
303 304 305 306 307 308 309 310 311
			if gray != nil {
				offset := (y - b.Min.Y) * gray.Stride
				copy(cr[0][1:], gray.Pix[offset:offset+b.Dx()])
			} else {
				for x := b.Min.X; x < b.Max.X; x++ {
					c := color.GrayModel.Convert(m.At(x, y)).(color.Gray)
					cr[0][i] = c.Y
					i++
				}
Mathieu Lonjaret's avatar
Mathieu Lonjaret committed
312
			}
Nigel Tao's avatar
Nigel Tao committed
313
		case cbTC8:
314 315
			// We have previously verified that the alpha value is fully opaque.
			cr0 := cr[0]
316 317 318 319 320 321 322 323
			stride, pix := 0, []byte(nil)
			if rgba != nil {
				stride, pix = rgba.Stride, rgba.Pix
			} else if nrgba != nil {
				stride, pix = nrgba.Stride, nrgba.Pix
			}
			if stride != 0 {
				j0 := (y - b.Min.Y) * stride
324 325
				j1 := j0 + b.Dx()*4
				for j := j0; j < j1; j += 4 {
326 327 328
					cr0[i+0] = pix[j+0]
					cr0[i+1] = pix[j+1]
					cr0[i+2] = pix[j+2]
329
					i += 3
330 331 332 333
				}
			} else {
				for x := b.Min.X; x < b.Max.X; x++ {
					r, g, b, _ := m.At(x, y).RGBA()
334 335 336 337
					cr0[i+0] = uint8(r >> 8)
					cr0[i+1] = uint8(g >> 8)
					cr0[i+2] = uint8(b >> 8)
					i += 3
338
				}
Nigel Tao's avatar
Nigel Tao committed
339
			}
Nigel Tao's avatar
Nigel Tao committed
340
		case cbP8:
341 342 343
			if paletted != nil {
				offset := (y - b.Min.Y) * paletted.Stride
				copy(cr[0][1:], paletted.Pix[offset:offset+b.Dx()])
344 345 346 347 348 349 350
			} else {
				pi := m.(image.PalettedImage)
				for x := b.Min.X; x < b.Max.X; x++ {
					cr[0][i] = pi.ColorIndexAt(x, y)
					i += 1
				}
			}
Nigel Tao's avatar
Nigel Tao committed
351
		case cbTCA8:
352 353 354 355 356 357 358 359 360 361 362 363 364
			if nrgba != nil {
				offset := (y - b.Min.Y) * nrgba.Stride
				copy(cr[0][1:], nrgba.Pix[offset:offset+b.Dx()*4])
			} else {
				// Convert from image.Image (which is alpha-premultiplied) to PNG's non-alpha-premultiplied.
				for x := b.Min.X; x < b.Max.X; x++ {
					c := color.NRGBAModel.Convert(m.At(x, y)).(color.NRGBA)
					cr[0][i+0] = c.R
					cr[0][i+1] = c.G
					cr[0][i+2] = c.B
					cr[0][i+3] = c.A
					i += 4
				}
Nigel Tao's avatar
Nigel Tao committed
365
			}
Nigel Tao's avatar
Nigel Tao committed
366 367
		case cbG16:
			for x := b.Min.X; x < b.Max.X; x++ {
368
				c := color.Gray16Model.Convert(m.At(x, y)).(color.Gray16)
369 370 371
				cr[0][i+0] = uint8(c.Y >> 8)
				cr[0][i+1] = uint8(c.Y)
				i += 2
Nigel Tao's avatar
Nigel Tao committed
372 373
			}
		case cbTC16:
374
			// We have previously verified that the alpha value is fully opaque.
Nigel Tao's avatar
Nigel Tao committed
375 376
			for x := b.Min.X; x < b.Max.X; x++ {
				r, g, b, _ := m.At(x, y).RGBA()
377 378 379 380 381 382 383
				cr[0][i+0] = uint8(r >> 8)
				cr[0][i+1] = uint8(r)
				cr[0][i+2] = uint8(g >> 8)
				cr[0][i+3] = uint8(g)
				cr[0][i+4] = uint8(b >> 8)
				cr[0][i+5] = uint8(b)
				i += 6
Nigel Tao's avatar
Nigel Tao committed
384 385 386 387
			}
		case cbTCA16:
			// Convert from image.Image (which is alpha-premultiplied) to PNG's non-alpha-premultiplied.
			for x := b.Min.X; x < b.Max.X; x++ {
388
				c := color.NRGBA64Model.Convert(m.At(x, y)).(color.NRGBA64)
389 390 391 392 393 394 395 396 397
				cr[0][i+0] = uint8(c.R >> 8)
				cr[0][i+1] = uint8(c.R)
				cr[0][i+2] = uint8(c.G >> 8)
				cr[0][i+3] = uint8(c.G)
				cr[0][i+4] = uint8(c.B >> 8)
				cr[0][i+5] = uint8(c.B)
				cr[0][i+6] = uint8(c.A >> 8)
				cr[0][i+7] = uint8(c.A)
				i += 8
Nigel Tao's avatar
Nigel Tao committed
398
			}
Nigel Tao's avatar
Nigel Tao committed
399 400 401
		}

		// Apply the filter.
402
		f := filter(&cr, pr, bpp)
Nigel Tao's avatar
Nigel Tao committed
403 404

		// Write the compressed bytes.
405
		if _, err := zw.Write(cr[f]); err != nil {
406
			return err
Nigel Tao's avatar
Nigel Tao committed
407 408 409
		}

		// The current row for y is the previous row for y+1.
410
		pr, cr[0] = cr[0], pr
Nigel Tao's avatar
Nigel Tao committed
411
	}
412
	return nil
Nigel Tao's avatar
Nigel Tao committed
413 414 415 416 417
}

// Write the actual image data to one or more IDAT chunks.
func (e *encoder) writeIDATs() {
	if e.err != nil {
418
		return
Nigel Tao's avatar
Nigel Tao committed
419
	}
420
	var bw *bufio.Writer
421
	bw = bufio.NewWriterSize(e, 1<<15)
Nigel Tao's avatar
Nigel Tao committed
422
	e.err = writeImage(bw, e.m, e.cb)
Nigel Tao's avatar
Nigel Tao committed
423
	if e.err != nil {
424
		return
Nigel Tao's avatar
Nigel Tao committed
425
	}
426
	e.err = bw.Flush()
Nigel Tao's avatar
Nigel Tao committed
427 428
}

429
func (e *encoder) writeIEND() { e.writeChunk(nil, "IEND") }
Nigel Tao's avatar
Nigel Tao committed
430

431 432
// Encode writes the Image m to w in PNG format. Any Image may be encoded, but
// images that are not image.NRGBA might be encoded lossily.
433
func Encode(w io.Writer, m image.Image) error {
434 435 436
	// Obviously, negative widths and heights are invalid. Furthermore, the PNG
	// spec section 11.2.2 says that zero is invalid. Excessively large images are
	// also rejected.
437
	mw, mh := int64(m.Bounds().Dx()), int64(m.Bounds().Dy())
Nigel Tao's avatar
Nigel Tao committed
438
	if mw <= 0 || mh <= 0 || mw >= 1<<32 || mh >= 1<<32 {
439
		return FormatError("invalid image size: " + strconv.FormatInt(mw, 10) + "x" + strconv.FormatInt(mh, 10))
Nigel Tao's avatar
Nigel Tao committed
440 441
	}

442 443 444
	var e encoder
	e.w = w
	e.m = m
445

446
	var pal color.Palette
447 448
	// cbP8 encoding needs PalettedImage's ColorIndexAt method.
	if _, ok := m.(image.PalettedImage); ok {
449
		pal, _ = m.ColorModel().(color.Palette)
450
	}
Nigel Tao's avatar
Nigel Tao committed
451
	if pal != nil {
Nigel Tao's avatar
Nigel Tao committed
452
		e.cb = cbP8
Mathieu Lonjaret's avatar
Mathieu Lonjaret committed
453 454
	} else {
		switch m.ColorModel() {
455
		case color.GrayModel:
Nigel Tao's avatar
Nigel Tao committed
456
			e.cb = cbG8
457
		case color.Gray16Model:
Nigel Tao's avatar
Nigel Tao committed
458
			e.cb = cbG16
459
		case color.RGBAModel, color.NRGBAModel, color.AlphaModel:
Nigel Tao's avatar
Nigel Tao committed
460 461 462 463 464
			if opaque(m) {
				e.cb = cbTC8
			} else {
				e.cb = cbTCA8
			}
Mathieu Lonjaret's avatar
Mathieu Lonjaret committed
465 466
		default:
			if opaque(m) {
Nigel Tao's avatar
Nigel Tao committed
467 468 469
				e.cb = cbTC16
			} else {
				e.cb = cbTCA16
Mathieu Lonjaret's avatar
Mathieu Lonjaret committed
470 471
			}
		}
Nigel Tao's avatar
Nigel Tao committed
472 473
	}

474 475
	_, e.err = io.WriteString(w, pngHeader)
	e.writeIHDR()
Nigel Tao's avatar
Nigel Tao committed
476
	if pal != nil {
477
		e.writePLTEAndTRNS(pal)
Nigel Tao's avatar
Nigel Tao committed
478
	}
479 480 481
	e.writeIDATs()
	e.writeIEND()
	return e.err
Nigel Tao's avatar
Nigel Tao committed
482
}