├── .github └── workflows │ └── python-publish.yml ├── .gitignore ├── README.md ├── README_CN.md ├── ffmpegcv ├── __init__.py ├── ffmpeg_noblock.py ├── ffmpeg_reader.py ├── ffmpeg_reader_camera.py ├── ffmpeg_reader_cuda.py ├── ffmpeg_reader_noblock.py ├── ffmpeg_reader_pannels.py ├── ffmpeg_reader_qsv.py ├── ffmpeg_reader_stream.py ├── ffmpeg_reader_stream_realtime.py ├── ffmpeg_writer.py ├── ffmpeg_writer_noblock.py ├── ffmpeg_writer_qsv.py ├── ffmpeg_writer_stream_realtime.py ├── stream_info.py ├── version.py └── video_info.py └── setup.py /.github/workflows/python-publish.yml: -------------------------------------------------------------------------------- 1 | # This workflow will upload a Python Package using Twine when a release is created 2 | # For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries 3 | 4 | # This workflow uses actions that are not certified by GitHub. 5 | # They are provided by a third-party and are governed by 6 | # separate terms of service, privacy policy, and support 7 | # documentation. 8 | 9 | name: Upload Python Package 10 | 11 | on: 12 | release: 13 | types: [published] 14 | 15 | permissions: 16 | contents: read 17 | 18 | jobs: 19 | deploy: 20 | 21 | runs-on: ubuntu-latest 22 | 23 | steps: 24 | - uses: actions/checkout@v3 25 | - name: Set up Python 26 | uses: actions/setup-python@v3 27 | with: 28 | python-version: '3.x' 29 | - name: Install dependencies 30 | run: | 31 | python -m pip install --upgrade pip 32 | pip install build 33 | - name: Build package 34 | run: python -m build 35 | - name: Publish package 36 | uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29 37 | with: 38 | skip-existing: true 39 | user: __token__ 40 | password: ${{ secrets.PYPI_API_TOKEN }} 41 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /build 2 | *.pyc 3 | /dist 4 | /ffmpegcv.egg-info -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # FFMPEGCV is an alternative to OPENCV for video reading&writing. 2 | ![Python versions](https://img.shields.io/badge/Python-3.6%2B-blue.svg) 3 | [![PyPI version](https://img.shields.io/pypi/v/ffmpegcv.svg?logo=pypi&logoColor=FFE873)](https://pypi.org/project/ffmpegcv/) 4 | [![PyPI downloads](https://img.shields.io/pypi/dm/ffmpegcv.svg)](https://pypistats.org/packages/ffmpegcv) 5 | ![Code size](https://shields.io/github/languages/code-size/chenxinfeng4/ffmpegcv 6 | ) 7 | ![Last Commit](https://shields.io/github/last-commit/chenxinfeng4/ffmpegcv) 8 | 9 | English Version | [中文版本](./README_CN.md) 10 | 11 | Here is the Python version of ffmpegcv. For the C++ version, please visit [FFMPEGCV-CPP](https://github.com/chenxinfeng4/ffmpegcv-cpp) 12 | 13 | The ffmpegcv provide Video Reader and Video Witer with ffmpeg backbone, which are faster and powerful than cv2. Integrating ffmpegcv into your deeplearning pipeline is very smooth. 14 | 15 | - The ffmpegcv is api **compatible** to open-cv. 16 | - The ffmpegcv can use **GPU accelerate** encoding and decoding*. 17 | - The ffmpegcv supports much more video **codecs** v.s. open-cv. 18 | - The ffmpegcv supports **RGB** & BGR & GRAY format as you like. 19 | - The ffmpegcv supports fp32 CHW & HWC format shortcut to CUDA memory. 20 | - The ffmpegcv supports **Stream reading** (IP Camera) in low latency. 21 | - The ffmpegcv supports ROI operations.You can **crop**, **resize** and **pad** the ROI. 22 | 23 | In all, ffmpegcv is just similar to opencv api. But it has more codecs and does't require opencv installed at all. It's great for deeplearning pipeline. 24 | 25 | 26 |

27 | 28 |

29 | 30 | ## Functions: 31 | - `VideoWriter`: Write a video file. 32 | - `VideoCapture`: Read a video file. 33 | - `VideoCaptureNV`: Read a video file by NVIDIA GPU. 34 | - `VideoCaptureQSV`: Read a video file by Intel QuickSync Video. 35 | - `VideoCaptureCAM`: Read a camera. 36 | - `VideoCaptureStream`: Read a RTP/RTSP/RTMP/HTTP stream. 37 | - `VideoCaptureStreamRT`: Read a RTSP stream (IP Camera) in real time low latency as possible. 38 | - `noblock`: Read/Write a video file in background using mulitprocssing. 39 | - `toCUDA`: Translate a video/stream as CHW/HWC-float32 format into CUDA device, >2x faster. 40 | 41 | ## Install 42 | You need to download ffmpeg before you can use ffmpegcv. 43 | ``` 44 | #1A. LINUX: sudo apt install ffmpeg 45 | #1B. MAC (No NVIDIA GPU): brew install ffmpeg 46 | #1C. WINDOWS: download ffmpeg and add to the path 47 | #1D. CONDA: conda install ffmpeg=6.0.0 #don't use the default 4.x.x version 48 | 49 | #2A. python 50 | pip install ffmpegcv #stable verison 51 | pip install git+https://github.com/chenxinfeng4/ffmpegcv #latest verison 52 | 53 | #2B. recommand only when you want advanced functions. See the toCUDA section 54 | pip install ffmpegcv[cuda] 55 | ``` 56 | 57 | ## When should choose `ffmpegcv` other than `opencv`: 58 | - The `opencv` is hard to install. The ffmpegcv only requires `numpy` and `FFmpeg`, works across Mac/Windows/Linux platforms. 59 | - The `opencv` packages too much image processing toolbox. You just want a simple video/camero IO with GPU accessible. 60 | - The `opencv` didn't support profiling `h264`/`h265` and other video writers. 61 | - You want to **crop**, **resize** and **pad** the video/camero ROI. 62 | - You are interested in deeplearning pipeline. 63 | ## Basic example 64 | Read a video by CPU, and rewrite it by GPU. 65 | ```python 66 | vidin = ffmpegcv.VideoCapture(vfile_in) 67 | vidout = ffmpegcv.VideoWriterNV(vfile_out, 'h264', vidin.fps) #NVIDIA-GPU 68 | 69 | with vidin, vidout: 70 | for frame in vidin: 71 | cv2.imshow('image', frame) 72 | vidout.write(frame) 73 | ``` 74 | 75 | Read the camera. 76 | ```python 77 | # by device ID 78 | cap = ffmpegcv.VideoCaptureCAM(0) 79 | # by device name 80 | cap = ffmpegcv.VideoCaptureCAM("Integrated Camera") 81 | ``` 82 | 83 | Deeplearning pipeline. 84 | ```python 85 | """ 86 | —————————— NVIDIA GPU accelerating ⤴⤴ ——————— 87 | | | 88 | V V 89 | video -> decode -> crop -> resize -> RGB -> CUDA:CHW float32 -> model 90 | """ 91 | cap = ffmpegcv.toCUDA( 92 | ffmpegcv.VideoCaptureNV(file, pix_fmt='nv12', resize=(W,H)), 93 | tensor_format='chw') 94 | 95 | for frame_CHW_cuda in cap: 96 | frame_CHW_cuda = (frame_CHW_cuda - mean) / std 97 | result = model(frame_CHW_cuda) 98 | ``` 99 | ## Cross platform 100 | 101 | The ffmpegcv is based on Python+FFmpeg, it can cross platform among `Windows, Linux, Mac, X86, Arm`systems. 102 | 103 | ## GPU Acceleration 104 | 105 | - Support **NVIDIA** card only, test in x86_64 only. 106 | - Works in **Windows**, **Linux** and **Anaconda**. 107 | - Works in the **Google Colab** notebook. 108 | - Infeasible in the **MacOS**. That ffmpeg didn't supports NVIDIA at all. 109 | 110 | > \* The ffmegcv GPU reader is a bit slower than CPU reader, but much faster when use ROI operations (crop, resize, pad). 111 | 112 | ## Codecs 113 | 114 | | Codecs | OpenCV-reader | ffmpegcv-cpu-r | gpu-r | OpenCV-writer | ffmpegcv-cpu-w | gpu-w | 115 | | ----------- | ------------- | ---------------- | ---- | ------------- | ---------------- | ---- | 116 | | h264 | √ | √ | √ | × | √ | √ | 117 | | h265 (hevc) | not sure | √ | √ | × | √ | √ | 118 | | mjpeg | √ | √ | × | √ | √ | × | 119 | | mpeg | √ | √ | × | √ | √ | × | 120 | | others | not sure | ffmpeg -decoders | × | not sure | ffmpeg -encoders | × | 121 | 122 | ## Benchmark 123 | *On the way...（maybe never）* 124 | 125 | 126 | ## Video Reader 127 | --- 128 | The ffmpegcv is just similar to opencv in api. 129 | ```python 130 | # open cv 131 | import cv2 132 | cap = cv2.VideoCapture(file) 133 | while True: 134 | ret, frame = cap.read() 135 | if not ret: 136 | break 137 | pass 138 | 139 | # ffmpegcv 140 | import ffmpegcv 141 | cap = ffmpegcv.VideoCapture(file) 142 | while True: 143 | ret, frame = cap.read() 144 | if not ret: 145 | break 146 | pass 147 | cap.release() 148 | 149 | # alternative 150 | cap = ffmpegcv.VideoCapture(file) 151 | nframe = len(cap) 152 | for frame in cap: 153 | pass 154 | cap.release() 155 | 156 | # more pythonic, recommand 157 | with ffmpegcv.VideoCapture(file) as cap: 158 | nframe = len(cap) 159 | for iframe, frame in enumerate(cap): 160 | if iframe>100: break 161 | pass 162 | ``` 163 | 164 | Use GPU to accelerate decoding. It depends on the video codes. 165 | h264_nvcuvid, hevc_nvcuvid .... 166 | ```python 167 | cap_cpu = ffmpegcv.VideoCapture(file) 168 | cap_gpu0 = ffmpegcv.VideoCaptureNV(file) #NVIDIA GPU0 169 | cap_gpu1 = ffmpegcv.VideoCaptureNV(file, gpu=1) #NVIDIA GPU1 170 | cap_qsv = ffmpegcv.VideoCaptureQSV(file) #Intel QSV, experimental 171 | ``` 172 | 173 | Use `rgb24` instead of `bgr24`. The `gray` version would be more efficient. 174 | ```python 175 | cap = ffmpegcv.VideoCapture(file, pix_fmt='rgb24') #rgb24, bgr24, gray 176 | ret, frame = cap.read() 177 | plt.imshow(frame) 178 | ``` 179 | 180 | ### ROI Operations 181 | You can crop, resize and pad the video. These ROI operation is `ffmpegcv-GPU` > `ffmpegcv-CPU` >> `opencv`. 182 | 183 | **Crop** video, which will be much faster than read the whole canvas. The top-left corner is (0, 0). 184 | ```python 185 | cap = ffmpegcv.VideoCapture(file, crop_xywh=(0, 0, 640, 480)) 186 | ``` 187 | 188 | **Resize** the video to the given size. 189 | ```python 190 | cap = ffmpegcv.VideoCapture(file, resize=(640, 480)) 191 | ``` 192 | 193 | **Resize** and keep the aspect ratio with black border **padding**. 194 | ```python 195 | cap = ffmpegcv.VideoCapture(file, resize=(640, 480), resize_keepratio=True) 196 | ``` 197 | 198 | **Crop** and then **resize** the video. 199 | ```python 200 | cap = ffmpegcv.VideoCapture(file, crop_xywh=(0, 0, 640, 480), resize=(512, 512)) 201 | ``` 202 | 203 | ## toCUDA device 204 | --- 205 | The ffmpegcv can translate the video/stream from HWC-uint8 cpu to CHW-float32 in CUDA device. It significantly reduce your cpu load, and get >2x faster than your manually convertion. 206 | 207 | Prepare your environment. The cuda environment is required. The `pycuda` package is required. The `pytorch` package is non-essential. 208 | > nvcc --version # check you've installed NVIDIA CUDA Compiler. Already installed if 209 | > you've installed Tensorflow-gpu or Pytorch-gpu 210 | > 211 | > pip install ffmpegcv[cuda] #auto install pycuda 212 | 213 | ```python 214 | # Read a video file to CUDA device, original 215 | cap = ffmpegcv.VideoCaptureNV(file, pix_fmt='rgb24') 216 | ret, frame_HWC_CPU = cap.read() 217 | frame_CHW_CUDA = torch.from_numpy(frame_HWC_CPU).permute(2, 0, 1).cuda().contiguous().float() # 120fps, 1200% CPU load 218 | 219 | # speed up 220 | cap = toCUDA(ffmpegcv.VideoCapture(file, pix_fmt='yuv420p')) #pix_fmt: 'yuv420p' or 'nv12' only 221 | cap = toCUDA(ffmpegcv.VideoCaptureNV(file, pix_fmt='nv12')) #'nv12' is better for gpu 222 | cap = toCUDA(vid, tensor_format='chw') #tensor format:'chw'(default) or 'hwc', fp32 precision 223 | cap = toCUDA(vid, gpu=1) #choose gpu 224 | 225 | # read to the cuda device 226 | ret, frame_CHW_pycuda = cap.read() #380fps, 200% CPU load, dtype is [pycuda array] 227 | ret, frame_CHW_pycudamem = cap.read_cudamem() #dtype is [pycuda mem_alloc] 228 | ret, frame_CHW_CUDA = cap.read_torch() #dtype is [pytorch tensor] 229 | ret, _ = cap.read_torch(frame_CHW_CUDA) #no copy, but need to specify the output memory 230 | 231 | frame_CHW_pycuda[:] = (frame_CHW_pycuda - mean) / std #normalize 232 | ``` 233 | 234 | How can `toCUDA` make it faster in your deeplearning pipeline than `opencv` or `ffmpeg`? 235 | > 1. The opencv/ffmpeg uses the cpu to convert video pix_fmt from original YUV to RGB24, which is slow. The ffmpegcv use the cuda to accelerate pix_fmt convertion. 236 | > 2. Use `yuv420p` or `nv12` can save the cpu load and reduce the memory copy from CPU to GPU. 237 | > 3. The ffmpeg stores the image as HWC format. The ffmpegcv can use HWC & CHW format to accelerate the video reading. 238 | 239 | ## Video Writer 240 | --- 241 | 242 | ```python 243 | # cv2 244 | out = cv2.VideoWriter('outpy.avi', 245 | cv2.VideoWriter_fourcc('M','J','P','G'), 246 | 10, 247 | (w, h)) 248 | out.write(frame1) 249 | out.write(frame2) 250 | out.release() 251 | 252 | # ffmpegcv, default codec is 'h264' in cpu 'h265' in gpu. 253 | # frameSize is decided by the size of the first frame. 254 | # use the 'mp4/mkv' instead of 'avi' to avoid the codec outdated. 255 | out = ffmpegcv.VideoWriter('outpy.mp4', None, 10) 256 | out.write(frame1) 257 | out.write(frame2) 258 | out.release() 259 | 260 | # more pythonic 261 | with ffmpegcv.VideoWriter('outpy.mp4', None, 10) as out: 262 | out.write(frame1) 263 | out.write(frame2) 264 | ``` 265 | 266 | 267 | Use GPU to accelerate encoding. Such as h264_nvenc, hevc_nvenc. 268 | ```python 269 | out_cpu = ffmpegcv.VideoWriter('outpy.mp4', None, 10) 270 | out_gpu0 = ffmpegcv.VideoWriterNV('outpy.mp4', 'h264', 10) #NVIDIA GPU0 271 | out_gpu1 = ffmpegcv.VideoWriterNV('outpy.mp4', 'hevc', 10, gpu=1) #NVIDIA GPU1 272 | out_qsv = ffmpegcv.VideoWriterQSV('outpy.mp4', 'h264', 10) #Intel QSV, experimental 273 | 274 | ``` 275 | 276 | Input image is rgb24 instead of bgr24 277 | ```python 278 | out = ffmpegcv.VideoWriter('outpy.mp4', None, 10, pix_fmt='rgb24') 279 | ``` 280 | 281 | Resize the video 282 | ```python 283 | out_resz = ffmpegcv.VideoWriter('outpy.mp4', None, 10, resize=(640, 480)) #Resize 284 | ``` 285 | ## Video Reader and Writer 286 | --- 287 | ```python 288 | import ffmpegcv 289 | vfile_in = 'A.mp4' 290 | vfile_out = 'A_h264.mp4' 291 | vidin = ffmpegcv.VideoCapture(vfile_in) 292 | vidout = ffmpegcv.VideoWriter(vfile_out, None, vidin.fps) 293 | 294 | with vidin, vidout: 295 | for frame in vidin: 296 | vidout.write(frame) 297 | ``` 298 | 299 | ## Camera Reader 300 | --- 301 | **Experimental feature**. The ffmpegcv offers Camera reader. Which is consistent with VideoFiler reader. 302 | 303 | - The `VideoCaptureCAM` aims to support ROI operations. The Opencv will be general fascinating than ffmpegcv in camera read. **I recommand the opencv in most camera reading case**. 304 | - The ffmpegcv can use name to retrieve the camera device. Use `ffmpegcv.VideoCaptureCAM("Integrated Camera")` is readable than `cv2.VideoCaptureCAM(0)`. 305 | - The `VideoCaptureCAM` will be laggy and dropping frames if your post-process takes long time. The VideoCaptureCAM will buffer the recent frames. 306 | - The `VideoCaptureCAM` is continously working on background even if you didn't read it. **Please release it in time**. 307 | - Works perfect in Windows, not-perfect in Linux and macOS. 308 | 309 | ```python 310 | import cv2 311 | cap = cv2.VideoCapture(0) 312 | while True: 313 | ret, frame = cap.read() 314 | cv2.imshow('frame', frame) 315 | if cv2.waitKey(1) & 0xFF == ord('q'): 316 | break 317 | cap.release() 318 | 319 | # ffmpegcv, in Windows&Linux 320 | import ffmpegcv 321 | cap = ffmpegcv.VideoCaptureCAM(0) 322 | while True: 323 | ret, frame = cap.read() 324 | cv2.imshow('frame', frame) 325 | if cv2.waitKey(1) & 0xFF == ord('q'): 326 | break 327 | cap.release() 328 | 329 | # ffmpegcv use by camera name, in Windows&Linux 330 | cap = ffmpegcv.VideoCaptureCAM("Integrated Camera") 331 | 332 | # ffmpegcv use camera path if multiple cameras conflict 333 | cap = ffmpegcv.VideoCaptureCAM('@device_pnp_\\\\?\\usb#vid_2304&' 334 | 'pid_oot#media#0001#{65e8773d-8f56-11d0-a3b9-00a0c9223196}' 335 | '\\global') 336 | 337 | # ffmpegcv use camera with ROI operations 338 | cap = ffmpegcv.VideoCaptureCAM("Integrated Camera", crop_xywh=(0, 0, 640, 480), resize=(512, 512), resize_keepratio=True) 339 | 340 | 341 | ``` 342 | 343 | **List all camera devices** 344 | ```python 345 | from ffmpegcv.ffmpeg_reader_camera import query_camera_devices 346 | 347 | devices = query_camera_devices() 348 | print(devices) 349 | ``` 350 | >{0: ('Integrated Camera', '@device_pnp_\\\\?\\usb#vid_2304&pid_oot#media#0001#{65e8773d-8f56-11d0-a3b9-00a0c9223196}\\global'), 351 | 1: ('OBS Virtual Camera', '@device_sw_{860BB310-5D01-11D0-BD3B-00A0C911CE86}\\{A3FCE0F5-3493-419F-958A-ABA1250EC20B}')} 352 | 353 | 354 | **Set the camera resolution, fps, vcodec/pixel-format** 355 | 356 | ```python 357 | from ffmpegcv.ffmpeg_reader_camera import query_camera_options 358 | 359 | options = query_camera_options(0) # or query_camera_options("Integrated Camera") 360 | print(options) 361 | cap = ffmpegcv.VideoCaptureCAM(0, **options[-1]) 362 | ``` 363 | >[{'camcodec': 'mjpeg', 'campix_fmt': None, 'camsize_wh': (1280, 720), 'camfps': 60.0002}, {'camcodec': 'mjpeg', 'campix_fmt': None, 'camsize_wh': (640, 480), 'camfps': 60.0002}, {'camcodec': 'mjpeg', 'campix_fmt': None, 'camsize_wh': (1920, 1080), 'camfps': 60.0002}, {'camcodec': None, 'campix_fmt': 'yuyv422', 'camsize_wh': (1280, 720), 'camfps': 10}, {'camcodec': None, 'campix_fmt': 'yuyv422', 'camsize_wh': (640, 480), 'camfps': 30}, {'camcodec': None, 'campix_fmt': 'yuyv422', 'camsize_wh': (1920, 1080), 'camfps': 5}] 364 | 365 | **Known issues** 366 | 1. The VideoCaptureCAM didn't give a smooth experience in macOS. You must specify all the camera parameters. And the query_camera_options woun't give any suggestion. That's because the `ffmpeg` cannot list device options using mac native `avfoundation`. 367 | ```python 368 | # The macOS requires full argument. 369 | cap = ffmpegcv.VideoCaptureCAM('FaceTime HD Camera', camsize_wh=(1280,720), camfps=30, campix_fmt='nv12') 370 | ``` 371 | 372 | 2. The VideoCaptureCAM cann't list the FPS in linux. Because the `ffmpeg` cound't query the device's FPS using linux native `v4l2` module. However, it's just OK to let the FPS empty. 373 | 374 | 375 | ## Stream Reader (Live streaming, RTSP IP cameras) 376 | **Experimental feature**. The ffmpegcv offers Stream reader. Which is consistent with VideoFiler reader, and more similiar to the camera. 377 | Becareful when using it. 378 | 379 | - Support `RTSP`, `RTP`, `RTMP`, `HTTP`, `HTTPS` streams. 380 | - The `VideoCaptureStream` will be laggy and dropping frames if your post-process takes long time. The VideoCaptureCAM will buffer the recent frames. 381 | - The `VideoCaptureStreamRT` is continously working on background even if you didn't read it. **Please release it in time**. 382 | - **Low latency** RTSP IP camera reader. Batter than opencv. 383 | - *It's still experimental*. Recommand you to use opencv. 384 | 385 | 386 | ```python 387 | # opencv 388 | import cv2 389 | stream_url = 'http://devimages.apple.com.edgekey.net/streaming/examples/bipbop_4x3/gear2/prog_index.m3u8' 390 | cap = cv2.VideoCapture(stream_url, cv2.CAP_FFMPEG) 391 | 392 | if not cap.isOpened(): 393 | print('Cannot open the stream') 394 | exit(-1) 395 | 396 | while True: 397 | ret, frame = cap.read() 398 | if not ret: 399 | break 400 | pass 401 | 402 | # ffmpegcv 403 | import ffmpegcv 404 | cap = ffmpegcv.VideoCaptureStream(stream_url) 405 | while True: 406 | ret, frame = cap.read() 407 | if not ret: 408 | break 409 | pass 410 | 411 | # ffmpegcv, IP Camera Low-latency 412 | # e.g. HIK Vision IP Camera, `101` Main camera stream, `102` the second 413 | stream_url = 'rtsp://admin:PASSWD@192.168.1.xxx:8554/Streaming/Channels/102' 414 | cap = ffmpegcv.VideoCaptureStreamRT(stream_url) # Low latency & recent buffered 415 | cap = ffmpegcv.ReadLiveLast(ffmpegcv.VideoCaptureStreamRT, stream_url) #no buffer 416 | while True: 417 | ret, frame = cap.read() 418 | if not ret: 419 | break 420 | pass 421 | ``` 422 | 423 | ## Noblock 424 | A proxy to automatic prepare frames in backgroud, which does not block when reading&writing current frame (multiprocessing). This make your python program more efficient in CPU usage. Up to 2x boost. 425 | 426 | > ffmpegcv.VideoCapture(*args) -> ffmpegcv.noblock(ffmpegcv.VideoCapture, *args) 427 | > 428 | > ffmpegcv.VideoWriter(*args) -> ffmpegcv.noblock(ffmpegcv.VideoWriter, *args) 429 | ```python 430 | #Proxy any VideoCapture&VideoWriter args and kargs 431 | vid_noblock = ffmpegcv.noblock(ffmpegcv.VideoCapture, vfile, pix_fmt='rbg24') 432 | 433 | # this is fast 434 | def cpu_tense(): time.sleep(0.01) 435 | for _ in tqdm.trange(1000): 436 | ret, img = vid_noblock.read() #current img is already buffered, take no time 437 | cpu_tense() #meanwhile, the next img is buffering in background 438 | 439 | # this is slow 440 | vid = ffmpegcv.VideoCapture(vfile, pix_fmt='rbg24') 441 | for _ in tqdm.trange(1000): 442 | ret, img = vid.read() #this read will block cpu, take time 443 | cpu_tense() 444 | ``` -------------------------------------------------------------------------------- /README_CN.md: -------------------------------------------------------------------------------- 1 | # FFMPEGCV 读写视频，替代 OPENCV. 2 | ![Python versions](https://img.shields.io/badge/Python-3.6%2B-blue.svg) 3 | [![PyPI version](https://img.shields.io/pypi/v/ffmpegcv.svg?logo=pypi&logoColor=FFE873)](https://pypi.org/project/ffmpegcv/) 4 | [![PyPI downloads](https://img.shields.io/pypi/dm/ffmpegcv.svg)](https://pypistats.org/packages/ffmpegcv) 5 | ![Code size](https://shields.io/github/languages/code-size/chenxinfeng4/ffmpegcv 6 | ) 7 | ![Last Commit](https://shields.io/github/last-commit/chenxinfeng4/ffmpegcv) 8 | 9 | [English Version](./README.md) | 中文版本 10 | 11 | ffmpegcv提供了基于ffmpeg的视频读取器和视频编写器，比cv2更快和更强大。适合深度学习的视频处理。 12 | 13 | - ffmpegcv与open-cv具有**兼容**的API。 14 | - ffmpegcv可以使用**GPU加速**编码和解码。 15 | - ffmpegcv支持比open-cv更多的**视频编码器**。 16 | - ffmpegcv原生支持**RGB**/BGR/灰度像素格式。 17 | - ffmpegcv支持网络**流视频读取** (网线监控相机)。 18 | - ffmpegcv支持ROI（感兴趣区域）操作，可以对ROI进行**裁剪**、**调整大小**和**填充**。 19 | 总之，ffmpegcv与opencv的API非常相似。但它具有更多的编码器，并且不需要安装opencv。 20 | - ffmpegcv支持导出图像帧到CUDA设备。 21 | 22 |

