From cffdef368439dcfdc884c6e2fdc55dc150282b62 Mon Sep 17 00:00:00 2001
From: David Plowman <david.plowman@raspberrypi.com>
Date: Mon, 1 Jul 2024 12:07:43 +0100
Subject: [PATCH] Allow quality_check.py test to use a file for bitrate checks

The file $ASSET_DIR/quality_check.mp4, if it exists, is injected over
the camera images to make the bitrate quality checks repeatable.

Signed-off-by: David Plowman <david.plowman@raspberrypi.com>
---
 tests/quality_check.py | 47 +++++++++++++++++++++++++++++++++++++-----
 1 file changed, 42 insertions(+), 5 deletions(-)
 mode change 100644 => 100755 tests/quality_check.py

diff --git a/tests/quality_check.py b/tests/quality_check.py
old mode 100644
new mode 100755
index e368c16d..9b2f664b
--- a/tests/quality_check.py
+++ b/tests/quality_check.py
@@ -1,7 +1,12 @@
+#!/usr/bin/python3
+
 import io
+import os
 import time
 
-from picamera2 import Picamera2
+import cv2
+
+from picamera2 import MappedArray, Picamera2
 from picamera2.encoders import H264Encoder, JpegEncoder, MJPEGEncoder, Quality
 from picamera2.outputs import FileOutput
 
@@ -10,7 +15,38 @@
 # between very low/high is so great that the behaviour should be reliable enough.
 
 
+# If ASSET_DIR is set, we can inject this file instead of using camera images.
+FILENAME = "quality_check.mp4"
+FRAMES = []
+COUNTER = 0
+SIZE = (640, 360)
+if os.environ.get('ASSET_DIR') is not None:
+    filename = os.path.join(os.path.expanduser(os.environ['ASSET_DIR']), FILENAME)
+    if os.path.isfile(filename):
+        print("Using file", filename)
+        cap = cv2.VideoCapture(filename)
+        while len(FRAMES) < 100:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            if frame.shape != (SIZE[1], SIZE[0], 3):
+                frame = cv2.resize(frame, dsize=SIZE)
+            FRAMES.append(frame)
+else:
+    print("Using camera")
+
+
+def callback(request):
+    if FRAMES:
+        global COUNTER
+        with MappedArray(request, 'main') as m:
+            m.array[...] = FRAMES[COUNTER]
+        COUNTER = (COUNTER + 1) % len(FRAMES)
+
+
 def do_encode(encoder, quality):
+    global COUNTER
+    COUNTER = 0
     data = io.BytesIO()
     picam2.start_encoder(encoder, output=FileOutput(data), name='main', quality=quality)
     time.sleep(time_seconds)
@@ -19,25 +55,26 @@ def do_encode(encoder, quality):
 
 
 picam2 = Picamera2()
-config = picam2.create_video_configuration({'format': 'RGB888', 'size': (640, 360)}, lores={'size': (640, 360)})
+config = picam2.create_video_configuration({'format': 'RGB888', 'size': SIZE}, lores={'size': SIZE})
 picam2.configure(config)
 picam2.start()
+picam2.pre_callback = callback
 time_seconds = 5
 
 low_quality = do_encode(MJPEGEncoder(), Quality.VERY_LOW)
 high_quality = do_encode(MJPEGEncoder(), Quality.VERY_HIGH)
 print("MJPEGEncoder: low quality", low_quality, "high quality", high_quality)
-if (low_quality > high_quality):
+if (1.5 * low_quality > high_quality):
     print("Error: MJPEGEncoder file sizes not as expected")
 
 low_quality = do_encode(H264Encoder(), Quality.VERY_LOW)
 high_quality = do_encode(H264Encoder(), Quality.VERY_HIGH)
 print("H264Encoder: low quality", low_quality, "high quality", high_quality)
-if (low_quality > high_quality):
+if (1.5 * low_quality > high_quality):
     print("Error: H264Encoder file sizes not as expected")
 
 low_quality = do_encode(JpegEncoder(), Quality.VERY_LOW)
 high_quality = do_encode(JpegEncoder(), Quality.VERY_HIGH)
 print("JpegEncoder: low quality", low_quality, "high quality", high_quality)
-if (low_quality > high_quality):
+if (1.5 * low_quality > high_quality):
     print("Error: JpegEncoder file sizes not as expected")