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