writer.go

     1  // Copyright 2013 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 gif
     6  
     7  import (
     8  	"bufio"
     9  	"bytes"
    10  	"compress/lzw"
    11  	"errors"
    12  	"image"
    13  	"image/color"
    14  	"image/color/palette"
    15  	"image/draw"
    16  	"io"
    17  )
    18  
    19  // Graphic control extension fields.
    20  const (
    21  	gcLabel     = 0xF9
    22  	gcBlockSize = 0x04
    23  )
    24  
    25  var log2Lookup = [8]int{2, 4, 8, 16, 32, 64, 128, 256}
    26  
    27  func log2(x int) int {
    28  	for i, v := range log2Lookup {
    29  		if x <= v {
    30  			return i
    31  		}
    32  	}
    33  	return -1
    34  }
    35  
    36  // Little-endian.
    37  func writeUint16(b []uint8, u uint16) {
    38  	b[0] = uint8(u)
    39  	b[1] = uint8(u >> 8)
    40  }
    41  
    42  // writer is a buffered writer.
    43  type writer interface {
    44  	Flush() error
    45  	io.Writer
    46  	io.ByteWriter
    47  }
    48  
    49  // encoder encodes an image to the GIF format.
    50  type encoder struct {
    51  	// w is the writer to write to. err is the first error encountered during
    52  	// writing. All attempted writes after the first error become no-ops.
    53  	w   writer
    54  	err error
    55  	// g is a reference to the data that is being encoded.
    56  	g GIF
    57  	// globalCT is the size in bytes of the global color table.
    58  	globalCT int
    59  	// buf is a scratch buffer. It must be at least 256 for the blockWriter.
    60  	buf              [256]byte
    61  	globalColorTable [3 * 256]byte
    62  	localColorTable  [3 * 256]byte
    63  }
    64  
    65  // blockWriter writes the block structure of GIF image data, which
    66  // comprises (n, (n bytes)) blocks, with 1 <= n <= 255. It is the
    67  // writer given to the LZW encoder, which is thus immune to the
    68  // blocking.
    69  type blockWriter struct {
    70  	e *encoder
    71  }
    72  
    73  func (b blockWriter) setup() {
    74  	b.e.buf[0] = 0
    75  }
    76  
    77  func (b blockWriter) Flush() error {
    78  	return b.e.err
    79  }
    80  
    81  func (b blockWriter) WriteByte(c byte) error {
    82  	if b.e.err != nil {
    83  		return b.e.err
    84  	}
    85  
    86  	// Append c to buffered sub-block.
    87  	b.e.buf[0]++
    88  	b.e.buf[b.e.buf[0]] = c
    89  	if b.e.buf[0] < 255 {
    90  		return nil
    91  	}
    92  
    93  	// Flush block
    94  	b.e.write(b.e.buf[:256])
    95  	b.e.buf[0] = 0
    96  	return b.e.err
    97  }
    98  
    99  // blockWriter must be an io.Writer for lzw.NewWriter, but this is never
   100  // actually called.
   101  func (b blockWriter) Write(data []byte) (int, error) {
   102  	for i, c := range data {
   103  		if err := b.WriteByte(c); err != nil {
   104  			return i, err
   105  		}
   106  	}
   107  	return len(data), nil
   108  }
   109  
   110  func (b blockWriter) close() {
   111  	// Write the block terminator (0x00), either by itself, or along with a
   112  	// pending sub-block.
   113  	if b.e.buf[0] == 0 {
   114  		b.e.writeByte(0)
   115  	} else {
   116  		n := uint(b.e.buf[0])
   117  		b.e.buf[n+1] = 0
   118  		b.e.write(b.e.buf[:n+2])
   119  	}
   120  	b.e.flush()
   121  }
   122  
   123  func (e *encoder) flush() {
   124  	if e.err != nil {
   125  		return
   126  	}
   127  	e.err = e.w.Flush()
   128  }
   129  
   130  func (e *encoder) write(p []byte) {
   131  	if e.err != nil {
   132  		return
   133  	}
   134  	_, e.err = e.w.Write(p)
   135  }
   136  
   137  func (e *encoder) writeByte(b byte) {
   138  	if e.err != nil {
   139  		return
   140  	}
   141  	e.err = e.w.WriteByte(b)
   142  }
   143  
   144  func (e *encoder) writeHeader() {
   145  	if e.err != nil {
   146  		return
   147  	}
   148  	_, e.err = io.WriteString(e.w, "GIF89a")
   149  	if e.err != nil {
   150  		return
   151  	}
   152  
   153  	// Logical screen width and height.
   154  	writeUint16(e.buf[0:2], uint16(e.g.Config.Width))
   155  	writeUint16(e.buf[2:4], uint16(e.g.Config.Height))
   156  	e.write(e.buf[:4])
   157  
   158  	if p, ok := e.g.Config.ColorModel.(color.Palette); ok && len(p) > 0 {
   159  		paddedSize := log2(len(p)) // Size of Global Color Table: 2^(1+n).
   160  		e.buf[0] = fColorTable | uint8(paddedSize)
   161  		e.buf[1] = e.g.BackgroundIndex
   162  		e.buf[2] = 0x00 // Pixel Aspect Ratio.
   163  		e.write(e.buf[:3])
   164  		var err error
   165  		e.globalCT, err = encodeColorTable(e.globalColorTable[:], p, paddedSize)
   166  		if err != nil && e.err == nil {
   167  			e.err = err
   168  			return
   169  		}
   170  		e.write(e.globalColorTable[:e.globalCT])
   171  	} else {
   172  		// All frames have a local color table, so a global color table
   173  		// is not needed.
   174  		e.buf[0] = 0x00
   175  		e.buf[1] = 0x00 // Background Color Index.
   176  		e.buf[2] = 0x00 // Pixel Aspect Ratio.
   177  		e.write(e.buf[:3])
   178  	}
   179  
   180  	// Add animation info if necessary.
   181  	if len(e.g.Image) > 1 && e.g.LoopCount >= 0 {
   182  		e.buf[0] = 0x21 // Extension Introducer.
   183  		e.buf[1] = 0xff // Application Label.
   184  		e.buf[2] = 0x0b // Block Size.
   185  		e.write(e.buf[:3])
   186  		_, err := io.WriteString(e.w, "NETSCAPE2.0") // Application Identifier.
   187  		if err != nil && e.err == nil {
   188  			e.err = err
   189  			return
   190  		}
   191  		e.buf[0] = 0x03 // Block Size.
   192  		e.buf[1] = 0x01 // Sub-block Index.
   193  		writeUint16(e.buf[2:4], uint16(e.g.LoopCount))
   194  		e.buf[4] = 0x00 // Block Terminator.
   195  		e.write(e.buf[:5])
   196  	}
   197  }
   198  
   199  func encodeColorTable(dst []byte, p color.Palette, size int) (int, error) {
   200  	if uint(size) >= uint(len(log2Lookup)) {
   201  		return 0, errors.New("gif: cannot encode color table with more than 256 entries")
   202  	}
   203  	for i, c := range p {
   204  		if c == nil {
   205  			return 0, errors.New("gif: cannot encode color table with nil entries")
   206  		}
   207  		var r, g, b uint8
   208  		// It is most likely that the palette is full of color.RGBAs, so they
   209  		// get a fast path.
   210  		if rgba, ok := c.(color.RGBA); ok {
   211  			r, g, b = rgba.R, rgba.G, rgba.B
   212  		} else {
   213  			rr, gg, bb, _ := c.RGBA()
   214  			r, g, b = uint8(rr>>8), uint8(gg>>8), uint8(bb>>8)
   215  		}
   216  		dst[3*i+0] = r
   217  		dst[3*i+1] = g
   218  		dst[3*i+2] = b
   219  	}
   220  	n := log2Lookup[size]
   221  	if n > len(p) {
   222  		// Pad with black.
   223  		fill := dst[3*len(p) : 3*n]
   224  		for i := range fill {
   225  			fill[i] = 0
   226  		}
   227  	}
   228  	return 3 * n, nil
   229  }
   230  
   231  func (e *encoder) colorTablesMatch(localLen, transparentIndex int) bool {
   232  	localSize := 3 * localLen
   233  	if transparentIndex >= 0 {
   234  		trOff := 3 * transparentIndex
   235  		return bytes.Equal(e.globalColorTable[:trOff], e.localColorTable[:trOff]) &&
   236  			bytes.Equal(e.globalColorTable[trOff+3:localSize], e.localColorTable[trOff+3:localSize])
   237  	}
   238  	return bytes.Equal(e.globalColorTable[:localSize], e.localColorTable[:localSize])
   239  }
   240  
   241  func (e *encoder) writeImageBlock(pm *image.Paletted, delay int, disposal byte) {
   242  	if e.err != nil {
   243  		return
   244  	}
   245  
   246  	if len(pm.Palette) == 0 {
   247  		e.err = errors.New("gif: cannot encode image block with empty palette")
   248  		return
   249  	}
   250  
   251  	b := pm.Bounds()
   252  	if b.Min.X < 0 || b.Max.X >= 1<<16 || b.Min.Y < 0 || b.Max.Y >= 1<<16 {
   253  		e.err = errors.New("gif: image block is too large to encode")
   254  		return
   255  	}
   256  	if !b.In(image.Rectangle{Max: image.Point{e.g.Config.Width, e.g.Config.Height}}) {
   257  		e.err = errors.New("gif: image block is out of bounds")
   258  		return
   259  	}
   260  
   261  	transparentIndex := -1
   262  	for i, c := range pm.Palette {
   263  		if c == nil {
   264  			e.err = errors.New("gif: cannot encode color table with nil entries")
   265  			return
   266  		}
   267  		if _, _, _, a := c.RGBA(); a == 0 {
   268  			transparentIndex = i
   269  			break
   270  		}
   271  	}
   272  
   273  	if delay > 0 || disposal != 0 || transparentIndex != -1 {
   274  		e.buf[0] = sExtension  // Extension Introducer.
   275  		e.buf[1] = gcLabel     // Graphic Control Label.
   276  		e.buf[2] = gcBlockSize // Block Size.
   277  		if transparentIndex != -1 {
   278  			e.buf[3] = 0x01 | disposal<<2
   279  		} else {
   280  			e.buf[3] = 0x00 | disposal<<2
   281  		}
   282  		writeUint16(e.buf[4:6], uint16(delay)) // Delay Time (1/100ths of a second)
   283  
   284  		// Transparent color index.
   285  		if transparentIndex != -1 {
   286  			e.buf[6] = uint8(transparentIndex)
   287  		} else {
   288  			e.buf[6] = 0x00
   289  		}
   290  		e.buf[7] = 0x00 // Block Terminator.
   291  		e.write(e.buf[:8])
   292  	}
   293  	e.buf[0] = sImageDescriptor
   294  	writeUint16(e.buf[1:3], uint16(b.Min.X))
   295  	writeUint16(e.buf[3:5], uint16(b.Min.Y))
   296  	writeUint16(e.buf[5:7], uint16(b.Dx()))
   297  	writeUint16(e.buf[7:9], uint16(b.Dy()))
   298  	e.write(e.buf[:9])
   299  
   300  	// To determine whether or not this frame's palette is the same as the
   301  	// global palette, we can check a couple things. First, do they actually
   302  	// point to the same []color.Color? If so, they are equal so long as the
   303  	// frame's palette is not longer than the global palette...
   304  	paddedSize := log2(len(pm.Palette)) // Size of Local Color Table: 2^(1+n).
   305  	if gp, ok := e.g.Config.ColorModel.(color.Palette); ok && len(pm.Palette) <= len(gp) && &gp[0] == &pm.Palette[0] {
   306  		e.writeByte(0) // Use the global color table.
   307  	} else {
   308  		ct, err := encodeColorTable(e.localColorTable[:], pm.Palette, paddedSize)
   309  		if err != nil {
   310  			if e.err == nil {
   311  				e.err = err
   312  			}
   313  			return
   314  		}
   315  		// This frame's palette is not the very same slice as the global
   316  		// palette, but it might be a copy, possibly with one value turned into
   317  		// transparency by DecodeAll.
   318  		if ct <= e.globalCT && e.colorTablesMatch(len(pm.Palette), transparentIndex) {
   319  			e.writeByte(0) // Use the global color table.
   320  		} else {
   321  			// Use a local color table.
   322  			e.writeByte(fColorTable | uint8(paddedSize))
   323  			e.write(e.localColorTable[:ct])
   324  		}
   325  	}
   326  
   327  	litWidth := paddedSize + 1
   328  	if litWidth < 2 {
   329  		litWidth = 2
   330  	}
   331  	e.writeByte(uint8(litWidth)) // LZW Minimum Code Size.
   332  
   333  	bw := blockWriter{e: e}
   334  	bw.setup()
   335  	lzww := lzw.NewWriter(bw, lzw.LSB, litWidth)
   336  	if dx := b.Dx(); dx == pm.Stride {
   337  		_, e.err = lzww.Write(pm.Pix[:dx*b.Dy()])
   338  		if e.err != nil {
   339  			lzww.Close()
   340  			return
   341  		}
   342  	} else {
   343  		for i, y := 0, b.Min.Y; y < b.Max.Y; i, y = i+pm.Stride, y+1 {
   344  			_, e.err = lzww.Write(pm.Pix[i : i+dx])
   345  			if e.err != nil {
   346  				lzww.Close()
   347  				return
   348  			}
   349  		}
   350  	}
   351  	lzww.Close() // flush to bw
   352  	bw.close()   // flush to e.w
   353  }
   354  
   355  // Options are the encoding parameters.
   356  type Options struct {
   357  	// NumColors is the maximum number of colors used in the image.
   358  	// It ranges from 1 to 256.
   359  	NumColors int
   360  
   361  	// Quantizer is used to produce a palette with size NumColors.
   362  	// palette.Plan9 is used in place of a nil Quantizer.
   363  	Quantizer draw.Quantizer
   364  
   365  	// Drawer is used to convert the source image to the desired palette.
   366  	// draw.FloydSteinberg is used in place of a nil Drawer.
   367  	Drawer draw.Drawer
   368  }
   369  
   370  // EncodeAll writes the images in g to w in GIF format with the
   371  // given loop count and delay between frames.
   372  func EncodeAll(w io.Writer, g *GIF) error {
   373  	if len(g.Image) == 0 {
   374  		return errors.New("gif: must provide at least one image")
   375  	}
   376  
   377  	if len(g.Image) != len(g.Delay) {
   378  		return errors.New("gif: mismatched image and delay lengths")
   379  	}
   380  
   381  	e := encoder{g: *g}
   382  	// The GIF.Disposal, GIF.Config and GIF.BackgroundIndex fields were added
   383  	// in Go 1.5. Valid Go 1.4 code, such as when the Disposal field is omitted
   384  	// in a GIF struct literal, should still produce valid GIFs.
   385  	if e.g.Disposal != nil && len(e.g.Image) != len(e.g.Disposal) {
   386  		return errors.New("gif: mismatched image and disposal lengths")
   387  	}
   388  	if e.g.Config == (image.Config{}) {
   389  		p := g.Image[0].Bounds().Max
   390  		e.g.Config.Width = p.X
   391  		e.g.Config.Height = p.Y
   392  	} else if e.g.Config.ColorModel != nil {
   393  		if _, ok := e.g.Config.ColorModel.(color.Palette); !ok {
   394  			return errors.New("gif: GIF color model must be a color.Palette")
   395  		}
   396  	}
   397  
   398  	if ww, ok := w.(writer); ok {
   399  		e.w = ww
   400  	} else {
   401  		e.w = bufio.NewWriter(w)
   402  	}
   403  
   404  	e.writeHeader()
   405  	for i, pm := range g.Image {
   406  		disposal := uint8(0)
   407  		if g.Disposal != nil {
   408  			disposal = g.Disposal[i]
   409  		}
   410  		e.writeImageBlock(pm, g.Delay[i], disposal)
   411  	}
   412  	e.writeByte(sTrailer)
   413  	e.flush()
   414  	return e.err
   415  }
   416  
   417  // Encode writes the Image m to w in GIF format.
   418  func Encode(w io.Writer, m image.Image, o *Options) error {
   419  	// Check for bounds and size restrictions.
   420  	b := m.Bounds()
   421  	if b.Dx() >= 1<<16 || b.Dy() >= 1<<16 {
   422  		return errors.New("gif: image is too large to encode")
   423  	}
   424  
   425  	opts := Options{}
   426  	if o != nil {
   427  		opts = *o
   428  	}
   429  	if opts.NumColors < 1 || 256 < opts.NumColors {
   430  		opts.NumColors = 256
   431  	}
   432  	if opts.Drawer == nil {
   433  		opts.Drawer = draw.FloydSteinberg
   434  	}
   435  
   436  	pm, _ := m.(*image.Paletted)
   437  	if pm == nil {
   438  		if cp, ok := m.ColorModel().(color.Palette); ok {
   439  			pm = image.NewPaletted(b, cp)
   440  			for y := b.Min.Y; y < b.Max.Y; y++ {
   441  				for x := b.Min.X; x < b.Max.X; x++ {
   442  					pm.Set(x, y, cp.Convert(m.At(x, y)))
   443  				}
   444  			}
   445  		}
   446  	}
   447  	if pm == nil || len(pm.Palette) > opts.NumColors {
   448  		// Set pm to be a palettedized copy of m, including its bounds, which
   449  		// might not start at (0, 0).
   450  		//
   451  		// TODO: Pick a better sub-sample of the Plan 9 palette.
   452  		pm = image.NewPaletted(b, palette.Plan9[:opts.NumColors])
   453  		if opts.Quantizer != nil {
   454  			pm.Palette = opts.Quantizer.Quantize(make(color.Palette, 0, opts.NumColors), m)
   455  		}
   456  		opts.Drawer.Draw(pm, b, m, b.Min)
   457  	}
   458  
   459  	// When calling Encode instead of EncodeAll, the single-frame image is
   460  	// translated such that its top-left corner is (0, 0), so that the single
   461  	// frame completely fills the overall GIF's bounds.
   462  	if pm.Rect.Min != (image.Point{}) {
   463  		dup := *pm
   464  		dup.Rect = dup.Rect.Sub(dup.Rect.Min)
   465  		pm = &dup
   466  	}
   467  
   468  	return EncodeAll(w, &GIF{
   469  		Image: []*image.Paletted{pm},
   470  		Delay: []int{0},
   471  		Config: image.Config{
   472  			ColorModel: pm.Palette,
   473  			Width:      b.Dx(),
   474  			Height:     b.Dy(),
   475  		},
   476  	})
   477  }
   478
View as plain text