mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
033ab9da1b
Summary: Adding a simple video data layer which allows to read video data from frames, videos and output 5D tensor. It also allows multiple labels. The current implementation is based on ffmpeg Differential Revision: D4801798 fbshipit-source-id: 46448e9c65fb055c2d71855447383a33ade0e444
33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
## @package tools
|
|
# Module caffe2.python.helpers.tools
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
|
|
def image_input(
|
|
model, blob_in, blob_out, order="NCHW", use_gpu_transform=False, **kwargs
|
|
):
|
|
if order == "NCHW":
|
|
if (use_gpu_transform):
|
|
kwargs['use_gpu_transform'] = 1 if use_gpu_transform else 0
|
|
# GPU transform will handle NHWC -> NCHW
|
|
data, label = model.net.ImageInput(
|
|
blob_in, [blob_out[0], blob_out[1]], **kwargs
|
|
)
|
|
pass
|
|
else:
|
|
data, label = model.net.ImageInput(
|
|
blob_in, [blob_out[0] + '_nhwc', blob_out[1]], **kwargs
|
|
)
|
|
data = model.net.NHWC2NCHW(data, blob_out[0])
|
|
else:
|
|
data, label = model.net.ImageInput(blob_in, blob_out, **kwargs)
|
|
return data, label
|
|
|
|
|
|
def video_input(model, blob_in, blob_out, **kwargs):
|
|
data, label = model.net.VideoInput(blob_in, blob_out, **kwargs)
|
|
return data, label
|