Skip to content

Commit 9fca773

Browse files
ArcaneNibblemarcan
ArcaneNibble
authored and
marcan
committed
jpeg: Implement decoding into planar formats
Signed-off-by: R <[email protected]>
1 parent 3a18075 commit 9fca773

1 file changed

Lines changed: 184 additions & 33 deletions

File tree

‎proxyclient/experiments/jpeg.py‎

Lines changed: 184 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,26 @@ def divroundup(val, div):
1818
return (val + div - 1) // div
1919

2020

21+
def yuv2rgb(y, u, v):
22+
y -= 16
23+
u -= 128
24+
v -= 128
25+
26+
y /= 255
27+
u /= 255
28+
v /= 255
29+
30+
r = y + 1.13983 * v
31+
g = y - 0.39465 * u - 0.58060 * v
32+
b = y + 2.03211 * u
33+
34+
r = min(255, max(0, int(r * 255)))
35+
g = min(255, max(0, int(g * 255)))
36+
b = min(255, max(0, int(b * 255)))
37+
38+
return (r, g, b)
39+
40+
2141
ap = argparse.ArgumentParser(description='JPEG block experiment')
2242
ap.add_argument("--jpeg", dest='which_jpeg', type=str, default='jpeg0',
2343
help='which JPEG instance (jpeg0/jpeg1)')
@@ -46,6 +66,9 @@ def divroundup(val, div):
4666
'RGB565',
4767
'YUV422-CbYCrY',
4868
'YUV422-YCbYCr',
69+
'YUV422-planar',
70+
'YUV420-planar',
71+
'YUV444-planar',
4972
]
5073
pixfmt = args.decode_pixelfmt
5174

@@ -128,14 +151,39 @@ def divroundup(val, div):
128151
BYTESPP = 4
129152
elif pixfmt in ['RGB565', 'YUV422-CbYCrY', 'YUV422-YCbYCr']:
130153
BYTESPP = 2
154+
elif pixfmt in ['YUV422-planar', 'YUV420-planar', 'YUV444-planar']:
155+
BYTESPP = 1
131156
else:
132157
assert False
133158
surface_stride = surface_W * BYTESPP
159+
surface_sz = surface_stride*surface_H
160+
161+
if pixfmt == 'YUV422-planar':
162+
P1_MULW = 1 # FIXME UGLY
163+
P1_DIVW = 1
164+
P1_DIVH = 1
165+
elif pixfmt == 'YUV420-planar':
166+
P1_MULW = 1
167+
P1_DIVW = 1
168+
P1_DIVH = 2
169+
elif pixfmt == 'YUV444-planar':
170+
P1_MULW = 2
171+
P1_DIVW = 1
172+
P1_DIVH = 1
173+
if pixfmt in ['YUV422-planar', 'YUV420-planar', 'YUV444-planar']:
174+
surface_P1_W = surface_W * P1_MULW // P1_DIVW
175+
surface_P1_H = surface_H // P1_DIVH
176+
surface_P1_stride = surface_P1_W
177+
surface_P1_off = surface_sz
178+
surface_sz += surface_P1_stride*surface_P1_H
179+
else:
180+
surface_P1_stride = 0
181+
surface_P1_off = 0
134182

135183
input_mem_sz = align_up(len(jpeg_data))
136184
print(f"Using size {input_mem_sz:08X} for JPEG data")
137185

138-
output_mem_sz = align_up(surface_stride*surface_H)
186+
output_mem_sz = align_up(surface_sz)
139187
print(f"Using size {output_mem_sz:08X} for output image")
140188
else:
141189
assert False
@@ -330,13 +378,23 @@ def set_default_regs(param1=0):
330378
jpeg.DECODE_PIXEL_FORMAT.set(FORMAT=E_DECODE_PIXEL_FORMAT.RGB565)
331379
elif pixfmt == 'YUV422-CbYCrY' or pixfmt == 'YUV422-YCbYCr':
332380
jpeg.DECODE_PIXEL_FORMAT.set(FORMAT=E_DECODE_PIXEL_FORMAT.YUV422_linear)
381+
elif pixfmt == 'YUV422-planar':
382+
jpeg.DECODE_PIXEL_FORMAT.set(FORMAT=E_DECODE_PIXEL_FORMAT.YUV422_planar)
383+
elif pixfmt == 'YUV420-planar':
384+
jpeg.DECODE_PIXEL_FORMAT.set(FORMAT=E_DECODE_PIXEL_FORMAT.YUV420_planar)
385+
elif pixfmt == 'YUV444-planar':
386+
jpeg.DECODE_PIXEL_FORMAT.set(FORMAT=E_DECODE_PIXEL_FORMAT.YUV444_planar)
333387
else:
334388
assert False
335389

