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 7from time import monotonic 8 9from xpra.log import Logger 10log = Logger("decoder", "jpeg") 11 12from xpra.util import envbool, reverse_dict 13from xpra.codecs.image_wrapper import ImageWrapper 14from libc.stdint cimport uintptr_t 15from xpra.buffers.membuf cimport getbuf, MemBuf #pylint: disable=syntax-error 16from libc.string cimport memset #pylint: disable=syntax-error 17 18from libc.stdint cimport uint8_t 19 20LOG_PERF = envbool("XPRA_JPEG_LOG_PERF", False) 21 22ctypedef int TJSAMP 23ctypedef int TJPF 24ctypedef int TJCS 25 26cdef extern from "Python.h": 27 int PyObject_GetBuffer(object obj, Py_buffer *view, int flags) 28 void PyBuffer_Release(Py_buffer *view) 29 int PyBUF_ANY_CONTIGUOUS 30 31cdef extern from "turbojpeg.h": 32 TJSAMP TJSAMP_444 33 TJSAMP TJSAMP_422 34 TJSAMP TJSAMP_420 35 TJSAMP TJSAMP_GRAY 36 TJSAMP TJSAMP_440 37 TJSAMP TJSAMP_411 38 39 TJPF TJPF_RGB 40 TJPF TJPF_BGR 41 TJPF TJPF_RGBX 42 TJPF TJPF_BGRX 43 TJPF TJPF_XBGR 44 TJPF TJPF_XRGB 45 TJPF TJPF_GRAY 46 TJPF TJPF_RGBA 47 TJPF TJPF_BGRA 48 TJPF TJPF_ABGR 49 TJPF TJPF_ARGB 50 TJPF TJPF_CMYK 51 52 TJCS TJCS_RGB 53 TJCS TJCS_YCbCr 54 TJCS TJCS_GRAY 55 TJCS TJCS_CMYK 56 TJCS TJCS_YCCK 57 58 int TJFLAG_BOTTOMUP 59 int TJFLAG_FASTUPSAMPLE 60 int TJFLAG_FASTDCT 61 int TJFLAG_ACCURATEDCT 62 63 ctypedef void* tjhandle 64 tjhandle tjInitDecompress() 65 int tjDecompressHeader3(tjhandle handle, 66 const unsigned char *jpegBuf, unsigned long jpegSize, int *width, 67 int *height, int *jpegSubsamp, int *jpegColorspace) 68 int tjDecompress2(tjhandle handle, 69 const unsigned char *jpegBuf, unsigned long jpegSize, unsigned char *dstBuf, 70 int width, int pitch, int height, int pixelFormat, int flags) nogil 71 int tjDecompressToYUVPlanes(tjhandle handle, 72 const unsigned char *jpegBuf, unsigned long jpegSize, 73 unsigned char **dstPlanes, int width, int *strides, int height, int flags) nogil 74 int tjDestroy(tjhandle handle) 75 char* tjGetErrorStr() 76 77 int tjPlaneWidth(int componentID, int width, int subsamp) 78 unsigned long tjBufSizeYUV2(int width, int pad, int height, int subsamp) 79 unsigned long tjPlaneSizeYUV(int componentID, int width, int stride, int height, int subsamp) 80 int tjPlaneWidth(int componentID, int width, int subsamp) 81 int tjPlaneHeight(int componentID, int height, int subsamp) 82 83 84TJSAMP_STR = { 85 TJSAMP_444 : "444", 86 TJSAMP_422 : "422", 87 TJSAMP_420 : "420", 88 TJSAMP_GRAY : "GRAY", 89 TJSAMP_440 : "440", 90 TJSAMP_411 : "411", 91 } 92 93TJCS_STR = { 94 TJCS_RGB : "RGB", 95 TJCS_YCbCr : "YCbCr", 96 TJCS_GRAY : "GRAY", 97 TJCS_CMYK : "CMYK", 98 TJCS_YCCK : "YCCK", 99 } 100 101TJPF_STR = { 102 TJPF_RGB : "RGB", 103 TJPF_BGR : "BGR", 104 TJPF_RGBX : "RGBX", 105 TJPF_BGRX : "BGRX", 106 TJPF_XBGR : "XBGR", 107 TJPF_XRGB : "XRGB", 108 TJPF_GRAY : "GRAY", 109 TJPF_RGBA : "RGBA", 110 TJPF_BGRA : "BGRA", 111 TJPF_ABGR : "ABGR", 112 TJPF_ARGB : "ARGB", 113 TJPF_CMYK : "CMYK", 114 } 115TJPF_VAL = reverse_dict(TJPF_STR) 116 117 118def get_version(): 119 return (1, 0) 120 121def get_encodings(): 122 return ("jpeg", "jpega") 123 124 125cdef inline int roundup(int n, int m): 126 return (n + m - 1) & ~(m - 1) 127 128 129def get_error_str(): 130 cdef char *err = tjGetErrorStr() 131 return str(err) 132 133def decompress_to_yuv(data, unsigned char nplanes=3): 134 cdef tjhandle decompressor = tjInitDecompress() 135 if decompressor==NULL: 136 raise Exception("failed to instantiate a JPEG decompressor") 137 138 cdef Py_buffer py_buf 139 if PyObject_GetBuffer(data, &py_buf, PyBUF_ANY_CONTIGUOUS): 140 tjDestroy(decompressor) 141 raise Exception("failed to read compressed data from %s" % type(data)) 142 143 cdef int r 144 def close(): 145 PyBuffer_Release(&py_buf) 146 r = tjDestroy(decompressor) 147 if r: 148 log.error("Error: failed to destroy the JPEG decompressor, code %i:", r) 149 log.error(" %s", get_error_str()) 150 151 cdef int w, h, subsamp, cs 152 r = tjDecompressHeader3(decompressor, 153 <const unsigned char *> py_buf.buf, py_buf.len, 154 &w, &h, &subsamp, &cs) 155 if r: 156 err = get_error_str() 157 close() 158 raise Exception("failed to decompress JPEG header: %s" % err) 159 160 subsamp_str = TJSAMP_STR.get(subsamp, subsamp) 161 assert subsamp in (TJSAMP_444, TJSAMP_422, TJSAMP_420, TJSAMP_GRAY), "unsupported JPEG colour subsampling: %s" % subsamp_str 162 log("jpeg.decompress_to_yuv size: %4ix%-4i, subsampling=%-4s, colorspace=%s", 163 w, h, subsamp_str, TJCS_STR.get(cs, cs)) 164 if nplanes==3: 165 pixel_format = "YUV%sP" % subsamp_str 166 elif nplanes==1: 167 pixel_format = "YUV400P" 168 else: 169 close() 170 raise ValueError("invalid number of planes: %i" % nplanes) 171 #allocate YUV buffers: 172 cdef unsigned long plane_size 173 cdef unsigned char *planes[3] 174 cdef int strides[3] 175 cdef int i, stride 176 cdef MemBuf membuf 177 cdef MemBuf empty 178 cdef int flags = 0 179 pystrides = [] 180 pyplanes = [] 181 cdef unsigned long total_size = 0 182 cdef double start, elapsed 183 for i in range(3): 184 strides[i] = 0 185 planes[i] = NULL 186 try: 187 for i in range(nplanes): 188 stride = tjPlaneWidth(i, w, subsamp) 189 if stride<=0: 190 if subsamp!=TJSAMP_GRAY or i==0: 191 raise ValueError("cannot get size for plane %r for mode %r" % ("YUV"[i], subsamp_str)) 192 stride = roundup(w//2, 4) 193 plane_size = stride * roundup(h, 2)//2 194 if i==1: 195 #allocate empty U and V planes: 196 empty = getbuf(plane_size) 197 memset(<void *> empty.get_mem(), 128, plane_size) 198 pixel_format = "YUV420P" 199 membuf = empty 200 else: 201 stride = roundup(stride, 4) 202 strides[i] = stride 203 plane_size = tjPlaneSizeYUV(i, w, stride, h, subsamp) 204 membuf = getbuf(plane_size) #add padding? 205 planes[i] = <unsigned char*> membuf.get_mem() 206 total_size += plane_size 207 #python objects for each plane: 208 pystrides.append(stride) 209 pyplanes.append(memoryview(membuf)) 210 #log("jpeg strides: %s, plane sizes=%s", pystrides, [int(plane_sizes[i]) for i in range(3)]) 211 start = monotonic() 212 with nogil: 213 r = tjDecompressToYUVPlanes(decompressor, 214 <const unsigned char*> py_buf.buf, py_buf.len, 215 planes, w, strides, h, flags) 216 if r: 217 raise Exception("failed to decompress %s JPEG data to YUV: %s" % (subsamp_str, get_error_str())) 218 finally: 219 close() 220 if LOG_PERF: 221 elapsed = monotonic()-start 222 log("decompress jpeg to %s: %4i MB/s (%9i bytes in %2.1fms)", 223 pixel_format, total_size/elapsed//1024//1024, total_size, 1000*elapsed) 224 return ImageWrapper(0, 0, w, h, pyplanes, pixel_format, 24, pystrides, ImageWrapper.PLANAR_3) 225 226 227def decompress_to_rgb(rgb_format, data, unsigned long alpha_offset=0): 228 assert rgb_format in TJPF_VAL 229 cdef TJPF pixel_format = TJPF_VAL[rgb_format] 230 231 cdef tjhandle decompressor = tjInitDecompress() 232 if decompressor==NULL: 233 raise Exception("failed to instantiate a JPEG decompressor") 234 235 cdef Py_buffer py_buf 236 if PyObject_GetBuffer(data, &py_buf, PyBUF_ANY_CONTIGUOUS): 237 tjDestroy(decompressor) 238 raise Exception("failed to read compressed data from %s" % type(data)) 239 240 cdef int r 241 def close(): 242 PyBuffer_Release(&py_buf) 243 r = tjDestroy(decompressor) 244 if r: 245 log.error("Error: failed to destroy the JPEG decompressor, code %i:", r) 246 log.error(" %s", get_error_str()) 247 248 cdef int w, h, subsamp, cs 249 cdef uintptr_t buf = <uintptr_t> py_buf.buf 250 cdef unsigned long buf_size = py_buf.len 251 if alpha_offset>0: 252 buf_size = alpha_offset 253 log("decompressing buffer at %#x of size %i", buf, buf_size) 254 r = tjDecompressHeader3(decompressor, 255 <const unsigned char *> buf, buf_size, 256 &w, &h, &subsamp, &cs) 257 if r: 258 err = get_error_str() 259 close() 260 raise Exception("failed to decompress JPEG header: %s" % err) 261 subsamp_str = TJSAMP_STR.get(subsamp, subsamp) 262 log("jpeg.decompress_to_rgb: size=%4ix%-4i, subsampling=%3s, colorspace=%s", 263 w, h, subsamp_str, TJCS_STR.get(cs, cs)) 264 cdef int flags = 0 #TJFLAG_BOTTOMUP 265 cdef double elapsed 266 #TODO: add padding and rounding? 267 cdef double start = monotonic() 268 cdef int stride = w*4 269 cdef unsigned long size = stride*h 270 cdef MemBuf membuf = getbuf(size) 271 cdef unsigned char *dst_buf = <unsigned char*> membuf.get_mem() 272 with nogil: 273 r = tjDecompress2(decompressor, 274 <const unsigned char *> buf, buf_size, dst_buf, 275 w, stride, h, pixel_format, flags) 276 if r: 277 close() 278 raise Exception("failed to decompress %s JPEG data to %s: %s" % (subsamp_str, rgb_format, get_error_str())) 279 #deal with alpha channel if there is one: 280 cdef int aw, ah 281 cdef unsigned char *planes[3] 282 cdef int strides[3] 283 cdef MemBuf alpha 284 cdef unsigned long alpha_size 285 cdef int x, y, alpha_stride 286 cdef unsigned char* alpha_plane 287 cdef char alpha_index 288 if alpha_offset: 289 alpha_index = rgb_format.find("A") 290 assert alpha_index>=0, "no 'A' in %s" % rgb_format 291 assert len(rgb_format)==4, "unsupported rgb format for alpha: %s" % rgb_format 292 assert <unsigned long> py_buf.len>alpha_offset, "alpha offset is beyond the end of the compressed buffer" 293 buf = (<uintptr_t> py_buf.buf) + alpha_offset 294 buf_len = py_buf.len - alpha_offset 295 r = tjDecompressHeader3(decompressor, 296 <const unsigned char *> buf, buf_len, 297 &aw, &ah, &subsamp, &cs) 298 assert aw==w and ah==h, "alpha plane dimensions %ix%i don't match main image %ix%i" % ( 299 aw, ah, w, h) 300 subsamp_str = TJSAMP_STR.get(subsamp, subsamp) 301 log("found alpha plane %r at %#x size %i", subsamp_str, buf, buf_len) 302 assert subsamp==TJSAMP_GRAY, "unsupported JPEG alpha subsampling: %s" % subsamp_str 303 for i in range(3): 304 strides[i] = 0 305 planes[i] = NULL 306 alpha_stride = tjPlaneWidth(0, w, subsamp) 307 strides[0] = alpha_stride 308 alpha_size = tjPlaneSizeYUV(0, w, alpha_stride, h, subsamp) 309 alpha = getbuf(alpha_size) 310 alpha_plane = <unsigned char*> alpha.get_mem() 311 planes[0] = alpha_plane 312 with nogil: 313 r = tjDecompressToYUVPlanes(decompressor, 314 <const unsigned char*> buf, buf_len, 315 planes, w, strides, h, flags) 316 if r: 317 close() 318 raise Exception("failed to decompress %s JPEG alpha data: %s" % (subsamp_str, get_error_str())) 319 #merge alpha into rgb buffer: 320 for y in range(h): 321 for x in range(w): 322 dst_buf[y*stride+x*4+alpha_index] = alpha_plane[y*alpha_stride+x] 323 close() 324 if LOG_PERF: 325 elapsed = monotonic()-start 326 log("decompress jpeg to %s: %4i MB/s (%9i bytes in %2.1fms)", 327 rgb_format, size/elapsed//1024//1024, size, 1000*elapsed) 328 return ImageWrapper(0, 0, w, h, memoryview(membuf), rgb_format, 24, stride, ImageWrapper.PACKED) 329 330 331def selftest(full=False): 332 try: 333 log("jpeg selftest") 334 import binascii 335 data = binascii.unhexlify("ffd8ffe000104a46494600010101004800480000fffe00134372656174656420776974682047494d50ffdb0043000302020302020303030304030304050805050404050a070706080c0a0c0c0b0a0b0b0d0e12100d0e110e0b0b1016101113141515150c0f171816141812141514ffdb00430103040405040509050509140d0b0d1414141414141414141414141414141414141414141414141414141414141414141414141414141414141414141414141414ffc20011080010001003011100021101031101ffc4001500010100000000000000000000000000000008ffc40014010100000000000000000000000000000000ffda000c03010002100310000001aa4007ffc40014100100000000000000000000000000000020ffda00080101000105021fffc40014110100000000000000000000000000000020ffda0008010301013f011fffc40014110100000000000000000000000000000020ffda0008010201013f011fffc40014100100000000000000000000000000000020ffda0008010100063f021fffc40014100100000000000000000000000000000020ffda0008010100013f211fffda000c03010002000300000010924fffc40014110100000000000000000000000000000020ffda0008010301013f101fffc40014110100000000000000000000000000000020ffda0008010201013f101fffc40014100100000000000000000000000000000020ffda0008010100013f101fffd9") 336 def test_rgbx(*args): 337 return decompress_to_rgb("RGBX", *args) 338 for fn in (decompress_to_yuv, test_rgbx): 339 img = fn(data) 340 log("%s(%i bytes)=%s", fn, len(data), img) 341 if full: 342 try: 343 v = decompress_to_yuv(data[:len(data)//2]) 344 assert v is not None 345 except: 346 pass 347 else: 348 raise Exception("should not be able to decompress incomplete data, but got %s" % v) 349 finally: 350 pass 351