我的文件格式
def improveUSM(src, m, sigma, threshold, amount): # USM算法
blur = cv2.GaussianBlur(src, (m, m), sigma)
diff = src-blur
mask = diff < threshold
dst = cv2.addWeighted(src, 1+amount, blur, -1*amount, 0)
mask = mask.astype(np.int16)
mask2 = cv2.GaussianBlur(mask, (m, m), sigma)
res = src*mask2+dst*(1-mask2)
return res # 返回对Y通道做USM之后的结果
screenLevels = 255.0
def yuv_import(filename,dims,numfrm,startfrm): # 提取YUV通道
fp=open(filename,'rb')
blk_size = int(prod(dims)*3/2) #每一帧的长度
fp.seek(blk_size*startfrm,0) # 移动指针
Y=[]
U=[]
V=[]
d00=dims[0]//2
d01=dims[1]//2
Yt=zeros((dims[0],dims[1]),uint8,'C')
Ut=zeros((d00,d01),uint8,'C')
Vt=zeros((d00,d01),uint8,'C')
for i in range(numfrm): # 帧数
for m in range(dims[0]): # 长度
for n in range(dims[1]): #宽度
Yt[m,n]=ord(fp.read(1))
for m in range(d00):
for n in range(d01):
Ut[m,n]=ord(fp.read(1))
for m in range(d00):
for n in range(d01):
Vt[m,n]=ord(fp.read(1))
Y=Y+[Yt]
U=U+[Ut]
V=V+[Vt]
fp.close()
return (Y,U,V)
def merge_yuv(new_filename, data, Y): # 合并YUV三通道
f = open(new_filename, 'ab+') #追加写入
f.write(Y.astype('uint8')) #强制类型转换
f.write(data[1][0].astype('uint8'))
f.write(data[2][0].astype('uint8'))
f.close()
return f
if __name__ == '__main__':
width = 1920
from tqdm import tqdm
import numpy as np
root = os.getcwd() + '/ori_yuv/'
for files in os.listdir(root):
# 计算出每一个YUV序列的帧数
height = int(files.split("_")[1].split("x")[1]) #这里的height是我文件名字本身带的,可以split取到
path = os.path.join(root, files)
fp = open(path, 'rb')
framesize = height * width * 3 // 2 # 一帧图像所含的像素个数
fp.seek(0, 2) # 设置文件指针到文件流的尾部
ps = fp.tell() # 当前文件指针位置
numfrm = ps // framesize # 计算输出帧数
usm_path = './usm_yuv/' + '{}_usm.yuv'.format(files.split('.')[0])
if os.path.exists(usm_path): # 如果usm_yuv文件已经存在则删除,防止重复追加
# print(usm_path)
os.remove(usm_path)
for i in tqdm(range(numfrm)):
data = yuv_import(root + '{}.yuv'.format(files.split('.')[0]), (height, width), 1, i)
new_data_Y = improveUSM(data[0][0], m=5, sigma=0, threshold=1.0, amount=0.5)
merge_yuv(usm_path, data, new_data_Y)
aomenc ffmepg编解码
def parse_param_code(ori_yuv, usm_yuv, keyint):
ori_yuv_name = ori_yuv.split('_')[0]
ori_yuv_length = ori_yuv.split('_')[1]
usm_yuv_name = usm_yuv.split('_')[0]
usm_yuv_length = usm_yuv.split('_')[1]
ori_param_bs = ori_yuv_name + "_" + ori_yuv_length + "_" + str(keyint) + ".bit"
usm_param_bs = usm_yuv_name + "_" + usm_yuv_length + "_" + str(keyint) + ".bit"
param_csv = ori_yuv_name + "_" + str(keyint) + ".csv"
new_root = os.getcwd()
# 跨行字符串拼接
ori_yuv_command = "aomenc --width=1920 --height={} --fps={}/1".format(ori_yuv_length.split('x')[1], ori_yuv.split('.')[0].split('_')[2]) + \
" --limit=128 --end-usage=q --passes=2 --threads=1 --verbose --psnr --ivf --sb-size=128 --kf-min-dist=128 --kf-max-dist=128 " + \
"--tune=ssim --test-decode=fatal --cpu-used=2 --cq-level={}".format(keyint) +\
" --output={}/ori_yuv_bit/{}_1920x{}_{}.bit ".format(new_root, ori_yuv_name, ori_yuv_length.split('x')[1], keyint) + \
root + ori_yuv + " >{}/ori_yuv_bit/{}-{}bit.log 2>&1".format(new_root, ori_yuv_name, keyint)
usm_yuv_command = "aomenc --width=1920 --height={} --fps={}/1".format(usm_yuv_length.split('x')[1], ori_yuv.split('.')[0].split('_')[2]) + \
" --limit=128 --end-usage=q --passes=2 --threads=1 --verbose --psnr --ivf --sb-size=128 --kf-min-dist=128 --kf-max-dist=128 " + \
"--tune=ssim --test-decode=fatal --cpu-used=2 --cq-level={}".format(keyint) + \
" --output={}/usm_yuv_bit/{}_1920x{}_{}.bit ".format(new_root, usm_yuv_name, ori_yuv_length.split('x')[1], keyint) + \
usm_root + usm_yuv + " >{}/usm_yuv_bit/{}-{}-usm-bit.log 2>&1".format(new_root, usm_yuv_name, keyint)
if(not os.path.isfile(param_csv)):
# 编码
print(ori_yuv_command)
encoder_ori = subprocess.Popen(ori_yuv_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
print(usm_yuv_command)
encoder_usm = subprocess.Popen(usm_yuv_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
encoder_ori.wait()
encoder_usm.wait()
# FFmpeg解码
ori_param_dec = new_root + "/ori_yuv_decoder/" + "dec_" + ori_yuv_name + "_" + "1920x{}_".format(ori_yuv_length.split('x')[1]) +str(keyint) + ".yuv"
ori_ffmpegstr = "ffmpeg -y -i " + new_root + "/ori_yuv_bit/" + ori_param_bs + " " + ori_param_dec
print(ori_ffmpegstr)
decoder1 = subprocess.Popen(ori_ffmpegstr, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
decoder1.wait()
usm_param_dec = new_root + "/usm_yuv_decoder/" + "dec_" + ori_yuv_name + "_" + "1920x{}_".format(ori_yuv_length.split('x')[1]) + str(keyint) + ".yuv"
usm_ffmpegstr = "ffmpeg -y -i " + new_root + "/usm_yuv_bit/" + usm_param_bs + " " + usm_param_dec
print(usm_ffmpegstr)
decoder2 = subprocess.Popen(usm_ffmpegstr, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
decoder2.wait()
print('-------------------Done------------------------')