add scenario encode to measure the quality with vmaf

Author: Pengfei Qu

CMakeLists.txt

@@ -22,7 +22,7 @@ if (NOT DEFINED SCENARIO)
     set(SCENARIO "transcode")
 endif()
 
-if (SCENARIO STREQUAL "transcode")
+if (NOT (SCENARIO STREQUAL "cdn"))
     set(NLIVES "0")
 endif()
 

common/ffmpegcmd.py

@@ -190,13 +190,14 @@ def _dash(self):
     def _tc(self):
         cmd_1 = []
         params = self._tc_params
+        stream_name = self._target.split("/")[-1].split(".")[0]
         for item in self._renditions:
             width = item[0]
             height = item[1]
             v_bitrate = self._to_kps(item[2])
             a_bitrate = self._to_kps(item[3])
             maxrate = self._to_kps(item[2] * self._max_bitrate_ratio)
-            name= self._target+"/"+self._codec_type+"_"+str(height)+"p."+self._streaming_type if self._streaming_type == "mp4" else self._target+"_"+self._codec_type+str(height)+"p"
+            name= self._target+"/"+stream_name+self._codec_type+"_"+str(height)+"p."+self._streaming_type if self._streaming_type == "mp4" else self._target+"_"+self._codec_type+str(height)+"p"
 
             if self._acc_type == "vaapi":
                 cmd_1 += ["-vf", "scale_vaapi=w="+str(width)+":"+"h="+str(height)+":format=nv12", "-c:v", self._codec]

content-provider/build.sh

@@ -21,6 +21,7 @@ clips=(
 case "$(cat /proc/1/sched | head -n 1)" in
 *build.sh*)
     cd /mnt
+    mkdir -p /mnt/raw
     for clip in "${clips[@]}"; do
         clip_name="$(echo $clip | cut -f5 -d/).mp4"
         if test ! -f "archive/$clip_name"; then
@@ -40,6 +41,12 @@ case "$(cat /proc/1/sched | head -n 1)" in
             ffmpeg -i "archive/$clip_name" -vf "thumbnail,scale=640:360" -frames:v 1 -y "archive/$clip_name".png
         fi
     done
+    for clip in `find archive -name "*.mp4" -print`; do
+        clip_name="${clip/*\//}"
+        if test ! -f "raw/$clip_name".yuv; then
+            ffmpeg -i "archive/$clip_name" -vcodec rawvideo -an -frames:v 600 -y "raw/$clip_name".yuv
+        fi
+    done
     wait
     ;;
 *) 

xcode-server/main.py

@@ -25,9 +25,12 @@
 DASH_ROOT = "/var/www/video/dash"
 HLS_ROOT = "/var/www/video/hls"
 MP4_ROOT = "/var/www/video/mp4"
+CACHE_RAW_YUV_ROOT = "/var/www/video/rawyuv"
+CACHE_DECODED_YUV_ROOT = "/var/www/video/decodedyuv"
 
 HW_ACC_TYPE=os.getenv("HW_ACC_TYPE","sw")
 HW_DEVICE=os.getenv("HW_DEVICE",None)
+SCENARIO=os.getenv("SCENARIO","transcode")
 
 fps_regex = re.compile(
             r"\s*frame=\s*(?P<frame_count>\d+)\s*fps=\s*(?P<fps>\d+\.?\d*).*"
@@ -54,7 +57,7 @@ def get_fps(next_line,start_time):
         return {"fps":round(fps,1), "speed":round(speed,3), "frames":frame_count, "start":round(start_time,3), "duration":round(now-start_time,3), "end":round(now,3), "status": "active"}
     return {}
 
-def execute(idx, name, cmd):
+def execute(idx, name, cmd,kafka=True):
     p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=1, universal_newlines=True)
     p.poll()
     start_time=time.time()
@@ -84,6 +87,46 @@ def execute(idx, name, cmd):
         print("Exception: {}".format(e))
     return p.returncode
 
