1/*
2Copyright 2014 The Perkeep Authors
3
4Licensed under the Apache License, Version 2.0 (the "License");
5you may not use this file except in compliance with the License.
6You may obtain a copy of the License at
7
8     http://www.apache.org/licenses/LICENSE-2.0
9
10Unless required by applicable law or agreed to in writing, software
11distributed under the License is distributed on an "AS IS" BASIS,
12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13See the License for the specific language governing permissions and
14limitations under the License.
15*/
16
17// Package fastjpeg uses djpeg(1), from the Independent JPEG Group's
18// (www.ijg.org) jpeg package, to quickly down-sample images on load.  It can
19// sample images by a factor of 1, 2, 4 or 8.
20// This reduces the amount of data that must be decompressed into memory when
21// the full resolution image isn't required, i.e. in the case of generating
22// thumbnails.
23package fastjpeg // import "perkeep.org/internal/images/fastjpeg"
24
25import (
26	"bytes"
27	"errors"
28	"expvar"
29	"fmt"
30	"image"
31	"image/color"
32	_ "image/jpeg"
33	"io"
34	"log"
35	"os"
36	"os/exec"
37	"strconv"
38	"sync"
39
40	"perkeep.org/pkg/buildinfo"
41
42	"go4.org/readerutil"
43)
44
45var (
46	ErrDjpegNotFound = errors.New("fastjpeg: djpeg not found in path")
47)
48
49// DjpegFailedError wraps errors returned when calling djpeg and handling its
50// response.  Used for type asserting and retrying with other jpeg decoders,
51// i.e. the standard library's jpeg.Decode.
52type DjpegFailedError struct {
53	Err error
54}
55
56func (dfe DjpegFailedError) Error() string {
57	return dfe.Err.Error()
58}
59
60// TODO(wathiede): do we need to conditionally add ".exe" on Windows? I have
61// no access to test on Windows.
62const djpegBin = "djpeg"
63
64var (
65	checkAvailability sync.Once
66	available         bool
67)
68
69var (
70	djpegSuccessVar = expvar.NewInt("fastjpeg-djpeg-success")
71	djpegFailureVar = expvar.NewInt("fastjpeg-djpeg-failure")
72	// Bytes read from djpeg subprocess
73	djpegBytesReadVar = expvar.NewInt("fastjpeg-djpeg-bytes-read")
74	// Bytes written to djpeg subprocess
75	djpegBytesWrittenVar = expvar.NewInt("fastjpeg-djpeg-bytes-written")
76)
77
78func Available() bool {
79	checkAvailability.Do(func() {
80		if ok, _ := strconv.ParseBool(os.Getenv("CAMLI_DISABLE_DJPEG")); ok {
81			log.Println("images/fastjpeg: CAMLI_DISABLE_DJPEG set in environment; disabling fastjpeg.")
82			return
83		}
84
85		if p, err := exec.LookPath(djpegBin); p != "" && err == nil {
86			available = true
87			log.Printf("images/fastjpeg: fastjpeg enabled with %s.", p)
88		}
89		if !available {
90			// TODO: also use Docker. https://github.com/perkeep/perkeep/issues/1106.
91			log.Printf("images/fastjpeg: %s not found in PATH, disabling fastjpeg.", djpegBin)
92		}
93	})
94
95	return available
96}
97
98func init() {
99	buildinfo.RegisterDjpegStatusFunc(djpegStatus)
100}
101
102func djpegStatus() string {
103	// TODO: more info: its path, whether it works, its version, etc.
104	if Available() {
105		return "djpeg available"
106	}
107	return "djpeg optimization unavailable"
108}
109
110func readPNM(buf *bytes.Buffer) (image.Image, error) {
111	var imgType, w, h int
112	nTokens, err := fmt.Fscanf(buf, "P%d\n%d %d\n255\n", &imgType, &w, &h)
113	if err != nil {
114		return nil, err
115	}
116	if nTokens != 3 {
117		hdr := buf.Bytes()
118		if len(hdr) > 100 {
119			hdr = hdr[:100]
120		}
121		return nil, fmt.Errorf("fastjpeg: Invalid PNM header: %q", hdr)
122	}
123
124	switch imgType {
125	case 5: // Gray
126		src := buf.Bytes()
127		if len(src) != w*h {
128			return nil, fmt.Errorf("fastjpeg: grayscale source buffer not sized w*h")
129		}
130		im := &image.Gray{
131			Pix:    src,
132			Stride: w,
133			Rect:   image.Rect(0, 0, w, h),
134		}
135		return im, nil
136	case 6: // RGB
137		src := buf.Bytes()
138		if len(src) != w*h*3 {
139			return nil, fmt.Errorf("fastjpeg: RGB source buffer not sized w*h*3")
140		}
141		im := image.NewRGBA(image.Rect(0, 0, w, h))
142		dst := im.Pix
143		for i := 0; i < len(src)/3; i++ {
144			dst[4*i+0] = src[3*i+0] // R
145			dst[4*i+1] = src[3*i+1] // G
146			dst[4*i+2] = src[3*i+2] // B
147			dst[4*i+3] = 255        // Alpha
148		}
149		return im, nil
150	default:
151		return nil, fmt.Errorf("fastjpeg: Unsupported PNM type P%d", imgType)
152	}
153}
154
155// Factor returns the sample factor DecodeSample should use to generate a
156// sampled image greater than or equal to sw x sh pixels given a source image
157// of w x h pixels.
158func Factor(w, h, sw, sh int) int {
159	switch {
160	case w>>3 >= sw && h>>3 >= sh:
161		return 8
162	case w>>2 >= sw && h>>2 >= sh:
163		return 4
164	case w>>1 >= sw && h>>1 >= sh:
165		return 2
166	}
167	return 1
168}
169
170// DecodeDownsample decodes JPEG data in r, down-sampling it by factor.
171// If djpeg is not found, err is ErrDjpegNotFound and r is not read from.
172// If the execution of djpeg, or decoding the resulting PNM fails, error will
173// be of type DjpegFailedError.
174func DecodeDownsample(r io.Reader, factor int) (image.Image, error) {
175	if !Available() {
176		return nil, ErrDjpegNotFound
177	}
178	switch factor {
179	case 1, 2, 4, 8:
180	default:
181		return nil, fmt.Errorf("fastjpeg: unsupported sample factor %d", factor)
182	}
183
184	buf := new(bytes.Buffer)
185	tr := io.TeeReader(r, buf)
186	ic, format, err := image.DecodeConfig(tr)
187	if err != nil {
188		return nil, err
189	}
190	if format != "jpeg" {
191		return nil, fmt.Errorf("fastjpeg: Unsupported format %q", format)
192	}
193	var bpp int
194	switch ic.ColorModel {
195	case color.YCbCrModel:
196		bpp = 4 // JPEG will decode to RGB, and we'll expand inplace to RGBA.
197	case color.GrayModel:
198		bpp = 1
199	default:
200		return nil, fmt.Errorf("fastjpeg: Unsupported thumnbnail color model %T", ic.ColorModel)
201	}
202	args := []string{djpegBin, "-scale", fmt.Sprintf("1/%d", factor)}
203	cmd := exec.Command(args[0], args[1:]...)
204	cmd.Stdin = readerutil.NewStatsReader(djpegBytesWrittenVar, io.MultiReader(buf, r))
205
206	// Allocate space for the RGBA / Gray pixel data plus some extra for PNM
207	// header info.  Explicitly allocate all the memory upfront to prevent
208	// many smaller allocations.
209	pixSize := ic.Width*ic.Height*bpp/factor/factor + 128
210	w := bytes.NewBuffer(make([]byte, 0, pixSize))
211	cmd.Stdout = w
212
213	stderrW := new(bytes.Buffer)
214	cmd.Stderr = stderrW
215	if err := cmd.Run(); err != nil {
216		djpegFailureVar.Add(1)
217		return nil, DjpegFailedError{Err: fmt.Errorf("%v: %s", err, stderrW)}
218	}
219	djpegSuccessVar.Add(1)
220	djpegBytesReadVar.Add(int64(w.Len()))
221	m, err := readPNM(w)
222	if err != nil {
223		return m, DjpegFailedError{Err: err}
224	}
225	return m, nil
226}
227