336-
jpeg.PX_USE_PLANE1 = 0
390+
if pixfmt in ['YUV422-planar', 'YUV420-planar', 'YUV444-planar']:
391+
jpeg.PX_USE_PLANE1 = 1
392+
jpeg.PX_PLANE1_WIDTH = jpeg_W * P1_MULW // P1_DIVW // decode_scale - 1
393+
jpeg.PX_PLANE1_HEIGHT = jpeg_H // P1_DIVH // decode_scale - 1
394+
else:
395+
jpeg.PX_USE_PLANE1 = 0
337396
jpeg.PX_PLANE0_WIDTH = jpeg_W*BYTESPP // decode_scale - 1
338397
jpeg.PX_PLANE0_HEIGHT = jpeg_H // decode_scale - 1
339-
# TODO P1
340398
jpeg.TIMEOUT = 266000000
341399

342400
jpeg.REG_0x94 = 0x1f
@@ -348,15 +406,16 @@ def set_default_regs(param1=0):
348406
jpeg_W - divroundup(jpeg_W, macroblock_W)*macroblock_W + macroblock_W
349407
bot_edge_px = \
350408
jpeg_H - divroundup(jpeg_H, macroblock_H)*macroblock_H + macroblock_H
351-
# XXX changing this does not seem to do anything
409+
# XXX changing this does not seem to do anything.
410+
# Does it possibly affect scaling down?
352411
jpeg.RIGHT_EDGE_PIXELS.val = right_edge_px
353412
jpeg.BOTTOM_EDGE_PIXELS.val = bot_edge_px
354413
jpeg.RIGHT_EDGE_SAMPLES.val = right_edge_px // (macroblock_W // 8)
355414
jpeg.BOTTOM_EDGE_SAMPLES.val = bot_edge_px // (macroblock_H // 8)
356415