+def decode_yuv(yuv_path,in_stream_name,nframes=None):
+    try:
+        yuv_name=yuv_path+"/"+in_stream_name.split("/")[-1].replace(".mp4",".yuv")
+        if not os.path.exists(path): makedirs(yuv_path)
+        if os.path.exists(yuv_name): return
+        if nframes:
+            cmd = ["ffmpeg", "-hide_banner", "-i",in_stream_name, "-vcodec rawvideo", "-an","-frames:v", str(nframes),"-y",yuv_name]
+        else:
+            cmd = ["ffmpeg", "-hide_banner", "-i",in_stream_name, "-vcodec rawvideo", "-an","-y",yuv_name]
+        print(cmd, flush=True)
+        execute(5001,"decoded",cmd,kafka=False)
+    except Exception as e:
+        print("Exception: {}".format(e))
+        pass
+    return yuv_name
+
+def measure_quality_vamf(idx,name, raw_mp4_path, target_mp4_path,width,height, nframes=100):
+    vmaf_score=None
+    model_path="/home/models/"
+    try:
+        if width >=1920 and height >=1080:
+            model_name=model_path+"vmaf_4k_v0.6.1.json"
+        else:
+            model_name=model_path+"vmaf_v0.6.1.json"
+        log_path=target_mp4_path+".json"
+        framerate=24
+        cmd = ["ffmpeg", "-r", str(framerate), "-i",raw_mp4_path, "-r", str(framerate), "-i", target_mp4_path, "-lavfi", "[0:v]trim=end_frame={},scale={}:{}:flags=bicubic,setpts=PTS-STARTPTS[reference];[1:v]trim=end_frame={},setpts=PTS-STARTPTS[distorted];[distorted][reference]libvmaf=log_fmt=json:log_path={}:model_path={}".format(nframes,width,height,nframes,log_path,model_name),"-f", "null", "-"]
+        print(cmd,flush=True)
+        execute(str(idx+1000),name,cmd,kafka=False)
+        with open(log_path) as f:
+            obj = json.load(f)
+            sinfo={"id": str(idx+1000), "stream":name}
+            vmaf_score=float(obj["pooled_metrics"]["vmaf"]["mean"])
+            sinfo.update({"vmaf":vmaf_score})
+            producer.send(KAFKA_WORKLOAD_TOPIC, json.dumps(sinfo))
+
+    except Exception as e:
+        print("Exception: {}".format(e))
+    return vmaf_score
+
 def process_stream_vods(msg):
     stream_name=msg["name"]
     stream_type=msg["output"]["type"]
@@ -110,7 +153,8 @@ def process_stream_vods(msg):
 
     if zk.process_start():
         try:
-            cmd = FFMpegCmd(ARCHIVE_ROOT+"/"+stream_name, target_root+"/"+stream_name, stream_type, params=stream_parameters, acc_type=HW_ACC_TYPE, loop=loop, device=HW_DEVICE).cmd()
+            input_stream=ARCHIVE_ROOT+"/"+stream_name
+            cmd = FFMpegCmd(input_stream, target_root+"/"+stream_name, stream_type, params=stream_parameters, acc_type=HW_ACC_TYPE, loop=loop, device=HW_DEVICE).cmd()
             if cmd:
                 print(cmd, flush=True)
                 r = execute(idx, stream_name, cmd)
@@ -130,7 +174,7 @@ def process_stream_lives(msg):
     stream_type=msg["output"]["type"]
     target=msg["output"]["target"]
     loop= msg["loop"]
-    idx=msg["idx"] if "idx" in msg.keys() else int(random.random()*10000)
+    idx=int(msg["idx"]) if "idx" in msg.keys() else int(random.random()*10000)
     stream=stream_type+"/"+stream_name
 
     if not isfile(ARCHIVE_ROOT+"/"+stream_name):
@@ -156,13 +200,19 @@ def process_stream_lives(msg):
 
     if zk.process_start():
         try:
-            cmd = FFMpegCmd(ARCHIVE_ROOT+"/"+stream_name, target_name, stream_type, params=stream_parameters, acc_type=HW_ACC_TYPE, loop=loop, device=HW_DEVICE).cmd()
+            input_stream = ARCHIVE_ROOT+"/"+stream_name
+            cmd = FFMpegCmd(input_stream, target_name, stream_type, params=stream_parameters, acc_type=HW_ACC_TYPE, loop=loop, device=HW_DEVICE).cmd()
 
             if cmd:
                 print(cmd, flush=True)
                 r = execute(idx, stream_name, cmd)
                 if r:
                     raise Exception("status code: "+str(r))
+                if SCENARIO == "encode":
+                    width=stream_parameters["renditions"][0][0]
+                    height=stream_parameters["renditions"][0][1]
+                    mp4_file=cmd[-1]
+                    measure_quality_vamf(idx,stream_name,raw_mp4_path=input_stream,target_mp4_path=mp4_file,width=width,height=height)
                 zk.process_end()
         except:
             print(traceback.format_exc(), flush=True)