diff --git a/core/cmd/dms/commands_screenshot.go b/core/cmd/dms/commands_screenshot.go index 0ab482328..a24709048 100644 --- a/core/cmd/dms/commands_screenshot.go +++ b/core/cmd/dms/commands_screenshot.go @@ -381,7 +381,7 @@ func copyImageToClipboard(buf *screenshot.ShmBuffer, format screenshot.Format, q } default: mimeType = "image/png" - if err := screenshot.EncodePNGTagged(&data, screenshot.BufferToImageDeep(buf, pixelFormat), cicp); err != nil { + if err := screenshot.EncodeBufferPNG(&data, buf, pixelFormat, cicp); err != nil { return err } } @@ -394,7 +394,7 @@ func writeImageToStdout(buf *screenshot.ShmBuffer, format screenshot.Format, qua case screenshot.FormatJPEG: return screenshot.EncodeJPEG(os.Stdout, screenshot.BufferToImageWithFormat(buf, pixelFormat), quality) default: - return screenshot.EncodePNGTagged(os.Stdout, screenshot.BufferToImageDeep(buf, pixelFormat), cicp) + return screenshot.EncodeBufferPNG(os.Stdout, buf, pixelFormat, cicp) } } diff --git a/core/internal/screenshot/cicp.go b/core/internal/screenshot/cicp.go index dab98ca1d..ca2f7a6c9 100644 --- a/core/internal/screenshot/cicp.go +++ b/core/internal/screenshot/cicp.go @@ -3,8 +3,6 @@ package screenshot import ( "encoding/binary" "hash/crc32" - "image" - "io" "github.com/AvengeMedia/DankMaterialShell/core/internal/proto/wp_color_management" ) @@ -77,6 +75,13 @@ func cicpFromNamed(primaries, tf uint32) *CICP { return &CICP{Primaries: p, TF: t} } +func (c *CICP) chunks() [][]byte { + if c == nil { + return nil + } + return [][]byte{c.chunk()} +} + func (c *CICP) chunk() []byte { payload := []byte{'c', 'I', 'C', 'P', c.Primaries, c.TF, 0, 1} // matrix=RGB, full range chunk := binary.BigEndian.AppendUint32(make([]byte, 0, 16), 4) @@ -84,13 +89,6 @@ func (c *CICP) chunk() []byte { return binary.BigEndian.AppendUint32(chunk, crc32.ChecksumIEEE(payload)) } -func EncodePNGTagged(w io.Writer, img image.Image, cicp *CICP) error { - if cicp == nil { - return encodePNG(w, img) - } - return encodePNG(w, img, cicp.chunk()) -} - // outputCICP queries the output's image description via wp_color_management_v1. // Returns nil unless the output uses a wide-gamut or HDR encoding. func (s *Screenshoter) outputCICP(output *WaylandOutput) *CICP { diff --git a/core/internal/screenshot/encode.go b/core/internal/screenshot/encode.go index 33fd5fd1b..f6723ccf5 100644 --- a/core/internal/screenshot/encode.go +++ b/core/internal/screenshot/encode.go @@ -47,39 +47,6 @@ func tenBitToRGBA(buf *ShmBuffer, format uint32) *image.RGBA { return img } -// BufferToImageDeep preserves 10-bit sources as 16-bit-per-channel; others decode to 8-bit RGBA. -func BufferToImageDeep(buf *ShmBuffer, format uint32) image.Image { - if !PixelFormat(format).Is10Bit() { - return BufferToImageWithFormat(buf, format) - } - - img := image.NewNRGBA64(image.Rect(0, 0, buf.Width, buf.Height)) - data := buf.Data() - swapRB := tenBitSwapRB(format) - - for y := 0; y < buf.Height; y++ { - srcOff := y * buf.Stride - dstOff := y * img.Stride - for x := 0; x < buf.Width; x++ { - si := srcOff + x*4 - di := dstOff + x*8 - if si+4 > len(data) || di+8 > len(img.Pix) { - continue - } - v := binary.LittleEndian.Uint32(data[si:]) - c0, c1, c2 := v&0x3FF, (v>>10)&0x3FF, (v>>20)&0x3FF - if swapRB { - c0, c2 = c2, c0 - } - binary.BigEndian.PutUint16(img.Pix[di+0:], uint16(c0<<6|c0>>4)) - binary.BigEndian.PutUint16(img.Pix[di+2:], uint16(c1<<6|c1>>4)) - binary.BigEndian.PutUint16(img.Pix[di+4:], uint16(c2<<6|c2>>4)) - binary.BigEndian.PutUint16(img.Pix[di+6:], 0xFFFF) - } - } - return img -} - func BufferToImageWithFormat(buf *ShmBuffer, format uint32) *image.RGBA { if PixelFormat(format).Is10Bit() { return tenBitToRGBA(buf, format) @@ -172,10 +139,6 @@ func ImageToBuffer(img image.Image) (*ShmBuffer, error) { return buf, nil } -func EncodePNG(w io.Writer, img image.Image) error { - return encodePNG(w, img) -} - func EncodeJPEG(w io.Writer, img image.Image, quality int) error { return jpeg.Encode(w, img, &jpeg.Options{Quality: quality}) } @@ -251,6 +214,6 @@ func WriteToFileWithFormat(buf *ShmBuffer, path string, format Format, quality i case FormatPPM: return EncodePPM(f, BufferToImageWithFormat(buf, pixelFormat)) default: - return EncodePNGTagged(f, BufferToImageDeep(buf, pixelFormat), cicp) + return EncodeBufferPNG(f, buf, pixelFormat, cicp) } } diff --git a/core/internal/screenshot/encode_test.go b/core/internal/screenshot/encode_test.go index e37e52eb0..4b335454a 100644 --- a/core/internal/screenshot/encode_test.go +++ b/core/internal/screenshot/encode_test.go @@ -2,7 +2,6 @@ package screenshot import ( "encoding/binary" - "image" "testing" ) @@ -44,27 +43,3 @@ func TestBufferToImageWithFormat10Bit(t *testing.T) { }) } } - -func TestBufferToImageDeep10Bit(t *testing.T) { - const r10, g10, b10 = uint32(1023), uint32(512), uint32(0) - buf := newTenBitBuffer(t, FormatXRGB2101010, r10<<20|g10<<10|b10) - - img, ok := BufferToImageDeep(buf, uint32(FormatXRGB2101010)).(*image.NRGBA64) - if !ok { - t.Fatal("expected *image.NRGBA64 for 10-bit source") - } - - want := [4]uint16{0xFFFF, 512<<6 | 512>>4, 0, 0xFFFF} - for i, w := range want { - if got := binary.BigEndian.Uint16(img.Pix[i*2:]); got != w { - t.Errorf("channel %d = %d, want %d", i, got, w) - } - } -} - -func TestBufferToImageDeep8BitFallback(t *testing.T) { - buf := newTenBitBuffer(t, FormatXRGB8888, 0) - if _, ok := BufferToImageDeep(buf, uint32(FormatXRGB8888)).(*image.RGBA); !ok { - t.Error("expected *image.RGBA for 8-bit source") - } -} diff --git a/core/internal/screenshot/png.go b/core/internal/screenshot/png.go index 0d6b8d90f..ce9b37dc3 100644 --- a/core/internal/screenshot/png.go +++ b/core/internal/screenshot/png.go @@ -25,6 +25,7 @@ const ( adlerBase = 65521 ) +// pngSource reads rows in place from pix; pack converts a row to PNG channel order, nil when it already is. type pngSource struct { pix []byte stride int @@ -32,71 +33,107 @@ type pngSource struct { depth byte colorType byte srcBpp int - dstBpp int + pack func(dst, src []byte) } -func pngSourceOf(pix []byte, stride, width, height int, depth byte, opaque bool) pngSource { - s := pngSource{pix: pix, stride: stride, width: width, height: height, depth: depth} - s.srcBpp = 4 * int(depth) / 8 - s.dstBpp, s.colorType = s.srcBpp, pngColorRGBA - if opaque { - s.dstBpp, s.colorType = 3*int(depth)/8, pngColorRGB +func (s pngSource) dstBpp() int { + channels := 4 + if s.colorType == pngColorRGB { + channels = 3 + } + return channels * int(s.depth) / 8 +} + +func packRGBX(dst, src []byte) { + for x := range len(dst) / 3 { + d, p := dst[x*3:x*3+3], src[x*4:x*4+4] + d[0], d[1], d[2] = p[0], p[1], p[2] + } +} + +func packBGRX(dst, src []byte) { + for x := range len(dst) / 3 { + d, p := dst[x*3:x*3+3], src[x*4:x*4+4] + d[0], d[1], d[2] = p[2], p[1], p[0] + } +} + +func packRGBX16(dst, src []byte) { + for x := range len(dst) / 6 { + copy(dst[x*6:x*6+6], src[x*8:x*8+6]) + } +} + +func pack2101010(swapRB bool) func(dst, src []byte) { + return func(dst, src []byte) { + for x := range len(dst) / 6 { + v := binary.LittleEndian.Uint32(src[x*4:]) + c0, c1, c2 := v&0x3FF, v>>10&0x3FF, v>>20&0x3FF + if swapRB { + c0, c2 = c2, c0 + } + d := dst[x*6 : x*6+6] + binary.BigEndian.PutUint16(d[0:], uint16(c0<<6|c0>>4)) + binary.BigEndian.PutUint16(d[2:], uint16(c1<<6|c1>>4)) + binary.BigEndian.PutUint16(d[4:], uint16(c2<<6|c2>>4)) + } } - return s } func newPNGSource(img image.Image) pngSource { b := img.Bounds() + s := pngSource{width: b.Dx(), height: b.Dy(), colorType: pngColorRGBA} switch src := img.(type) { case *image.RGBA: - if src.Opaque() { - return pngSourceOf(src.Pix, src.Stride, b.Dx(), b.Dy(), 8, true) + if !src.Opaque() { + break } + s.pix, s.stride, s.depth, s.srcBpp = src.Pix, src.Stride, 8, 4 + s.colorType, s.pack = pngColorRGB, packRGBX + return s case *image.NRGBA: - return pngSourceOf(src.Pix, src.Stride, b.Dx(), b.Dy(), 8, src.Opaque()) + s.pix, s.stride, s.depth, s.srcBpp = src.Pix, src.Stride, 8, 4 + if src.Opaque() { + s.colorType, s.pack = pngColorRGB, packRGBX + } + return s case *image.NRGBA64: - return pngSourceOf(src.Pix, src.Stride, b.Dx(), b.Dy(), 16, src.Opaque()) + s.pix, s.stride, s.depth, s.srcBpp = src.Pix, src.Stride, 16, 8 + if src.Opaque() { + s.colorType, s.pack = pngColorRGB, packRGBX16 + } + return s } - nrgba := image.NewNRGBA(image.Rect(0, 0, b.Dx(), b.Dy())) + nrgba := image.NewNRGBA(image.Rect(0, 0, s.width, s.height)) draw.Draw(nrgba, nrgba.Rect, img, b.Min, draw.Src) return newPNGSource(nrgba) } -func (s pngSource) row(y int, scratch []byte) []byte { - src := s.pix[y*s.stride : y*s.stride+s.width*s.srcBpp] - switch s.dstBpp { - case s.srcBpp: - return src - case 3: - for x := 0; x < s.width; x++ { - d, p := scratch[x*3:x*3+3], src[x*4:x*4+4] - d[0], d[1], d[2] = p[0], p[1], p[2] - } +func newShmPNGSource(buf *ShmBuffer, format uint32) pngSource { + s := pngSource{pix: buf.Data(), stride: buf.Stride, width: buf.Width, height: buf.Height, depth: 8, colorType: pngColorRGB, srcBpp: 4} + switch f := PixelFormat(format); { + case f.Is10Bit(): + s.depth, s.pack = 16, pack2101010(tenBitSwapRB(format)) + case f == FormatABGR8888, f == FormatXBGR8888: + s.pack = packRGBX default: - for x := 0; x < s.width; x++ { - copy(scratch[x*s.dstBpp:(x+1)*s.dstBpp], src[x*s.srcBpp:]) - } + s.pack = packBGRX } - return scratch + return s } -func (s pngSource) filterBand(y0, y1 int) []byte { - rowLen := s.width * s.dstBpp - out := make([]byte, (rowLen+1)*(y1-y0)) - scratch := make([]byte, rowLen) - bpp := s.dstBpp - for y := y0; y < y1; y++ { - line := out[(y-y0)*(rowLen+1):][:rowLen+1] - line[0] = pngFilterSub - row := s.row(y, scratch) - copy(line[1:1+bpp], row[:bpp]) - cur, prev, dst := row[bpp:], row[:rowLen-bpp], line[1+bpp:] - prev, dst = prev[:len(cur)], dst[:len(cur)] - for i := range cur { - dst[i] = cur[i] - prev[i] - } +func (s pngSource) filterRow(y, bpp int, packed, dst []byte) { + row := s.pix[y*s.stride:][:s.width*s.srcBpp] + if s.pack != nil { + s.pack(packed, row) + row = packed + } + copy(dst[:bpp], row[:bpp]) + cur, prev, out := row[bpp:], row[:len(row)-bpp], dst[bpp:] + prev, out = prev[:len(cur)], out[:len(cur)] + for i := range cur { + out[i] = cur[i] - prev[i] } - return out } type pngBand struct { @@ -106,21 +143,30 @@ type pngBand struct { } func (s pngSource) compressBand(y0, y1 int, final bool) (pngBand, error) { - raw := s.filterBand(y0, y1) + bpp := s.dstBpp() + rowLen := s.width * bpp + line := make([]byte, rowLen+1) + line[0] = pngFilterSub + packed := make([]byte, rowLen) var buf bytes.Buffer fw, err := flate.NewWriter(&buf, flate.BestSpeed) if err != nil { return pngBand{}, err } - if _, err := fw.Write(raw); err != nil { - return pngBand{}, err + sum := adler32.New() + for y := y0; y < y1; y++ { + s.filterRow(y, bpp, packed, line[1:]) + sum.Write(line) + if _, err := fw.Write(line); err != nil { + return pngBand{}, err + } } if final { err = fw.Close() } else { err = fw.Flush() } - return pngBand{data: buf.Bytes(), adler: adler32.Checksum(raw), size: len(raw)}, err + return pngBand{data: buf.Bytes(), adler: sum.Sum32(), size: len(line) * (y1 - y0)}, err } func adlerCombine(a, b uint32, lenB int) uint32 { @@ -175,21 +221,20 @@ func (s pngSource) ihdr() []byte { return b } -// encodePNG deflates row bands in parallel into a single zlib stream; extraChunks go right after IHDR. -func encodePNG(w io.Writer, img image.Image, extraChunks ...[]byte) error { - src := newPNGSource(img) - if src.width <= 0 || src.height <= 0 { +// encode deflates row bands in parallel into a single zlib stream; extraChunks go right after IHDR. +func (s pngSource) encode(w io.Writer, extraChunks ...[]byte) error { + if s.width <= 0 || s.height <= 0 { return errors.New("png: empty image") } - bands := min(runtime.GOMAXPROCS(0), max(1, src.height/pngMinBandRows)) + bands := min(runtime.GOMAXPROCS(0), max(1, s.height/pngMinBandRows)) results := make([]pngBand, bands) errs := make([]error, bands) var wg sync.WaitGroup for b := range bands { wg.Go(func() { - y0, y1 := src.height*b/bands, src.height*(b+1)/bands - results[b], errs[b] = src.compressBand(y0, y1, b == bands-1) + y0, y1 := s.height*b/bands, s.height*(b+1)/bands + results[b], errs[b] = s.compressBand(y0, y1, b == bands-1) }) } wg.Wait() @@ -209,7 +254,7 @@ func encodePNG(w io.Writer, img image.Image, extraChunks ...[]byte) error { if _, err := w.Write(pngSignature); err != nil { return err } - if err := writeChunk(w, "IHDR", src.ihdr()); err != nil { + if err := writeChunk(w, "IHDR", s.ihdr()); err != nil { return err } for _, c := range extraChunks { @@ -222,3 +267,16 @@ func encodePNG(w io.Writer, img image.Image, extraChunks ...[]byte) error { } return writeChunk(w, "IEND") } + +func EncodePNG(w io.Writer, img image.Image) error { + return newPNGSource(img).encode(w) +} + +func EncodePNGTagged(w io.Writer, img image.Image, cicp *CICP) error { + return newPNGSource(img).encode(w, cicp.chunks()...) +} + +// EncodeBufferPNG reads the capture mapping directly; 10-bit formats are kept as 16-bit channels. +func EncodeBufferPNG(w io.Writer, buf *ShmBuffer, format uint32, cicp *CICP) error { + return newShmPNGSource(buf, format).encode(w, cicp.chunks()...) +} diff --git a/core/internal/screenshot/png_test.go b/core/internal/screenshot/png_test.go index b4a95b742..7ebc23f4c 100644 --- a/core/internal/screenshot/png_test.go +++ b/core/internal/screenshot/png_test.go @@ -2,6 +2,7 @@ package screenshot import ( "bytes" + "encoding/binary" "hash/adler32" "image" "image/color" @@ -136,3 +137,58 @@ func TestAdlerCombine(t *testing.T) { t.Error("combine with empty prefix must be identity") } } + +func TestEncodeBufferPNG(t *testing.T) { + const w, h = 5, 3 + const r10, g10, b10 = uint32(1023), uint32(512), uint32(0) + cases := []struct { + name string + format PixelFormat + pixel uint32 + want color.RGBA64 + }{ + {"xrgb8888", FormatXRGB8888, 0x00FF8000, color.RGBA64{0xFFFF, 0x8080, 0, 0xFFFF}}, + {"argb8888-alpha-ignored", FormatARGB8888, 0x10FF8000, color.RGBA64{0xFFFF, 0x8080, 0, 0xFFFF}}, + {"xbgr8888", FormatXBGR8888, 0x000080FF, color.RGBA64{0xFFFF, 0x8080, 0, 0xFFFF}}, + {"xrgb2101010", FormatXRGB2101010, r10<<20 | g10<<10 | b10, color.RGBA64{0xFFFF, 512<<6 | 512>>4, 0, 0xFFFF}}, + {"xbgr2101010", FormatXBGR2101010, b10<<20 | g10<<10 | r10, color.RGBA64{0xFFFF, 512<<6 | 512>>4, 0, 0xFFFF}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + buf, err := CreateShmBuffer(w, h, w*4+8) + if err != nil { + t.Fatal(err) + } + defer buf.Close() + for y := range h { + for x := range w { + binary.LittleEndian.PutUint32(buf.Data()[y*buf.Stride+x*4:], tc.pixel) + } + } + + var out bytes.Buffer + if err := EncodeBufferPNG(&out, buf, uint32(tc.format), nil); err != nil { + t.Fatal(err) + } + data := out.Bytes() + img, err := png.Decode(bytes.NewReader(data)) + if err != nil { + t.Fatal(err) + } + if img.Bounds().Dx() != w || img.Bounds().Dy() != h { + t.Fatalf("bounds = %v", img.Bounds()) + } + got := color.RGBA64Model.Convert(img.At(w-1, h-1)).(color.RGBA64) + if got != tc.want { + t.Errorf("pixel = %v, want %v", got, tc.want) + } + wantDepth := byte(8) + if tc.format.Is10Bit() { + wantDepth = 16 + } + if depth := data[24]; depth != wantDepth { + t.Errorf("bit depth = %d, want %d", depth, wantDepth) + } + }) + } +}