23 | 24 |

25 | 26 | ## 功能： 27 | - `VideoWriter`：写入视频文件。 28 | - `VideoCapture`：读取视频文件。 29 | - `VideoCaptureNV`：使用NVIDIA GPU读取视频文件。 30 | - `VideoCaptureQSV`: 使用Intel集成显卡读取视频文件. 31 | - `VideoCaptureCAM`：读取摄像头。 32 | - `VideoCaptureStream`：读取RTP/RTSP/RTMP/HTTP流。 33 | - `VideoCaptureStreamRT`: 读取RTSP流 (网线监控相机)，实时、低延迟。 34 | - `noblock`：在后台读取视频文件（更快）,使用多进程。 35 | - `toCUDA`：将图像帧导出到CUDA设备，以 CHW/HWC-float32 格式存储，超过2倍性能提升。 36 | 37 | ## 安装 38 | 在使用ffmpegcv之前，您需要下载`ffmpeg`。 39 | ``` 40 | #1A. LINUX: sudo apt install ffmpeg 41 | #1B. MAC: brew install ffmpeg 42 | #1C. WINDOWS: 下载ffmpeg并添加至环境变量的路径中 43 | #1D. CONDA: conda install ffmpeg=6.0.0 44 | 45 | #2. python 46 | pip install ffmpegcv #stable verison 47 | pip install git+https://github.com/chenxinfeng4/ffmpegcv #latest verison 48 | ``` 49 | 50 | ## 何时选择 `ffmpegcv` 而不是 `opencv`： 51 | - 安装`opencv`比较困难。ffmpegcv仅需要`numpy`和`FFmpeg`，可以在Mac/Windows/Linux平台上工作。 52 | - `opencv`包含太多的图像处理工具箱，而您只是想使用带GPU支持的简单视频/摄像头输入输出操作。 53 | - `opencv`不支持`h264`/`h265`和其他视频编码器。 54 | - 您想对视频/摄像头的感兴趣区域（ROI）进行**裁剪**、**调整大小**和**填充**操作。 55 | 56 | 57 | ## 基本示例 58 | 通过CPU读取视频，并通过GPU重写视频。 59 | ```python 60 | vidin = ffmpegcv.VideoCapture(vfile_in) 61 | vidout = ffmpegcv.VideoWriterNV(vfile_out, 'h264', vidin.fps) #NVIDIA 显卡 62 | 63 | with vidin, vidout: 64 | for frame in vidin: 65 | cv2.imshow('image', frame) 66 | vidout.write(frame) 67 | ``` 68 | 69 | 读取摄像头。 70 | ```python 71 | # 通过设备ID 72 | cap = ffmpegcv.VideoCaptureCAM(0) 73 | # 通过设备名称 74 | cap = ffmpegcv.VideoCaptureCAM("Integrated Camera") 75 | ``` 76 | 77 | 深度学习流水线 78 | ```python 79 | """ 80 | —————————— NVIDIA GPU 加速 ⤴⤴ ——————— 81 | | | 82 | V V 83 | 视频 -> 解码器 -> 裁剪 -> 缩放 -> RGB -> CUDA:CHW float32 -> 模型 84 | """ 85 | cap = ffmpegcv.toCUDA( 86 | ffmpegcv.VideoCaptureNV(file, pix_fmt='nv12', resize=(W,H)), 87 | tensor_format='chw') 88 | 89 | for frame_CHW_cuda in cap: 90 | frame_CHW_cuda = (frame_CHW_cuda - mean) / std 91 | result = model(frame_CHW_cuda) 92 | ``` 93 | 94 | ## GPU加速 95 | - 仅支持NVIDIA显卡，在 x86_64 上测试。 96 | - 原生支持**Windows**, **Linux**, **Anaconda**。 97 | - 在**Google Colab**上顺利运行。 98 | - 在**MacOS**仅能使用CPU功能，上无法进行GPU加速，因为Mac根本就不支持NVIDIA。 99 | 100 | > 在CPU数量充足的条件下，GPU读取速度可能比CPU读取速度稍慢。在使用感兴趣区域（ROI）操作（裁剪、调整大小、填充）时，GPU优势更凸显。 101 | 102 | ## 编解码器 103 | 104 | | 编解码器 | OpenCV读取器 | ffmpegcv-CPU读取器 | GPU读取器 | OpenCV写入器 | ffmpegcv-CPU写入器 | GPU写入器 | 105 | | ----------- | ------------- | ---------------- | ---- | ------------- | ---------------- | ---- | 106 | | h264 | √ | √ | √ | × | √ | √ | 107 | | h265 (hevc) | 不确定 | √ | √ | × | √ | √ | 108 | | mjpeg | √ | √ | × | √ | √ | × | 109 | | mpeg | √ | √ | × | √ | √ | × | 110 | | 其他 | 不确定 | ffmpeg -decoders | × | 不确定 | ffmpeg -encoders | × | 111 | 112 | ## 基准测试 113 | *正在进行中...(遥遥无期)* 114 | 115 | 116 | ## 视频读取器 117 | --- 118 | ffmpegcv与opencv在API上非常类似。 119 | ```python 120 | # OpenCV 121 | import cv2 122 | cap = cv2.VideoCapture(file) 123 | while True: 124 | ret, frame = cap.read() 125 | if not ret: 126 | break 127 | pass 128 | 129 | # ffmpegcv 130 | import ffmpegcv 131 | cap = ffmpegcv.VideoCapture(file) 132 | while True: 133 | ret, frame = cap.read() 134 | if not ret: 135 | break 136 | pass 137 | cap.release() 138 | 139 | # 另一种写法 140 | cap = ffmpegcv.VideoCapture(file) 141 | nframe = len(cap) 142 | for frame in cap: 143 | pass 144 | cap.release() 145 | 146 | # 更加Pythonic的写法，推荐使用 147 | with ffmpegcv.VideoCapture(file) as cap: 148 | nframe = len(cap) 149 | for iframe, frame in enumerate(cap): 150 | if iframe>100: break 151 | pass 152 | ``` 153 | 154 | 使用GPU加速解码。具体取决于视频编码格式。 155 | h264_nvcuvid, hevc_nvcuvid .... 156 | ```python 157 | cap_cpu = ffmpegcv.VideoCapture(file) 158 | cap_gpu = ffmpegcv.VideoCapture(file, codec='h264_cuvid') # NVIDIA GPU0 159 | cap_gpu0 = ffmpegcv.VideoCaptureNV(file) # NVIDIA GPU0 160 | cap_gpu1 = ffmpegcv.VideoCaptureNV(file, gpu=1) # NVIDIA GPU1 161 | cap_qsv = ffmpegcv.VideoCaptureQSV(file) #Intel QSV, 测试中 162 | ``` 163 | 164 | 使用`rgb24`代替`bgr24`。`gray`版本会更高效。 165 | ```python 166 | cap = ffmpegcv.VideoCapture(file, pix_fmt='rgb24') # rgb24, bgr24, gray 167 | ret, frame = cap.read() 168 | plt.imshow(frame) 169 | ``` 170 | 171 | ### 感兴趣区域（ROI）操作 172 | 您可以对视频进行裁剪、调整大小和填充。这些ROI操作中，`ffmpegcv-GPU` > `ffmpegcv-CPU` >> `opencv` 在性能上。 173 | 174 | **裁剪**视频，比读取整个画布要快得多。 175 | ```python 176 | cap = ffmpegcv.VideoCapture(file, crop_xywh=(0, 0, 640, 480)) 177 | ``` 178 | 179 | 将视频调整为给定大小的**大小**。 180 | ```python 181 | cap = ffmpegcv.VideoCapture(file, resize=(640, 480)) 182 | ``` 183 | 184 | **调整大小**并保持宽高比，使用黑色边框进行**填充**。 185 | ```python 186 | cap = ffmpegcv.VideoCapture(file, resize=(640, 480), resize_keepratio=True) 187 | ``` 188 | 189 | 对视频进行**裁剪**，然后进行**调整大小**。左上角为坐标 (0,0)。 190 | ```python 191 | cap = ffmpegcv.VideoCapture(file, crop_xywh=(0, 0, 640, 480), resize=(512, 512)) 192 | ``` 193 | 194 | ## toCUDA 将图像帧快速导出到CUDA设备 195 | --- 196 | ffmpegcv 可以将 HWC-uint8 cpu 中的视频/流转换为 CUDA 设备中的 CHW-float32。它可以显著减少你的 CPU 负载，并比你的手动转换快 2 倍以上。 197 | 198 | 准备环境。你需要具备 cuda 环境，并且安装 pycuda 包。注意，pytorch 包是非必须的。 199 | > nvcc --version # 检查你是否已经安装了 NVIDIA CUDA 编译器 200 | > pip install pycuda # 安装 pycuda 201 | 202 | ```python 203 | # 读取视频到CUDA设备，加速前 204 | cap = ffmpegcv.VideoCaptureNV(file, pix_fmt='rgb24') 205 | ret, frame_HWC_CPU = cap.read() 206 | frame_CHW_CUDA = torch.from_numpy(frame_HWC_CPU).permute(2, 0, 1).cuda().contiguous().float() # 120fps, 1200% CPU 使用率 207 | 208 | # 加速后 209 | cap = toCUDA(ffmpegcv.VideoCapture(file, pix_fmt='yuv420p')) #必须设置, yuv420p 针对 cpu 210 | cap = toCUDA(ffmpegcv.VideoCaptureNV(file, pix_fmt='nv12')) #必须设置, nv12 针对 gpu 211 | cap = toCUDA(vid, tensor_format='chw') #tensor 格式:'chw'(默认) or 'hwc' 212 | cap = toCUDA(vid, gpu=1) #选择 gpu 213 | 214 | ret, frame_CHW_pycuda = cap.read() #380fps, 200% CPU load, [pycuda array] 215 | ret, frame_CHW_pycudamem = cap.read_cudamem() #same as [pycuda mem_alloc] 216 | ret, frame_CHW_CUDA = cap.read_torch() #same as [pytorch tensor] 217 | ret, _ = cap.read_torch(frame_CHW_CUDA) #不拷贝, 但需要提前分配内存 218 | 219 | frame_CHW_pycuda[:] = (frame_CHW_pycuda - mean) / std #归一化 220 | ``` 221 | 222 | 为什么在深度学习流水线中使用 toCUDA 会更快？ 223 | 224 | > 1. ffmpeg 使用 CPU 将视频像素格式从原始 YUV 转换为 RGB24，这个过程很慢。`toCUDA` 使用 cuda 加速像素格式转换。 225 | > 2. 使用 yuv420p 或 nv12 可以节省 CPU 负载并减少从 CPU 到 GPU 的内存复制。 226 | > 3. ffmpeg 将图像存储为 HWC 格式。ffmpegcv 可以使用 HWC 和 CHW 格式来加速视频存储。 227 | 228 | ## 视频写入器 229 | --- 230 | ```python 231 | # cv2 232 | out = cv2.VideoWriter('outpy.avi', 233 | cv2.VideoWriter_fourcc('M','J','P','G'), 234 | 10, 235 | (w, h)) 236 | out.write(frame1) 237 | out.write(frame2) 238 | out.release() 239 | 240 | # ffmpegcv，默认的编码器为'h264'在CPU上，'h265'在GPU上。 241 | # 帧大小由第一帧决定 242 | # 使用 'mp4/mkv' 来替代古老的 'avi' 格式 243 | out = ffmpegcv.VideoWriter('outpy.mp4', None, 10) 244 | out.write(frame1) 245 | out.write(frame2) 246 | out.release() 247 | 248 | # 更加Pythonic的写法 249 | with ffmpegcv.VideoWriter('outpy.mp4', None, 10) as out: 250 | out.write(frame1) 251 | out.write(frame2) 252 | ``` 253 | 254 | 使用GPU加速编码。例如h264_nvenc，hevc_nvenc。 255 | ```python 256 | out_cpu = ffmpegcv.VideoWriter('outpy.mp4', None, 10) 257 | out_gpu0 = ffmpegcv.VideoWriterNV('outpy.mp4', 'h264', 10) # NVIDIA GPU0 258 | out_gpu1 = ffmpegcv.VideoWriterNV('outpy.mp4', 'hevc', 10, gpu=1) # NVIDIA GPU1 259 | out_qsv = ffmpegcv.VideoWriterQSV('outpy.mp4', 'h264', 10) #Intel QSV, 测试中 260 | ``` 261 | 262 | 输入图像使用rgb24而不是bgr24。 263 | ```python 264 | out = ffmpegcv.VideoWriter('outpy.mp4', None, 10, pix_fmt='rgb24') 265 | ``` 266 | 267 | 缩放图像尺寸 268 | ```python 269 | out_resz = ffmpegcv.VideoWriter('outpy.mp4', None, 10, resize=(640, 480)) 270 | ``` 271 | 272 | ## 视频读取器和写入器 273 | --- 274 | ```python 275 | import ffmpegcv 276 | vfile_in = 'A.mp4' 277 | vfile_out = 'A_h264.mp4' 278 | vidin = ffmpegcv.VideoCapture(vfile_in) 279 | vidout = ffmpegcv.VideoWriter(vfile_out, None, vidin.fps) 280 | 281 | with vidin, vidout: 282 | for frame in vidin: 283 | vidout.write(frame) 284 | ``` 285 | 286 | ## 相机读取器 287 | --- 288 | **实验性功能**。ffmpegcv提供了相机读取器。与VideoCapture读取器一致。 289 | 290 | - VideoCaptureCAM旨在支持感兴趣区域（ROI）操作。在相机读取方面，Opencv比ffmpegcv更具吸引力。**对于大多数相机读取情况，我推荐使用Opencv**。 291 | - ffmpegcv可以使用名称检索相机设备，使用`ffmpegcv.VideoCaptureCAM("Integrated Camera")`比使用`cv2.VideoCaptureCAM(0)`更易读。 292 | - 如果后处理时间过长，VideoCaptureCAM将会出现卡顿和丢帧。VideoCaptureCAM会缓冲最近的帧。 293 | - 即使没有读取视频帧，VideoCaptureCAM也会在后台不断工作。**请及时释放资源**。 294 | - 在Windows上表现良好，在Linux和macOS上表现不完美。 295 | 296 | ```python 297 | import cv2 298 | cap = cv2.VideoCapture(0) 299 | while True: 300 | ret, frame = cap.read() 301 | cv2.imshow('frame', frame) 302 | if cv2.waitKey(1) & 0xFF == ord('q'): 303 | break 304 | cap.release() 305 | 306 | # ffmpegcv，在Windows和Linux上 307 | import ffmpegcv 308 | cap = ffmpegcv.VideoCaptureCAM(0) 309 | while True: 310 | ret, frame = cap.read() 311 | cv2.imshow('frame', frame) 312 | if cv2.waitKey(1) & 0xFF == ord('q'): 313 | break 314 | cap.release() 315 | 316 | # ffmpegcv 使用相机名称，在Windows和Linux上 317 | cap = ffmpegcv.VideoCaptureCAM("Integrated Camera") 318 | 319 | # ffmpegcv 使用相机路径（避免多个相机冲突） 320 | cap = ffmpegcv.VideoCaptureCAM('@device_pnp_\\\\?\\usb#vid_2304&' 321 | 'pid_oot#media#0001#{65e8773d-8f56-11d0-a3b9-00a0c9223196}' 322 | '\\global') 323 | 324 | # ffmpegcv 使用具有ROI操作的相机 325 | cap = ffmpegcv.VideoCaptureCAM("Integrated Camera", crop_xywh=(0, 0, 640, 480), resize=(512, 512), resize_keepratio=True) 326 | 327 | 328 | ``` 329 | 330 | **列出所有相机设备** 331 | ```python 332 | from ffmpegcv.ffmpeg_reader_camera import query_camera_devices 333 | 334 | devices = query_camera_devices() 335 | print(devices) 336 | ``` 337 | >{0: ('Integrated Camera', '@device_pnp_\\\\?\\usb#vid_2304&pid_oot#media#0001#{65e8773d-8f56-11d0-a3b9-00a0c9223196}\\global'), 338 | 1: ('OBS Virtual Camera', '@device_sw_{860BB310-5D01-11D0-BD3B-00A0C911CE86}\\{A3FCE0F5-3493-419F-958A-ABA1250EC20B}')} 339 | 340 | 341 | **设置相机的分辨率、帧率、视频编码/像素格式** 342 | 343 | ```python 344 | from ffmpegcv.ffmpeg_reader_camera import query_camera_options 345 | 346 | options = query_camera_options(0) # 或者 query_camera_options("Integrated Camera") 347 | print(options) 348 | cap = ffmpegcv.VideoCaptureCAM(0, **options[-1]) 349 | ``` 350 | >[{'camcodec': 'mjpeg', 'campix_fmt': None, 'camsize_wh': (1280, 720), 'camfps': 60.0002}, {'camcodec': 'mjpeg', 'campix_fmt': None, 'camsize_wh': (640, 480), 'camfps': 60.0002}, {'camcodec': 'mjpeg', 'campix_fmt': None, 'camsize_wh': (1920, 1080), 'camfps': 60.0002}, {'camcodec': None, 'campix_fmt': 'yuyv422', 'camsize_wh': (1280, 720), 'camfps': 10}, {'camcodec': None, 'campix_fmt': 'yuyv422', 'camsize_wh': (640, 480), 'camfps': 30}, {'camcodec': None, 'campix_fmt': 'yuyv422', 'camsize_wh': (1920, 1080), 'camfps': 5}] 351 | 352 | **已知问题** 353 | 1. VideoCaptureCAM在macOS上的体验不太流畅。你必须指定所有相机参数。而且query_camera_options不会给出任何建议。这是因为`ffmpeg`无法使用mac本机的`avfoundation`列出设备选项。 354 | ```python 355 | # macOS需要提供完整参数。 356 | cap = ffmpegcv.VideoCaptureCAM('FaceTime HD Camera', camsize_wh=(1280,720), camfps=30, campix_fmt='nv12') 357 | ``` 358 | 359 | 2. 在Linux上VideoCaptureCAM无法列出FPS，因为`ffmpeg`无法使用Linux本机的`v4l2`模块查询设备的FPS。不过，让FPS为空也没问题。 360 | 361 | ## 流读取器 （直播流，网络监控摄像头） 362 | **实验性功能**。ffmpegcv提供了流读取器，与VideoFile读取器一致，更类似于相机。 363 | 364 | - 支持`RTSP`、`RTP`、`RTMP`、`HTTP`、`HTTPS`流。 365 | - 如果后处理时间过长，VideoCaptureStream会出现卡顿和丢帧。VideoCaptureCAM会缓冲最近的帧。 366 | - 即使没有读取视频帧，VideoCaptureStream也会在后台不断工作。**请及时释放资源**。 367 | - 这仍然是实验性功能。建议您使用opencv。 368 | 369 | ```python 370 | # opencv 371 | import cv2 372 | stream_url = 'http://devimages.apple.com.edgekey.net/streaming/examples/bipbop_4x3/gear2/prog_index.m3u8' 373 | cap = cv2.VideoCapture(stream_url, cv2.CAP_FFMPEG) 374 | 375 | if not cap.isOpened(): 376 | print('无法打开流') 377 | exit(-1) 378 | 379 | while True: 380 | ret, frame = cap.read() 381 | if not ret: 382 | break 383 | pass 384 | 385 | # ffmpegcv 386 | import ffmpegcv 387 | cap = ffmpegcv.VideoCaptureStream(stream_url) 388 | while True: 389 | ret, frame = cap.read() 390 | if not ret: 391 | break 392 | pass 393 | 394 | # ffmpegcv, 网络监控摄像头 395 | # 例如 海康威视, `101` 主视频流, `102` 子视频流 396 | stream_url = 'rtsp://admin:PASSWD@192.168.1.xxx:8554/Streaming/Channels/102' 397 | cap = ffmpegcv.VideoCaptureStreamRT(stream_url) # 低延迟 & 缓存 398 | cap = ffmpegcv.ReadLiveLast(ffmpegcv.VideoCaptureStreamRT, stream_url) #不缓存 399 | while True: 400 | ret, frame = cap.read() 401 | if not ret: 402 | break 403 | pass 404 | ``` 405 | 406 | ## FFmpegReaderNoblock 407 | 更快的读写取视频。利用多进程在后台自动准备帧，这样在读写当前帧时不会阻塞。这使得您的Python程序在CPU使用方面更高效。带来最大翻倍效率提升。 408 | 409 | > ffmpegcv.VideoCapture(*args) -> ffmpegcv.noblock(ffmpegcv.VideoCapture, *args) 410 | > 411 | > ffmpegcv.VideoWriter(*args) -> ffmpegcv.noblock(ffmpegcv.VideoWriter, *args) 412 | 413 | ```python 414 | # 代理任何 VideoCapture&VideoWriter 的参数和kargs 415 | vid_noblock = ffmpegcv.noblock(ffmpegcv.VideoCapture, vfile, pix_fmt='rbg24') 416 | 417 | # 这很快 418 | def cpu_tense(): time.sleep(0.01) 419 | for _ in tqdm.trange(1000): 420 | ret, img = vid_noblock.read() #当前图像已经被缓冲，不会占用时间 421 | cpu_tense() #同时，下一帧在后台缓冲 422 | 423 | # 这很慢 424 | vid = ffmpegcv.VideoCapture(vfile, pix_fmt='rbg24') 425 | for _ in tqdm.trange(2000): 426 | ret, img = vid.read() #此读取将阻塞CPU，占用时间 427 | cpu_tense() 428 | ``` 429 | -------------------------------------------------------------------------------- /ffmpegcv/__init__.py: -------------------------------------------------------------------------------- 1 | from .ffmpeg_reader import FFmpegReader, FFmpegReaderNV 2 | from .ffmpeg_writer import FFmpegWriter, FFmpegWriterNV 3 | from .ffmpeg_reader_camera import FFmpegReaderCAM 4 | from .ffmpeg_reader_stream import FFmpegReaderStream 5 | from .ffmpeg_reader_stream_realtime import FFmpegReaderStreamRT, FFmpegReaderStreamRTNV 6 | from .ffmpeg_writer_stream_realtime import FFmpegWriterStreamRT 7 | from .ffmpeg_reader_qsv import FFmpegReaderQSV 8 | from .ffmpeg_writer_qsv import FFmpegWriterQSV 9 | from .ffmpeg_reader_pannels import FFmpegReaderPannels 10 | from .ffmpeg_noblock import noblock, ReadLiveLast 11 | from .video_info import get_num_NVIDIA_GPUs 12 | import shutil 13 | from subprocess import DEVNULL, check_output 14 | 15 | from .version import __version__ 16 | 17 | 18 | def _check(): 19 | if not shutil.which("ffmpeg") or not shutil.which("ffprobe"): 20 | raise RuntimeError( 21 | "The ffmpeg is not installed. \n\n" 22 | "Please install ffmpeg via:\n " 23 | "conda install ffmpeg" 24 | ) 25 | 26 | 27 | _check() 28 | 29 | _check_nvidia_init = None 30 | 31 | 32 | def _check_nvidia(): 33 | global _check_nvidia_init 34 | run = lambda x: check_output(x, shell=True, stderr=DEVNULL) 35 | if _check_nvidia_init is None: 36 | calling_output = run("ffmpeg -h encoder=hevc_nvenc") 37 | if "AVOptions" not in calling_output.decode("utf-8"): 38 | raise RuntimeError( 39 | "The ffmpeg is not compiled with NVENC support.\n\n" 40 | "Please re-compile ffmpeg following the instructions at:\n " 41 | "https://docs.nvidia.com/video-technologies/video-codec-sdk/ffmpeg-with-nvidia-gpu/" 42 | ) 43 | 44 | calling_output = run("ffmpeg -h decoder=hevc_cuvid") 45 | if "AVOptions" not in calling_output.decode("utf-8"): 46 | raise RuntimeError( 47 | "The ffmpeg is not compiled with NVENC support.\n\n" 48 | "Please re-compile ffmpeg following the instructions at:\n " 49 | "https://docs.nvidia.com/video-technologies/video-codec-sdk/ffmpeg-with-nvidia-gpu/" 50 | ) 51 | 52 | if get_num_NVIDIA_GPUs() == 0: 53 | raise RuntimeError( 54 | "No NVIDIA GPU found.\n\n" 55 | "Please use a NVIDIA GPU card listed at:\n " 56 | "https://developer.nvidia.com/video-encode-and-decode-gpu-support-matrix-new" 57 | ) 58 | 59 | _check_nvidia_init = True 60 | 61 | return True 62 | 63 | 64 | def VideoCapture( 65 | file, 66 | codec=None, 67 | pix_fmt="bgr24", 68 | crop_xywh=None, 69 | resize=None, 70 | resize_keepratio=True, 71 | resize_keepratioalign="center", 72 | ) -> FFmpegReader: 73 | """ 74 | Alternative to cv2.VideoCapture 75 | 76 | Parameters 77 | ---------- 78 | file : str 79 | Path to video file. 80 | codec : str 81 | Codec to use. Optional. Default is `None`. 82 | pix_fmt : str 83 | Pixel format. ['bgr24' | 'rgb24']. Optional. Default is 'bgr24'. 84 | crop_xywh : tuple 85 | Crop the frame. (x, y, width, height). Optional. Default is `None`. 86 | resize : tuple 87 | Resize the video to the given size. Optional. Default is `None`. 88 | resize_keepratio : bool 89 | Keep the aspect ratio and the border is black. Optional. Default is `True`. 90 | resize_keepratioalign : str 91 | Align the image to the `center`, `topleft`, `topright`, `bottomleft` or 92 | `bottomright`. Optional. Default is 'center'. 93 | 94 | Examples 95 | -------- 96 | opencv 97 | ``` 98 | cap = cv2.VideoCapture(file) 99 | while True: 100 | ret, frame = cap.read() 101 | if not ret: 102 | break 103 | pass 104 | ``` 105 | 106 | ffmpegcv 107 | ``` 108 | cap = ffmpegcv.VideoCapture(file) 109 | while True: 110 | ret, frame = cap.read() 111 | if not ret: 112 | break 113 | pass 114 | ``` 115 | 116 | Or use iterator 117 | ``` 118 | cap = ffmpegcv.VideoCapture(file) 119 | for frame in cap: 120 | pass 121 | counts = len(cap) 122 | ``` 123 | 124 | Use GPU to accelerate decoding 125 | ``` 126 | cap_cpu = ffmpegcv.VideoCapture(file) 127 | cap_gpu = ffmpegcv.VideoCaptureNV(file) 128 | ``` 129 | 130 | Use rgb24 instead of bgr24 131 | ``` 132 | cap = ffmpegcv.VideoCapture(file, pix_fmt='rgb24') 133 | ``` 134 | 135 | Crop video. 136 | ```python 137 | cap = ffmpegcv.VideoCapture(file, crop_xywh=(0, 0, 640, 480)) 138 | ``` 139 | 140 | Resize the video to the given size 141 | ``` 142 | cap = ffmpegcv.VideoCapture(file, resize=(640, 480)) 143 | ``` 144 | 145 | Resize and keep the aspect ratio with black border 146 | ``` 147 | cap = ffmpegcv.VideoCapture(file, resize=(640, 480), resize_keepratio=True) 148 | ``` 149 | 150 | Crop and then resize the video. 151 | ```python 152 | cap = ffmpegcv.VideoCapture(file, crop_xywh=(0, 0, 640, 480), resize=(512, 512)) 153 | ``` 154 | Author: Chenxinfeng 2022-04-16, cxf529125853@163.com 155 | """ 156 | return FFmpegReader.VideoReader( 157 | file, codec, pix_fmt, crop_xywh, resize, resize_keepratio, resize_keepratioalign 158 | ) 159 | 160 | 161 | VideoReader = VideoCapture 162 | 163 | 164 | def VideoWriter( 165 | file, codec=None, fps=30, pix_fmt="bgr24", bitrate=None, resize=None, preset=None 166 | ) -> FFmpegWriter: 167 | """ 168 | Alternative to cv2.VideoWriter 169 | 170 | Parameters 171 | ---------- 172 | file : str 173 | Path to video file. 174 | codec : str 175 | Codec to use. Optional. Default is `None` (x264). 176 | fps : number 177 | Frames per second. Optional. Default is 30. 178 | pix_fmt : str 179 | Pixel format of input. ['bgr24' | 'rgb24']. Optional. Default is 'bgr24'. 180 | bitrate : str 181 | Bitrate of output video. Optional. Default is `None`. 182 | resize : tuple 183 | Frame size of output. (width, height). Optional. Default is `None`. 184 | preset : str 185 | Preset of ffmpeg. Optional. Default is `None`. 186 | Examples 187 | -------- 188 | opencv 189 | ``` 190 | out = cv2.VideoWriter('outpy.avi', 191 | cv2.VideoWriter_fourcc('M','J','P','G'), 192 | 10, 193 | (w, h)) 194 | out.write(frame1) 195 | out.write(frame2) 196 | out.release() 197 | ``` 198 | 199 | ffmpegcv 200 | ``` 201 | out = ffmpegcv.VideoWriter('outpy.avi', None, 10) 202 | out.write(frame1) 203 | out.write(frame2) 204 | out.release() 205 | ``` 206 | 207 | frameSize is decided by the size of the first frame 208 | ``` 209 | out = ffmpegcv.VideoWriter('outpy.avi', None, 10) 210 | ``` 211 | 212 | Use GPU to accelerate encoding 213 | ``` 214 | out_cpu = ffmpegcv.VideoWriter('outpy.avi', None, 10) 215 | out_gpu = ffmpegcv.VideoWriter('outpy.avi', 'h264_nvenc', 10) 216 | ``` 217 | 218 | Use rgb24 instead of bgr24 219 | ``` 220 | out = ffmpegcv.VideoWriter('outpy.avi', None, 10, pix_fmt='rgb24') 221 | out.write(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) 222 | ``` 223 | 224 | Author: Chenxinfeng 2022-04-16, cxf529125853@163.com 225 | """ 226 | return FFmpegWriter.VideoWriter( 227 | file, codec, fps, pix_fmt, bitrate, resize, preset=preset 228 | ) 229 | 230 | 231 | def VideoCaptureNV( 232 | file, 233 | pix_fmt="bgr24", 234 | crop_xywh=None, 235 | resize=None, 236 | resize_keepratio=True, 237 | resize_keepratioalign="center", 238 | gpu=0, 239 | ) -> FFmpegReaderNV: 240 | """ 241 | `ffmpegcv.VideoCaptureNV` is a gpu version for `ffmpegcv.VideoCapture`. 242 | """ 243 | _check_nvidia() 244 | return FFmpegReaderNV.VideoReader( 245 | file, pix_fmt, crop_xywh, resize, resize_keepratio, resize_keepratioalign, gpu 246 | ) 247 | 248 | 249 | VideoReaderNV = VideoCaptureNV 250 | 251 | 252 | def VideoCaptureQSV( 253 | file, 254 | pix_fmt="bgr24", 255 | crop_xywh=None, 256 | resize=None, 257 | resize_keepratio=True, 258 | resize_keepratioalign="center", 259 | gpu=0, 260 | ) -> FFmpegReaderQSV: 261 | """ 262 | `ffmpegcv.VideoCaptureQSV` is a gpu version for `ffmpegcv.VideoCapture`. 263 | """ 264 | return FFmpegReaderQSV.VideoReader( 265 | file, pix_fmt, crop_xywh, resize, resize_keepratio, resize_keepratioalign, gpu 266 | ) 267 | 268 | 269 | VideoReaderQSV = VideoCaptureQSV 270 | 271 | 272 | def VideoWriterNV( 273 | file, 274 | codec=None, 275 | fps=30, 276 | pix_fmt="bgr24", 277 | gpu=0, 278 | bitrate=None, 279 | resize=None, 280 | preset=None, 281 | ) -> FFmpegWriterNV: 282 | """ 283 | `ffmpegcv.VideoWriterNV` is a gpu version for `ffmpegcv.VideoWriter`. 284 | """ 285 | _check_nvidia() 286 | return FFmpegWriterNV.VideoWriter( 287 | file, codec, fps, pix_fmt, gpu, bitrate, resize, preset=preset 288 | ) 289 | 290 | 291 | def VideoWriterQSV( 292 | file, 293 | codec=None, 294 | fps=30, 295 | pix_fmt="bgr24", 296 | gpu=0, 297 | bitrate=None, 298 | resize=None, 299 | preset=None, 300 | ) -> FFmpegWriterQSV: 301 | """ 302 | `ffmpegcv.VideoWriterQSV` is a gpu version for `ffmpegcv.VideoWriter`. 303 | """ 304 | return FFmpegWriterQSV.VideoWriter( 305 | file, codec, fps, pix_fmt, gpu, bitrate, resize, preset=preset 306 | ) 307 | 308 | 309 | def VideoWriterStreamRT( 310 | url, pix_fmt="bgr24", bitrate=None, resize=None, preset=None 311 | ) -> FFmpegWriterStreamRT: 312 | return FFmpegWriterStreamRT.VideoWriter( 313 | url, "libx264", pix_fmt, bitrate, resize, preset 314 | ) 315 | 316 | 317 | def VideoCaptureCAM( 318 | camname, 319 | pix_fmt="bgr24", 320 | crop_xywh=None, 321 | resize=None, 322 | resize_keepratio=True, 323 | resize_keepratioalign="center", 324 | camsize_wh=None, 325 | camfps=None, 326 | camcodec=None, 327 | campix_fmt=None, 328 | ) -> FFmpegReaderCAM: 329 | """ 330 | Alternative to cv2.VideoCapture 331 | 332 | Parameters 333 | ---------- 334 | file : see ffmpegcv.VideoReader 335 | codec : see ffmpegcv.VideoReader 336 | pix_fmt : see ffmpegcv.VideoReader 337 | crop_xywh : see ffmpegcv.VideoReader 338 | resize : see ffmpegcv.VideoReader 339 | resize_keepratio : see ffmpegcv.VideoReader 340 | resize_keepratioalign : see ffmpegcv.VideoReader 341 | camsize_wh: tuple or None 342 | Camera resolution (width, height). e.g (800, 600) 343 | camfps: float or None 344 | Camera framerate. e.g. 30. 345 | camcodec: str or None 346 | Camera codec. e.g. 'mjpeg' or 'h264'. 347 | campix_fmt: str or None 348 | Camera pixel format. e.g. 'rgb24' or 'yuv420p'. 349 | Just set one of `camcodec` or `campix_fmt`. 350 | Examples 351 | -------- 352 | opencv 353 | ``` 354 | cap = cv2.VideoCapture(0) 355 | while True: 356 | ret, frame = cap.read() 357 | if not ret: 358 | break 359 | pass 360 | ``` 361 | 362 | ffmpegcv 363 | ``` 364 | cap = ffmpegcv.VideoCaptureCAM(0) 365 | while True: 366 | ret, frame = cap.read() 367 | if not ret: 368 | break 369 | pass 370 | ``` 371 | 372 | Or use camera name 373 | ``` 374 | cap = ffmpegcv.VideoCaptureCAM("Integrated Camera") 375 | ``` 376 | 377 | Use full camera parameter 378 | ``` 379 | cap = ffmpegcv.VideoCaptureCAM('FaceTime HD Camera', 380 | camsize_wh = (1280,720), 381 | camfps = 30, 382 | campix_fmt = 'nv12') 383 | ``` 384 | 385 | Use camera with ROI operations 386 | ``` 387 | cap = ffmpegcv.VideoCaptureCAM("Integrated Camera", 388 | crop_xywh = (0, 0, 640, 480), 389 | resize = (512, 512), 390 | resize_keepratio = True) 391 | ``` 392 | Author: Chenxinfeng 2023-05-11, cxf529125853@163.com 393 | """ 394 | return FFmpegReaderCAM.VideoReader( 395 | camname, 396 | pix_fmt, 397 | crop_xywh, 398 | resize, 399 | resize_keepratio, 400 | resize_keepratioalign, 401 | camsize_wh=camsize_wh, 402 | camfps=camfps, 403 | camcodec=camcodec, 404 | campix_fmt=campix_fmt, 405 | ) 406 | 407 | 408 | VideoReaderCAM = VideoCaptureCAM 409 | 410 | 411 | def VideoCaptureStream( 412 | stream_url, 413 | codec=None, 414 | pix_fmt="bgr24", 415 | crop_xywh=None, 416 | resize=None, 417 | resize_keepratio=True, 418 | resize_keepratioalign="center", 419 | timeout=None, 420 | ) -> FFmpegReaderStream: 421 | """ 422 | Alternative to cv2.VideoCapture 423 | 424 | Parameters 425 | ---------- 426 | stream_url : RTSP, RTP, RTMP, HTTP, HTTPS url 427 | codec : see ffmpegcv.VideoReader 428 | pix_fmt : see ffmpegcv.VideoReader 429 | crop_xywh : see ffmpegcv.VideoReader 430 | resize : see ffmpegcv.VideoReader 431 | resize_keepratio : see ffmpegcv.VideoReader 432 | resize_keepratioalign : see ffmpegcv.VideoReader 433 | timeout : waits in seconds for stream video to connect 434 | 435 | Examples 436 | -------- 437 | opencv 438 | ``` 439 | stream_url = 'http://devimages.apple.com.edgekey.net/streaming/examples/bipbop_4x3/gear2/prog_index.m3u8' 440 | cap = cv2.VideoCapture(stream_url, cv2.CAP_FFMPEG) 441 | 442 | if not cap.isOpened(): 443 | print('Cannot open the stream') 444 | exit(-1) 445 | 446 | while True: 447 | ret, frame = cap.read() 448 | if not ret: 449 | break 450 | pass 451 | ``` 452 | 453 | ffmpegcv 454 | ``` 455 | cap = ffmpegcv.VideoCaptureStream(stream_url) 456 | while True: 457 | ret, frame = cap.read() 458 | if not ret: 459 | break 460 | pass 461 | ``` 462 | 463 | Author: Chenxinfeng 2023-05-31, cxf529125853@163.com 464 | """ 465 | return FFmpegReaderStream.VideoReader( 466 | stream_url, 467 | codec, 468 | pix_fmt, 469 | crop_xywh, 470 | resize, 471 | resize_keepratio, 472 | resize_keepratioalign, 473 | timeout, 474 | ) 475 | 476 | 477 | VideoReaderStream = VideoCaptureStream 478 | 479 | 480 | def VideoCaptureStreamRT( 481 | stream_url, 482 | codec=None, 483 | pix_fmt="bgr24", 484 | crop_xywh=None, 485 | resize=None, 486 | resize_keepratio=True, 487 | resize_keepratioalign="center", 488 | gpu=None, 489 | timeout=None, 490 | ) -> FFmpegReaderStreamRT: 491 | if gpu is None: 492 | return FFmpegReaderStreamRT.VideoReader( 493 | stream_url, 494 | codec, 495 | pix_fmt, 496 | crop_xywh, 497 | resize, 498 | resize_keepratio, 499 | resize_keepratioalign, 500 | timeout=timeout, 501 | ) 502 | else: 503 | return FFmpegReaderStreamRTNV.VideoReader( 504 | stream_url, 505 | codec, 506 | pix_fmt, 507 | crop_xywh, 508 | resize, 509 | resize_keepratio, 510 | resize_keepratioalign, 511 | gpu=gpu, 512 | timeout=timeout, 513 | ) 514 | 515 | 516 | VideoReaderStreamRT = VideoCaptureStreamRT 517 | 518 | 519 | def VideoCapturePannels( 520 | file: str, crop_xywh_l: list, codec=None, pix_fmt="bgr24", resize=None 521 | ): 522 | """ 523 | Alternative to cv2.VideoCapture 524 | 525 | Parameters 526 | ---------- 527 | file : str 528 | Path to video file. 529 | crop_xywh_l : list of crop_xywh 530 | Crop the frame. [(x0, y0, w0, h0), (x1, y1, w1, h1), ...]. 531 | codec : str 532 | Codec to use. Optional. Default is `None`. 533 | pix_fmt : str 534 | Pixel format. ['bgr24' | 'rgb24' | 'gray']. Optional. Default is 'bgr24'. 535 | resize : tuple as (w,h) 536 | Resize the pannels to identical size. Optional. Default is `None`. 537 | Does not keep the ratio. 538 | 539 | 540 | Examples 541 | -------- 542 | 543 | 544 | ffmpegcv 545 | ``` 546 | w,h = 1280,800 547 | cap = ffmpegcv.VideoCapturePannels(file, 548 | [[0,0,w,h], [w,0,w,h],[0,h,w,h], [w,h,w,h]]) 549 | while True: 550 | ret, frame_pannels = cap.read() 551 | if not ret: 552 | break 553 | print(frame_pannels.shape) # shape=(4,h,w,3) 554 | ``` 555 | 556 | Author: Chenxinfeng 2024-02-15, cxf529125853@163.com 557 | """ 558 | 559 | return FFmpegReaderPannels.VideoReader( 560 | file, 561 | crop_xywh_l, 562 | codec, 563 | pix_fmt, 564 | resize, 565 | ) 566 | 567 | 568 | VideoReaderPannels = VideoCapturePannels 569 | 570 | 571 | def toCUDA(vid: FFmpegReader, gpu: int = 0, tensor_format: str = "chw") -> FFmpegReader: 572 | """ 573 | Convert frames to CUDA tensor float32 in 'chw' or 'hwc' format. 574 | """ 575 | from ffmpegcv.ffmpeg_reader_cuda import FFmpegReaderCUDA 576 | 577 | return FFmpegReaderCUDA(vid, gpu, tensor_format) 578 | -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_noblock.py: -------------------------------------------------------------------------------- 1 | from .ffmpeg_reader_noblock import FFmpegReaderNoblock 2 | from .ffmpeg_writer_noblock import FFmpegWriterNoblock 3 | from typing import Callable 4 | import ffmpegcv 5 | import threading 6 | import numpy as np 7 | import queue 8 | 9 | 10 | def noblock(fun:Callable, *v_args, **v_kargs): 11 | readerfuns = (ffmpegcv.VideoCapture, ffmpegcv.VideoCaptureNV) 12 | writerfuns = (ffmpegcv.VideoWriter, ffmpegcv.VideoWriterNV) 13 | 14 | if fun in readerfuns: 15 | proxyfun = FFmpegReaderNoblock(fun, *v_args, **v_kargs) 16 | elif fun in writerfuns: 17 | proxyfun = FFmpegWriterNoblock(fun, *v_args, **v_kargs) 18 | else: 19 | raise ValueError('The function is not supported as a Reader or Writer') 20 | 21 | return proxyfun 22 | 23 | 24 | class ReadLiveLast(threading.Thread, ffmpegcv.FFmpegReader): 25 | def __init__(self, fun, *args, **kvargs): 26 | threading.Thread.__init__(self) 27 | ffmpegcv.FFmpegReader.__init__(self) 28 | self.vid = vid = fun(*args, **kvargs) 29 | props_name = ['width', 'height', 'fps', 'count', 'codec', 'ffmpeg_cmd', 30 | 'size', 'pix_fmt', 'out_numpy_shape', 'iframe', 31 | 'duration', 'origin_width', 'origin_height'] 32 | for name in props_name: 33 | setattr(self, name, getattr(vid, name, None)) 34 | 35 | self.img = np.zeros(self.out_numpy_shape, dtype=np.uint8) 36 | self.ret = True 37 | self._isopen = True 38 | self._q = queue.Queue(maxsize=1) # synchronize new frame 39 | self._lock = threading.Lock() 40 | self.start() 41 | 42 | def read(self): 43 | if self.ret: 44 | self._q.get() # if reading too freq, then wait until new frame 45 | self.iframe += 1 46 | return self.ret, self.img 47 | 48 | def release(self): 49 | with self._lock: 50 | self._isopen = False 51 | self.vid.release() 52 | 53 | def run(self): 54 | while True: 55 | with self._lock: 56 | if self._isopen: 57 | self.ret, self.img = self.vid.read() 58 | else: 59 | break 60 | if not self._q.full(): 61 | self._q.put(None) 62 | if not self.ret: 63 | break -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_reader.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pprint 3 | import warnings 4 | import os 5 | import sys 6 | import select 7 | from .video_info import ( 8 | run_async_reader as run_async, 9 | get_info, 10 | get_num_NVIDIA_GPUs, 11 | decoder_to_nvidia, 12 | release_process, 13 | ) 14 | 15 | 16 | def get_videofilter_cpu( 17 | originsize: list, 18 | pix_fmt: str, 19 | crop_xywh: list, 20 | resize: list, 21 | resize_keepratio: bool, 22 | resize_keepratioalign: str, 23 | ): 24 | """ 25 | ONGONING: common filter for video/cam/stream capture. 26 | """ 27 | assert pix_fmt in ["rgb24", "bgr24", "yuv420p", "yuvj420p", "nv12", "gray"] 28 | origin_width, origin_height = originsize 29 | if crop_xywh: 30 | crop_w, crop_h = crop_xywh[2:] 31 | if not all([n % 2 == 0] for n in crop_xywh): 32 | print("Warning 'crop_xywh' would be replaced into even numbers") 33 | crop_xywh = [int(n//2*2) for n in crop_xywh] 34 | assert crop_w <= origin_width and crop_h <= origin_height 35 | x, y, w, h = crop_xywh 36 | cropopt = f"crop={w}:{h}:{x}:{y}" 37 | else: 38 | crop_w, crop_h = origin_width, origin_height 39 | cropopt = "" 40 | 41 | crop_wh = (crop_w, crop_h) 42 | if resize is None or tuple(resize) == crop_wh: 43 | scaleopt = "" 44 | padopt = "" 45 | final_size_wh = crop_wh 46 | else: 47 | final_size_wh = (dst_width, dst_height) = resize 48 | assert all([n % 2 == 0] for n in resize), "'resize' must be even number" 49 | if not resize_keepratio: 50 | scaleopt = f"scale={dst_width}x{dst_height}" 51 | padopt = "" 52 | else: 53 | re_width, re_height = crop_w / (crop_h / dst_height), dst_height 54 | if re_width > dst_width: 55 | re_width, re_height = dst_width, crop_h / (crop_w / dst_width) 56 | re_width, re_height = int(re_width), int(re_height) 57 | scaleopt = f"scale={re_width}x{re_height}" 58 | if resize_keepratioalign is None: 59 | resize_keepratioalign = "center" 60 | paddings = { 61 | "center": ((dst_width - re_width) // 2, (dst_height - re_height) // 2,), 62 | "topleft": (0, 0), 63 | "topright": (dst_width - re_width, 0), 64 | "bottomleft": (0, dst_height - re_height), 65 | "bottomright": (dst_width - re_width, dst_height - re_height), 66 | } 67 | assert ( 68 | resize_keepratioalign in paddings 69 | ), 'resize_keepratioalign must be one of "center"(mmpose), "topleft"(mmdetection), "topright", "bottomleft", "bottomright"' 70 | xpading, ypading = paddings[resize_keepratioalign] 71 | padopt = f"pad={dst_width}:{dst_height}:{xpading}:{ypading}:black" 72 | 73 | pix_fmtopt = "extractplanes=y" if pix_fmt == "gray" else "" 74 | if any([cropopt, scaleopt, padopt, pix_fmtopt]): 75 | filterstr = ",".join(x for x in [cropopt, scaleopt, padopt, pix_fmtopt] if x) 76 | filteropt = f"-vf {filterstr}" 77 | else: 78 | filteropt = "" 79 | return crop_wh, final_size_wh, filteropt 80 | 81 | 82 | def get_videofilter_gpu( 83 | originsize: list, 84 | pix_fmt: str, 85 | crop_xywh: list, 86 | resize: list, 87 | resize_keepratio: bool, 88 | resize_keepratioalign: str, 89 | ): 90 | assert pix_fmt in ["rgb24", "bgr24", "yuv420p", "yuvj420p", "nv12", "gray"] 91 | origin_width, origin_height = originsize 92 | if crop_xywh: 93 | crop_w, crop_h = crop_xywh[2:] 94 | assert all([n % 2 == 0] for n in crop_xywh), "'crop_xywh' must be even number" 95 | assert crop_w <= origin_width and crop_h <= origin_height 96 | x, y, w, h = crop_xywh 97 | top, bottom, left, right = ( 98 | y, 99 | origin_height - (y + h), 100 | x, 101 | origin_width - (x + w), 102 | ) # crop length 103 | cropopt = f"-crop {top}x{bottom}x{left}x{right}" 104 | else: 105 | crop_w, crop_h = origin_width, origin_height 106 | cropopt = "" 107 | 108 | crop_wh = (crop_w, crop_h) 109 | filteropt = "" 110 | scaleopt = "" 111 | if resize is None or tuple(resize) == crop_wh: 112 | final_size_wh = crop_wh 113 | else: 114 | final_size_wh = (dst_width, dst_height) = resize 115 | assert all([n % 2 == 0] for n in resize), "'resize' must be even number" 116 | if not resize_keepratio: 117 | scaleopt = f"-resize {dst_width}x{dst_height}" 118 | else: 119 | re_width, re_height = crop_w / (crop_h / dst_height), dst_height 120 | if re_width > dst_width: 121 | re_width, re_height = dst_width, crop_h / (crop_w / dst_width) 122 | re_width, re_height = int(re_width), int(re_height) 123 | scaleopt = f"-resize {re_width}x{re_height}" 124 | if resize_keepratioalign is None: 125 | resize_keepratioalign = "center" 126 | paddings = { 127 | "center": ((dst_width - re_width) // 2, (dst_height - re_height) // 2,), 128 | "topleft": (0, 0), 129 | "topright": (dst_width - re_width, 0), 130 | "bottomleft": (0, dst_height - re_height), 131 | "bottomright": (dst_width - re_width, dst_height - re_height), 132 | } 133 | assert ( 134 | resize_keepratioalign in paddings 135 | ), 'resize_keepratioalign must be one of "center"(mmpose), "topleft"(mmdetection), "topright", "bottomleft", "bottomright"' 136 | xpading, ypading = paddings[resize_keepratioalign] 137 | padopt = f"pad={dst_width}:{dst_height}:{xpading}:{ypading}:black" 138 | filteropt = f"-vf {padopt}" 139 | 140 | if pix_fmt == "gray": 141 | if filteropt: 142 | filteropt = f"{filteropt},extractplanes=y" 143 | else: 144 | filteropt = f"-vf extractplanes=y" 145 | 146 | return crop_wh, final_size_wh, [cropopt, scaleopt, filteropt] 147 | 148 | 149 | def get_outnumpyshape(size_wh: list, pix_fmt: str) -> tuple: 150 | width, height = size_wh 151 | assert (not pix_fmt == "yuv420p") or ( 152 | height % 2 == 0 and width % 2 == 0 153 | ), "yuv420p must be even" 154 | out_numpy_shape = { 155 | "rgb24": (height, width, 3), 156 | "bgr24": (height, width, 3), 157 | "yuv420p": (int(height * 1.5), width), 158 | "yuvj420p": (int(height * 1.5), width), 159 | "nv12": (int(height * 1.5), width), 160 | "gray": (height, width, 1), 161 | }[pix_fmt] 162 | return out_numpy_shape 163 | 164 | 165 | class FFmpegReader: 166 | def __init__(self): 167 | self.filename:str = '' 168 | self.iframe:int = -1 169 | self.width:int = None 170 | self.height:int = None 171 | self.size = (None, None) 172 | self.waitInit:bool = True 173 | self.process = None 174 | self._isopen:bool = True 175 | self.debug:bool = False 176 | self.fps:float = None 177 | self.out_numpy_shape = (None, None, None) 178 | 179 | def __repr__(self): 180 | props = pprint.pformat(self.__dict__).replace("{", " ").replace("}", " ") 181 | return f"{self.__class__}\n" + props 182 | 183 | def __enter__(self): 184 | return self 185 | 186 | def __exit__(self, type, value, traceback): 187 | self.release() 188 | 189 | def __len__(self): 190 | return self.count 191 | 192 | def __iter__(self): 193 | return self 194 | 195 | def __next__(self): 196 | ret, img = self.read() 197 | if ret: 198 | return img 199 | else: 200 | raise StopIteration 201 | 202 | @staticmethod 203 | def VideoReader( 204 | filename, 205 | codec, 206 | pix_fmt, 207 | crop_xywh, 208 | resize, 209 | resize_keepratio, 210 | resize_keepratioalign, 211 | ): 212 | assert os.path.exists(filename) and os.path.isfile( 213 | filename 214 | ), f"{filename} not exists" 215 | 216 | vid = FFmpegReader() 217 | vid.filename = filename 218 | videoinfo = get_info(filename) 219 | vid.origin_width = videoinfo.width 220 | vid.origin_height = videoinfo.height 221 | vid.fps = videoinfo.fps 222 | vid.count = videoinfo.count 223 | vid.duration = videoinfo.duration 224 | vid.pix_fmt = pix_fmt 225 | vid.codec = videoinfo.codec 226 | 227 | if codec is not None: 228 | warnings.warn( 229 | "The 'codec' parameter is auto detected and will be removed " 230 | "in future versions. Please refrain from using this parameter.", 231 | DeprecationWarning 232 | ) 233 | 234 | ( 235 | (vid.crop_width, vid.crop_height), 236 | (vid.width, vid.height), 237 | filteropt, 238 | ) = get_videofilter_cpu( 239 | (vid.origin_width, vid.origin_height), 240 | pix_fmt, 241 | crop_xywh, 242 | resize, 243 | resize_keepratio, 244 | resize_keepratioalign, 245 | ) 246 | vid.size = (vid.width, vid.height) 247 | 248 | vid.ffmpeg_cmd = ( 249 | f"ffmpeg -loglevel error " 250 | f' -vcodec {vid.codec} -r {vid.fps} -i "{filename}" ' 251 | f" {filteropt} -pix_fmt {pix_fmt} -r {vid.fps} -f rawvideo pipe:" 252 | ) 253 | vid.out_numpy_shape = get_outnumpyshape(vid.size, pix_fmt) 254 | return vid 255 | 256 | def read(self): 257 | if self.waitInit: 258 | self.process = run_async(self.ffmpeg_cmd) 259 | self.waitInit = False 260 | 261 | in_bytes = self.process.stdout.read(np.prod(self.out_numpy_shape)) 262 | # check if ffmpeg process error 263 | # if self.process.stderr.readable(): 264 | # print('---a') 265 | # data = self.process.stderr.read() 266 | # sys.stderr.buffer.write(data) 267 | # print('---f') 268 | 269 | if not in_bytes: 270 | self.release() 271 | return False, None 272 | self.iframe += 1 273 | img = np.frombuffer(in_bytes, np.uint8).reshape(self.out_numpy_shape) 274 | 275 | return True, img 276 | 277 | def isOpened(self): 278 | return self._isopen 279 | 280 | def release(self): 281 | self._isopen = False 282 | release_process(self.process, forcekill=True) 283 | 284 | def close(self): 285 | return self.release() 286 | 287 | 288 | class FFmpegReaderNV(FFmpegReader): 289 | def _get_opts( 290 | vid, 291 | videoinfo, 292 | crop_xywh, 293 | resize, 294 | resize_keepratio, 295 | resize_keepratioalign, 296 | isgray, 297 | ): 298 | vid.origin_width = videoinfo.width 299 | vid.origin_height = videoinfo.height 300 | vid.fps = videoinfo.fps 301 | vid.count = videoinfo.count 302 | vid.duration = videoinfo.duration 303 | vid.width, vid.height = vid.origin_width, vid.origin_height 304 | vid.codec = videoinfo.codec 305 | assert vid.origin_height % 2 == 0, "height must be even" 306 | assert vid.origin_width % 2 == 0, "width must be even" 307 | if crop_xywh: 308 | crop_w, crop_h = crop_xywh[2:] 309 | vid.width, vid.height = crop_w, crop_h 310 | x, y, w, h = crop_xywh 311 | top, bottom, left, right = ( 312 | y, 313 | vid.origin_height - (y + h), 314 | x, 315 | vid.origin_width - (x + w), 316 | ) # crop length 317 | cropopt = f"-crop {top}x{bottom}x{left}x{right}" 318 | else: 319 | crop_w, crop_h = vid.origin_width, vid.origin_height 320 | cropopt = "" 321 | 322 | vid.crop_width, vid.crop_height = crop_w, crop_h 323 | 324 | if resize is None or tuple(resize) == (vid.crop_width, vid.crop_height): 325 | scaleopt = "" 326 | filteropt = "" 327 | else: 328 | vid.width, vid.height = dst_width, dst_height = resize 329 | if not resize_keepratio: 330 | scaleopt = f"-resize {dst_width}x{dst_height}" 331 | filteropt = "" 332 | else: 333 | re_width, re_height = crop_w / (crop_h / dst_height), dst_height 334 | if re_width > dst_width: 335 | re_width, re_height = dst_width, crop_h / (crop_w / dst_width) 336 | re_width, re_height = int(re_width), int(re_height) 337 | scaleopt = f"-resize {re_width}x{re_height}" 338 | if resize_keepratioalign is None: 339 | resize_keepratioalign = "center" 340 | paddings = { 341 | "center": ( 342 | (dst_width - re_width) // 2, 343 | (dst_height - re_height) // 2, 344 | ), 345 | "topleft": (0, 0), 346 | "topright": (dst_width - re_width, 0), 347 | "bottomleft": (0, dst_height - re_height), 348 | "bottomright": (dst_width - re_width, dst_height - re_height), 349 | } 350 | assert ( 351 | resize_keepratioalign in paddings 352 | ), 'resize_keepratioalign must be one of "center"(mmpose), "topleft"(mmdetection), "topright", "bottomleft", "bottomright"' 353 | xpading, ypading = paddings[resize_keepratioalign] 354 | padopt = f"pad={dst_width}:{dst_height}:{xpading}:{ypading}:black" 355 | filteropt = f"-vf {padopt}" 356 | 357 | if isgray: 358 | if filteropt: 359 | filteropt = f"{filteropt},extractplanes=y" 360 | else: 361 | filteropt = f"-vf extractplanes=y" 362 | 363 | vid.size = (vid.width, vid.height) 364 | return cropopt, scaleopt, filteropt 365 | 366 | @staticmethod 367 | def VideoReader( 368 | filename, 369 | pix_fmt, 370 | crop_xywh, 371 | resize, 372 | resize_keepratio, 373 | resize_keepratioalign, 374 | gpu, 375 | ): 376 | assert os.path.exists(filename) and os.path.isfile( 377 | filename 378 | ), f"{filename} not exists" 379 | assert pix_fmt in ["rgb24", "bgr24", "yuv420p", "yuvj420p", "nv12", "gray"] 380 | numGPU = get_num_NVIDIA_GPUs() 381 | assert numGPU > 0, "No GPU found" 382 | gpu = int(gpu) % numGPU if gpu is not None else 0 383 | assert ( 384 | resize is None or len(resize) == 2 385 | ), "resize must be a tuple of (width, height)" 386 | videoinfo = get_info(filename) 387 | vid = FFmpegReaderNV() 388 | vid.filename = filename 389 | isgray = pix_fmt == "gray" 390 | cropopt, scaleopt, filteropt = vid._get_opts( 391 | videoinfo, 392 | crop_xywh, 393 | resize, 394 | resize_keepratio, 395 | resize_keepratioalign, 396 | isgray, 397 | ) 398 | vid.codecNV = decoder_to_nvidia(vid.codec) 399 | 400 | vid.ffmpeg_cmd = ( 401 | f"ffmpeg -loglevel error -hwaccel cuda -hwaccel_device {gpu} " 402 | f' -vcodec {vid.codecNV} {cropopt} {scaleopt} -r {vid.fps} -i "{filename}" ' 403 | f" {filteropt} -pix_fmt {pix_fmt} -r {vid.fps} -f rawvideo pipe:" 404 | ) 405 | 406 | vid.pix_fmt = pix_fmt 407 | vid.out_numpy_shape = get_outnumpyshape(vid.size, pix_fmt) 408 | return vid 409 | -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_reader_camera.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pprint 3 | from .video_info import run_async, release_process 4 | import re 5 | import subprocess 6 | from threading import Thread 7 | from queue import Queue 8 | import sys 9 | import os 10 | from ffmpegcv.ffmpeg_reader import get_videofilter_cpu, get_outnumpyshape 11 | 12 | 13 | class platform: 14 | win = 0 15 | linux = 1 16 | mac = 2 17 | other = 3 18 | 19 | 20 | if sys.platform.startswith("linux"): 21 | this_os = platform.linux 22 | elif sys.platform.startswith("win32"): 23 | this_os = platform.win 24 | elif sys.platform.startswith("darwin"): 25 | this_os = platform.mac 26 | else: 27 | this_os = platform.other 28 | 29 | 30 | def _query_camera_divices_mac() -> dict: 31 | # run the command 'ffmpeg -f avfoundation -list_devices true -i "" ' 32 | command = 'ffmpeg -hide_banner -f avfoundation -list_devices true -i ""' 33 | process = subprocess.Popen( 34 | command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE 35 | ) 36 | stdout, stderr = process.communicate() 37 | 38 | # parse the output into a dictionary 39 | lines = stderr.decode("utf-8").split("AVFoundation audio devices:")[0].split("\n") 40 | id_device_map = dict() 41 | device_id_pattern = re.compile(r"\[[^\]]*?\] \[(\d*)\]") 42 | device_name_pattern = re.compile(r".*\] (.*)") 43 | for line in lines[1:-1]: 44 | device_id = int(re.search(device_id_pattern, line).group(1)) 45 | device_name = re.search(device_name_pattern, line).group(1) 46 | id_device_map[device_id] = (device_name, device_id) 47 | return id_device_map 48 | 49 | 50 | def _query_camera_divices_win() -> dict: 51 | command = "ffmpeg -hide_banner -list_devices true -f dshow -i dummy" 52 | process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) 53 | stdout, stderr = process.communicate() 54 | dshowliststr = stderr.decode("utf-8") 55 | dshowliststr = dshowliststr.split("DirectShow audio devices")[0] 56 | pattern = re.compile(r'\[*?\] *"([^"]*)"') 57 | matches = pattern.findall(dshowliststr) 58 | alternative_pattern = re.compile(r'Alternative name "(.*)"') 59 | alternative_names = alternative_pattern.findall(dshowliststr) 60 | assert len(matches) == len(alternative_names) 61 | id_device_map = { 62 | i: device for i, device in enumerate(zip(matches, alternative_names)) 63 | } 64 | if len(id_device_map) == 0: 65 | print("No camera divice found") 66 | return id_device_map 67 | 68 | 69 | def _query_camera_divices_linux() -> dict: 70 | "edit from https://github.com/p513817/python-get-cam-name/blob/master/get_cam_name.py" 71 | root = "/sys/class/video4linux" 72 | cam_info = [] 73 | 74 | for index in sorted([file for file in os.listdir(root)]): 75 | # Get Camera Name From /sys/class/video4linux//name 76 | real_index_file = os.path.realpath("/sys/class/video4linux/" + index + "/index") 77 | with open(real_index_file, "r") as name_file: 78 | _index = name_file.read().rstrip() 79 | if _index != "0": 80 | continue 81 | 82 | real_file = os.path.realpath("/sys/class/video4linux/" + index + "/name") 83 | with open(real_file, "r") as name_file: 84 | name = name_file.read().rstrip() 85 | name = name.split(":")[0] 86 | 87 | # Setup Each Camera and Index ( video* ) 88 | cam_info.append((name, "/dev/" + index)) 89 | 90 | id_device_map = {i: vname for i, vname in enumerate(cam_info)} 91 | return id_device_map 92 | 93 | 94 | def query_camera_devices(verbose_dict: bool = False) -> dict: 95 | result = { 96 | platform.linux: _query_camera_divices_linux, 97 | platform.mac: _query_camera_divices_mac, 98 | platform.win: _query_camera_divices_win, 99 | }[this_os]() 100 | if verbose_dict: 101 | dict_by_v0 = {v[0]: v for v in result.values()} 102 | dict_by_v1 = {v[1]: v for v in result.values()} 103 | result.update(dict_by_v0) 104 | result.update(dict_by_v1) 105 | 106 | return result 107 | 108 | 109 | def _query_camera_options_mac(cam_id_name) -> str: 110 | print( 111 | "\033[33m" 112 | + "FFmpeg& FFmpegcv CAN NOT query the camera options in MAC platform." 113 | + "\033[0m" 114 | ) 115 | print("Please find the proper parameter other way.") 116 | return [{"camsize_wh": None, "camfps": None}] 117 | 118 | 119 | def _query_camera_options_linux(cam_id_name) -> str: 120 | print( 121 | "\033[33m" 122 | + "FFmpeg& FFmpegcv CAN NOT query the camera FPS in Linux platform." 123 | + "\033[0m" 124 | ) 125 | print("Please find the proper parameter other way.") 126 | camname = query_camera_devices(verbose_dict=True)[cam_id_name][1] 127 | command = f'ffmpeg -hide_banner -f v4l2 -list_formats all -i "{camname}"' 128 | process = subprocess.Popen( 129 | command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE 130 | ) 131 | stdout, stderr = process.communicate() 132 | lines = stderr.decode("utf-8").split("\n") 133 | lines = [l for l in lines if "v4l2" in l] 134 | outlist = [] 135 | for line in lines: 136 | _, vcodec, *_, resolutions = line.split(":") 137 | vcodec = vcodec.strip() 138 | israw = "Raw" in line 139 | camcodec = None if israw else vcodec 140 | campix_fmt = vcodec if israw else None 141 | resolutions = resolutions.strip() 142 | 143 | camsize_wh_l = [tuple(map(int, r.split("x"))) for r in resolutions.split()] 144 | outlist.extend( 145 | [ 146 | { 147 | "camcodec": camcodec, 148 | "campix_fmt": campix_fmt, 149 | "camsize_wh": wh, 150 | "camfps": None, 151 | } 152 | for wh in camsize_wh_l 153 | ] 154 | ) 155 | return outlist 156 | 157 | 158 | def _query_camera_options_win(cam_id_name) -> str: 159 | if isinstance(cam_id_name, int): 160 | id_device_map = query_camera_devices() 161 | camname = id_device_map[cam_id_name][1] 162 | elif isinstance(cam_id_name, str): 163 | camname = cam_id_name 164 | else: 165 | raise ValueError("Not valid camname") 166 | command = f'ffmpeg -hide_banner -f dshow -list_options true -i video="{camname}"' 167 | process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) 168 | stdout, stderr = process.communicate() 169 | dshowliststr = stderr.decode("utf-8").replace("\r\n", "\n").replace("\r", "\n") 170 | dshowlist = [s for s in dshowliststr.split("\n") if "fps=" in s] 171 | from collections import OrderedDict 172 | 173 | unique_dshowlist = list(OrderedDict.fromkeys(dshowlist)) 174 | outlist = [] 175 | for text in unique_dshowlist: 176 | cam_options = dict() 177 | cam_options["camcodec"] = ( 178 | re.search(r"vcodec=(\w+)", text).group(1) if "vcodec" in text else None 179 | ) 180 | cam_options["campix_fmt"] = ( 181 | re.search(r"pixel_format=(\w+)", text).group(1) 182 | if "pixel_format" in text 183 | else None 184 | ) 185 | camsize_wh = re.search(r"min s=(\w+)", text).group(1) 186 | cam_options["camsize_wh"] = tuple(int(v) for v in camsize_wh.split("x")) 187 | camfps = float(re.findall(r"fps=([\d.]+)", text)[-1]) 188 | cam_options["camfps"] = int(camfps) if int(camfps) == camfps else camfps 189 | outlist.append(cam_options) 190 | return outlist 191 | 192 | 193 | def query_camera_options(cam_id_name) -> str: 194 | return { 195 | platform.linux: _query_camera_options_linux, 196 | platform.mac: _query_camera_options_mac, 197 | platform.win: _query_camera_options_win, 198 | }[this_os](cam_id_name) 199 | 200 | 201 | class ProducerThread(Thread): 202 | def __init__(self, vid, q): 203 | super(ProducerThread, self).__init__() 204 | self.vid = vid 205 | self.q = q 206 | 207 | def run(self): 208 | q = self.q 209 | while True: 210 | if not self.vid.isOpened(): 211 | break 212 | ret, img = self.vid.read_() 213 | 214 | if q.full(): 215 | q.get() # drop frames 216 | q.put((ret, img)) 217 | 218 | 219 | class FFmpegReaderCAM: 220 | def __init__(self): 221 | self.iframe = -1 222 | self._isopen = True 223 | 224 | def __repr__(self): 225 | props = pprint.pformat(self.__dict__).replace("{", " ").replace("}", " ") 226 | return f"{self.__class__}\n" + props 227 | 228 | def __enter__(self): 229 | return self 230 | 231 | def __exit__(self, type, value, traceback): 232 | self.release() 233 | 234 | def __iter__(self): 235 | return self 236 | 237 | def __next__(self): 238 | ret, img = self.read() 239 | if ret: 240 | return img 241 | else: 242 | raise StopIteration 243 | 244 | @staticmethod 245 | def VideoReader( 246 | cam_id_name, 247 | pix_fmt, 248 | crop_xywh, 249 | resize, 250 | resize_keepratio, 251 | resize_keepratioalign, 252 | camsize_wh=None, 253 | camfps=None, 254 | camcodec=None, 255 | campix_fmt=None, 256 | step=1, 257 | ): 258 | vid = FFmpegReaderCAM() 259 | if this_os == platform.mac: 260 | # use cam_id as the device marker 261 | if isinstance(cam_id_name, str): 262 | id_device_map = query_camera_devices() 263 | camname = cam_id_name 264 | id_device_map.update({v[0]: v for v in id_device_map.values()}) 265 | camid = id_device_map[cam_id_name][1] 266 | else: 267 | camname = None 268 | camid = cam_id_name 269 | elif this_os == platform.linux: 270 | id_device_map = query_camera_devices(verbose_dict=True) 271 | camname = id_device_map[cam_id_name][-1] 272 | camid = None 273 | else: 274 | if isinstance(cam_id_name, int): 275 | id_device_map = query_camera_devices() 276 | camname = id_device_map[cam_id_name][1] 277 | camid = cam_id_name 278 | else: 279 | camname = cam_id_name 280 | camid = None 281 | 282 | vid.camname = camname 283 | vid.camid = camid 284 | 285 | if camsize_wh is None: 286 | cam_options = query_camera_options(camname) 287 | resolutions = [c["camsize_wh"] for c in cam_options] 288 | camsize_wh = max(resolutions, key=lambda x: sum(x)) 289 | 290 | assert len(camsize_wh) == 2 291 | vid.origin_width, vid.origin_height = camsize_wh 292 | 293 | opt_camfps = f" -framerate {camfps} " if camfps else "" 294 | vid.camfps = camfps if camfps else None 295 | 296 | opt_camcodec_ = { 297 | platform.linux: "input_format", 298 | platform.mac: "", 299 | platform.win: "vcodec", 300 | }[this_os] 301 | opt_camcodec = f" -{opt_camcodec_} {camcodec} " if camcodec else "" 302 | vid.camcodec = camcodec if camcodec else None 303 | vid.pix_fmt = pix_fmt 304 | 305 | opt_campix_fmt_ = { 306 | platform.linux: "input_format", 307 | platform.mac: "pixel_format", 308 | platform.win: "pixel_format", 309 | }[this_os] 310 | opt_campix_fmt = f" -{opt_campix_fmt_} {campix_fmt} " if campix_fmt else "" 311 | vid.campix_fmt = campix_fmt if campix_fmt else None 312 | 313 | opt_camname = { 314 | platform.linux: f'"{camname}"', 315 | platform.win: f'video="{camname}"', 316 | platform.mac: f"{camid}:none", 317 | }[this_os] 318 | 319 | (vid.crop_width, vid.crop_height), (vid.width, vid.height), filteropt = get_videofilter_cpu( 320 | (vid.origin_width, vid.origin_height), pix_fmt, crop_xywh, resize, resize_keepratio, resize_keepratioalign) 321 | vid.size = (vid.width, vid.height) 322 | 323 | opt_driver_ = { 324 | platform.linux: "v4l2", 325 | platform.mac: "avfoundation", 326 | platform.win: "dshow", 327 | }[this_os] 328 | 329 | vid.ffmpeg_cmd = ( 330 | f"ffmpeg -loglevel warning " 331 | f" -f {opt_driver_} " 332 | f" -video_size {vid.origin_width}x{vid.origin_height} " 333 | f" {opt_camfps} {opt_camcodec} {opt_campix_fmt} " 334 | f" -i {opt_camname} " 335 | f" {filteropt} -pix_fmt {pix_fmt} -f rawvideo pipe:" 336 | ) 337 | 338 | vid.out_numpy_shape = get_outnumpyshape(vid.size, pix_fmt) 339 | vid.process = run_async(vid.ffmpeg_cmd) 340 | 341 | # producer 342 | assert step >= 1 and isinstance(step, int) 343 | vid.step = step 344 | vid.q = Queue(maxsize=30) 345 | producer = ProducerThread(vid, vid.q) 346 | producer.start() 347 | return vid 348 | 349 | def read_(self): 350 | for i in range(self.step): 351 | in_bytes = self.process.stdout.read(np.prod(self.out_numpy_shape)) 352 | if not in_bytes: 353 | self.release() 354 | return False, None 355 | 356 | self.iframe += 1 357 | img = None 358 | img = np.frombuffer(in_bytes, np.uint8).reshape(self.out_numpy_shape) 359 | return True, img 360 | 361 | def read(self): 362 | ret, img = self.q.get() 363 | return ret, img 364 | 365 | def isOpened(self): 366 | return self._isopen 367 | 368 | def release(self): 369 | self._isopen = False 370 | release_process(self.process) 371 | 372 | def close(self): 373 | return self.release() 374 | -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_reader_cuda.py: -------------------------------------------------------------------------------- 1 | import pycuda.driver as cuda 2 | from pycuda.driver import PointerHolderBase 3 | from pycuda.compiler import SourceModule 4 | from pycuda import gpuarray 5 | from ffmpegcv.ffmpeg_reader import FFmpegReader, FFmpegReaderNV 6 | import numpy as np 7 | from typing import Tuple 8 | 9 | 10 | cuda.init() 11 | 12 | mod_code = (""" 13 | __device__ void yuv_to_rgb(unsigned char &y, unsigned char &u, unsigned char &v, 14 | float &r, float &g, float &b) 15 | { 16 | // https://fourcc.org/fccyvrgb.php 17 | float Y_val = (float)1.164 * ((float)y - 16.0); 18 | float U_val = (float)u - 128.0; 19 | float V_val = (float)v - 128.0; 20 | r = Y_val + 1.596 * V_val; 21 | r = max(0.0, min(255.0, r)); // clamp(r, 0.0, 255.0); 22 | g = Y_val - 0.813 * V_val - 0.391 * U_val; 23 | g = max(0.0, min(255.0, g)); 24 | b = Y_val + 2.018 * U_val; 25 | b = max(0.0, min(255.0, b)); 26 | } 27 | 28 | __global__ void yuv420p_CHW_fp32(unsigned char *YUV420p, float *RGB24, int *width_, int *height_) 29 | { 30 | int width = *width_; int height = *height_; 31 | // Get the thread index 32 | int x = blockIdx.x * blockDim.x + threadIdx.x; 33 | int y = blockIdx.y * blockDim.y + threadIdx.y; 34 | 35 | // Check if we are within the bounds of the image 36 | if (x >= width || y >= height) 37 | return; 38 | 39 | // Get the Y, U, and V values for this pixel 40 | auto w_h = width * height; 41 | auto yW = y * width; 42 | auto out_ind = yW + x; 43 | unsigned char *Y = YUV420p; 44 | unsigned char *U = YUV420p + w_h; 45 | unsigned char *V = U + w_h/4; 46 | int delta = (y/2)*(width/2)+x/2; 47 | 48 | yuv_to_rgb(Y[out_ind], U[delta], V[delta], 49 | RGB24[out_ind], RGB24[out_ind + w_h], RGB24[out_ind + w_h*2]); 50 | } 51 | 52 | __global__ void yuv420p_HWC_fp32(unsigned char *YUV420p, float *RGB24, int *width_, int *height_) 53 | { 54 | int width = *width_; int height = *height_; 55 | // Get the thread index 56 | int x = blockIdx.x * blockDim.x + threadIdx.x; 57 | int y = blockIdx.y * blockDim.y + threadIdx.y; 58 | 59 | // Check if we are within the bounds of the image 60 | if (x >= width || y >= height) 61 | return; 62 | 63 | // Get the Y, U, and V values for this pixel 64 | auto w_h = width * height; 65 | auto yW = y * width; 66 | auto out_ind = yW + x; 67 | unsigned char *Y = YUV420p; 68 | unsigned char *U = YUV420p + w_h; 69 | unsigned char *V = U + w_h/4; 70 | int delta = (y/2)*(width/2)+x/2; 71 | auto ind = (yW + x)*3; 72 | 73 | yuv_to_rgb(Y[out_ind], U[delta], V[delta], 74 | RGB24[ind], RGB24[ind+1], RGB24[ind+2]); 75 | } 76 | 77 | __global__ void NV12_CHW_fp32(unsigned char *NV12, float *RGB24, int *width_, int *height_) 78 | { 79 | int width = *width_; int height = *height_; 80 | // Get the thread index 81 | int x = blockIdx.x * blockDim.x + threadIdx.x; 82 | int y = blockIdx.y * blockDim.y + threadIdx.y; 83 | 84 | // Check if we are within the bounds of the image 85 | if (x >= width || y >= height) 86 | return; 87 | 88 | // Get the Y, U, and V values for this pixel 89 | auto w_h = width * height; 90 | auto yW = y * width; 91 | auto out_ind = yW + x; 92 | unsigned char *Y = NV12 + out_ind; 93 | unsigned char *UV = NV12 + w_h + (y / 2) * width + (x / 2) * 2; 94 | yuv_to_rgb(Y[0], UV[0], UV[1], 95 | RGB24[out_ind], RGB24[out_ind + w_h], RGB24[out_ind + w_h*2]); 96 | } 97 | 98 | __global__ void NV12_HWC_fp32(unsigned char *NV12, float *RGB24, int *width_, int *height_) 99 | { 100 | int width = *width_; int height = *height_; 101 | // Get the thread index 102 | int x = blockIdx.x * blockDim.x + threadIdx.x; 103 | int y = blockIdx.y * blockDim.y + threadIdx.y; 104 | 105 | // Check if we are within the bounds of the image 106 | if (x >= width || y >= height) 107 | return; 108 | 109 | // Get the Y, U, and V values for this pixel 110 | auto w_h = width * height; 111 | auto yW = y * width; 112 | auto out_ind = yW + x; 113 | unsigned char *Y = NV12 + out_ind; 114 | unsigned char *UV = NV12 + w_h + (y / 2) * width + (x / 2) * 2; 115 | auto ind = (yW + x)*3; 116 | yuv_to_rgb(Y[0], UV[0], UV[1], 117 | RGB24[ind], RGB24[ind+1], RGB24[ind+2]); 118 | } 119 | """ 120 | ) 121 | 122 | 123 | def load_cuda_module(): 124 | mod = SourceModule(mod_code) 125 | converter = {('yuv420p', 'chw'): mod.get_function('yuv420p_CHW_fp32'), 126 | ('yuv420p', 'hwc'): mod.get_function('yuv420p_HWC_fp32'), 127 | ('nv12', 'chw'): mod.get_function('NV12_CHW_fp32'), 128 | ('nv12', 'hwc'): mod.get_function('NV12_HWC_fp32')} 129 | return converter 130 | 131 | 132 | class Holder(PointerHolderBase): 133 | def __init__(self, tensor): 134 | super().__init__() 135 | self.tensor = tensor 136 | self.gpudata = tensor.data_ptr() 137 | 138 | def get_pointer(self): 139 | return self.tensor.data_ptr() 140 | 141 | def __index__(self): 142 | return self.gpudata 143 | 144 | 145 | def tensor_to_gpuarray(tensor) -> gpuarray.GPUArray: 146 | '''Convert a :class:`torch.Tensor` to a :class:`pycuda.gpuarray.GPUArray`. The underlying 147 | storage will be shared, so that modifications to the array will reflect in the tensor object. 148 | Parameters 149 | ---------- 150 | tensor : torch.Tensor 151 | Returns 152 | ------- 153 | pycuda.gpuarray.GPUArray 154 | Raises 155 | ------ 156 | ValueError 157 | If the ``tensor`` does not live on the gpu 158 | ''' 159 | return gpuarray.GPUArray(tensor.shape, dtype=np.float32, gpudata=Holder(tensor)) 160 | 161 | 162 | class PycudaContext: 163 | def __init__(self, gpu=0): 164 | self.ctx = cuda.Device(gpu).make_context() 165 | 166 | def __enter__(self): 167 | if self.ctx is not None: 168 | self.ctx.push() 169 | return self 170 | 171 | def __exit__(self, *args, **kwargs): 172 | if self.ctx is not None: 173 | self.ctx.pop() 174 | 175 | def __del__(self): 176 | if self.ctx is not None: 177 | self.ctx.pop() 178 | 179 | 180 | class FFmpegReaderCUDA(FFmpegReader): 181 | def __init__(self, vid:FFmpegReader, gpu=0, tensor_format='hwc'): 182 | assert vid.pix_fmt in ['yuv420p', 'nv12'], 'Set pix_fmt to yuv420p or nv12. Auto convert to rgb in cuda.' 183 | assert tensor_format in ['hwc', 'chw'], 'tensor_format must be hwc or chw' 184 | if isinstance(vid, FFmpegReaderNV) and vid.pix_fmt != 'nv12': 185 | print('--Tips: please use VideoCaptureNV(..., pix_fmt="NV12") for better performance.') 186 | elif not isinstance(vid, FFmpegReaderNV) and vid.pix_fmt == 'nv12': 187 | print('--Tips: please use VideoCapture(..., pix_fmt="yuv420p") for better performance.') 188 | 189 | # work like normal FFmpegReaderObj 190 | props_name = ['width', 'height', 'fps', 'count', 'codec', 'ffmpeg_cmd', 191 | 'size', 'out_numpy_shape', 'iframe', 192 | 'duration', 'origin_width', 'origin_height'] 193 | for name in props_name: 194 | setattr(self, name, getattr(vid, name, None)) 195 | self.ctx = PycudaContext(gpu) 196 | self.pix_fmt = 'rgb24' 197 | self.vid = vid 198 | self.out_numpy_shape = (vid.height, vid.width, 3) if tensor_format == 'hwc' else (3, vid.height, vid.width) 199 | self.torch_device = f'cuda:{gpu}' 200 | self.block_size = (16, 16, 1) 201 | self.grid_size = ((self.width + self.block_size[0] - 1) // self.block_size[0], 202 | (self.height + self.block_size[1] - 1) // self.block_size[1]) 203 | self.process = None 204 | with self.ctx: 205 | self.converter = load_cuda_module()[(vid.pix_fmt, tensor_format)] 206 | 207 | def read(self, out_MAT:gpuarray.GPUArray=None) -> Tuple[bool, gpuarray.GPUArray]: 208 | self.waitInit = False 209 | ret, frame_yuv420p = self.vid.read() 210 | if not ret: 211 | return False, None 212 | 213 | with self.ctx: 214 | if out_MAT is None: 215 | out_MAT = gpuarray.empty(self.out_numpy_shape, dtype=np.float32) 216 | self.converter(cuda.In(frame_yuv420p), out_MAT, 217 | cuda.In(np.int32(self.width)), cuda.In(np.int32(self.height)), 218 | block=self.block_size, grid=self.grid_size) 219 | return True, out_MAT 220 | 221 | def read_cudamem(self, out_MAT:cuda.DeviceAllocation=None) -> Tuple[bool, cuda.DeviceAllocation]: 222 | self.waitInit = False 223 | ret, frame_yuv420p = self.vid.read() 224 | if not ret: 225 | return False, None 226 | 227 | with self.ctx: 228 | if out_MAT is None: 229 | out_MAT = cuda.mem_alloc(int(np.prod(self.out_numpy_shape) * 230 | np.dtype(np.float32).itemsize)) 231 | self.converter(cuda.In(frame_yuv420p), out_MAT, 232 | cuda.In(np.int32(self.width)), cuda.In(np.int32(self.height)), 233 | block=self.block_size, grid=self.grid_size) 234 | return True, out_MAT 235 | 236 | def read_torch(self, out_MAT=None): 237 | import torch 238 | self.waitInit = False 239 | ret, frame_yuv420p = self.vid.read() 240 | if not ret: 241 | return False, None 242 | 243 | with self.ctx: 244 | if out_MAT is None: 245 | out_MAT = torch.empty(self.out_numpy_shape, dtype=torch.float32, device=self.torch_device) 246 | tensor_proxy = tensor_to_gpuarray(out_MAT) 247 | self.converter(cuda.In(frame_yuv420p), tensor_proxy.gpudata, 248 | cuda.In(np.int32(self.width)), cuda.In(np.int32(self.height)), 249 | block=self.block_size, grid=self.grid_size) 250 | return True, out_MAT 251 | 252 | def release(self): 253 | self.vid.release() 254 | super().release() 255 | -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_reader_noblock.py: -------------------------------------------------------------------------------- 1 | import multiprocessing 2 | from multiprocessing import Queue 3 | import numpy as np 4 | from .ffmpeg_reader import FFmpegReader 5 | 6 | NFRAME = 10 7 | 8 | class FFmpegReaderNoblock(FFmpegReader): 9 | def __init__(self, 10 | vcap_fun, 11 | *vcap_args, **vcap_kwargs): 12 | vid:FFmpegReader = vcap_fun(*vcap_args, **vcap_kwargs) 13 | vid.release() 14 | 15 | # work like normal FFmpegReaderObj 16 | props_name = ['width', 'height', 'fps', 'count', 'codec', 'ffmpeg_cmd', 17 | 'size', 'pix_fmt', 'out_numpy_shape', 'iframe', 18 | 'duration', 'origin_width', 'origin_height'] 19 | for name in props_name: 20 | setattr(self, name, getattr(vid, name, None)) 21 | 22 | # 创建共享内存的NumPy数组 23 | shared_array = multiprocessing.Array('b', int(NFRAME*np.prod(self.out_numpy_shape))) 24 | 25 | # 将共享内存的NumPy数组转换为NumPy数组 26 | self.np_array = np.frombuffer(shared_array.get_obj(), dtype=np.uint8).reshape((NFRAME,*self.out_numpy_shape)) 27 | 28 | self.shared_array = shared_array 29 | self.vcap_args = vcap_args 30 | self.vcap_kwargs = vcap_kwargs 31 | self.q = Queue(maxsize=(NFRAME-2)) #buffer index, gluttonous snake NO biting its own tail 32 | self.vcap_fun = vcap_fun 33 | self.has_init = False 34 | self.process = None 35 | 36 | def read(self): 37 | if not self.has_init: 38 | self.has_init = True 39 | process = multiprocessing.Process(target=child_process, 40 | args=(self.shared_array, self.q, self.vcap_fun, 41 | self.vcap_args, self.vcap_kwargs)) 42 | process.start() 43 | self.process = process 44 | 45 | data_id = self.q.get() # 读取子进程写入的数据 46 | if data_id is None: 47 | return False, None 48 | else: 49 | self.iframe += 1 50 | return True, self.np_array[data_id] 51 | 52 | 53 | def child_process(shared_array, q:Queue, vcap_fun, vcap_args, vcap_kwargs): 54 | vid = vcap_fun(*vcap_args, **vcap_kwargs) 55 | np_array = np.frombuffer(shared_array.get_obj(), dtype=np.uint8).reshape((NFRAME,*vid.out_numpy_shape)) 56 | anything = True 57 | with vid: 58 | for i, img in enumerate(vid): 59 | iloop = i % NFRAME 60 | # 在子进程中修改共享内存的NumPy数组 61 | np_array[iloop] = img 62 | q.put(iloop) # 通知主进程已经写入了 63 | q.put(None) 64 | -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_reader_pannels.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | from ffmpegcv.ffmpeg_reader import FFmpegReader, get_outnumpyshape 4 | 5 | from ffmpegcv.video_info import ( 6 | get_info, 7 | run_async 8 | ) 9 | 10 | 11 | class FFmpegReaderPannels(FFmpegReader): 12 | @staticmethod 13 | def VideoReader( 14 | filename:str, 15 | crop_xywh_l:list, 16 | codec, 17 | pix_fmt='bgr24', 18 | resize=None 19 | ): 20 | assert os.path.exists(filename) and os.path.isfile( 21 | filename 22 | ), f"{filename} not exists" 23 | assert pix_fmt in ["rgb24", "bgr24", "yuv420p", "nv12", "gray"] 24 | vid = FFmpegReaderPannels() 25 | crop_xywh_l = np.array(crop_xywh_l) 26 | vid.crop_xywh_l = crop_xywh_l 27 | videoinfo = get_info(filename) 28 | vid.origin_width = videoinfo.width 29 | vid.origin_height = videoinfo.height 30 | vid.fps = videoinfo.fps 31 | vid.count = videoinfo.count 32 | vid.duration = videoinfo.duration 33 | vid.pix_fmt = pix_fmt 34 | vid.codec = codec if codec else videoinfo.codec 35 | 36 | vid.crop_width_l = crop_xywh_l[:,2] 37 | vid.crop_height_l = crop_xywh_l[:,3] 38 | vid.size_l = crop_xywh_l[:,2:][:,::-1] 39 | vid.npannel = len(crop_xywh_l) 40 | vid.out_numpy_shape_l = [get_outnumpyshape(s[::-1], pix_fmt) for s in vid.size_l] 41 | if len(set(vid.crop_width_l)) == len(set(vid.crop_height_l)) == 1: 42 | vid.is_pannel_similar = True 43 | vid.crop_width = vid.crop_width_l[0] 44 | vid.crop_height = vid.crop_height_l[0] 45 | vid.size = vid.size_l[0] 46 | vid.out_numpy_shape = (vid.npannel, *vid.out_numpy_shape_l[0]) 47 | else: 48 | vid.is_pannel_similar = False 49 | vid.crop_width = vid.crop_height = vid.size = None 50 | vid.out_numpy_shape = (np.sum(np.prod(s) for s in vid.out_numpy_shape_l),) 51 | 52 | VINSRCs =''.join(f'[VSRC{i}]' for i in range(vid.npannel)) 53 | pix_fmtopt = ',extractplanes=y' if pix_fmt=='gray' else '' 54 | CROPs = ';'.join(f'[VSRC{i}]crop={w}:{h}:{x}:{y}{pix_fmtopt}[VPANEL{i}]' 55 | for i, (x,y,w,h) in enumerate(vid.crop_xywh_l)) 56 | filteropt = f' -filter_complex "split={vid.npannel}{VINSRCs};{CROPs}"' 57 | outmaps = ''.join(f' -map [VPANEL{i}] -pix_fmt {pix_fmt} -r {vid.fps} -f rawvideo pipe:' 58 | for i in range(vid.npannel)) 59 | 60 | vid.ffmpeg_cmd = ( 61 | f"ffmpeg -loglevel warning " 62 | f' -r {vid.fps} -i "{filename}" ' 63 | f" {filteropt} {outmaps}" 64 | ) 65 | return vid 66 | 67 | def read(self): 68 | if self.waitInit: 69 | self.process = run_async(self.ffmpeg_cmd) 70 | self.waitInit = False 71 | 72 | in_bytes = self.process.stdout.read(np.prod(self.out_numpy_shape)) 73 | if not in_bytes: 74 | self.release() 75 | return False, None 76 | self.iframe += 1 77 | img0 = np.frombuffer(in_bytes, np.uint8) 78 | if self.is_pannel_similar: 79 | img = img0.reshape(self.out_numpy_shape) 80 | else: 81 | img = [] 82 | for out_numpy_shape in self.out_numpy_shape_l: 83 | nbuff = np.prod(out_numpy_shape) 84 | img.append(img0[:nbuff].reshape(out_numpy_shape)) 85 | img0 = img0[nbuff:] 86 | 87 | return True, img 88 | -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_reader_qsv.py: -------------------------------------------------------------------------------- 1 | import os 2 | from ffmpegcv.ffmpeg_reader import FFmpegReader, get_videofilter_cpu, get_outnumpyshape 3 | from .video_info import ( 4 | get_info, 5 | get_num_QSV_GPUs, 6 | decoder_to_qsv, 7 | ) 8 | 9 | 10 | class FFmpegReaderQSV(FFmpegReader): 11 | @staticmethod 12 | def VideoReader( 13 | filename, 14 | pix_fmt, 15 | crop_xywh, 16 | resize, 17 | resize_keepratio, 18 | resize_keepratioalign, 19 | gpu, 20 | ): 21 | """ 22 | TODO: 1. only 1 gpu is recognized 23 | TODO: 2. 'crop_xywh', 'resize*' are not supported yet 24 | """ 25 | assert os.path.exists(filename) and os.path.isfile( 26 | filename 27 | ), f"{filename} not exists" 28 | assert gpu is None or gpu == 0, "Cannot use multiple QSV gpu yet." 29 | numGPU = get_num_QSV_GPUs() 30 | assert numGPU > 0, "No GPU found" 31 | gpu = int(gpu) % numGPU if gpu is not None else 0 32 | assert ( 33 | resize is None or len(resize) == 2 34 | ), "resize must be a tuple of (width, height)" 35 | 36 | vid = FFmpegReaderQSV() 37 | videoinfo = get_info(filename) 38 | vid.origin_width = videoinfo.width 39 | vid.origin_height = videoinfo.height 40 | vid.fps = videoinfo.fps 41 | vid.count = videoinfo.count 42 | vid.duration = videoinfo.duration 43 | vid.codecQSV = decoder_to_qsv(videoinfo.codec) 44 | vid.pix_fmt = pix_fmt 45 | 46 | ( 47 | (vid.crop_width, vid.crop_height), 48 | (vid.width, vid.height), 49 | filteropt, 50 | ) = get_videofilter_cpu( 51 | (vid.origin_width, vid.origin_height), 52 | pix_fmt, 53 | crop_xywh, 54 | resize, 55 | resize_keepratio, 56 | resize_keepratioalign, 57 | ) 58 | vid.size = (vid.width, vid.height) 59 | 60 | vid.ffmpeg_cmd = ( 61 | f"ffmpeg -loglevel warning " 62 | f' -vcodec {vid.codecQSV} -r {vid.fps} -i "{filename}" ' 63 | f" {filteropt} -pix_fmt {pix_fmt} -r {vid.fps} -f rawvideo pipe:" 64 | ) 65 | 66 | vid.out_numpy_shape = get_outnumpyshape(vid.size, pix_fmt) 67 | return vid 68 | -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_reader_stream.py: -------------------------------------------------------------------------------- 1 | from .video_info import run_async 2 | from queue import Queue 3 | from ffmpegcv.stream_info import get_info 4 | from ffmpegcv.ffmpeg_reader_camera import FFmpegReaderCAM, ProducerThread 5 | from ffmpegcv.ffmpeg_reader import ( 6 | get_videofilter_cpu, 7 | get_outnumpyshape, 8 | get_videofilter_gpu, 9 | get_num_NVIDIA_GPUs, 10 | decoder_to_nvidia, 11 | ) 12 | 13 | 14 | class FFmpegReaderStream(FFmpegReaderCAM): 15 | def __init__(self): 16 | super().__init__() 17 | 18 | @staticmethod 19 | def VideoReader( 20 | stream_url, 21 | codec, 22 | pix_fmt, 23 | crop_xywh, 24 | resize, 25 | resize_keepratio, 26 | resize_keepratioalign, 27 | timeout, 28 | ): 29 | 30 | vid = FFmpegReaderStream() 31 | videoinfo = get_info(stream_url, timeout) 32 | vid.origin_width = videoinfo.width 33 | vid.origin_height = videoinfo.height 34 | vid.fps = videoinfo.fps 35 | vid.codec = codec if codec else videoinfo.codec 36 | vid.count = videoinfo.count 37 | vid.duration = videoinfo.duration 38 | vid.pix_fmt = pix_fmt 39 | 40 | ( 41 | (vid.crop_width, vid.crop_height), 42 | (vid.width, vid.height), 43 | filteropt, 44 | ) = get_videofilter_cpu( 45 | (vid.origin_width, vid.origin_height), 46 | pix_fmt, 47 | crop_xywh, 48 | resize, 49 | resize_keepratio, 50 | resize_keepratioalign, 51 | ) 52 | vid.size = (vid.width, vid.height) 53 | 54 | rtsp_opt = '' if not stream_url.startswith('rtsp://') else '-rtsp_flags prefer_tcp -pkt_size 736 ' 55 | vid.ffmpeg_cmd = ( 56 | f"ffmpeg -loglevel warning " 57 | f" {rtsp_opt} " 58 | f" -vcodec {vid.codec} -i {stream_url} " 59 | f" {filteropt} -pix_fmt {pix_fmt} -f rawvideo pipe:" 60 | ) 61 | vid.out_numpy_shape = get_outnumpyshape(vid.size, pix_fmt) 62 | vid.process = run_async(vid.ffmpeg_cmd) 63 | 64 | vid.isopened = True 65 | 66 | # producer 67 | vid.step = 1 68 | vid.q = Queue(maxsize=30) 69 | producer = ProducerThread(vid, vid.q) 70 | producer.start() 71 | return vid 72 | 73 | 74 | class FFmpegReaderStreamNV(FFmpegReaderCAM): 75 | def __init__(self): 76 | super().__init__() 77 | 78 | @staticmethod 79 | def VideoReader( 80 | stream_url, 81 | codec, 82 | pix_fmt, 83 | crop_xywh, 84 | resize, 85 | resize_keepratio, 86 | resize_keepratioalign, 87 | gpu, 88 | timeout, 89 | ): 90 | numGPU = get_num_NVIDIA_GPUs() 91 | vid = FFmpegReaderStreamNV() 92 | videoinfo = get_info(stream_url, timeout) 93 | vid.origin_width = videoinfo.width 94 | vid.origin_height = videoinfo.height 95 | vid.fps = videoinfo.fps 96 | vid.codec = codec if codec else videoinfo.codec 97 | vid.codecNV = decoder_to_nvidia(vid.codec) 98 | vid.count = videoinfo.count 99 | vid.duration = videoinfo.duration 100 | vid.pix_fmt = pix_fmt 101 | 102 | ( 103 | (vid.crop_width, vid.crop_height), 104 | (vid.width, vid.height), 105 | (cropopt, scaleopt, filteropt), 106 | ) = get_videofilter_gpu( 107 | (vid.origin_width, vid.origin_height), 108 | pix_fmt, 109 | crop_xywh, 110 | resize, 111 | resize_keepratio, 112 | resize_keepratioalign, 113 | ) 114 | vid.size = (vid.width, vid.height) 115 | 116 | rtsp_opt = '' if not stream_url.startswith('rtsp://') else '-rtsp_flags prefer_tcp -pkt_size 736 ' 117 | vid.ffmpeg_cmd = ( 118 | f"ffmpeg -loglevel warning -hwaccel cuda -hwaccel_device {gpu} " 119 | f" {rtsp_opt} " 120 | f' -vcodec {vid.codecNV} {cropopt} {scaleopt} -i "{stream_url}" ' 121 | f" {filteropt} -pix_fmt {pix_fmt} -f rawvideo pipe:" 122 | ) 123 | vid.out_numpy_shape = get_outnumpyshape(vid.size, pix_fmt) 124 | vid.process = run_async(vid.ffmpeg_cmd) 125 | 126 | vid.isopened = True 127 | 128 | # producer 129 | vid.step = 1 130 | vid.q = Queue(maxsize=30) 131 | producer = ProducerThread(vid, vid.q) 132 | producer.start() 133 | return vid 134 | -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_reader_stream_realtime.py: -------------------------------------------------------------------------------- 1 | from ffmpegcv.ffmpeg_reader import ( 2 | FFmpegReader, 3 | get_videofilter_cpu, 4 | get_outnumpyshape, 5 | get_videofilter_gpu, 6 | get_num_NVIDIA_GPUs, 7 | decoder_to_nvidia, 8 | ) 9 | from ffmpegcv.stream_info import get_info 10 | 11 | 12 | class FFmpegReaderStreamRT(FFmpegReader): 13 | def __init__(self): 14 | super().__init__() 15 | 16 | @staticmethod 17 | def VideoReader( 18 | stream_url, 19 | codec, 20 | pix_fmt, 21 | crop_xywh, 22 | resize, 23 | resize_keepratio, 24 | resize_keepratioalign, 25 | timeout, 26 | ): 27 | vid = FFmpegReaderStreamRT() 28 | videoinfo = get_info(stream_url, timeout) 29 | vid.origin_width = videoinfo.width 30 | vid.origin_height = videoinfo.height 31 | vid.fps = videoinfo.fps 32 | vid.codec = codec if codec else videoinfo.codec 33 | vid.count = videoinfo.count 34 | vid.duration = videoinfo.duration 35 | vid.pix_fmt = pix_fmt 36 | 37 | ( 38 | (vid.crop_width, vid.crop_height), 39 | (vid.width, vid.height), 40 | filteropt, 41 | ) = get_videofilter_cpu( 42 | (vid.origin_width, vid.origin_height), 43 | pix_fmt, 44 | crop_xywh, 45 | resize, 46 | resize_keepratio, 47 | resize_keepratioalign, 48 | ) 49 | vid.size = (vid.width, vid.height) 50 | 51 | rtsp_opt = '' if not stream_url.startswith('rtsp://') else '-rtsp_flags prefer_tcp -pkt_size 736 ' 52 | vid.ffmpeg_cmd = ( 53 | f"ffmpeg -loglevel error " 54 | f" {rtsp_opt} " 55 | "-fflags nobuffer -flags low_delay -strict experimental " 56 | f" -vcodec {vid.codec} -i {stream_url}" 57 | f" {filteropt} -pix_fmt {pix_fmt} -f rawvideo pipe:" 58 | ) 59 | 60 | vid.out_numpy_shape = get_outnumpyshape(vid.size, pix_fmt) 61 | return vid 62 | 63 | 64 | class FFmpegReaderStreamRTNV(FFmpegReader): 65 | def __init__(self): 66 | super().__init__() 67 | 68 | @staticmethod 69 | def VideoReader( 70 | stream_url, 71 | codec, 72 | pix_fmt, 73 | crop_xywh, 74 | resize, 75 | resize_keepratio, 76 | resize_keepratioalign, 77 | gpu, 78 | timeout, 79 | ): 80 | vid = FFmpegReaderStreamRTNV() 81 | videoinfo = get_info(stream_url, timeout) 82 | vid.origin_width = videoinfo.width 83 | vid.origin_height = videoinfo.height 84 | vid.fps = videoinfo.fps 85 | vid.codec = codec if codec else videoinfo.codec 86 | vid.codecNV = decoder_to_nvidia(vid.codec) 87 | vid.count = videoinfo.count 88 | vid.duration = videoinfo.duration 89 | vid.pix_fmt = pix_fmt 90 | 91 | numGPU = get_num_NVIDIA_GPUs() 92 | assert numGPU > 0, "No GPU found" 93 | gpu = int(gpu) % numGPU if gpu is not None else 0 94 | 95 | ( 96 | (vid.crop_width, vid.crop_height), 97 | (vid.width, vid.height), 98 | (cropopt, scaleopt, filteropt), 99 | ) = get_videofilter_gpu( 100 | (vid.origin_width, vid.origin_height), 101 | pix_fmt, 102 | crop_xywh, 103 | resize, 104 | resize_keepratio, 105 | resize_keepratioalign, 106 | ) 107 | vid.size = (vid.width, vid.height) 108 | 109 | rtsp_opt = "-rtsp_transport tcp " if stream_url.startswith("rtsp://") else "" 110 | vid.ffmpeg_cmd = ( 111 | f"ffmpeg -loglevel error -hwaccel cuda -hwaccel_device {gpu} " 112 | f" {rtsp_opt} " 113 | f"{cropopt} {scaleopt} " 114 | "-fflags nobuffer -flags low_delay -strict experimental " 115 | f' -vcodec {vid.codecNV} -i "{stream_url}" ' 116 | f" {filteropt} -pix_fmt {pix_fmt} -f rawvideo pipe:" 117 | ) 118 | 119 | vid.out_numpy_shape = get_outnumpyshape(vid.size, pix_fmt) 120 | return vid 121 | -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_writer.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import warnings 3 | import pprint 4 | import select 5 | import sys 6 | from .video_info import run_async, release_process_writer, get_num_NVIDIA_GPUs 7 | 8 | 9 | IN_COLAB = "google.colab" in sys.modules 10 | 11 | 12 | class FFmpegWriter: 13 | def __init__(self): 14 | self.iframe = -1 15 | self.size = None 16 | self.width, self.height = None, None 17 | self.waitInit = True 18 | self._isopen = True 19 | 20 | def __enter__(self): 21 | return self 22 | 23 | def __exit__(self, type, value, traceback): 24 | self.release() 25 | 26 | def __del__(self): 27 | self.release() 28 | 29 | def __repr__(self): 30 | props = pprint.pformat(self.__dict__).replace("{", " ").replace("}", " ") 31 | return f"{self.__class__}\n" + props 32 | 33 | @staticmethod 34 | def VideoWriter( 35 | filename, codec, fps, pix_fmt, bitrate=None, resize=None, preset=None 36 | ): 37 | if codec is None: 38 | codec = "h264" 39 | elif not isinstance(codec, str): 40 | codec = "h264" 41 | warnings.simplefilter( 42 | """ 43 | Codec should be a string. Eg `h264`, `h264_nvenc`. 44 | You may used CV2.VideoWriter_fourcc, which will be ignored. 45 | """ 46 | ) 47 | assert resize is None or len(resize) == 2 48 | 49 | vid = FFmpegWriter() 50 | vid.fps = fps 51 | vid.codec, vid.pix_fmt, vid.filename = codec, pix_fmt, filename 52 | vid.bitrate = bitrate 53 | vid.resize = resize 54 | vid.preset = preset 55 | return vid 56 | 57 | def _init_video_stream(self): 58 | bitrate_str = f"-b:v {self.bitrate} " if self.bitrate else "" 59 | rtsp_str = f"-f rtsp" if self.filename.startswith("rtsp://") else "" 60 | filter_str = ( 61 | "" 62 | if self.resize == self.size 63 | else f"-vf scale={self.resize[0]}:{self.resize[1]}" 64 | ) 65 | target_pix_fmt = getattr(self, "target_pix_fmt", "yuv420p") 66 | preset_str = f"-preset {self.preset} " if self.preset else "" 67 | 68 | self.ffmpeg_cmd = ( 69 | f"ffmpeg -y -loglevel error " 70 | f"-f rawvideo -pix_fmt {self.pix_fmt} -s {self.width}x{self.height} -r {self.fps} -i pipe: " 71 | f"{bitrate_str} " 72 | f"-r {self.fps} -c:v {self.codec} " 73 | f"{preset_str}" 74 | f"{filter_str} {rtsp_str} " 75 | f'-pix_fmt {target_pix_fmt} "{self.filename}"' 76 | ) 77 | self.process = run_async(self.ffmpeg_cmd) 78 | 79 | def write(self, img: np.ndarray): 80 | if self.waitInit: 81 | if self.pix_fmt in ("nv12", "yuv420p", "yuvj420p"): 82 | height_15, width = img.shape[:2] 83 | assert width % 2 == 0 and height_15 * 2 % 3 == 0 84 | height = int(height_15 / 1.5) 85 | else: 86 | height, width = img.shape[:2] 87 | self.width, self.height = width, height 88 | self.in_numpy_shape = img.shape 89 | self.size = (width, height) 90 | self.resize = self.size if self.resize is None else tuple(self.resize) 91 | self._init_video_stream() 92 | self.waitInit = False 93 | 94 | self.iframe += 1 95 | assert self.in_numpy_shape == img.shape 96 | img = img.astype(np.uint8).tobytes() 97 | self.process.stdin.write(img) 98 | 99 | stderrreadable, _, _ = select.select([self.process.stderr], [], [], 0) 100 | if stderrreadable: 101 | data = self.process.stderr.read(1024) 102 | sys.stderr.buffer.write(data) 103 | 104 | def isOpened(self): 105 | return self._isopen 106 | 107 | def release(self): 108 | self._isopen = False 109 | if hasattr(self, "process"): 110 | release_process_writer(self.process) 111 | 112 | def close(self): 113 | return self.release() 114 | 115 | 116 | class FFmpegWriterNV(FFmpegWriter): 117 | @staticmethod 118 | def VideoWriter( 119 | filename, codec, fps, pix_fmt, gpu, bitrate=None, resize=None, preset=None 120 | ): 121 | numGPU = get_num_NVIDIA_GPUs() 122 | assert numGPU 123 | gpu = int(gpu) % numGPU if gpu is not None else 0 124 | if codec is None: 125 | codec = "hevc_nvenc" 126 | elif not isinstance(codec, str): 127 | codec = "hevc_nvenc" 128 | warnings.simplefilter( 129 | """ 130 | Codec should be a string. Eg `h264`, `h264_nvenc`. 131 | You may used CV2.VideoWriter_fourcc, which will be ignored. 132 | """ 133 | ) 134 | elif codec.endswith("_nvenc"): 135 | codec = codec 136 | else: 137 | codec = codec + "_nvenc" 138 | assert codec in [ 139 | "hevc_nvenc", 140 | "h264_nvenc", 141 | ], "codec should be `hevc_nvenc` or `h264_nvenc`" 142 | assert resize is None or len(resize) == 2 143 | 144 | vid = FFmpegWriterNV() 145 | vid.fps = fps 146 | vid.codec, vid.pix_fmt, vid.filename = codec, pix_fmt, filename 147 | vid.gpu = gpu 148 | vid.bitrate = bitrate 149 | vid.resize = resize 150 | vid.preset = preset if preset is not None else ("default" if IN_COLAB else "fast") 151 | return vid 152 | 153 | def _init_video_stream(self): 154 | bitrate_str = f"-b:v {self.bitrate} " if self.bitrate else "" 155 | rtsp_str = f"-f rtsp" if self.filename.startswith("rtsp://") else "" 156 | filter_str = ( 157 | "" 158 | if self.resize == self.size 159 | else f"-vf scale={self.resize[0]}:{self.resize[1]}" 160 | ) 161 | self.ffmpeg_cmd = ( 162 | f"ffmpeg -y -loglevel error " 163 | f"-f rawvideo -pix_fmt {self.pix_fmt} -s {self.width}x{self.height} -r {self.fps} -i pipe: " 164 | f"-preset {self.preset} {bitrate_str} " 165 | f"-r {self.fps} -gpu {self.gpu} -c:v {self.codec} " 166 | f"{filter_str} {rtsp_str} " 167 | f'-pix_fmt yuv420p "{self.filename}"' 168 | ) 169 | self.process = run_async(self.ffmpeg_cmd) 170 | -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_writer_noblock.py: -------------------------------------------------------------------------------- 1 | from multiprocessing import Queue, Process, Array 2 | import numpy as np 3 | from .ffmpeg_writer import FFmpegWriter 4 | 5 | NFRAME = 10 6 | 7 | class FFmpegWriterNoblock(FFmpegWriter): 8 | def __init__(self, 9 | vwriter_fun, 10 | *vwriter_args, **vwriter_kwargs): 11 | super().__init__() 12 | vid:FFmpegWriter = vwriter_fun(*vwriter_args, **vwriter_kwargs) 13 | vid.release() 14 | 15 | props_name = ['width', 'height', 'fps', 'codec', 'pix_fmt', 16 | 'filename', 'size', 'bitrate'] 17 | for name in props_name: 18 | setattr(self, name, getattr(vid, name, None)) 19 | 20 | self.vwriter_fun = vwriter_fun 21 | self.vwriter_args = vwriter_args 22 | self.vwriter_kwargs = vwriter_kwargs 23 | self.q = Queue(maxsize=(NFRAME-2)) #buffer index, gluttonous snake NO biting its own tail 24 | self.waitInit = True 25 | self.process = None 26 | 27 | def write(self, img:np.ndarray): 28 | if self.waitInit: 29 | if self.size is None: 30 | self.size = (img.shape[1], img.shape[0]) 31 | else: 32 | assert tuple(self.size) == (img.shape[1], img.shape[0]) 33 | self.in_numpy_shape = img.shape 34 | self._init_share_array() 35 | process = Process(target=child_process, 36 | args=(self.shared_array, self.q, self.in_numpy_shape, 37 | self.vwriter_fun, self.vwriter_args, self.vwriter_kwargs)) 38 | process.start() 39 | self.process = process 40 | self.width, self.height = self.size 41 | self.waitInit = False 42 | 43 | self.iframe += 1 44 | data_id = self.iframe % NFRAME 45 | self.np_array[data_id] = img 46 | self.q.put(data_id) 47 | 48 | def _init_share_array(self): 49 | self.shared_array = Array('b', int(NFRAME*np.prod(self.in_numpy_shape))) 50 | self.np_array = np.frombuffer(self.shared_array.get_obj(), dtype=np.uint8).reshape((NFRAME,*self.in_numpy_shape)) 51 | 52 | def release(self): 53 | if self.process is not None and self.process.is_alive(): 54 | self.q.put(None) 55 | self.process.join() 56 | 57 | 58 | def child_process(shared_array, q:Queue, in_numpy_shape, vwriter_fun, vwriter_args, vwriter_kwargs): 59 | vid = vwriter_fun(*vwriter_args, **vwriter_kwargs) 60 | np_array = np.frombuffer(shared_array.get_obj(), dtype=np.uint8).reshape((NFRAME,*in_numpy_shape)) 61 | with vid: 62 | while True: 63 | data_id = q.get() 64 | if data_id is None: 65 | break 66 | else: 67 | img = np_array[data_id] 68 | vid.write(img) 69 | -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_writer_qsv.py: -------------------------------------------------------------------------------- 1 | import warnings 2 | from ffmpegcv.ffmpeg_writer import FFmpegWriter 3 | from .video_info import ( 4 | run_async, 5 | get_num_QSV_GPUs, 6 | decoder_to_qsv, 7 | ) 8 | 9 | 10 | class FFmpegWriterQSV(FFmpegWriter): 11 | @staticmethod 12 | def VideoWriter(filename, codec, fps, pix_fmt, gpu, bitrate=None, resize=None, preset=None): 13 | assert gpu is None or gpu == 0, 'Cannot use multiple QSV gpu yet.' 14 | numGPU = get_num_QSV_GPUs() 15 | assert numGPU 16 | gpu = int(gpu) % numGPU if gpu is not None else 0 17 | if codec is None: 18 | codec = "hevc_qsv" 19 | elif not isinstance(codec, str): 20 | codec = "hevc_qsv" 21 | warnings.simplefilter( 22 | """ 23 | Codec should be a string. Eg `h264`, `hevc`. 24 | You may used CV2.VideoWriter_fourcc, which will be ignored. 25 | """ 26 | ) 27 | else: 28 | codec = decoder_to_qsv(codec) 29 | assert resize is None or len(resize) == 2 30 | 31 | vid = FFmpegWriterQSV() 32 | vid.fps = fps 33 | vid.codec, vid.pix_fmt, vid.filename = codec, pix_fmt, filename 34 | vid.gpu = gpu 35 | vid.bitrate = bitrate 36 | vid.resize = resize 37 | vid.preset = preset 38 | return vid 39 | -------------------------------------------------------------------------------- /ffmpegcv/ffmpeg_writer_stream_realtime.py: -------------------------------------------------------------------------------- 1 | from .video_info import run_async 2 | from ffmpegcv.ffmpeg_writer import FFmpegWriter 3 | 4 | 5 | class FFmpegWriterStreamRT(FFmpegWriter): 6 | @staticmethod 7 | def VideoWriter( 8 | filename: str, codec, pix_fmt, bitrate=None, resize=None, preset=None 9 | ) -> FFmpegWriter: 10 | assert codec in ["h264", "libx264", "x264", "mpeg4"] 11 | assert pix_fmt in ["bgr24", "rgb24", "gray"] 12 | # assert filename.startswith('rtmp://'), 'currently only support rtmp' 13 | assert resize is None or len(resize) == 2 14 | vid = FFmpegWriterStreamRT() 15 | vid.filename = filename 16 | vid.codec = codec 17 | vid.pix_fmt = pix_fmt 18 | vid.bitrate = bitrate 19 | vid.resize = resize 20 | if preset is not None: 21 | print("Preset is auto configured in FFmpegWriterStreamRT") 22 | vid.preset = "ultrafast" 23 | return vid 24 | 25 | def _init_video_stream(self): 26 | bitrate_str = f"-b:v {self.bitrate} " if self.bitrate else "" 27 | rtsp_str = f"-f rtsp" if self.filename.startswith("rtsp://") else "" 28 | filter_str = ( 29 | "" 30 | if self.resize == self.size 31 | else f"-vf scale={self.resize[0]}:{self.resize[1]}" 32 | ) 33 | self.ffmpeg_cmd = ( 34 | f"ffmpeg -loglevel warning " 35 | f"-f rawvideo -pix_fmt {self.pix_fmt} -s {self.width}x{self.height} -i pipe: " 36 | f"{bitrate_str} -f flv -rtsp_transport tcp " 37 | f" -tune zerolatency -preset {self.preset} " 38 | f"{filter_str} {rtsp_str} " 39 | f' -c:v {self.codec} -g 50 -pix_fmt yuv420p "{self.filename}"' 40 | ) 41 | self.process = run_async(self.ffmpeg_cmd) 42 | -------------------------------------------------------------------------------- /ffmpegcv/stream_info.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | from collections import namedtuple 3 | import json 4 | import shlex 5 | 6 | 7 | def get_info(stream_url, timeout=None, duration_ms: int = 100): 8 | rtsp_opt = '' if not stream_url.startswith('rtsp://') else '-rtsp_flags prefer_tcp -pkt_size 736 ' 9 | analyze_duration = f'-analyzeduration {duration_ms * 1000}' 10 | cmd = (f'ffprobe -v quiet -print_format json=compact=1 {rtsp_opt} {analyze_duration} ' 11 | f'-select_streams v:0 -show_format -show_streams "{stream_url}"') 12 | output = subprocess.check_output(shlex.split(cmd), shell=False, timeout=timeout) 13 | data: dict = json.loads(output) 14 | vinfo: dict = data['streams'][0] 15 | 16 | StreamInfo = namedtuple( 17 | "StreamInfo", ["width", "height", "fps", "count", "codec", "duration"] 18 | ) 19 | outinfo = dict() 20 | outinfo["width"] = int(vinfo["width"]) 21 | outinfo["height"] = int(vinfo["height"]) 22 | outinfo["fps"] = eval(vinfo.get("avg_frame_rate", vinfo["r_frame_rate"])) 23 | outinfo["count"] = None 24 | outinfo["codec"] = vinfo["codec_name"] 25 | outinfo["duration"] = None 26 | streaminfo = StreamInfo(**outinfo) 27 | 28 | return streaminfo 29 | 30 | 31 | if __name__ == "__main__": 32 | stream_url = "http://devimages.apple.com.edgekey.net/streaming/examples/bipbop_4x3/gear2/prog_index.m3u8" 33 | streaminfo = get_info(stream_url) 34 | print(streaminfo) 35 | -------------------------------------------------------------------------------- /ffmpegcv/version.py: -------------------------------------------------------------------------------- 1 | __version__='0.3.18' -------------------------------------------------------------------------------- /ffmpegcv/video_info.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | from subprocess import Popen, PIPE 3 | import re 4 | from collections import namedtuple 5 | import json 6 | import shlex 7 | import platform 8 | 9 | scan_the_whole = {"mkv", "flv", "ts"} # scan the whole file to the count, slow 10 | 11 | _is_windows = platform.system() == "Windows" 12 | _inited_get_num_NVIDIA_GPUs = False 13 | _inited_get_num_QSV_GPUs = False 14 | _num_NVIDIA_GPUs = -1 15 | _num_QSV_GPUs = -1 16 | 17 | 18 | def get_info(video: str): 19 | do_scan_the_whole = video.split(".")[-1] in scan_the_whole 20 | 21 | def ffprobe_info_(do_scan_the_whole): 22 | use_count_packets = '-count_packets' if do_scan_the_whole else '' 23 | cmd = 'ffprobe -v quiet -print_format json=compact=1 -select_streams v:0 {} -show_streams "{}"'.format( 24 | use_count_packets, video) 25 | 26 | output = subprocess.check_output(shlex.split(cmd), shell=False) 27 | data: dict = json.loads(output) 28 | vinfo: dict = data['streams'][0] 29 | return vinfo 30 | 31 | vinfo = ffprobe_info_(do_scan_the_whole) 32 | 33 | if "nb_frames" not in vinfo: 34 | do_scan_the_whole = True 35 | vinfo = ffprobe_info_(do_scan_the_whole) 36 | 37 | # VideoInfo = namedtuple( 38 | # "VideoInfo", ["width", "height", "fps", "count", "codec", "duration", "pix_fmt"] 39 | # ) 40 | VideoInfo = namedtuple( 41 | "VideoInfo", ["width", "height", "fps", "count", "codec", "duration"] 42 | ) 43 | outinfo = dict() 44 | outinfo["width"] = int(vinfo["width"]) 45 | outinfo["height"] = int(vinfo["height"]) 46 | outinfo["fps"] = eval(vinfo["r_frame_rate"]) 47 | outinfo["count"] = int( 48 | vinfo["nb_read_packets" if do_scan_the_whole else "nb_frames"] 49 | ) # nb_read_packets | nb_frames 50 | outinfo["codec"] = vinfo["codec_name"] 51 | # outinfo['pix_fmt'] = vinfo['pix_fmt'] 52 | 53 | outinfo["duration"] = ( 54 | float(vinfo["duration"]) 55 | if "duration" in vinfo 56 | else outinfo["count"] / outinfo["fps"] 57 | ) 58 | videoinfo = VideoInfo(**outinfo) 59 | 60 | return videoinfo 61 | 62 | 63 | def get_info_precise(video: str): 64 | videoinfo = get_info(video) 65 | cmd = ( 66 | "ffprobe -v error -select_streams v:0 -show_entries frame=pts_time " 67 | f' -of default=noprint_wrappers=1:nokey=1 -read_intervals 0%+#1,99999% "{video}"' 68 | ) 69 | output = subprocess.check_output( 70 | shlex.split(cmd), shell=False, stderr=subprocess.DEVNULL 71 | ) 72 | pts_start, *_, pts_end = output.decode().split() 73 | pts_start, pts_end = float(pts_start), float(pts_end) 74 | videoinfod = videoinfo._asdict() 75 | duration_ = pts_end - pts_start 76 | videoinfod["fps"] = round((videoinfo.count - 1) / duration_, 3) 77 | videoinfod["duration"] = round(duration_ + 1 / videoinfod["fps"], 3) 78 | videoinfo_precise = videoinfo.__class__(*videoinfod.values()) 79 | return videoinfo_precise 80 | 81 | 82 | def get_num_NVIDIA_GPUs(): 83 | global _num_NVIDIA_GPUs, _inited_get_num_NVIDIA_GPUs 84 | if not _inited_get_num_NVIDIA_GPUs: 85 | cmd = "ffmpeg -f lavfi -i nullsrc -c:v h264_nvenc -gpu list -f null -" 86 | p = Popen(cmd.split(), shell=False, stdin=PIPE, stdout=PIPE, stderr=PIPE) 87 | stdout, stderr = p.communicate(b"") 88 | p.stdin.close() 89 | p.stdout.close() 90 | p.terminate() 91 | pattern = re.compile(r"GPU #\d+ - < ") 92 | nv_info = pattern.findall(stderr.decode()) 93 | _num_NVIDIA_GPUs = len(nv_info) 94 | _inited_get_num_NVIDIA_GPUs = True 95 | return _num_NVIDIA_GPUs 96 | 97 | 98 | def get_num_QSV_GPUs(): 99 | global _num_QSV_GPUs, _inited_get_num_QSV_GPUs 100 | if not _inited_get_num_QSV_GPUs: 101 | cmd = "ffmpeg -hide_banner -f qsv -h encoder=h264_qsv" 102 | p = Popen(cmd.split(), shell=False, stdin=PIPE, stdout=PIPE, stderr=PIPE) 103 | stdout, stderr = p.communicate(b"") 104 | _num_QSV_GPUs = 1 if len(stdout) > 50 else 0 105 | _inited_get_num_QSV_GPUs = True 106 | return _num_QSV_GPUs 107 | 108 | 109 | def encoder_to_nvidia(codec): 110 | codec_map = {"h264": "h264_nvenc", "hevc": "hevc_nvenc"} 111 | 112 | if codec in codec_map: 113 | return codec_map[codec] 114 | elif codec in codec_map.values(): 115 | return codec 116 | else: 117 | raise Exception("No NV codec found for %s" % codec) 118 | 119 | 120 | def encoder_to_qsv(codec): 121 | codec_map = { 122 | "h264": "h264_qsv", 123 | "hevc": "hevc_qsv", 124 | "mjpeg": "mjpeg_qsv", 125 | "mpeg2video": "mpeg2_qsv", 126 | "vp9": "vp9_qsv", 127 | } 128 | 129 | if codec in codec_map: 130 | return codec_map[codec] 131 | elif codec in codec_map.values(): 132 | return codec 133 | else: 134 | raise Exception("No QSV codec found for %s" % codec) 135 | 136 | 137 | def decoder_to_nvidia(codec): 138 | codec_map = { 139 | "av1": "av1_cuvid", 140 | "h264": "h264_cuvid", 141 | "x264": "h264_cuvid", 142 | "hevc": "hevc_cuvid", 143 | "x265": "hevc_cuvid", 144 | "h265": "hevc_cuvid", 145 | "mjpeg": "mjpeg_cuvid", 146 | "mpeg1video": "mpeg1_cuvid", 147 | "mpeg2video": "mpeg2_cuvid", 148 | "mpeg4": "mpeg4_cuvid", 149 | "vp1": "vp1_cuvid", 150 | "vp8": "vp8_cuvid", 151 | "vp9": "vp9_cuvid", 152 | } 153 | 154 | if codec in codec_map: 155 | return codec_map[codec] 156 | elif codec in codec_map.values(): 157 | return codec 158 | else: 159 | raise Exception("No NV codec found for %s" % codec) 160 | 161 | 162 | def decoder_to_qsv(codec): 163 | codec_map = { 164 | "av1": "av1_qsv", 165 | "h264": "h264_qsv", 166 | "hevc": "hevc_qsv", 167 | "mjpeg": "mjpeg_qsv", 168 | "mpeg2video": "mpeg2_qsv", 169 | "vc1": "vc1_qsv", 170 | "vp8": "vp8_qsv", 171 | "vp9": "vp9_qsv", 172 | } 173 | 174 | if codec in codec_map: 175 | return codec_map[codec] 176 | elif codec in codec_map.values(): 177 | return codec 178 | else: 179 | raise Exception("No QSV codec found for %s" % codec) 180 | 181 | 182 | def run_async(args): 183 | bufsize = -1 184 | if isinstance(args, str): 185 | args = shlex.split(args) 186 | return Popen( 187 | args, 188 | stdin=PIPE, 189 | stdout=PIPE, 190 | stderr=PIPE, 191 | shell=False, 192 | bufsize=bufsize, 193 | ) 194 | 195 | def run_async_reader(args): 196 | bufsize = -1 197 | if isinstance(args, str): 198 | args = shlex.split(args) 199 | 200 | return Popen( 201 | args, 202 | stdin=None, 203 | stdout=PIPE, 204 | stderr=subprocess.DEVNULL, 205 | shell=False, 206 | bufsize=bufsize, 207 | ) 208 | 209 | 210 | def release_process(process: Popen, forcekill=False): 211 | if hasattr(process, "stdin") and process.stdin is not None: 212 | process.stdin.close() 213 | if hasattr(process, "stdout") and process.stdout is not None: 214 | process.stdout.close() 215 | if hasattr(process, "stderr") and process.stderr is not None: 216 | process.stderr.close() 217 | if forcekill and hasattr(process, "terminate") and not _is_windows: 218 | process.terminate() 219 | if forcekill and hasattr(process, "wait"): 220 | process.wait() 221 | 222 | 223 | def release_process_writer(process: Popen): 224 | if hasattr(process, "stdin"): 225 | process.stdin.close() 226 | if hasattr(process, "stdout"): 227 | process.stdout.close() 228 | if hasattr(process, "wait"): 229 | process.wait() 230 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # setup.py 2 | from setuptools import setup, find_packages 3 | from pathlib import Path 4 | 5 | this_directory = Path(__file__).parent 6 | long_description = (this_directory / "README.md").read_text("utf-8") 7 | 8 | setup( 9 | name="ffmpegcv", # 应用名 10 | version="0.3.18", # 版本号 11 | packages=find_packages(include=["ffmpegcv*"]), # 包括在安装包内的 Python 包 12 | author="chenxf", 13 | author_email="cxf529125853@163.com", 14 | url="https://github.com/chenxinfeng4/ffmpegcv", 15 | long_description=long_description, 16 | long_description_content_type="text/markdown", 17 | # 添加依赖项 18 | python_requires=">=3.6", 19 | install_requires=[ 20 | "numpy", 21 | ], 22 | extras_require={"cuda": ["pycuda"]}, # 定义一个名为cuda的可选依赖项，并指定pycuda 23 | ) 24 | --------------------------------------------------------------------------------