8000 Detect: Add NVidia model decoder data · kkroening/ffmpeg-python@d084c38 · GitHub
[go: up one dir, main page]
Skip to content 8000

Commit d084c38

Browse files
rpattersonrpatterson
committed
Detect: Add NVidia model decoder data
1 parent 769d45b commit d084c38

File tree

3 files changed

+2291
-35
lines changed

3 files changed

+2291
-35
lines changed

‎examples/get_detect_data.py

Lines changed: 94 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,14 @@
44
"""
55

66
import sys
7+
import math
78
import json
89

910
import requests
1011
import pandas
1112

13+
from ffmpeg import _detect
14+
1215
PLATFORM_TO_PY = {
1316
'Apple': 'Darwin',
1417
}
@@ -31,6 +34,10 @@
3134
'VideoToolbox': 'videotoolbox',
3235
}
3336

37+
NVIDIA_GPU_MATRIX_URL = (
38+
'https://developer.nvidia.com/video-encode-decode-gpu-support-matrix')
39+
NVIDIA_LINE_SUFFIXES = {'geforce': ['gtx titan', 'gtx', 'gt', 'rtx']}
40+
NVIDIA_CODEC_COLUMN_PREFIXES = {'h.264': 'h264', 'h.265': 'hevc'}
3441

3542

3643
def get_hwaccel_data():
@@ -58,11 +65,97 @@ def get_hwaccel_data():
5865
return hwaccels
5966

6067

68+
def get_nvidia_data():
69+
"""
70+
Download the NVIDIA GPU support matrix to detection data.
71+
"""
72+
response = requests.get(NVIDIA_GPU_MATRIX_URL)
73+
tables = pandas.read_html(response.content)
74+
(
75+
nvenc_recent, nvenc_consumer, nvenc_workstation, nvenc_virt,
76+
nvdec_recent, nvdec_consumer, nvdec_workstation, nvdec_virt) = tables
77+
nvidia = dict(lines=[], model_lines={}, boards={})
78+
79+
# Compile aggregate data needed to parse individual rows
80+
for nvenc_table in (
81+
nvenc_recent, nvenc_consumer, nvenc_workstation, nvenc_virt):
82+
for board in nvenc_table['BOARD']:
83+
line = board.replace('\xa0', ' ').split(None, 1)[0].lower()
84+
if line not in nvidia['lines']:
85+
nvidia['lines'].append(line)
86+
for line, line_suffixes in NVIDIA_LINE_SUFFIXES.items():
87+
for line_suffix in reversed(line_suffixes):
88+
nvidia['lines'].insert(0, ' '.join((line, line_suffix)))
89+
90+
for nvenc_table in (
91+
nvenc_recent, nvenc_consumer, nvenc_workstation, nvenc_virt):
92+
for nvenc_row_idx, nvenc_row in nvenc_table.iterrows():
93+
nvenc_row_values = {
94+
idx: cell for idx, cell in enumerate(nvenc_row[1:]) if (
95+
cell and
96+
not (isinstance(cell, float) and math.isnan(cell)))}
97+
if not nvenc_row_values:
98+
# Divider row
99+
continue
100+
101+
# Assemble the data for this row to use for each model or range
102+
model_data = {}
103+
for key, value in nvenc_row.items():
104+
if value in {'YES', 'NO'}:
105+
model_data[key] = value == 'YES'
106+
else:
107+
model_data[key] = value
108+
model_data['BOARD'] = model_data['BOARD'].replace(
109+
'\xa0', ' ')
110+
# Add keys for the data for the ffmpeg codec names for fast lookup
111+
for codec_prefix, codec in NVIDIA_CODEC_COLUMN_PREFIXES.items():
112+
for column_idx, column in enumerate(nvenc_row.keys()):
113+
if column.lower().startswith(codec_prefix):
114+
model_data[codec] = nvenc_row[column_idx] == 'YES'
115+
break
116+
nvidia['boards'][model_data['BOARD']] = model_data
117+
118+
_detect._parse_models(
119+
model_lines=nvidia['lines'],
120+
boards=model_data['BOARD'].lower(),
121+
model_data=model_data['BOARD'],
122+
model_lines_data=nvidia['model_lines'])
123+
124+
# Clean up some annoying clashes between the titan model line and GeForce
125+
# GTX model numbers
126+
for model_line, model_line_suffixes in NVIDIA_LINE_SUFFIXES.items():
127+
models_data = nvidia['model_lines'][model_line]['models']
128+
for model_num in models_data:
129+
for model_line_suffix in model_line_suffixes:
130+
if model_num.startswith(model_line_suffix + ' '):
131+
models_data[model_num[
132+
len(model_line_suffix + ' '):]] = models_data.pop(
133+
model_num)
134+
for titan_model_num in {'black', 'xp'}:
135+
nvidia['model_lines']['geforce gtx']['models'][
136+
'titan ' + titan_model_num] = nvidia['model_lines'][
137+
'titan']['models'].pop(titan_model_num)
138+
for titan_model_num in list(nvidia['model_lines'][
139+
'geforce gtx titan']['models'].keys()):
140+
nvidia['model_lines']['geforce gtx']['models'][
141+
'titan ' + titan_model_num] = nvidia['model_lines'][
142+
'geforce gtx titan']['models'].pop(titan_model_num)
143+
nvidia['model_lines']['geforce gtx']['models']['titan'] = nvidia[
144+
'model_lines']['geforce gtx']['models']['titan black']
145+
del nvidia['model_lines']['geforce gtx']['models']['titan ']
146+
del nvidia['model_lines']['geforce gtx titan']
147+
148+
return nvidia
149+
150+
61151
def main():
62152
"""
63153
Download ffmpeg detection data.
64154
"""
65-
data = dict(hwaccels=get_hwaccel_data())
155+
data = dict(
156+
hwaccels=get_hwaccel_data(),
157+
nvidia=get_nvidia_data(),
158+
)
66159
json.dump(data, sys.stdout, indent=2)
67160

68161

‎ffmpeg/_detect.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -275,9 +275,13 @@ def _parse_models(
275275
model_range_separator)
276276
if len(model_range_parameters) > 1:
277277
# This is a range of models
278-
model_line_data['model_ranges'].append(
279-
[model_range, model_data])
280-
models.remove(model_range)
278+
if model_range in model_line_data['model_ranges']:
279+
model_line_data['model_ranges'][
280+
model_line_data['model_ranges'].index(
281+
model_range)] = model_data
282+
else:
283+
model_line_data['model_ranges'].append(
284+
[model_range, model_data])
281285
break
282286
else:
283287
model_line_data['models'][model_range] = model_data

0 commit comments

Comments
 (0)
0