Skip to content

Commit 3a18075

Browse files
ArcaneNibblemarcan
authored andcommitted
jpeg: Decode linear YUV
Signed-off-by: R <[email protected]>
1 parent 4c8807f commit 3a18075

1 file changed

Lines changed: 116 additions & 27 deletions

File tree

proxyclient/experiments/jpeg.py

Lines changed: 116 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@ def divroundup(val, div):
4444
'RGBA',
4545
'BGRA',
4646
'RGB565',
47+
'YUV422-CbYCrY',
48+
'YUV422-YCbYCr',
4749
]
4850
pixfmt = args.decode_pixelfmt
4951

@@ -124,7 +126,7 @@ def divroundup(val, div):
124126
surface_H = divroundup(jpeg_H // decode_scale, macroblock_H) * macroblock_H
125127
if pixfmt in ['RGBA', 'BGRA']:
126128
BYTESPP = 4
127-
elif pixfmt == 'RGB565':
129+
elif pixfmt in ['RGB565', 'YUV422-CbYCrY', 'YUV422-YCbYCr']:
128130
BYTESPP = 2
129131
else:
130132
assert False
@@ -326,6 +328,10 @@ def set_default_regs(param1=0):
326328
jpeg.DECODE_PIXEL_FORMAT.set(FORMAT=E_DECODE_PIXEL_FORMAT.RGBA8888)
327329
elif pixfmt == 'RGB565':
328330
jpeg.DECODE_PIXEL_FORMAT.set(FORMAT=E_DECODE_PIXEL_FORMAT.RGB565)
331+
elif pixfmt == 'YUV422-CbYCrY' or pixfmt == 'YUV422-YCbYCr':
332+
jpeg.DECODE_PIXEL_FORMAT.set(FORMAT=E_DECODE_PIXEL_FORMAT.YUV422_linear)
333+
else:
334+
assert False
329335

330336
jpeg.PX_USE_PLANE1 = 0
331337
jpeg.PX_PLANE0_WIDTH = jpeg_W*BYTESPP // decode_scale - 1
@@ -349,7 +355,11 @@ def set_default_regs(param1=0):
349355
jpeg.BOTTOM_EDGE_SAMPLES.val = bot_edge_px // (macroblock_H // 8)
350356

351357
jpeg.PX_TILES_H = divroundup(jpeg_H, macroblock_H)
352-
jpeg.PX_TILES_W = divroundup(jpeg_W // decode_scale, macroblock_W)
358+
# FIXME explain this
359+
if pixfmt in ['RGBA', 'BGRA', 'RGB565']:
360+
jpeg.PX_TILES_W = divroundup(jpeg_W // decode_scale, macroblock_W)
361+
else:
362+
jpeg.PX_TILES_W = divroundup(jpeg_W // decode_scale, max(macroblock_W, 16))
353363
if pixfmt == 'RGBA' or pixfmt == 'BGRA':
354364
if jpeg_MODE == '444' or jpeg_MODE == '400':
355365
jpeg.PX_PLANE0_TILING_H = 4
@@ -396,15 +406,50 @@ def set_default_regs(param1=0):
396406
jpeg.PX_PLANE1_TILING_V = 0
397407
else:
398408
assert False
409+
elif pixfmt == 'YUV422-CbYCrY' or pixfmt == 'YUV422-YCbYCr':
410+
if jpeg_MODE == '444' or jpeg_MODE == '400':
411+
jpeg.PX_PLANE0_TILING_H = 4
412+
jpeg.PX_PLANE0_TILING_V = 8 // decode_scale
413+
jpeg.PX_PLANE1_TILING_H = 1
414+
jpeg.PX_PLANE1_TILING_V = 1
415+
elif jpeg_MODE == '422':
416+
jpeg.PX_PLANE0_TILING_H = 4
417+
jpeg.PX_PLANE0_TILING_V = 8 // decode_scale
418+
jpeg.PX_PLANE1_TILING_H = 1
419+
jpeg.PX_PLANE1_TILING_V = 1
420+
elif jpeg_MODE == '420':
421+
jpeg.PX_PLANE0_TILING_H = 4
422+
jpeg.PX_PLANE0_TILING_V = 16 // decode_scale
423+
jpeg.PX_PLANE1_TILING_H = 0
424+
jpeg.PX_PLANE1_TILING_V = 0
425+
elif jpeg_MODE == '411':
426+
jpeg.PX_PLANE0_TILING_H = 8
427+
jpeg.PX_PLANE0_TILING_V = 8 // decode_scale
428+
jpeg.PX_PLANE1_TILING_H = 0
429+
jpeg.PX_PLANE1_TILING_V = 0
430+
else:
431+
assert False
432+
else:
433+
assert False
434+
435+
if pixfmt in ['RGBA', 'BGRA', 'RGB565']:
436+
if jpeg_MODE in ['422', '420']:
437+
jpeg.CHROMA_DOUBLE_H = 1
399438

400-
if jpeg_MODE in ['422', '420']: # TODO
401-
jpeg.CHROMA_DOUBLE_H = 1
439+
if jpeg_MODE == '411':
440+
jpeg.CHROMA_QUADRUPLE_H = 1
402441

403-
if jpeg_MODE == '411': # TODO
404-
jpeg.CHROMA_QUADRUPLE_H = 1
442+
if jpeg_MODE == '420':
443+
jpeg.CHROMA_DOUBLE_V = 1
444+
elif pixfmt in ["YUV422-CbYCrY", "YUV422-YCbYCr"]:
445+
if jpeg_MODE == '444':
446+
jpeg.CHROMA_HALVE_H_TYPE1 = 1
405447

406-
if jpeg_MODE == '420': # TODO
407-
jpeg.CHROMA_DOUBLE_V = 1
448+
if jpeg_MODE == '411':
449+
jpeg.CHROMA_DOUBLE_H = 1
450+
451+
if jpeg_MODE == '420':
452+
jpeg.CHROMA_DOUBLE_V = 1
408453

409454
jpeg.MATRIX_MULT[0].val = 0x100
410455
jpeg.MATRIX_MULT[1].val = 0x0
@@ -420,6 +465,7 @@ def set_default_regs(param1=0):
420465

421466
jpeg.RGBA_ALPHA = args.decode_rgba_alpha
422467
jpeg.RGBA_ORDER = pixfmt == "RGBA"
468+
jpeg.YUV422_ORDER = pixfmt == "YUV422-YCbYCr"
423469

424470
if decode_scale == 1:
425471
jpeg.SCALE_FACTOR.set(SCALE=E_SCALE.DIV1)
@@ -473,23 +519,66 @@ def set_default_regs(param1=0):
473519
with Image.new(
474520
mode='RGBA',
475521
size=(jpeg_W // decode_scale, jpeg_H // decode_scale)) as im:
476-
for y in range(jpeg_H // decode_scale):
477-
for x in range(jpeg_W // decode_scale):
478-
block = output_data[
479-
y*surface_stride + x*BYTESPP:
480-
y*surface_stride + (x+1)*BYTESPP]
481-
482-
if pixfmt == "RGBA":
483-
r, g, b, a = block
484-
elif pixfmt == "BGRA":
485-
b, g, r, a = block
486-
elif pixfmt == "RGB565":
487-
rgb = struct.unpack("<H", block)[0]
488-
b = (rgb & 0b11111) << 3
489-
g = ((rgb >> 5) & 0b111111) << 2
490-
r = ((rgb >> 11) & 0b11111) << 3
491-
a = 255
492-
else:
493-
assert False
494-
im.putpixel((x, y), (r, g, b, a))
522+
if pixfmt in ["RGBA", "BGRA", "RGB565"]:
523+
for y in range(jpeg_H // decode_scale):
524+
for x in range(jpeg_W // decode_scale):
525+
block = output_data[
526+
y*surface_stride + x*BYTESPP:
527+
y*surface_stride + (x+1)*BYTESPP]
528+
529+
if pixfmt == "RGBA":
530+
r, g, b, a = block
531+
elif pixfmt == "BGRA":
532+
b, g, r, a = block
533+
elif pixfmt == "RGB565":
534+
rgb = struct.unpack("<H", block)[0]
535+
b = (rgb & 0b11111) << 3
536+
g = ((rgb >> 5) & 0b111111) << 2
537+
r = ((rgb >> 11) & 0b11111) << 3
538+
a = 255
539+
else:
540+
assert False
541+
im.putpixel((x, y), (r, g, b, a))
542+
elif pixfmt in ["YUV422-CbYCrY", "YUV422-YCbYCr"]:
543+
for y in range(jpeg_H // decode_scale):
544+
for x in range(0, jpeg_W // decode_scale, 2):
545+
block = output_data[
546+
y*surface_stride + x*BYTESPP:
547+
y*surface_stride + (x+2)*BYTESPP]
548+
549+
if pixfmt == "YUV422-CbYCrY":
550+
cb, y0, cr, y1 = block
551+
elif pixfmt == "YUV422-YCbYCr":
552+
y0, cb, y1, cr = block
553+
554+
y0 -= 16
555+
y1 -= 16
556+
cb -= 128
557+
cr -= 128
558+
559+
cb /= 255
560+
y0 /= 255
561+
cr /= 255
562+
y1 /= 255
563+
564+
r0 = y0 + 1.13983 * cr
565+
g0 = y0 - 0.39465 * cb - 0.58060 * cr
566+
b0 = y0 + 2.03211 * cb
567+
r1 = y1 + 1.13983 * cr
568+
g1 = y1 - 0.39465 * cb - 0.58060 * cr
569+
b1 = y1 + 2.03211 * cb
570+
571+
r0 = min(255, max(0, int(r0 * 255)))
572+
g0 = min(255, max(0, int(g0 * 255)))
573+
b0 = min(255, max(0, int(b0 * 255)))
574+
r1 = min(255, max(0, int(r1 * 255)))
575+
g1 = min(255, max(0, int(g1 * 255)))
576+
b1 = min(255, max(0, int(b1 * 255)))
577+
578+
im.putpixel((x, y), (r0, g0, b0, 255))
579+
# XXX this really needs some fixing
580+
if x+1 < jpeg_W // decode_scale:
581+
im.putpixel((x+1, y), (r1, g1, b1, 255))
582+
else:
583+
assert False
495584
im.save(args.output)

0 commit comments

Comments
 (0)