357416
jpeg.PX_TILES_H = divroundup(jpeg_H, macroblock_H)
358417
# FIXME explain this
359-
if pixfmt in ['RGBA', 'BGRA', 'RGB565']:
418+
if pixfmt in ['RGBA', 'BGRA', 'RGB565', 'YUV444-planar']:
360419
jpeg.PX_TILES_W = divroundup(jpeg_W // decode_scale, macroblock_W)
361420
else:
362421
jpeg.PX_TILES_W = divroundup(jpeg_W // decode_scale, max(macroblock_W, 16))
@@ -429,10 +488,82 @@ def set_default_regs(param1=0):
429488
jpeg.PX_PLANE1_TILING_V = 0
430489
else:
431490
assert False
491+
elif pixfmt == 'YUV422-planar':
492+
if jpeg_MODE == '444' or jpeg_MODE == '400':
493+
jpeg.PX_PLANE0_TILING_H = 2
494+
jpeg.PX_PLANE0_TILING_V = 8 // decode_scale
495+
jpeg.PX_PLANE1_TILING_H = 2
496+
jpeg.PX_PLANE1_TILING_V = 8 // decode_scale
497+
elif jpeg_MODE == '422':
498+
jpeg.PX_PLANE0_TILING_H = 2
499+
jpeg.PX_PLANE0_TILING_V = 8 // decode_scale
500+
jpeg.PX_PLANE1_TILING_H = 2
501+
jpeg.PX_PLANE1_TILING_V = 8 // decode_scale
502+
elif jpeg_MODE == '420':
503+
jpeg.PX_PLANE0_TILING_H = 2
504+
jpeg.PX_PLANE0_TILING_V = 16 // decode_scale
505+
jpeg.PX_PLANE1_TILING_H = 2
506+
jpeg.PX_PLANE1_TILING_V = 16 // decode_scale
507+
elif jpeg_MODE == '411':
508+
jpeg.PX_PLANE0_TILING_H = 4
509+
jpeg.PX_PLANE0_TILING_V = 8 // decode_scale
510+
jpeg.PX_PLANE1_TILING_H = 4
511+
jpeg.PX_PLANE1_TILING_V = 8 // decode_scale
512+
else:
513+
assert False
514+
elif pixfmt == 'YUV420-planar':
515+
if jpeg_MODE == '444' or jpeg_MODE == '400':
516+
jpeg.PX_PLANE0_TILING_H = 2
517+
jpeg.PX_PLANE0_TILING_V = 8 // decode_scale
518+
jpeg.PX_PLANE1_TILING_H = 2
519+
jpeg.PX_PLANE1_TILING_V = 4 // decode_scale
520+
elif jpeg_MODE == '422':
521+
jpeg.PX_PLANE0_TILING_H = 2
522+
jpeg.PX_PLANE0_TILING_V = 8 // decode_scale
523+
jpeg.PX_PLANE1_TILING_H = 2
524+
jpeg.PX_PLANE1_TILING_V = 4 // decode_scale
525+
elif jpeg_MODE == '420':
526+
jpeg.PX_PLANE0_TILING_H = 2
527+
jpeg.PX_PLANE0_TILING_V = 16 // decode_scale
528+
jpeg.PX_PLANE1_TILING_H = 2
529+
jpeg.PX_PLANE1_TILING_V = 8 // decode_scale
530+
elif jpeg_MODE == '411':
531+
jpeg.PX_PLANE0_TILING_H = 4
532+
jpeg.PX_PLANE0_TILING_V = 8 // decode_scale
533+
jpeg.PX_PLANE1_TILING_H = 4
534+
jpeg.PX_PLANE1_TILING_V = 4 // decode_scale
535+
else:
536+
assert False
537+
elif pixfmt == 'YUV444-planar':
538+
if jpeg_MODE == '444' or jpeg_MODE == '400':
539+
jpeg.PX_PLANE0_TILING_H = 1
540+
jpeg.PX_PLANE0_TILING_V = 8 // decode_scale
541+
jpeg.PX_PLANE1_TILING_H = 2
542+
jpeg.PX_PLANE1_TILING_V = 8 // decode_scale
543+
elif jpeg_MODE == '422':
544+
# The driver doesn't use this, but guessing seems to be fine?
545+
jpeg.PX_PLANE0_TILING_H = 2
546+
jpeg.PX_PLANE0_TILING_V = 8 // decode_scale
547+
jpeg.PX_PLANE1_TILING_H = 4
548+
jpeg.PX_PLANE1_TILING_V = 8 // decode_scale
549+
elif jpeg_MODE == '420':
550+
# The driver doesn't use this, but guessing seems to be fine?
551+
jpeg.PX_PLANE0_TILING_H = 2
552+
jpeg.PX_PLANE0_TILING_V = 16 // decode_scale
553+
jpeg.PX_PLANE1_TILING_H = 4
554+
jpeg.PX_PLANE1_TILING_V = 16 // decode_scale
555+
elif jpeg_MODE == '411':
556+
# The driver doesn't use this, but guessing seems to be fine?
557+
jpeg.PX_PLANE0_TILING_H = 4
558+
jpeg.PX_PLANE0_TILING_V = 8 // decode_scale
559+
jpeg.PX_PLANE1_TILING_H = 8
560+
jpeg.PX_PLANE1_TILING_V = 8 // decode_scale
561+
else:
562+
assert False
432563
else:
433564
assert False
434565

435-
if pixfmt in ['RGBA', 'BGRA', 'RGB565']:
566+
if pixfmt in ['RGBA', 'BGRA', 'RGB565', 'YUV444-planar']:
436567
if jpeg_MODE in ['422', '420']:
437568
jpeg.CHROMA_DOUBLE_H = 1
438569

@@ -441,7 +572,7 @@ def set_default_regs(param1=0):
441572

442573
if jpeg_MODE == '420':
443574
jpeg.CHROMA_DOUBLE_V = 1
444-
elif pixfmt in ["YUV422-CbYCrY", "YUV422-YCbYCr"]:
575+
elif pixfmt in ["YUV422-CbYCrY", "YUV422-YCbYCr", "YUV422-planar"]:
445576
if jpeg_MODE == '444':
446577
jpeg.CHROMA_HALVE_H_TYPE1 = 1
447578

@@ -450,6 +581,17 @@ def set_default_regs(param1=0):
450581

451582
if jpeg_MODE == '420':
452583
jpeg.CHROMA_DOUBLE_V = 1
584+
elif pixfmt in ["YUV420-planar"]:
585+
if jpeg_MODE == '444':
586+
jpeg.CHROMA_HALVE_H_TYPE1 = 1
587+
588+
if jpeg_MODE in ['444', '422', '411']:
589+
jpeg.CHROMA_HALVE_V_TYPE1 = 1
590+
591+
if jpeg_MODE == '411':
592+
jpeg.CHROMA_DOUBLE_H = 1
593+
else:
594+
assert False
453595

454596
jpeg.MATRIX_MULT[0].val = 0x100
455597
jpeg.MATRIX_MULT[1].val = 0x0
@@ -482,11 +624,10 @@ def set_default_regs(param1=0):
482624
jpeg.INPUT_START2 = 0xdeadbeef
483625
jpeg.INPUT_END = input_buf_iova + input_mem_sz
484626
jpeg.OUTPUT_START1 = output_buf_iova
485-
# jpeg.OUTPUT_START2 = output_buf_iova + jpeg_W * 4 # HACK
486-
jpeg.OUTPUT_START2 = 0xdeadbeef
627+
jpeg.OUTPUT_START2 = output_buf_iova + surface_P1_off
487628
jpeg.OUTPUT_END = output_buf_iova + output_mem_sz
488629
jpeg.PX_PLANE0_STRIDE = surface_stride
489-
# jpeg.PX_PLANE1_STRIDE = output_W * 4 # HACK
630+
jpeg.PX_PLANE1_STRIDE = surface_P1_stride
490631

491632
jpeg.REG_0x1ac = 0x0
492633
jpeg.REG_0x1b0 = 0x0
@@ -516,6 +657,7 @@ def set_default_regs(param1=0):
516657
with open(args.raw_output, 'wb') as f:
517658
f.write(output_data)
518659

660+
# Just for demonstration purposes, wrangle everything back into RGB
519661
with Image.new(
520662
mode='RGBA',
521663
size=(jpeg_W // decode_scale, jpeg_H // decode_scale)) as im:
@@ -551,34 +693,43 @@ def set_default_regs(param1=0):
551693
elif pixfmt == "YUV422-YCbYCr":
552694
y0, cb, y1, cr = block
553695

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)))
696+
r0, g0, b0 = yuv2rgb(y0, cb, cr)
697+
r1, g1, b1 = yuv2rgb(y1, cb, cr)
577698

578699
im.putpixel((x, y), (r0, g0, b0, 255))
579700
# XXX this really needs some fixing
580701
if x+1 < jpeg_W // decode_scale:
581702
im.putpixel((x+1, y), (r1, g1, b1, 255))
703+
elif pixfmt == "YUV422-planar":
704+
for y in range(jpeg_H // decode_scale):
705+
for x in range(jpeg_W // decode_scale):
706+
y_ = output_data[y*surface_stride + x]
707+
cb = output_data[surface_P1_off + y*surface_P1_stride + x&~1]
708+
cr = output_data[surface_P1_off + y*surface_P1_stride + (x&~1)+1]
709+
710+
r, g, b = yuv2rgb(y_, cb, cr)
711+
712+
im.putpixel((x, y), (r, g, b, 255))
713+
elif pixfmt == "YUV420-planar":
714+
for y in range(jpeg_H // decode_scale):
715+
for x in range(jpeg_W // decode_scale):
716+
y_ = output_data[y*surface_stride + x]
717+
cb = output_data[surface_P1_off + (y//2)*surface_P1_stride + x&~1]
718+
cr = output_data[surface_P1_off + (y//2)*surface_P1_stride + (x&~1)+1]
719+
720+
r, g, b = yuv2rgb(y_, cb, cr)
721+
722+
im.putpixel((x, y), (r, g, b, 255))
723+
elif pixfmt == "YUV444-planar":
724+
for y in range(jpeg_H // decode_scale):
725+
for x in range(jpeg_W // decode_scale):
726+
y_ = output_data[y*surface_stride + x]
727+
cb = output_data[surface_P1_off + y*surface_P1_stride + x*2]
728+
cr = output_data[surface_P1_off + y*surface_P1_stride + x*2+1]
729+
730+
r, g, b = yuv2rgb(y_, cb, cr)
731+
732+
im.putpixel((x, y), (r, g, b, 255))
582733
else:
583734
assert False
584735
im.save(args.output)

0 commit comments

Comments
 (0)