image/jpeg: encode *image.Gray as grayscale JPEGs.

Fixes #8201.

LGTM=nigeltao
R=nigeltao
CC=golang-codereviews
https://golang.org/cl/105990046

src/pkg/image/jpeg/writer.go

@@ -312,32 +312,44 @@ func (e *encoder) writeDQT() {
 }
 
 // writeSOF0 writes the Start Of Frame (Baseline) marker.
-func (e *encoder) writeSOF0(size image.Point) {
-	const markerlen = 8 + 3*nColorComponent
+func (e *encoder) writeSOF0(size image.Point, nComponent int) {
+	markerlen := 8 + 3*nComponent
 	e.writeMarkerHeader(sof0Marker, markerlen)
 	e.buf[0] = 8 // 8-bit color.
 	e.buf[1] = uint8(size.Y >> 8)
 	e.buf[2] = uint8(size.Y & 0xff)
 	e.buf[3] = uint8(size.X >> 8)
 	e.buf[4] = uint8(size.X & 0xff)
-	e.buf[5] = nColorComponent
-	for i := 0; i < nColorComponent; i++ {
-		e.buf[3*i+6] = uint8(i + 1)
-		// We use 4:2:0 chroma subsampling.
-		e.buf[3*i+7] = "\x22\x11\x11"[i]
-		e.buf[3*i+8] = "\x00\x01\x01"[i]
+	e.buf[5] = uint8(nComponent)
+	if nComponent == 1 {
+		e.buf[6] = 1
+		// No subsampling for grayscale image.
+		e.buf[7] = 0x11
+		e.buf[8] = 0x00
+	} else {
+		for i := 0; i < nComponent; i++ {
+			e.buf[3*i+6] = uint8(i + 1)
+			// We use 4:2:0 chroma subsampling.
+			e.buf[3*i+7] = "\x22\x11\x11"[i]
+			e.buf[3*i+8] = "\x00\x01\x01"[i]
+		}
 	}
-	e.write(e.buf[:3*(nColorComponent-1)+9])
+	e.write(e.buf[:3*(nComponent-1)+9])
 }
 
 // writeDHT writes the Define Huffman Table marker.
-func (e *encoder) writeDHT() {
+func (e *encoder) writeDHT(nComponent int) {
 	markerlen := 2
-	for _, s := range theHuffmanSpec {
+	specs := theHuffmanSpec[:]
+	if nComponent == 1 {
+		// Drop the Chrominance tables.
+		specs = specs[:2]
+	}
+	for _, s := range specs {
 		markerlen += 1 + 16 + len(s.value)
 	}
 	e.writeMarkerHeader(dhtMarker, markerlen)
-	for i, s := range theHuffmanSpec {
+	for i, s := range specs {
 		e.writeByte("\x00\x10\x01\x11"[i])
 		e.write(s.count[:])
 		e.write(s.value)
@@ -345,8 +357,8 @@ func (e *encoder) writeDHT() {
 }
 
 // writeBlock writes a block of pixel data using the given quantization table,
-// returning the post-quantized DC value of the DCT-transformed block.
-// b is in natural (not zig-zag) order.
+// returning the post-quantized DC value of the DCT-transformed block. b is in
+// natural (not zig-zag) order.
 func (e *encoder) writeBlock(b *block, q quantIndex, prevDC int32) int32 {
 	fdct(b)
 	// Emit the DC delta.
@@ -390,6 +402,20 @@ func toYCbCr(m image.Image, p image.Point, yBlock, cbBlock, crBlock *block) {
 	}
 }
 
+// grayToY stores the 8x8 region of m whose top-left corner is p in yBlock.
+func grayToY(m *image.Gray, p image.Point, yBlock *block) {
+	b := m.Bounds()
+	xmax := b.Max.X - 1
+	ymax := b.Max.Y - 1
+	pix := m.Pix
+	for j := 0; j < 8; j++ {
+		for i := 0; i < 8; i++ {
+			idx := m.PixOffset(min(p.X+i, xmax), min(p.Y+j, ymax))
+			yBlock[8*j+i] = int32(pix[idx])
+		}
+	}
+}
+
 // rgbaToYCbCr is a specialized version of toYCbCr for image.RGBA images.
 func rgbaToYCbCr(m *image.RGBA, p image.Point, yBlock, cbBlock, crBlock *block) {
 	b := m.Bounds()
@@ -430,7 +456,18 @@ func scale(dst *block, src *[4]block) {
 	}
 }
 
-// sosHeader is the SOS marker "\xff\xda" followed by 12 bytes:
+// sosHeaderY is the SOS marker "\xff\xda" followed by 8 bytes:
+//	- the marker length "\x00\x08",
+//	- the number of components "\x01",
+//	- component 1 uses DC table 0 and AC table 0 "\x01\x00",
+//	- the bytes "\x00\x3f\x00". Section B.2.3 of the spec says that for
+//	  sequential DCTs, those bytes (8-bit Ss, 8-bit Se, 4-bit Ah, 4-bit Al)
+//	  should be 0x00, 0x3f, 0x00<<4 | 0x00.
+var sosHeaderY = []byte{
+	0xff, 0xda, 0x00, 0x08, 0x01, 0x01, 0x00, 0x00, 0x3f, 0x00,
+}
+
+// sosHeaderYCbCr is the SOS marker "\xff\xda" followed by 12 bytes:
 //	- the marker length "\x00\x0c",
 //	- the number of components "\x03",
 //	- component 1 uses DC table 0 and AC table 0 "\x01\x00",
@@ -439,14 +476,19 @@ func scale(dst *block, src *[4]block) {
 //	- the bytes "\x00\x3f\x00". Section B.2.3 of the spec says that for
 //	  sequential DCTs, those bytes (8-bit Ss, 8-bit Se, 4-bit Ah, 4-bit Al)
 //	  should be 0x00, 0x3f, 0x00<<4 | 0x00.
-var sosHeader = []byte{
+var sosHeaderYCbCr = []byte{
 	0xff, 0xda, 0x00, 0x0c, 0x03, 0x01, 0x00, 0x02,
 	0x11, 0x03, 0x11, 0x00, 0x3f, 0x00,
 }
 
 // writeSOS writes the StartOfScan marker.
 func (e *encoder) writeSOS(m image.Image) {
-	e.write(sosHeader)
+	switch m.(type) {
+	case *image.Gray:
+		e.write(sosHeaderY)
+	default:
+		e.write(sosHeaderYCbCr)
+	}
 	var (
 		// Scratch buffers to hold the YCbCr values.
 		// The blocks are in natural (not zig-zag) order.
@@ -456,24 +498,36 @@ func (e *encoder) writeSOS(m image.Image) {
 		prevDCY, prevDCCb, prevDCCr int32
 	)
 	bounds := m.Bounds()
-	rgba, _ := m.(*image.RGBA)
-	for y := bounds.Min.Y; y < bounds.Max.Y; y += 16 {
-		for x := bounds.Min.X; x < bounds.Max.X; x += 16 {
-			for i := 0; i < 4; i++ {
-				xOff := (i & 1) * 8
-				yOff := (i & 2) * 4
-				p := image.Pt(x+xOff, y+yOff)
-				if rgba != nil {
-					rgbaToYCbCr(rgba, p, &b, &cb[i], &cr[i])
-				} else {
-					toYCbCr(m, p, &b, &cb[i], &cr[i])
-				}
+	switch m := m.(type) {
+	// TODO(wathiede): switch on m.ColorModel() instead of type.
+	case *image.Gray:
+		for y := bounds.Min.Y; y < bounds.Max.Y; y += 8 {
+			for x := bounds.Min.X; x < bounds.Max.X; x += 8 {
+				p := image.Pt(x, y)
+				grayToY(m, p, &b)
 				prevDCY = e.writeBlock(&b, 0, prevDCY)
 			}
-			scale(&b, &cb)
-			prevDCCb = e.writeBlock(&b, 1, prevDCCb)
-			scale(&b, &cr)
-			prevDCCr = e.writeBlock(&b, 1, prevDCCr)
+		}
+	default:
+		rgba, _ := m.(*image.RGBA)
+		for y := bounds.Min.Y; y < bounds.Max.Y; y += 16 {
+			for x := bounds.Min.X; x < bounds.Max.X; x += 16 {
+				for i := 0; i < 4; i++ {
+					xOff := (i & 1) * 8
+					yOff := (i & 2) * 4
+					p := image.Pt(x+xOff, y+yOff)
+					if rgba != nil {
+						rgbaToYCbCr(rgba, p, &b, &cb[i], &cr[i])
+					} else {
+						toYCbCr(m, p, &b, &cb[i], &cr[i])
+					}
+					prevDCY = e.writeBlock(&b, 0, prevDCY)
+				}
+				scale(&b, &cb)
+				prevDCCb = e.writeBlock(&b, 1, prevDCCb)
+				scale(&b, &cr)
+				prevDCCr = e.writeBlock(&b, 1, prevDCCr)
+			}
 		}
 	}
 	// Pad the last byte with 1's.
@@ -532,6 +586,13 @@ func Encode(w io.Writer, m image.Image, o *Options) error {
 			e.quant[i][j] = uint8(x)
 		}
 	}
+	// Compute number of components based on input image type.
+	nComponent := 3
+	switch m.(type) {
+	// TODO(wathiede): switch on m.ColorModel() instead of type.
+	case *image.Gray:
+		nComponent = 1
+	}
 	// Write the Start Of Image marker.
 	e.buf[0] = 0xff
 	e.buf[1] = 0xd8
@@ -539,9 +600,9 @@ func Encode(w io.Writer, m image.Image, o *Options) error {
 	// Write the quantization tables.
 	e.writeDQT()
 	// Write the image dimensions.
-	e.writeSOF0(b.Size())
+	e.writeSOF0(b.Size(), nComponent)
 	// Write the Huffman tables.
-	e.writeDHT()
+	e.writeDHT(nComponent)
 	// Write the image data.
 	e.writeSOS(m)
 	// Write the End Of Image marker.

src/pkg/image/jpeg/writer_test.go

@@ -160,6 +160,34 @@ func TestWriter(t *testing.T) {
 	}
 }
 
+// TestWriteGrayscale tests that a grayscale images survives a round-trip
+// through encode/decode cycle.
+func TestWriteGrayscale(t *testing.T) {
+	m0 := image.NewGray(image.Rect(0, 0, 32, 32))
+	for i := range m0.Pix {
+		m0.Pix[i] = uint8(i)
+	}
+	var buf bytes.Buffer
+	if err := Encode(&buf, m0, nil); err != nil {
+		t.Fatal(err)
+	}
+	m1, err := Decode(&buf)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if m0.Bounds() != m1.Bounds() {
+		t.Fatalf("bounds differ: %v and %v", m0.Bounds(), m1.Bounds())
+	}
+	if _, ok := m1.(*image.Gray); !ok {
+		t.Errorf("got %T, want *image.Gray", m1)
+	}
+	// Compare the average delta to the tolerance level.
+	want := int64(2 << 8)
+	if got := averageDelta(m0, m1); got > want {
+		t.Errorf("average delta too high; got %d, want <= %d", got, want)
+	}
+}
+
 // averageDelta returns the average delta in RGB space. The two images must
 // have the same bounds.
 func averageDelta(m0, m1 image.Image) int64 {