1# This file is part of Xpra.
2# Copyright (C) 2017-2021 Antoine Martin <antoine@xpra.org>
3# Xpra is released under the terms of the GNU GPL v2, or, at your option, any
4# later version. See the file COPYING for details.
5
6#cython: wraparound=False
7
8import time
9
10from xpra.util import envbool, typedict
11from xpra.log import Logger
12log = Logger("encoder", "jpeg")
13
14from libc.stdint cimport uintptr_t
15from xpra.buffers.membuf cimport makebuf, MemBuf, buffer_context    #pylint: disable=syntax-error
16
17from xpra.codecs.codec_constants import get_subsampling_divs
18from xpra.net.compression import Compressed
19from xpra.util import csv
20from xpra.os_util import bytestostr
21
22cdef int SAVE_TO_FILE = envbool("XPRA_SAVE_TO_FILE")
23
24
25ctypedef int TJSAMP
26ctypedef int TJPF
27
28cdef extern from "math.h":
29    double sqrt(double arg) nogil
30    double round(double arg) nogil
31
32cdef extern from "turbojpeg.h":
33    TJSAMP  TJSAMP_444
34    TJSAMP  TJSAMP_422
35    TJSAMP  TJSAMP_420
36    TJSAMP  TJSAMP_GRAY
37    TJSAMP  TJSAMP_440
38    TJSAMP  TJSAMP_411
39
40    TJPF    TJPF_RGB
41    TJPF    TJPF_BGR
42    TJPF    TJPF_RGBX
43    TJPF    TJPF_BGRX
44    TJPF    TJPF_XBGR
45    TJPF    TJPF_XRGB
46    TJPF    TJPF_GRAY
47    TJPF    TJPF_RGBA
48    TJPF    TJPF_BGRA
49    TJPF    TJPF_ABGR
50    TJPF    TJPF_ARGB
51    TJPF    TJPF_CMYK
52
53    int TJFLAG_BOTTOMUP
54    int TJFLAG_FASTUPSAMPLE
55    int TJFLAG_FASTDCT
56    int TJFLAG_ACCURATEDCT
57
58    ctypedef void* tjhandle
59    tjhandle tjInitCompress()
60    int tjDestroy(tjhandle handle)
61    char* tjGetErrorStr()
62    #unsigned long tjBufSize(int width, int height, int jpegSubsamp)
63    int tjCompress2(tjhandle handle, const unsigned char *srcBuf,
64                    int width, int pitch, int height, int pixelFormat, unsigned char **jpegBuf,
65                    unsigned long *jpegSize, int jpegSubsamp, int jpegQual, int flags) nogil
66
67    int tjCompressFromYUVPlanes(tjhandle handle,
68                    const unsigned char **srcPlanes,
69                    int width, const int *strides,
70                    int height, int subsamp,
71                    unsigned char **jpegBuf,
72                    unsigned long *jpegSize, int jpegQual,
73                    int flags) nogil
74
75TJPF_VAL = {
76    "RGB"   : TJPF_RGB,
77    "BGR"   : TJPF_BGR,
78    "RGBX"  : TJPF_RGBX,
79    "BGRX"  : TJPF_BGRX,
80    "XBGR"  : TJPF_XBGR,
81    "XRGB"  : TJPF_XRGB,
82    "GRAY"  : TJPF_GRAY,
83    "RGBA"  : TJPF_RGBA,
84    "BGRA"  : TJPF_BGRA,
85    "ABGR"  : TJPF_ABGR,
86    "ARGB"  : TJPF_ARGB,
87    "CMYK"  : TJPF_CMYK,
88    }
89TJSAMP_STR = {
90    TJSAMP_444  : "444",
91    TJSAMP_422  : "422",
92    TJSAMP_420  : "420",
93    TJSAMP_GRAY : "GRAY",
94    TJSAMP_440  : "440",
95    TJSAMP_411  : "411",
96    }
97
98
99def get_version():
100    return (2, 0)
101
102def get_type():
103    return "jpeg"
104
105def get_info():
106    return {"version"   : get_version()}
107
108def get_encodings():
109    return ("jpeg", "jpega")
110
111def init_module():
112    log("jpeg.init_module()")
113
114def cleanup_module():
115    log("jpeg.cleanup_module()")
116
117def get_input_colorspaces(encoding):
118    if encoding=="jpeg":
119        return ("BGRX", "RGBX", "XBGR", "XRGB", "RGB", "BGR", "YUV420P", "YUV422P", "YUV444P")
120    assert encoding=="jpega"
121    return ("BGRA", "RGBA", )
122
123def get_output_colorspaces(encoding, input_colorspace):
124    assert encoding in get_encodings()
125    assert input_colorspace in get_input_colorspaces(encoding)
126    return (input_colorspace, )
127
128def get_spec(encoding, colorspace):
129    assert encoding in ("jpeg", "jpega")
130    assert colorspace in get_input_colorspaces(encoding)
131    from xpra.codecs.codec_constants import video_spec
132    width_mask=0xFFFF
133    height_mask=0xFFFF
134    if colorspace in ("YUV420P", "YUV422P"):
135        width_mask=0xFFFE
136    if colorspace in ("YUV420P", ):
137        height_mask=0xFFFE
138    return video_spec(encoding, input_colorspace=colorspace, output_colorspaces=(colorspace, ), has_lossless_mode=False,
139                      codec_class=Encoder, codec_type="jpeg",
140                      setup_cost=0, cpu_cost=100, gpu_cost=0,
141                      min_w=16, min_h=16, max_w=16*1024, max_h=16*1024,
142                      can_scale=False,
143                      score_boost=-50,
144                      width_mask=width_mask, height_mask=height_mask)
145
146
147cdef inline int norm_quality(int quality) nogil:
148    if quality<=0:
149        return 0
150    if quality>=100:
151        return 100
152    return <int> round(sqrt(<double> quality)*10)
153
154cdef class Encoder:
155    cdef tjhandle compressor
156    cdef int width
157    cdef int height
158    cdef object src_format
159    cdef int quality
160    cdef int grayscale
161    cdef long frames
162    cdef object __weakref__
163
164    def __init__(self):
165        self.width = self.height = self.quality = self.frames = 0
166        self.compressor = tjInitCompress()
167        if self.compressor==NULL:
168            raise Exception("Error: failed to instantiate a JPEG compressor")
169
170    def init_context(self, encoding, width : int, height : int, src_format, options : typedict):
171        assert encoding=="jpeg"
172        assert src_format in get_input_colorspaces(encoding)
173        scaled_width = options.intget("scaled-width", width)
174        scaled_height = options.intget("scaled-height", height)
175        assert scaled_width==width and scaled_height==height, "jpeg encoder does not handle scaling"
176        self.width = width
177        self.height = height
178        self.src_format = src_format
179        self.grayscale = options.boolget("grayscale")
180        self.quality = options.intget("quality", 50)
181
182    def is_ready(self):
183        return self.compressor!=NULL
184
185    def is_closed(self):
186        return self.compressor==NULL
187
188    def clean(self):
189        self.width = self.height = self.quality = 0
190        r = tjDestroy(self.compressor)
191        self.compressor = NULL
192        if r:
193            log.error("Error: failed to destroy the JPEG compressor, code %i:", r)
194            log.error(" %s", get_error_str())
195
196    def get_encoding(self):
197        return "jpeg"
198
199    def get_width(self):
200        return self.width
201
202    def get_height(self):
203        return self.height
204
205    def get_type(self):
206        return "jpeg"
207
208    def get_src_format(self):
209        return self.src_format
210
211    def get_info(self) -> dict:
212        info = get_info()
213        info.update({
214            "frames"        : int(self.frames),
215            "width"         : self.width,
216            "height"        : self.height,
217            "quality"       : self.quality,
218            "grayscale"     : bool(self.grayscale),
219            })
220        return info
221
222    def compress_image(self, image, options=None):
223        options = options or {}
224        quality = options.get("quality", -1)
225        if quality>0:
226            self.quality = quality
227        else:
228            quality = self.quality
229        pfstr = image.get_pixel_format()
230        if pfstr in ("YUV420P", "YUV422P", "YUV444P"):
231            cdata = encode_yuv(self.compressor, image, quality, self.grayscale)
232        else:
233            cdata = encode_rgb(self.compressor, image, quality, self.grayscale)
234        if not cdata:
235            return None
236        self.frames += 1
237        return memoryview(cdata), {}
238
239
240def get_error_str():
241    cdef char *err = tjGetErrorStr()
242    return bytestostr(err)
243
244JPEG_INPUT_FORMATS = ("RGB", "RGBX", "BGRX", "XBGR", "XRGB", )
245JPEGA_INPUT_FORMATS = ("RGBA", "BGRA", "ABGR", "ARGB")
246
247def encode(coding, image, options=None):
248    assert coding in ("jpeg", "jpega")
249    options = options or {}
250    rgb_format = image.get_pixel_format()
251    if coding=="jpega" and rgb_format.find("A")<0:
252        #why did we select 'jpega' then!?
253        coding = "jpeg"
254    cdef int quality = options.get("quality", 50)
255    cdef int grayscale = options.get("grayscale", 0)
256    cdef int width = image.get_width()
257    cdef int height = image.get_height()
258    cdef int scaled_width = options.get("scaled-width", width)
259    cdef int scaled_height = options.get("scaled-height", height)
260    cdef char resize = scaled_width!=width or scaled_height!=height
261    log("encode%s", (coding, image, options))
262    input_formats = JPEG_INPUT_FORMATS if coding=="jpeg" else JPEGA_INPUT_FORMATS
263    if rgb_format not in input_formats or resize and len(rgb_format)!=4:
264        from xpra.codecs.argb.argb import argb_swap         #@UnresolvedImport
265        if not argb_swap(image, input_formats):
266            log("jpeg: argb_swap failed to convert %s to a suitable format: %s" % (
267                rgb_format, input_formats))
268        log("jpeg converted %s to %s", rgb_format, image)
269
270    if resize:
271        from xpra.codecs.argb.scale import scale_image
272        image = scale_image(image, scaled_width, scaled_height)
273        log("jpeg scaled image: %s", image)
274
275    client_options = {
276        "quality"   : quality
277        }
278    cdef tjhandle compressor = tjInitCompress()
279    if compressor==NULL:
280        log.error("Error: failed to instantiate a JPEG compressor")
281        return None
282    cdef int r
283    try:
284        cdata = encode_rgb(compressor, image, quality, grayscale)
285        if not cdata:
286            return None
287        now = time.time()
288        if SAVE_TO_FILE:    # pragma: no cover
289            filename = "./%s.jpeg" % (now, )
290            with open(filename, "wb") as f:
291                f.write(cdata)
292            log.info("saved %i bytes to %s", len(cdata), filename)
293        bpp = 24
294        if coding=="jpega":
295            from xpra.codecs.argb.argb import alpha  #@UnresolvedImport
296            a = alpha(image)
297            planes = (a, )
298            rowstrides = (image.get_rowstride()//4, )
299            adata = do_encode_yuv(compressor, "YUV400P", planes,
300                                  width, height, rowstrides,
301                                  quality, TJSAMP_GRAY)
302            if SAVE_TO_FILE:    # pragma: no cover
303                filename = "./%s.jpega" % (now, )
304                with open(filename, "wb") as f:
305                    f.write(adata)
306                log.info("saved %i bytes to %s", len(adata), filename)
307            client_options["alpha-offset"] = len(cdata)
308            cdata = memoryview(cdata).tobytes()+memoryview(adata).tobytes()
309            bpp = 32
310        return coding, Compressed(coding, memoryview(cdata), False), client_options, width, height, 0, bpp
311    finally:
312        r = tjDestroy(compressor)
313        if r:
314            log.error("Error: failed to destroy the JPEG compressor, code %i:", r)
315            log.error(" %s", get_error_str())
316
317cdef TJSAMP get_subsamp(int quality):
318    if quality<60:
319        return TJSAMP_420
320    elif quality<80:
321        return TJSAMP_422
322    return TJSAMP_444
323
324cdef encode_rgb(tjhandle compressor, image, int quality, int grayscale=0):
325    pfstr = image.get_pixel_format()
326    pf = TJPF_VAL.get(pfstr)
327    if pf is None:
328        raise Exception("invalid pixel format %s" % pfstr)
329    cdef TJPF tjpf = pf
330    cdef TJSAMP subsamp
331    if grayscale:
332        subsamp = TJSAMP_GRAY
333    else:
334        subsamp = get_subsamp(quality)
335    cdef int width = image.get_width()
336    cdef int height = image.get_height()
337    cdef int stride = image.get_rowstride()
338    pixels = image.get_pixels()
339    return do_encode_rgb(compressor, pfstr, pixels,
340                         width, height, stride,
341                         quality, tjpf, subsamp)
342
343cdef do_encode_rgb(tjhandle compressor, pfstr, pixels,
344                   int width, int height, int stride,
345                   int quality, TJPF tjpf, TJSAMP subsamp):
346    cdef int flags = 0
347    cdef unsigned char *out = NULL
348    cdef unsigned long out_size = 0
349    cdef int r = -1
350    cdef const unsigned char *src
351    log("jpeg.encode_rgb with subsampling=%s for pixel format=%s with quality=%s",
352        TJSAMP_STR.get(subsamp, subsamp), pfstr, quality)
353    with buffer_context(pixels) as bc:
354        assert len(bc)>=stride*height, "%s buffer is too small: %i bytes, %ix%i=%i bytes required" % (
355            pfstr, len(bc), stride, height, stride*height)
356        src = <const unsigned char *> (<uintptr_t> int(bc))
357        if src==NULL:
358            raise ValueError("missing pixel buffer address from context %s" % bc)
359        with nogil:
360            r = tjCompress2(compressor, src,
361                            width, stride, height, tjpf,
362                            &out, &out_size, subsamp, norm_quality(quality), flags)
363    if r!=0:
364        log.error("Error: failed to compress jpeg image, code %i:", r)
365        log.error(" %s", get_error_str())
366        log.error(" width=%i, stride=%i, height=%i", width, stride, height)
367        log.error(" quality=%i (from %i), flags=%x", norm_quality(quality), quality, flags)
368        log.error(" pixel format=%s", pfstr)
369        return None
370    assert out_size>0 and out!=NULL, "jpeg compression produced no data"
371    return makebuf(out, out_size)
372
373cdef encode_yuv(tjhandle compressor, image, int quality, int grayscale=0):
374    pfstr = image.get_pixel_format()
375    assert pfstr in ("YUV420P", "YUV422P"), "invalid yuv pixel format %s" % pfstr
376    cdef TJSAMP subsamp
377    if grayscale:
378        subsamp = TJSAMP_GRAY
379    elif pfstr=="YUV420P":
380        subsamp = TJSAMP_420
381    elif pfstr=="YUV422P":
382        subsamp = TJSAMP_422
383    elif pfstr=="YUV444P":
384        subsamp = TJSAMP_444
385    else:
386        raise ValueError("invalid yuv pixel format %s" % pfstr)
387    cdef int width = image.get_width()
388    cdef int height = image.get_height()
389    rowstrides = image.get_rowstride()
390    planes = image.get_pixels()
391    return do_encode_yuv(compressor, pfstr, planes,
392                         width, height, rowstrides,
393                         quality, subsamp)
394
395cdef do_encode_yuv(tjhandle compressor, pfstr, planes,
396                   int width, int height, rowstrides,
397                   int quality, TJSAMP subsamp):
398    cdef int flags = 0
399    cdef unsigned char *out = NULL
400    cdef unsigned long out_size = 0
401    cdef int r = -1
402    cdef int strides[3]
403    cdef const unsigned char *src[3]
404    divs = get_subsampling_divs(pfstr)
405    for i in range(3):
406        src[i] = NULL
407        strides[i] = 0
408    for i, (xdiv, ydiv) in enumerate(divs):
409        assert rowstrides[i]>=width//xdiv, "stride %i is too small for width %i of plane %s" % (
410            rowstrides[i], width//xdiv, "YUV"[i])
411        strides[i] = rowstrides[i]
412    contexts = []
413    try:
414        for i, (xdiv, ydiv) in enumerate(divs):
415            bc = buffer_context(planes[i])
416            bc.__enter__()
417            contexts.append(bc)
418            if len(bc)<strides[i]*height//ydiv:
419                raise ValueError("plane %r is only %i bytes, %ix%i=%i bytes required" % (
420                "YUV"[i], len(bc), strides[i], height, strides[i]*height//ydiv))
421            src[i] = <const unsigned char *> (<uintptr_t> int(bc))
422            if src[i]==NULL:
423                raise ValueError("missing plane %s from context %s" % ("YUV"[i], bc))
424        log("jpeg.encode_yuv with subsampling=%s for pixel format=%s with quality=%s",
425            TJSAMP_STR.get(subsamp, subsamp), pfstr, quality)
426        with nogil:
427            r = tjCompressFromYUVPlanes(compressor,
428                                        src,
429                                        width, <const int*> strides,
430                                        height, subsamp,
431                                        &out, &out_size, norm_quality(quality), flags)
432        if r!=0:
433            log.error("Error: failed to compress jpeg image, code %i:", r)
434            log.error(" %s", get_error_str())
435            log.error(" width=%i, strides=%s, height=%i", width, rowstrides, height)
436            log.error(" quality=%i (from %i), flags=%x", norm_quality(quality), quality, flags)
437            log.error(" pixel format=%s, subsampling=%s", pfstr, TJSAMP_STR.get(subsamp, subsamp))
438            log.error(" planes: %s", csv(<uintptr_t> src[i] for i in range(3)))
439            return None
440    finally:
441        for bc in contexts:
442            bc.__exit__()
443    assert out_size>0 and out!=NULL, "jpeg compression produced no data"
444    return makebuf(out, out_size)
445
446
447def selftest(full=False):
448    log("jpeg selftest")
449    from xpra.codecs.codec_checks import make_test_image
450    img = make_test_image("BGRA", 32, 32)
451    for q in (0, 50, 100):
452        v = encode("jpeg", img, {"quality" : q})
453        assert v, "encode output was empty!"
454