├── .gitignore ├── .style.yapf ├── .travis.yml ├── Dockerfile ├── README.md ├── draft ├── scannertools_pytorch │ ├── generate_trace.py │ ├── scannertools_pytorch │ │ ├── __init__.py │ │ └── resnet.py │ ├── scannertools_pytorch_cpp │ │ ├── CMakeLists.txt │ │ ├── pytorch.cpp │ │ └── scannertools_pytorch.proto │ ├── setup.cfg │ ├── setup.py │ └── tests │ │ └── test_all.py └── scannertools_tfcpp │ ├── scannertools_tfcpp │ └── __init__.py │ ├── scannertools_tfcpp_cpp │ ├── CMakeLists.txt │ ├── scannertools_tfcpp.proto │ └── tfcpp.cpp │ ├── setup.cfg │ ├── setup.py │ └── tests │ └── test_all.py ├── scannertools ├── README.md ├── scannertools │ ├── __init__.py │ ├── caffe2.py │ ├── face_detection.py │ ├── face_embedding.py │ ├── gender_detection.py │ ├── imgproc │ │ └── __init__.py │ ├── maskrcnn_detection.py │ ├── misc │ │ └── __init__.py │ ├── net_descriptor.py │ ├── object_detection.py │ ├── old │ │ ├── __init__.py │ │ ├── audio.py │ │ ├── bboxes.py │ │ ├── clothing_detection.py │ │ ├── cpp_ops │ │ │ ├── CMakeLists.txt │ │ │ ├── flow_histogram_kernel_cpu.cpp │ │ │ ├── imgproc.cpp │ │ │ └── imgproc.proto │ │ ├── face_landmark_detection.py │ │ ├── hairstyle_detection.py │ │ ├── histograms.py │ │ ├── imgproc.py │ │ ├── optical_flow.py │ │ ├── pipeline.py │ │ ├── pipeline2.py │ │ ├── pipeline3.py │ │ ├── pose_detection.py │ │ ├── prelude.py │ │ ├── test_audio.py │ │ ├── test_caption.py │ │ ├── transcript_alignment.py │ │ └── video.py │ ├── shot_detection.py │ ├── storage │ │ ├── __init__.py │ │ ├── audio.py │ │ ├── caption.py │ │ ├── files.py │ │ └── python.py │ ├── tensorflow.py │ ├── tf_vis_utils.py │ ├── torch.py │ ├── tracker.py │ ├── types.py │ └── vis.py ├── scannertools_cpp │ ├── CMakeLists.txt │ ├── imgproc │ │ ├── CMakeLists.txt │ │ ├── blur_kernel_cpu.cpp │ │ ├── convert_color_kernel.cpp │ │ ├── frame_difference_kernel_cpu.cpp │ │ ├── histogram_kernel_cpu.cpp │ │ ├── histogram_kernel_gpu.cpp │ │ ├── image_decoder_kernel_cpu.cpp │ │ ├── image_decoder_kernel_gpu.cpp │ │ ├── montage_kernel_cpu.cpp │ │ ├── montage_kernel_gpu.cpp │ │ ├── optical_flow_kernel_cpu.cpp │ │ ├── optical_flow_kernel_gpu.cpp │ │ ├── resize_kernel.cpp │ │ └── scannertools_imgproc.proto │ ├── misc │ │ ├── CMakeLists.txt │ │ ├── discard_kernel.cpp │ │ ├── info_from_frame_kernel.cpp │ │ └── pass_kernel.cpp │ └── storage │ │ ├── CMakeLists.txt │ │ ├── audio_source.cpp │ │ ├── captions_source.cpp │ │ ├── files_sink.cpp │ │ ├── files_source.cpp │ │ ├── packed_file_source.cpp │ │ ├── python_source.cpp │ │ ├── scannertools_storage.proto │ │ └── srtparser.h ├── setup.cfg ├── setup.py └── tests │ └── test_all.py ├── scannertools_caffe ├── scannertools_caffe │ ├── __init__.py │ └── pose_detection.py ├── scannertools_caffe_cpp │ ├── CMakeLists.txt │ ├── caffe_input_kernel.cpp │ ├── caffe_input_kernel.h │ ├── caffe_input_kernel_cpu.cpp │ ├── caffe_input_kernel_gpu.cpp │ ├── caffe_input_transformer_base.h │ ├── caffe_input_transformer_cpu.cpp │ ├── caffe_input_transformer_gpu.cpp │ ├── caffe_kernel.cpp │ ├── caffe_kernel.h │ ├── caffe_kernel_cpu.cpp │ ├── caffe_kernel_gpu.cpp │ ├── cmake │ │ ├── FindCaffe.cmake │ │ └── FindOpenPose.cmake │ ├── cpm2_input_kernel_gpu.cpp │ ├── cpm2_kernel.cpp │ ├── cpm2_output_kernel_cpu.cpp │ ├── facenet_input_kernel_cpu.cpp │ ├── facenet_input_kernel_gpu.cpp │ ├── facenet_kernel.cpp │ ├── facenet_output_kernel_cpu.cpp │ ├── faster_rcnn_kernel.cpp │ ├── faster_rcnn_output_kernel_cpu.cpp │ ├── openpose_kernel.cpp │ ├── scannertools_caffe.proto │ └── yolo_output_kernel_cpu.cpp ├── setup.cfg ├── setup.py └── tests │ └── test_all.py ├── scannertools_infra ├── scannertools_infra │ ├── __init__.py │ └── tests.py └── setup.py ├── scannertools_sql ├── scannertools_sql │ ├── __init__.py │ └── storage.py ├── scannertools_sql_cpp │ ├── CMakeLists.txt │ ├── scannertools_sql.proto │ ├── sql.cpp │ ├── sql.h │ ├── sql_sink.cpp │ └── sql_source.cpp ├── setup.cfg ├── setup.py └── tests │ └── test_all.py └── scripts ├── clean-all.sh ├── install-all.sh ├── test-all.sh ├── travis-build.sh └── travis-publish.sh /.gitignore: -------------------------------------------------------------------------------- 1 | .cache 2 | **/__pycache__ 3 | .gdb_history 4 | *.pyc 5 | doc/_build 6 | doc/source/ 7 | build 8 | dist 9 | *.egg-info 10 | *.swp 11 | **/build 12 | **/.eggs 13 | **/yapf*.py 14 | -------------------------------------------------------------------------------- /.style.yapf: -------------------------------------------------------------------------------- 1 | [style] 2 | based_on_style = pep8 3 | column_limit = 100 -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | python: 3 | - '3.5' 4 | sudo: required 5 | services: 6 | - docker 7 | env: 8 | global: 9 | - DOCKER_REPO=scannerresearch/scannertools 10 | - DOCKER_EMAIL=wcrichto@cs.stanford.edu 11 | - DOCKER_USER=wcrichto 12 | - secure: M+N4gEsYsiF+zfFBSG7bvpmKNOff+c0uX1dUZ9/rFoBro7Lftsq3FvGuLMvfb9gzuNcWH9vSKf9fmUI6apeBxgVqFIzn/LiNhbdrVzg8P2IRtZ2GSDOmzyp0HDzL/bET8FncQLccKCG3Il3QmGstBTJ3cpxJqJ1O60CsnKGX+UlYr51yz53hPBETbgLsEFTNFXVifvgrKJumfBqjx2QdqkYTUQO6JjYbcUo3rlgqA6eSluCY/DgSB0xzaznB8XbdJkIi7B98E/zmxM6uJKAYXNR/Srv+jbcryjW9GXldYHbQad06O1bKK12UPnrCKIVVg9iHd/Mx7S9eqeueGOMAfozIAcGG+oOV/5PIcvVlUeyQY8ATHx7GFYvmegUqoFIX3Ph5o4v12997rmzhaLag0qigOtZ6O3/Dg79BMbFyTdYz7LzHlmZ64ST3D/MRK76Hl73g4xSnUBNah3wysljPLcWn2WG+fAoRJ8ArwApJN85UWj+s9OPFjapxfGBNykYOoSTqEKspK+pbCUrPWY6yEDtTjUoTnyfKLSxQiCw/n4Hnxk7NqTAwLga441fymljgk3svdBTSnx00RP50SIRjm0HdGhO9XOWM04ba48nGHnLlpIlqPmFNaqFXlwePXu7UZOMYzJmxfksj5XlX4HZA8G3LoUBaD8GJm3D0z/6jkrg= 13 | - secure: 2Vons/Eaclk+6pBpkQ3phmAvXkMre5M+XwlyR6NN3dKlD7N+xXyUytkODUBRuISp/9R82gd0GI9H/ccY/yE8EhuNBjKbN4ernmC9NdFfS7RYSfgbdUH7J6LT9+FPeFchlsqv139fZfLkxC+MZf6derSZb0AeP67Ibn2mo4oh6YNflWD/gm7CKerA7in4eYIkPUubjtVyfSsQPRBTYQTdxXDXi3XOr3gnAgVwGcu9dS02n7AxVUdxGOJyEbZsvCjwE1s0h5BrezeUQ9EDJM07bqmI+yZxzY7ITxvBuIE+rcwc1dOLaNq3MAhVHML52Yjo9OKjto40UPvfAd8IFTmnWy2L3+y06XVj84fZvuatDLwwY77SzLaPwQ1iemmNT/NqjErr9BUQ8Rkz0cM8WqnaYHmcTbfU6A1pHGsfW1/A8fa99c7kmaYgYJCPhEWAHC9uonO+37y5DUnOxd4aW2bv5C+cfYCDkhErJluQyyGcaZINya7aZFi4L0U6iuAw/6tJDHvrSZiC5BFI6raxrNanoIoRe4YoPD7us4aVuGNzrjm0SDz9a5mjngGqCBkqIuYIKccmP9EjjUS+gudXtSsZj8Utw04BT8WczGOG9GhZAGJ4JDkUbAlpZFrQ/eiqUg1O009Y7EQpjyYY7253dhc8ohfkfbJ+DaEvnYh1y4Jnp94= 14 | matrix: 15 | - TAG=cpu 16 | - TAG=gpu-9.0-cudnn7 17 | - TAG=gpu-9.1-cudnn7 18 | - TAG=gpu-10.1-cudnn7 19 | script: "travis_retry ./scripts/travis-build.sh" 20 | after_success: "travis_retry ./scripts/travis-publish.sh" 21 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | ARG tag 2 | FROM scannerresearch/scanner:${tag}-latest 3 | ARG tag2 4 | 5 | WORKDIR /opt 6 | 7 | # Fixes travis pip failure 8 | RUN rm /usr/share/python-wheels/urllib3-1.13.1-py2.py3-none-any.whl && \ 9 | rm /usr/share/python-wheels/requests-2.9.1-py2.py3-none-any.whl && \ 10 | pip3 install requests[security] --upgrade -v 11 | RUN pip3 install face-alignment scipy pysrt 12 | RUN if [ "$tag2" = "cpu" ]; then pip3 install tensorflow==1.12.0; else pip3 install tensorflow-gpu==1.12.0; fi 13 | RUN git clone https://github.com/scanner-research/facenet && \ 14 | git clone https://github.com/scanner-research/rude-carnie 15 | ENV PYTHONPATH /opt/facenet/src:/opt/rude-carnie:$PYTHONPATH 16 | 17 | # pytorch (specific version for maskRCNN) 18 | RUN pip3 install torchvision==0.3.0 torch==1.1.0 19 | 20 | # Install PyTorch Detection 21 | RUN if [ "$tag2" = "cpu" ]; then FORCE_CUDA="0"; else FORCE_CUDA="1"; fi 22 | ENV FORCE_CUDA=${FORCE_CUDA} 23 | RUN git clone https://github.com/facebookresearch/maskrcnn-benchmark.git \ 24 | && cd maskrcnn-benchmark \ 25 | && python3 setup.py build develop 26 | ENV PYTHONPATH /opt/maskrcnn-benchmark:$PYTHONPATH 27 | 28 | RUN apt-get update && apt-get install -y jq 29 | 30 | RUN echo "deb http://packages.cloud.google.com/apt cloud-sdk-xenial main" | \ 31 | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \ 32 | curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \ 33 | apt-get update && apt-get install -y google-cloud-sdk kubectl 34 | 35 | # https://github.com/keras-team/keras/issues/9567#issuecomment-370887563 36 | RUN if [ "$tag2" != "cpu" ]; then \ 37 | apt-get update && apt-get install -y --allow-downgrades --allow-change-held-packages --no-install-recommends \ 38 | libcudnn7=7.0.5.15-1+cuda9.0 \ 39 | libcudnn7-dev=7.0.5.15-1+cuda9.0 && \ 40 | rm -rf /var/lib/apt/lists/*; \ 41 | fi 42 | 43 | COPY . scannertools 44 | RUN cd scannertools && pip3 install --upgrade setuptools && ./scripts/install-all.sh 45 | 46 | WORKDIR /app 47 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # scannertools: video processing toolkit   [![Build Status](https://travis-ci.org/scanner-research/scannertools.svg?branch=master)](https://travis-ci.org/scanner-research/scannertools) 2 | 3 | Scannertools is a Python library of easy-to-use, off-the-shelf pipelines written using the [Scanner](https://github.com/scanner-research/scanner/) video processing engine. Scannertools provides implementations of: 4 | 5 | * [Object detection](http://scanner.run/api/scannertools.html#object-detection) 6 | * [Face detection](http://scanner.run/api/scannertools.html#face-detection) 7 | * [Face embedding](http://scanner.run/api/scannertools.html#face-embedding) 8 | * [Gender detection](http://scanner.run/api/scannertools.html#gender-detection) 9 | * [Pose detection](http://scanner.run/api/scannertools.html#pose-detection) 10 | * [Optical flow](http://scanner.run/api/scannertools.html#optical-flow) 11 | * [Shot detection](http://scanner.run/api/scannertools.html#shot-detection) 12 | 13 | See the documentation on [scanner.run](http://scanner.run/api.html#scannertools-the-scanner-standard-library) for more details. 14 | 15 | ## Installation 16 | 17 | You must have Scanner and all of its dependencies installed. See our [installation guide](http://scanner.run/guide/getting-started.html). 18 | 19 | Each subdirectory prefixed with "scannertools" is a module containing Python and C++ Scanner ops. The modules are separated because each expects different system dependencies: 20 | 21 | * [scannertools](https://github.com/scanner-research/scannertools/tree/master/scannertools): no additional dependencies beyond Scanner. 22 | * [scannertools.face_detection](https://github.com/scanner-research/scannertools/blob/master/scannertools/scannertools/face_detection.py) and [scannertools.face_embedding](https://github.com/scanner-research/scannertools/blob/master/scannertools/scannertools/face_embedding.py): depends on [Facenet](https://github.com/scanner-research/facenet) 23 | * [scannertools.gender_detection](https://github.com/scanner-research/scannertools/blob/master/scannertools/scannertools/gender_detection.py): depends on [rude-carnie](https://github.com/dpressel/rude-carnie) 24 | * [scannertools.object_detection](https://github.com/scanner-research/scannertools/blob/master/scannertools/scannertools/object_detection.py): depends on [TensorFlow](https://www.tensorflow.org/) 25 | * [scannertools_caffe](https://github.com/scanner-research/scannertools/tree/master/scannertools_caffe): depends on [Caffe](http://caffe.berkeleyvision.org/installation.html). 26 | * [scannertools_sql](https://github.com/scanner-research/scannertools/tree/master/scannertools_sql): depends on [pqxx](https://github.com/jtv/libpqxx). 27 | 28 | The [scannertools_infra](https://github.com/scanner-research/scannertools/tree/master/scannertools_infra) package contains build and test infrastructure for each of the scannertools submodules. 29 | 30 | To install the optional dependencies that scannertools use, first clone the dependency repo and then add the source to the python pyath. 31 | 32 | ### From pip 33 | 34 | We'll be uploading pip packages soon. In the meantime, follow the install from source direction. 35 | 36 | ### From source 37 | 38 | First clone the repository and install the infrastructure. 39 | 40 | ``` 41 | git clone https://github.com/scanner-research/scannertools 42 | cd scannertools 43 | cd scannertools_infra 44 | pip3 install --user -e . 45 | ``` 46 | 47 | Then, for each submodule you want to install, go into the the subdirectory and run pip, e.g.: 48 | 49 | ``` 50 | cd scannertools_caffe 51 | pip3 install --user -e . 52 | ``` 53 | 54 | To build the ops with GPU compatibility, pass the path to your CUDA installation: 55 | 56 | ``` 57 | pip3 install --install-option="--build-cuda=/usr/local/cuda" --user -e . 58 | ``` 59 | 60 | ## Usage 61 | 62 | See the documentation on [scanner.run](http://scanner.run/api.html#scannertools-the-scanner-standard-library) for usage. 63 | -------------------------------------------------------------------------------- /draft/scannertools_pytorch/generate_trace.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision.models as models 3 | 4 | with torch.no_grad(): 5 | resnet18 = models.resnet18(pretrained=True) 6 | resnet18.eval() 7 | 8 | # use an example input to trace the operations of the model 9 | example_input = torch.rand(1, 3, 224, 224) # 224 is the least input size, depends on the dataset you use 10 | 11 | script_module = torch.jit.trace(resnet18.cuda(), example_input.cuda()) 12 | script_module.save('script_module.pt') 13 | -------------------------------------------------------------------------------- /draft/scannertools_pytorch/scannertools_pytorch/__init__.py: -------------------------------------------------------------------------------- 1 | from scannertools_infra import _register_module 2 | 3 | _register_module(__file__, "scannertools_pytorch") 4 | -------------------------------------------------------------------------------- /draft/scannertools_pytorch/scannertools_pytorch/resnet.py: -------------------------------------------------------------------------------- 1 | import scannerpy as sp 2 | from scannertools.torch import TorchKernel 3 | import torchvision.models as models 4 | import torchvision.transforms as transforms 5 | from scannerpy.types import NumpyArrayFloat32 6 | from typing import Sequence 7 | import torch 8 | import numpy as np 9 | from timeit import default_timer as now 10 | 11 | @sp.register_python_op(batch=2, device_sets=[(sp.DeviceType.CPU, 0), (sp.DeviceType.GPU, 1)]) 12 | class Resnet(TorchKernel): 13 | def build_model(self): 14 | self._mu = torch.tensor([0.485, 0.456, 0.406]).unsqueeze(1).unsqueeze(2) 15 | self._sigma = torch.tensor([0.229, 0.224, 0.225]).unsqueeze(1).unsqueeze(2) 16 | if not self.cpu_only: 17 | self._mu = self._mu.cuda() 18 | self._sigma = self._sigma.cuda() 19 | 20 | return models.resnet18(pretrained=True) 21 | 22 | def execute(self, frame: Sequence[sp.FrameType]) -> Sequence[NumpyArrayFloat32]: 23 | batch_size = len(frame) 24 | 25 | start = now() 26 | batch_tensor = torch.from_numpy( 27 | np.moveaxis(np.concatenate(np.expand_dims(frame, axis=0), axis=0), 3, 1)) \ 28 | .type(torch.FloatTensor) 29 | 30 | if not self.cpu_only: 31 | start = now() 32 | batch_tensor = batch_tensor.cuda() 33 | #print('Transfer to device: {:.3f}'.format(now() - start)) 34 | 35 | batch_tensor /= 255.0 36 | 37 | batch_tensor -= self._mu 38 | batch_tensor /= self._sigma 39 | #print('Transform: {:.3f}'.format(now() - start)) 40 | 41 | with torch.no_grad(): 42 | start = now() 43 | output = self.model.forward(batch_tensor) 44 | #print('Forward: {:.3f}'.format(now() - start)) 45 | 46 | if not self.cpu_only: 47 | start = now() 48 | output = output.cpu() 49 | #print('Transfer from device: {:.3f}'.format(now() - start)) 50 | 51 | import sys 52 | sys.stdout.flush() 53 | 54 | return [output[i, :].numpy() for i in range(batch_size)] 55 | -------------------------------------------------------------------------------- /draft/scannertools_pytorch/scannertools_pytorch_cpp/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.2.0 FATAL_ERROR) 2 | 3 | execute_process( 4 | OUTPUT_VARIABLE SCANNER_CMAKE_PATH 5 | COMMAND python3 -c "import scannerpy.build_flags as b; b.print_cmake()") 6 | include(${SCANNER_CMAKE_PATH}) 7 | 8 | set(SOURCES pytorch.cpp) 9 | 10 | build_op( 11 | LIB_NAME scannertools_pytorch 12 | CPP_SRCS ${SOURCES} 13 | BUILD_CUDA ${BUILD_CUDA}) 14 | #PROTO_SRC scannertools_pytorch.proto) 15 | 16 | execute_process( 17 | OUTPUT_VARIABLE PYTORCH_PATH 18 | COMMAND python3 -c "import torch, os, sys; sys.stdout.write(os.path.dirname(torch.__file__))") 19 | find_package(Torch REQUIRED PATHS "${PYTORCH_PATH}/share/cmake/Torch/") 20 | 21 | #set(LIBTF_PATH /home/will/libtensorflow) 22 | 23 | set(LIBRARIES ${TORCH_LIBRARIES}) 24 | set(INCLUDES ${TORCH_INCLUDE_DIRS}) 25 | 26 | target_include_directories(scannertools_pytorch PUBLIC ${INCLUDES}) 27 | target_link_libraries(scannertools_pytorch PUBLIC ${LIBRARIES}) 28 | set_property(TARGET scannertools_pytorch PROPERTY CXX_STANDARD 11) 29 | -------------------------------------------------------------------------------- /draft/scannertools_pytorch/scannertools_pytorch_cpp/pytorch.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/types.pb.h" 4 | #include "scanner/util/memory.h" 5 | #include "scanner/util/serialize.h" 6 | 7 | #include 8 | #include 9 | 10 | namespace scanner { 11 | 12 | float mu[] = {0.485, 0.456, 0.406}; 13 | float sigma[] = {0.229, 0.224, 0.225}; 14 | 15 | class PyTorchKernel : public BatchedKernel { 16 | public: 17 | PyTorchKernel(const KernelConfig &config) 18 | : BatchedKernel(config), device_(config.devices[0]) {} 19 | 20 | void setup_with_resources(proto::Result *result) override { 21 | module_ = 22 | torch::jit::load("/home/will/scannertools/scannertools_pytorch/script_module.pt"); 23 | auto torchdevice = device_.type == DeviceType::GPU 24 | ? at::device({at::kCUDA, device_.id}) : at::device(at::kCPU); 25 | 26 | if (device_.type == DeviceType::GPU) { 27 | LOG_IF(FATAL, !torch::cuda::is_available()) 28 | << "PyTorch kernel supposed to use GPU, but Torch CUDA API is not available"; 29 | module_->to({at::kCUDA, device_.id}); 30 | } 31 | 32 | mu_t_ = torch::from_blob(mu, {3}).unsqueeze(1).unsqueeze(2); 33 | sigma_t_ = torch::from_blob(sigma, {3}).unsqueeze(1).unsqueeze(2); 34 | if (device_.type == DeviceType::GPU) { 35 | mu_t_ = mu_t_.to(torchdevice); 36 | sigma_t_ = sigma_t_.to(torchdevice); 37 | } 38 | 39 | at::Tensor dummy = torch::zeros({1, 3, 224, 224}); 40 | dummy = dummy.cuda(); 41 | module_->forward({dummy}); 42 | 43 | result->set_success(true); 44 | } 45 | 46 | void execute(const BatchedElements &input_columns, 47 | BatchedElements &output_columns) override { 48 | size_t batch_size = input_columns[0].size(); 49 | const Frame* frame = input_columns[0][0].as_const_frame(); 50 | LOG_IF(FATAL, frame->shape[0] != 224 || frame->shape[1] != 224) 51 | << "Frame must be 224 x 224, was " << frame->shape[0] << " x " << frame->shape[1]; 52 | 53 | if (batch_buffer_ == nullptr) { 54 | batch_buffer_ = new_buffer(device_, batch_size * frame->size()); 55 | } 56 | 57 | { 58 | ProfileBlock _block(profiler_, "make_batch"); 59 | for (size_t i = 0; i < batch_size; ++i) { 60 | memcpy_buffer(&batch_buffer_[i * frame->size()], device_, 61 | input_columns[0][i].as_const_frame()->data, device_, 62 | frame->size()); 63 | } 64 | } 65 | 66 | auto torchdevice = device_.type == DeviceType::GPU 67 | ? at::device({at::kCUDA, device_.id}) : at::device(at::kCPU); 68 | 69 | auto opts = torchdevice.dtype(at::kByte); 70 | at::Tensor input_tensor = 71 | torch::from_blob( 72 | batch_buffer_, 73 | {batch_size, frame->shape[0], frame->shape[1], frame->shape[2]}, 74 | opts); 75 | 76 | at::Tensor normalized_tensor; 77 | { 78 | ProfileBlock _block(profiler_, "preprocess"); 79 | // convert to batch x channel x height x width and normalize to [0, 1] 80 | normalized_tensor = input_tensor.permute({0, 3, 1, 2}).to(at::kFloat) / 255.0; 81 | 82 | // normalize with mean/stddev. can't use torch::data::transforms::Normalize b/c it's only 83 | // implemented for CPU 84 | // https://caffe2.ai/doxygen-c/html/torch_2csrc_2api_2include_2torch_2data_2transforms_2tensor_8h_source.html 85 | normalized_tensor -= mu_t_; 86 | normalized_tensor /= sigma_t_; 87 | } 88 | 89 | auto start = now(); 90 | at::Tensor output = module_->forward({normalized_tensor}).toTensor(); 91 | profiler_->add_interval("forward", start, now()); 92 | 93 | { 94 | ProfileBlock _block(profiler_, "copy_out"); 95 | u8* output_buf = new_block_buffer(device_, output.nbytes(), batch_size); 96 | memcpy_buffer(output_buf, device_, (u8*) output.data_ptr(), device_, output.nbytes()); 97 | 98 | size_t elem_size = output.nbytes() / batch_size; 99 | for (size_t i = 0; i < batch_size; ++i) { 100 | insert_element(output_columns[0], &output_buf[i * elem_size], elem_size); 101 | } 102 | } 103 | } 104 | 105 | private: 106 | DeviceHandle device_; 107 | u8* batch_buffer_; 108 | std::shared_ptr module_; 109 | at::Tensor mu_t_; 110 | at::Tensor sigma_t_; 111 | }; 112 | 113 | REGISTER_OP(PyTorch).frame_input("frame").output("bytes", ColumnType::Bytes); 114 | 115 | REGISTER_KERNEL(PyTorch, PyTorchKernel) 116 | .device(DeviceType::CPU) 117 | .batch() 118 | .num_devices(BaseKernel::UnlimitedDevices); 119 | 120 | REGISTER_KERNEL(PyTorch, PyTorchKernel) 121 | .device(DeviceType::GPU) 122 | .batch() 123 | .num_devices(1); 124 | 125 | } // namespace scanner 126 | -------------------------------------------------------------------------------- /draft/scannertools_pytorch/scannertools_pytorch_cpp/scannertools_pytorch.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package scanner.proto; -------------------------------------------------------------------------------- /draft/scannertools_pytorch/setup.cfg: -------------------------------------------------------------------------------- 1 | [aliases] 2 | test=pytest 3 | 4 | [tool:pytest] 5 | addopts = -vvs 6 | testpaths = tests -------------------------------------------------------------------------------- /draft/scannertools_pytorch/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from scannertools_infra import CMakeExtension, CMakeBuild, CudaInstallCommand, CudaDevelopCommand 3 | 4 | if __name__ == "__main__": 5 | setup( 6 | name='scannertools_pytorch', 7 | version='0.2.15', 8 | description='Video analytics toolkit', 9 | url='http://github.com/scanner-research/scannertools', 10 | author='Will Crichton', 11 | author_email='wcrichto@cs.stanford.edu', 12 | license='Apache 2.0', 13 | packages=['scannertools_pytorch'], 14 | setup_requires=['pytest-runner', 'scannertools_infra'], 15 | tests_require=['pytest', 'scannertools_infra'], 16 | cmdclass=dict(build_ext=CMakeBuild, install=CudaInstallCommand, develop=CudaDevelopCommand), 17 | ext_modules=[CMakeExtension('scannertools_pytorch', 'scannertools_pytorch_cpp')], 18 | zip_safe=False) 19 | -------------------------------------------------------------------------------- /draft/scannertools_pytorch/tests/test_all.py: -------------------------------------------------------------------------------- 1 | from scannerpy import NamedVideoStream, CacheMode, NamedStream, PerfParams, DeviceType 2 | from scannertools_infra.tests import sc 3 | import scannertools_pytorch 4 | import scannertools.imgproc 5 | import scannertools_pytorch.resnet 6 | from timeit import default_timer as now 7 | 8 | def run(sc, op, name, device): 9 | vid = NamedVideoStream(sc, 'test3') 10 | inp = sc.io.Input([vid]) 11 | f = sc.streams.Gather(inp, [list(range(10000))]) 12 | res = sc.ops.Resize(frame=inp, width=[224], height=[224], device=device, batch=100) 13 | tf = op(frame=res, batch=100, device=device) 14 | out = NamedStream(sc, 'qq') 15 | outp = sc.io.Output(tf, [out]) 16 | 17 | s = now() 18 | sc.run( 19 | outp, 20 | PerfParams.manual(500, 10000, pipeline_instances_per_node=1), 21 | cache_mode=CacheMode.Overwrite, 22 | gpu_pool='8G') 23 | sc.table('qq').profiler().write_trace('{}.tar.gz'.format(name)) 24 | print('{:.1f}s'.format(now() - s)) 25 | 26 | def test_tf(sc): 27 | sc.ingest_videos([('test3', 'test_clip.mp4')]) 28 | device = DeviceType.GPU 29 | #run(sc, sc.ops.Resnet, 'torch_python', device) 30 | run(sc, sc.ops.PyTorch, 'torch_cpp', device) 31 | -------------------------------------------------------------------------------- /draft/scannertools_tfcpp/scannertools_tfcpp/__init__.py: -------------------------------------------------------------------------------- 1 | from scannertools_infra import _register_module 2 | 3 | _register_module(__file__, "scannertools_tfcpp") 4 | -------------------------------------------------------------------------------- /draft/scannertools_tfcpp/scannertools_tfcpp_cpp/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.2.0 FATAL_ERROR) 2 | 3 | execute_process( 4 | OUTPUT_VARIABLE SCANNER_CMAKE_PATH 5 | COMMAND python3 -c "import scannerpy.build_flags as b; b.print_cmake()") 6 | include(${SCANNER_CMAKE_PATH}) 7 | 8 | set(SOURCES tfcpp.cpp) 9 | 10 | build_op( 11 | LIB_NAME scannertools_tfcpp 12 | CPP_SRCS ${SOURCES} 13 | BUILD_CUDA ${BUILD_CUDA}) 14 | #PROTO_SRC scannertools_tfcpp.proto) 15 | 16 | set(LIBTF_PATH /home/will/libtensorflow) 17 | 18 | set(LIBRARIES ${LIBTF_PATH}/lib/libtensorflow.so) 19 | set(INCLUDES ${LIBTF_PATH}/include) 20 | 21 | target_include_directories(scannertools_tfcpp PUBLIC ${INCLUDES}) 22 | target_link_libraries(scannertools_tfcpp PUBLIC ${LIBRARIES}) 23 | -------------------------------------------------------------------------------- /draft/scannertools_tfcpp/scannertools_tfcpp_cpp/scannertools_tfcpp.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package scanner.proto; -------------------------------------------------------------------------------- /draft/scannertools_tfcpp/scannertools_tfcpp_cpp/tfcpp.cpp: -------------------------------------------------------------------------------- 1 | #include "tensorflow/c/c_api.h" 2 | 3 | #include "scanner/api/kernel.h" 4 | #include "scanner/api/op.h" 5 | #include "scanner/types.pb.h" 6 | #include "scanner/util/memory.h" 7 | #include "scanner/util/serialize.h" 8 | 9 | namespace scanner { 10 | 11 | // https://stackoverflow.com/questions/41688217/how-to-load-a-graph-with-tensorflow-so-and-c-api-h-in-c-language 12 | void free_buffer(void *data, size_t length) { free(data); } 13 | 14 | TF_Buffer *read_file(const char *file) { 15 | FILE *f = fopen(file, "rb"); 16 | fseek(f, 0, SEEK_END); 17 | long fsize = ftell(f); 18 | fseek(f, 0, SEEK_SET); // same as rewind(f); 19 | 20 | void *data = malloc(fsize); 21 | fread(data, fsize, 1, f); 22 | fclose(f); 23 | 24 | TF_Buffer *buf = TF_NewBuffer(); 25 | buf->data = data; 26 | buf->length = fsize; 27 | buf->data_deallocator = free_buffer; 28 | return buf; 29 | } 30 | 31 | void free_tensor(void *data, size_t len, void *arg) {} 32 | 33 | #define TF_CHECK_RETURN(STATUS) \ 34 | { \ 35 | if (TF_GetCode(STATUS) != TF_OK) { \ 36 | result->set_success(false); \ 37 | std::string msg(TF_Message(STATUS)); \ 38 | result->set_msg(msg); \ 39 | return; \ 40 | } \ 41 | } 42 | 43 | #define TF_CHECK_FAIL(STATUS) \ 44 | { \ 45 | if (TF_GetCode(STATUS) != TF_OK) { \ 46 | LOG(FATAL) << TF_Message(STATUS); \ 47 | } \ 48 | } 49 | 50 | class TFCPPKernel : public BatchedKernel { 51 | public: 52 | TFCPPKernel(const KernelConfig &config) 53 | : BatchedKernel(config), device_(config.devices[0]) {} 54 | 55 | void setup_with_resources(proto::Result *result) override { 56 | std::string graph_path = 57 | "/tmp/ssd_mobilenet_v1_coco_2017_11_17/frozen_inference_graph.pb"; 58 | TF_Buffer *graph_def = read_file(graph_path.c_str()); 59 | graph_ = TF_NewGraph(); 60 | 61 | status_ = TF_NewStatus(); 62 | TF_ImportGraphDefOptions *opts = TF_NewImportGraphDefOptions(); 63 | TF_GraphImportGraphDef(graph_, graph_def, opts, status_); 64 | TF_DeleteImportGraphDefOptions(opts); 65 | TF_CHECK_RETURN(status_); 66 | TF_DeleteBuffer(graph_def); 67 | 68 | TF_SessionOptions *sess_opts = TF_NewSessionOptions(); 69 | 70 | session_ = TF_NewSession(graph_, sess_opts, status_); 71 | TF_CHECK_RETURN(status_); 72 | 73 | TF_DeleteSessionOptions(sess_opts); 74 | 75 | result->set_success(true); 76 | } 77 | 78 | ~TFCPPKernel() { 79 | TF_DeleteSession(session_, status_); 80 | TF_CHECK_FAIL(status_); 81 | TF_DeleteGraph(graph_); 82 | TF_DeleteStatus(status_); 83 | } 84 | 85 | TF_DataType frame_type_to_tf_type(FrameType type) { 86 | if (type == FrameType::U8) { 87 | return TF_UINT8; 88 | } else if (type == FrameType::F32) { 89 | return TF_FLOAT; 90 | } else { 91 | LOG(FATAL) << "Can't convert frame type " << type << " to TF_DataType"; 92 | } 93 | } 94 | 95 | void execute(const BatchedElements &input_columns, 96 | BatchedElements &output_columns) override { 97 | 98 | TF_Output input_op = {TF_GraphOperationByName(graph_, "image_tensor"), 0}; 99 | LOG_IF(FATAL, input_op.oper == nullptr) << "Failed to find input op"; 100 | 101 | TF_Output output_op = {TF_GraphOperationByName(graph_, "detection_boxes"), 102 | 0}; 103 | LOG_IF(FATAL, output_op.oper == nullptr) << "Failed to find output op"; 104 | 105 | const Frame *frame = input_columns[0][0].as_const_frame(); 106 | int64_t shape[FRAME_DIMS + 1]; 107 | size_t batch_size = input_columns[0].size(); 108 | shape[0] = batch_size; 109 | for (size_t i = 0; i < FRAME_DIMS; ++i) { 110 | shape[i + 1] = frame->shape[i]; 111 | } 112 | 113 | u8* batch_buffer = new u8[frame->size() * batch_size]; 114 | size_t frame_size = frame->size(); 115 | for (size_t i = 0; i < batch_size; ++i) { 116 | std::memcpy(&batch_buffer[frame_size * i], input_columns[0][i].as_const_frame()->data, 117 | frame_size); 118 | } 119 | 120 | TF_Tensor *input_tensor = TF_NewTensor( 121 | frame_type_to_tf_type(frame->type), shape, FRAME_DIMS + 1, batch_buffer, 122 | frame->size() * batch_size, free_tensor, nullptr); 123 | TF_Tensor *output_tensor = nullptr; 124 | 125 | TF_SessionRun(session_, nullptr, &input_op, &input_tensor, 1, &output_op, 126 | &output_tensor, 1, nullptr, 0, nullptr, status_); 127 | TF_CHECK_FAIL(status_); 128 | 129 | delete batch_buffer; 130 | 131 | float *all_bboxes = (float *)TF_TensorData(output_tensor); 132 | int num_bboxes = TF_Dim(output_tensor, 1); 133 | int bbox_dim = TF_Dim(output_tensor, 2); 134 | for (size_t i = 0; i < batch_size; ++i) { 135 | std::vector bboxes; 136 | float *frame_bboxes = &all_bboxes[i * num_bboxes * bbox_dim]; 137 | for (size_t j = 0; j < num_bboxes; ++j) { 138 | proto::BoundingBox bbox; 139 | float *frame_bbox = &frame_bboxes[j * bbox_dim]; 140 | bbox.set_x1(frame_bboxes[1]); 141 | bbox.set_y1(frame_bboxes[0]); 142 | bbox.set_x2(frame_bboxes[3]); 143 | bbox.set_y2(frame_bboxes[2]); 144 | bboxes.push_back(bbox); 145 | } 146 | 147 | u8 *buffer; 148 | size_t size; 149 | serialize_bbox_vector(bboxes, buffer, size); 150 | 151 | insert_element(output_columns[0], buffer, size); 152 | } 153 | } 154 | 155 | private: 156 | DeviceHandle device_; 157 | TF_Graph *graph_; 158 | TF_Status *status_; 159 | TF_Session *session_; 160 | }; 161 | 162 | REGISTER_OP(TFC).frame_input("frame").output("bytes", ColumnType::Bytes); 163 | 164 | REGISTER_KERNEL(TFC, TFCPPKernel) 165 | .device(DeviceType::CPU) 166 | .batch() 167 | .num_devices(BaseKernel::UnlimitedDevices); 168 | 169 | REGISTER_KERNEL(TFC, TFCPPKernel) 170 | .device(DeviceType::GPU) 171 | .batch() 172 | .num_devices(1); 173 | 174 | } // namespace scanner 175 | -------------------------------------------------------------------------------- /draft/scannertools_tfcpp/setup.cfg: -------------------------------------------------------------------------------- 1 | [aliases] 2 | test=pytest 3 | 4 | [tool:pytest] 5 | addopts = -vvs 6 | testpaths = tests -------------------------------------------------------------------------------- /draft/scannertools_tfcpp/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from scannertools_infra import CMakeExtension, CMakeBuild, CudaInstallCommand, CudaDevelopCommand 3 | 4 | if __name__ == "__main__": 5 | setup( 6 | name='scannertools_tfcpp', 7 | version='0.2.15', 8 | description='Video analytics toolkit', 9 | url='http://github.com/scanner-research/scannertools', 10 | author='Will Crichton', 11 | author_email='wcrichto@cs.stanford.edu', 12 | license='Apache 2.0', 13 | packages=['scannertools_tfcpp'], 14 | setup_requires=['pytest-runner', 'scannertools_infra'], 15 | tests_require=['pytest', 'scannertools_infra'], 16 | cmdclass=dict(build_ext=CMakeBuild, install=CudaInstallCommand, develop=CudaDevelopCommand), 17 | ext_modules=[CMakeExtension('scannertools_tfcpp', 'scannertools_tfcpp_cpp')], 18 | zip_safe=False) 19 | -------------------------------------------------------------------------------- /draft/scannertools_tfcpp/tests/test_all.py: -------------------------------------------------------------------------------- 1 | from scannerpy import NamedVideoStream, CacheMode, NamedStream, PerfParams, DeviceType 2 | from scannertools_infra.tests import sc 3 | import scannertools_tfcpp 4 | import scannertools.object_detection 5 | from timeit import default_timer as now 6 | 7 | def run(sc, op, name): 8 | vid = NamedVideoStream(sc, 'test1') 9 | inp = sc.io.Input([vid]) 10 | #f = sc.streams.Gather(inp, [list(range(1000))]) 11 | tf = op(frame=inp, batch=100, device=DeviceType.CPU) 12 | out = NamedStream(sc, 'qq') 13 | outp = sc.io.Output(tf, [out]) 14 | 15 | s = now() 16 | sc.run(outp, PerfParams.estimate(), cache_mode=CacheMode.Overwrite, pipeline_instances_per_node=1) 17 | sc.table('qq').profiler().write_trace('{}.trace'.format(name)) 18 | print('{:.1f}s'.format(now() - s)) 19 | 20 | def test_tf(sc): 21 | run(sc, sc.ops.DetectObjects, 'python') 22 | run(sc, sc.ops.TFC, 'cpp') 23 | -------------------------------------------------------------------------------- /scannertools/README.md: -------------------------------------------------------------------------------- 1 | # scannertools: core standard library for Scanner 2 | 3 | See [scanner.run](http://scanner.run/api/scannertools.html) for a description of the operations in this library. 4 | 5 | ## Installation 6 | 7 | ### Dependencies 8 | 9 | Requirements: 10 | * OpenCV >= 3.4.0 11 | * pybind11 >= 2.2.3 12 | * ffmeg >= 3.3.1 13 | * tensorflow >= 1.11.0 14 | 15 | #### Ubuntu 16.04 16 | 17 | ``` 18 | sudo apt-get install -y ffmpeg pybind11-dev 19 | pip3 install tensorflow 20 | ``` 21 | 22 | See the [OpenCV docs](https://docs.opencv.org/3.4.1/d2/de6/tutorial_py_setup_in_ubuntu.html) for installing OpenCV. 23 | 24 | #### OS X 25 | 26 | ``` 27 | brew install ffmpeg opencv@3 pybind11 28 | pip3 install tensorflow 29 | ``` 30 | 31 | ### Library 32 | 33 | #### From pypi 34 | 35 | Coming soon, build from source for now. 36 | 37 | #### From source 38 | 39 | ``` 40 | git clone https://github.com/scanner-research/scannertools 41 | cd scannertools/scannertools_infra 42 | pip3 install -e . 43 | cd ../scannertools 44 | pip3 install -e . 45 | ``` 46 | 47 | To build with GPU support, instead run: 48 | 49 | ``` 50 | pip3 install --install-option="--build-cuda=/usr/local/cuda" -e . 51 | ``` 52 | -------------------------------------------------------------------------------- /scannertools/scannertools/__init__.py: -------------------------------------------------------------------------------- 1 | from . import face_detection 2 | from . import face_embedding 3 | from . import gender_detection 4 | from . import object_detection 5 | from . import shot_detection 6 | -------------------------------------------------------------------------------- /scannertools/scannertools/caffe2.py: -------------------------------------------------------------------------------- 1 | from ..kernel import Kernel 2 | from scannerpy import DeviceType 3 | 4 | from caffe2.python import workspace 5 | 6 | 7 | class Caffe2Kernel(Kernel): 8 | def __init__(self, config): 9 | workspace.GlobalInit(['caffe2', '--caffe2_log_level=0']) 10 | self.config = config 11 | self.protobufs = config.protobufs 12 | self.graph = self.build_graph() 13 | 14 | def close(self): 15 | del self.graph 16 | 17 | def build_graph(self): 18 | raise NotImplementedError 19 | 20 | def execute(self): 21 | raise NotImplementedError 22 | -------------------------------------------------------------------------------- /scannertools/scannertools/face_detection.py: -------------------------------------------------------------------------------- 1 | from scannerpy import DeviceType, FrameType, register_python_op, protobufs 2 | from scannerpy.types import BboxList 3 | from .tensorflow import TensorFlowKernel 4 | from typing import Sequence 5 | import os.path 6 | import numpy as np 7 | 8 | @register_python_op(device_sets=[[DeviceType.CPU, 0], [DeviceType.GPU, 1]], batch=5) 9 | class MTCNNDetectFaces(TensorFlowKernel): 10 | def build_graph(self): 11 | import tensorflow as tf 12 | self.pnet = None 13 | self.g = tf.Graph() 14 | self._g_default = self.g.as_default() 15 | return self.g 16 | 17 | def execute(self, frame: Sequence[FrameType]) -> Sequence[BboxList]: 18 | import align 19 | import align.detect_face 20 | 21 | if self.pnet is None: 22 | with self.g.as_default(): 23 | with self.sess.as_default(): 24 | print('Loading model...') 25 | self.pnet, self.rnet, self.onet = align.detect_face.create_mtcnn( 26 | self.sess, os.path.dirname(align.__file__)) 27 | print('Model loaded!') 28 | 29 | threshold = [0.45, 0.6, 0.7] 30 | factor = 0.709 31 | vmargin = 0.2582651235637604 32 | hmargin = 0.3449094129917718 33 | out_size = 160 34 | detection_window_size_ratio = .2 35 | 36 | imgs = frame 37 | #print(('Face detect on {} frames'.format(len(imgs)))) 38 | detections = align.detect_face.bulk_detect_face( 39 | imgs, detection_window_size_ratio, self.pnet, self.rnet, self.onet, threshold, factor) 40 | 41 | batch_faces = [] 42 | for img, bounding_boxes in zip(imgs, detections): 43 | if bounding_boxes == None: 44 | batch_faces.append([]) 45 | continue 46 | frame_faces = [] 47 | bounding_boxes = bounding_boxes[0] 48 | num_faces = bounding_boxes.shape[0] 49 | for i in range(num_faces): 50 | confidence = bounding_boxes[i][4] 51 | if confidence < .1: 52 | continue 53 | 54 | img_size = np.asarray(img.shape)[0:2] 55 | det = np.squeeze(bounding_boxes[i][0:5]) 56 | vmargin_pix = int((det[2] - det[0]) * vmargin) 57 | hmargin_pix = int((det[3] - det[1]) * hmargin) 58 | frame_faces.append( 59 | protobufs.BoundingBox( 60 | x1=np.maximum(det[0] - hmargin_pix / 2, 0) / img_size[1], 61 | y1=np.maximum(det[1] - vmargin_pix / 2, 0) / img_size[0], 62 | x2=np.minimum(det[2] + hmargin_pix / 2, img_size[1]) / img_size[1], 63 | y2=np.minimum(det[3] + vmargin_pix / 2, img_size[0]) / img_size[0], 64 | score=det[4])) 65 | 66 | batch_faces.append(frame_faces) 67 | 68 | return batch_faces 69 | -------------------------------------------------------------------------------- /scannertools/scannertools/face_embedding.py: -------------------------------------------------------------------------------- 1 | from scannerpy import FrameType, DeviceType 2 | import scannerpy 3 | from scannerpy.util import download_temp_file, temp_directory 4 | from .tensorflow import TensorFlowKernel 5 | import os 6 | import numpy as np 7 | from scannerpy.types import BboxList, UniformList, NumpyArrayFloat32 8 | from typing import Sequence 9 | 10 | MODEL_FILE = 'https://storage.googleapis.com/esper/models/facenet/20170512-110547.tar.gz' 11 | 12 | EMBEDDING_SIZE = 128 13 | FacenetEmbeddings = UniformList( 14 | 'FacenetEmbeddings', NumpyArrayFloat32, EMBEDDING_SIZE * np.dtype(np.float32).itemsize) 15 | 16 | 17 | @scannerpy.register_python_op(device_sets=[[DeviceType.CPU, 0], [DeviceType.GPU, 1]], batch=5) 18 | class EmbedFaces(TensorFlowKernel): 19 | def __init__(self, config, minibatch = 5): 20 | self._minibatch = 5 21 | TensorFlowKernel.__init__(self, config) 22 | 23 | def build_graph(self): 24 | import tensorflow as tf 25 | self.images_placeholder = None 26 | self.g = tf.Graph() 27 | self._g_default = self.g.as_default() 28 | self._model_dir = os.path.join(temp_directory(), '20170512-110547') 29 | return self.g 30 | 31 | def fetch_resources(self): 32 | download_temp_file(MODEL_FILE, untar=True) 33 | 34 | def execute(self, frame: Sequence[FrameType], bboxes: Sequence[BboxList]) -> Sequence[FacenetEmbeddings]: 35 | import facenet 36 | import cv2 37 | import tensorflow as tf 38 | 39 | if self.images_placeholder is None: 40 | print('Loading model...') 41 | with self.g.as_default(): 42 | with self.sess.as_default(): 43 | model_path = self._model_dir 44 | meta_file, ckpt_file = facenet.get_model_filenames(model_path) 45 | saver = tf.train.import_meta_graph(os.path.join(model_path, meta_file)) 46 | saver.restore(self.sess, os.path.join(model_path, ckpt_file)) 47 | 48 | self.images_placeholder = tf.get_default_graph().get_tensor_by_name('input:0') 49 | self.embeddings = tf.get_default_graph().get_tensor_by_name('embeddings:0') 50 | self.phase_train_placeholder = tf.get_default_graph().get_tensor_by_name( 51 | 'phase_train:0') 52 | print('Model loaded!') 53 | 54 | [h, w] = frame[0].shape[:2] 55 | 56 | out_size = 160 57 | outputs = b'' 58 | cleaned_images = [] 59 | source_indices = [] 60 | output_embs = [[None for _ in l] for l in bboxes] 61 | for i, frame_bboxes in enumerate(bboxes): 62 | for j, bbox in enumerate(frame_bboxes): 63 | # NOTE: if using output of mtcnn, not-normalized, so removing de-normalization factors here 64 | face_img = frame[i][int(bbox.y1 * h):int(bbox.y2 * h), int(bbox.x1 * w):int(bbox.x2 * w)] 65 | [fh, fw] = face_img.shape[:2] 66 | if fh == 0 or fw == 0: 67 | output_embs[i][j] = np.zeros(128, dtype=np.float32) 68 | else: 69 | face_img = cv2.resize(face_img, (out_size, out_size)) 70 | face_img = facenet.prewhiten(face_img) 71 | cleaned_images.append(face_img) 72 | source_indices.append((i, j)) 73 | 74 | for k in range(0, len(cleaned_images), self._minibatch): 75 | embs = self.sess.run( 76 | self.embeddings, 77 | feed_dict={ 78 | self.images_placeholder: cleaned_images[k:k+self._minibatch], 79 | self.phase_train_placeholder: False 80 | }) 81 | 82 | for emb, (i, j) in zip(embs, source_indices[k:k+self._minibatch]): 83 | output_embs[i][j] = emb 84 | 85 | for l in output_embs: 86 | for e in l: 87 | assert e is not None 88 | 89 | return output_embs 90 | -------------------------------------------------------------------------------- /scannertools/scannertools/gender_detection.py: -------------------------------------------------------------------------------- 1 | from scannerpy.util import download_temp_file, temp_directory 2 | from scannerpy.types import BboxList 3 | from scannerpy import FrameType 4 | from typing import Any 5 | import scannerpy 6 | import cv2 7 | import pickle 8 | import os 9 | 10 | MODEL_FILE = 'https://storage.googleapis.com/esper/models/rude-carnie/21936.tar.gz' 11 | 12 | 13 | @scannerpy.register_python_op() 14 | class DetectGender(scannerpy.Kernel): 15 | def fetch_resources(self): 16 | download_temp_file(MODEL_FILE, untar=True) 17 | 18 | def setup_with_resources(self): 19 | from carnie_helper import RudeCarnie 20 | self.rc = RudeCarnie(model_dir=os.path.join(temp_directory(), '21936')) 21 | 22 | def execute(self, frame: FrameType, bboxes: BboxList) -> Any: 23 | frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) 24 | [h, w] = frame.shape[:2] 25 | frames = [ 26 | frame[int(bbox.y1 * h):int(bbox.y2 * h), 27 | int(bbox.x1 * w):int(bbox.x2 * w)] for bbox in bboxes 28 | ] 29 | genders = self.rc.get_gender_batch(frames) 30 | return genders 31 | -------------------------------------------------------------------------------- /scannertools/scannertools/imgproc/__init__.py: -------------------------------------------------------------------------------- 1 | from scannertools_infra import _register_module 2 | import os 3 | _register_module(os.path.join(__file__, '..'), "scannertools_imgproc") 4 | -------------------------------------------------------------------------------- /scannertools/scannertools/misc/__init__.py: -------------------------------------------------------------------------------- 1 | from scannertools_infra import _register_module 2 | import os 3 | _register_module(os.path.join(__file__, '..'), "scannertools_misc") 4 | -------------------------------------------------------------------------------- /scannertools/scannertools/net_descriptor.py: -------------------------------------------------------------------------------- 1 | 2 | import toml 3 | 4 | from scannerpy.common import * 5 | 6 | class NetDescriptor(object): 7 | def __init__(self, db): 8 | self._descriptor = db.protobufs.NetDescriptor() 9 | self._descriptor.input_width = -1 10 | self._descriptor.input_height = -1 11 | self._descriptor.pad_mod = -1 12 | 13 | def _val(self, dct, key, default): 14 | if key in dct: 15 | return dct[key] 16 | else: 17 | return default 18 | 19 | @property 20 | def model_path(self): 21 | return self._descriptor.model_path 22 | 23 | @model_path.setter 24 | def model_path(self, value): 25 | self._descriptor.model_path = value 26 | 27 | @property 28 | def model_weights_path(self): 29 | return self._descriptor.model_weights_path 30 | 31 | @model_weights_path.setter 32 | def model_weights_path(self, value): 33 | self._descriptor.model_weights_path = value 34 | 35 | @property 36 | def input_layer_names(self): 37 | return self._descriptor.input_layer_names[:] 38 | 39 | @input_layer_names.setter 40 | def input_layer_names(self, value): 41 | del self._descriptor.input_layer_names[:] 42 | self._descriptor.input_layer_names.extend(value) 43 | 44 | @property 45 | def output_layer_names(self): 46 | return self._descriptor.output_layer_names[:] 47 | 48 | @output_layer_names.setter 49 | def output_layer_names(self, value): 50 | del self._descriptor.output_layer_names[:] 51 | self._descriptor.output_layer_names.extend(value) 52 | 53 | @property 54 | def input_width(self): 55 | return self._descriptor.input_width 56 | 57 | @input_width.setter 58 | def input_width(self, value): 59 | self._descriptor.input_width = value 60 | 61 | @property 62 | def input_height(self): 63 | return self._descriptor.input_height 64 | 65 | @input_width.setter 66 | def input_height(self, value): 67 | self._descriptor.input_height = value 68 | 69 | @property 70 | def normalize(self): 71 | return self._descriptor.normalize 72 | 73 | @normalize.setter 74 | def normalize(self, value): 75 | self._descriptor.normalize = value 76 | 77 | @property 78 | def preserve_aspect_ratio(self): 79 | return self._descriptor.preserve_aspect_ratio 80 | 81 | @preserve_aspect_ratio.setter 82 | def normalize(self, value): 83 | self._descriptor.preserve_aspect_ratio = value 84 | 85 | @property 86 | def transpose(self): 87 | return self._descriptor.transpose 88 | 89 | @transpose.setter 90 | def transpose(self, value): 91 | self._descriptor.transpose = value 92 | 93 | @property 94 | def pad_mod(self): 95 | return self._descriptor.pad_mod 96 | 97 | @pad_mod.setter 98 | def pad_mod(self, value): 99 | self._descriptor.pad_mod = value 100 | 101 | @property 102 | def uses_python(self): 103 | return self._descriptor.uses_python 104 | 105 | @uses_python.setter 106 | def uses_python(self, value): 107 | self._descriptor.uses_python = value 108 | 109 | @property 110 | def mean_colors(self): 111 | return self._descriptor.mean_colors 112 | 113 | @uses_python.setter 114 | def mean_colors(self, value): 115 | del self._descriptor.mean_colors[:] 116 | self._descriptor.mean_colors.extend(value) 117 | 118 | @classmethod 119 | def from_file(cls, db, path): 120 | self = cls(db) 121 | with open(path) as f: 122 | args = toml.loads(f.read()) 123 | 124 | d = self._descriptor 125 | net = args['net'] 126 | d.model_path = net['model'] 127 | d.model_weights_path = net['weights'] 128 | d.input_layer_names.extend(net['input_layers']) 129 | d.output_layer_names.extend(net['output_layers']) 130 | d.input_width = self._val(net, 'input_width', -1) 131 | d.input_height = self._val(net, 'input_height', -1) 132 | d.normalize = self._val(net, 'normalize', False) 133 | d.preserve_aspect_ratio = self._val(net, 'preserve_aspect_ratio', False) 134 | d.transpose = self._val(net, 'tranpose', False) 135 | d.pad_mod = self._val(net, 'pad_mod', -1) 136 | d.uses_python = self._val(net, 'uses_python', False) 137 | 138 | mean = args['mean-image'] 139 | if 'colors' in mean: 140 | order = net['input']['channel_ordering'] 141 | for color in order: 142 | d.mean_colors.append(mean['colors'][color]) 143 | elif 'image' in mean: 144 | d.mean_width = mean['width'] 145 | d.mean_height = mean['height'] 146 | # TODO: load mean binaryproto 147 | raise ScannerException('TODO') 148 | 149 | return self 150 | 151 | def as_proto(self): 152 | return self._descriptor 153 | -------------------------------------------------------------------------------- /scannertools/scannertools/object_detection.py: -------------------------------------------------------------------------------- 1 | from scannerpy.util import download_temp_file, temp_directory 2 | from .tensorflow import TensorFlowKernel 3 | from scannerpy import FrameType, DeviceType, protobufs 4 | from scannerpy.types import BboxList 5 | import os 6 | import tarfile 7 | import numpy as np 8 | import scannerpy 9 | from typing import Sequence 10 | 11 | SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) 12 | 13 | MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17' 14 | MODEL_FILE = MODEL_NAME + '.tar.gz' 15 | DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/' 16 | 17 | LABEL_URL = 'https://storage.googleapis.com/scanner-data/public/mscoco_label_map.pbtxt' 18 | 19 | GRAPH_PATH = os.path.join( 20 | temp_directory(), 'ssd_mobilenet_v1_coco_2017_11_17', 21 | 'frozen_inference_graph.pb') 22 | 23 | 24 | @scannerpy.register_python_op(device_sets=[[DeviceType.CPU, 0], [DeviceType.GPU, 1]], batch=5) 25 | class DetectObjects(TensorFlowKernel): 26 | def build_graph(self): 27 | import tensorflow as tf 28 | dnn = tf.Graph() 29 | with dnn.as_default(): 30 | od_graph_def = tf.GraphDef() 31 | with tf.gfile.GFile(GRAPH_PATH, 'rb') as fid: 32 | serialized_graph = fid.read() 33 | od_graph_def.ParseFromString(serialized_graph) 34 | tf.import_graph_def(od_graph_def, name='') 35 | return dnn 36 | 37 | def fetch_resources(self): 38 | model_tar_path = download_temp_file(DOWNLOAD_BASE + MODEL_FILE) 39 | with tarfile.open(model_tar_path) as f: 40 | f.extractall(temp_directory()) 41 | download_temp_file(LABEL_URL) 42 | 43 | # Evaluate object detection DNN model on a frame 44 | # Return bounding box position, class and score 45 | def execute(self, frame: Sequence[FrameType]) -> Sequence[BboxList]: 46 | image_tensor = self.graph.get_tensor_by_name('image_tensor:0') 47 | boxes = self.graph.get_tensor_by_name('detection_boxes:0') 48 | scores = self.graph.get_tensor_by_name('detection_scores:0') 49 | classes = self.graph.get_tensor_by_name('detection_classes:0') 50 | with self.graph.as_default(): 51 | (boxes, scores, classes) = self.sess.run( 52 | [boxes, scores, classes], feed_dict={image_tensor: np.concatenate(np.expand_dims(frame, axis=0), axis=0)}) 53 | 54 | bboxes = [ 55 | [ 56 | protobufs.BoundingBox( 57 | x1=box[1], y1=box[0], x2=box[3], y2=box[2], score=score, label=cls) 58 | for (box, score, cls) in zip( 59 | boxes[i, :, :].reshape(100, 4), scores[i, :].reshape(100, 1), classes[i, :].reshape(100, 1)) 60 | ] 61 | for i in range(len(frame)) 62 | ] 63 | 64 | return bboxes 65 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/__init__.py: -------------------------------------------------------------------------------- 1 | from .prelude import WithMany, init_storage, sample_video, imwrite, BoundOp, Pipeline, tile, DataSource 2 | from .video import Video 3 | from . import pose_detection 4 | from . import shot_detection 5 | from . import object_detection 6 | from . import gender_detection 7 | from . import face_detection 8 | from . import face_embedding 9 | from . import optical_flow 10 | from . import clothing_detection 11 | from . import imgproc 12 | from . import histograms 13 | from . import face_landmark_detection 14 | from . import tf_vis_utils 15 | from . import vis 16 | from . import bboxes 17 | from . import kube 18 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/audio.py: -------------------------------------------------------------------------------- 1 | from .prelude import * 2 | 3 | 4 | class AudioSource(DataSource): 5 | def __init__(self, video, frame_size=1.0, duration=None): 6 | self._video = video 7 | self._frame_size = frame_size 8 | self._duration = duration 9 | 10 | def scanner_source(self, db): 11 | return db.sources.Audio(frame_size=self._frame_size) 12 | 13 | def scanner_args(self, db): 14 | return { 15 | 'path': self._video.path(), 16 | 'duration': self._duration if self._duration is not None else 0.0 17 | } 18 | 19 | 20 | class CaptionSource(DataSource): 21 | def __init__(self, captions, max_time, window_size=10.0): 22 | self._captions = captions 23 | self._window_size = window_size 24 | self._max_time = max_time 25 | 26 | def scanner_source(self, db): 27 | return db.sources.Captions(window_size=self._window_size) 28 | 29 | def scanner_args(self, db): 30 | return {'path': self._captions, 'max_time': self._max_time} 31 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/bboxes.py: -------------------------------------------------------------------------------- 1 | import scannerpy 2 | import scannerpy.stdlib.readers as readers 3 | import scannerpy.stdlib.writers as writers 4 | import scannerpy.stdlib.bboxes as bboxes 5 | from scannerpy.stdlib.util import default 6 | 7 | 8 | @scannerpy.register_python_op() 9 | class BboxNMS(scannerpy.Kernel): 10 | def __init__(self, config): 11 | self._threshold = default(config.args, 'threshold', 0.3) 12 | self._config = config 13 | 14 | def execute(self, *input_columns) -> bytes: 15 | bboxes_list = [] 16 | for c in input_columns: 17 | bboxes_list += readers.bboxes(c, self._config.protobufs) 18 | 19 | nmsed_bboxes = bboxes.nms(bboxes_list, self._threshold) 20 | return writers.bboxes(nmsed_bboxes, self._config.protobufs) 21 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/cpp_ops/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # To build your custom op, you can either use our CMake convenience methods 2 | # or do it the slightly harder way in normal Make (see the Makefile). 3 | 4 | cmake_minimum_required(VERSION 3.2.0 FATAL_ERROR) 5 | 6 | execute_process( 7 | OUTPUT_VARIABLE SCANNER_CMAKE_PATH 8 | COMMAND python3 -c "import scannerpy.stdlib.build_flags as b; b.print_cmake()") 9 | include(${SCANNER_CMAKE_PATH}) 10 | 11 | # build_op will create a shared library called lib${LIB_NAME}.so that builds 12 | # from all of the CPP_SRCS. You can also optionally specify a PROTO_SRC that 13 | # points to a Protobuf file and will generate the C++ and Python bindings. 14 | build_op( 15 | LIB_NAME imgproc_op 16 | CPP_SRCS imgproc.cpp flow_histogram_kernel_cpu.cpp 17 | PROTO_SRC imgproc.proto) 18 | 19 | # The library specified in build_op is a normal CMake target, so you can use all 20 | # the normal CMake functions with it. 21 | find_package(OpenCV REQUIRED COMPONENTS core imgproc) 22 | include_directories(SYSTEM ${OpenCV_INCLUDE_DIRS}) 23 | target_link_libraries(imgproc_op PUBLIC "${OpenCV_LIBRARIES}") 24 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/cpp_ops/flow_histogram_kernel_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/memory.h" 4 | #include "scanner/util/opencv.h" 5 | #include 6 | 7 | namespace scanner { 8 | namespace { 9 | const i32 BINS = 64; 10 | } 11 | 12 | class FlowHistogramKernelCPU : public BatchedKernel { 13 | public: 14 | FlowHistogramKernelCPU(const KernelConfig& config) 15 | : BatchedKernel(config), device_(config.devices[0]) {} 16 | 17 | void execute(const BatchedElements& input_columns, 18 | BatchedElements& output_columns) override { 19 | auto& frame_col = input_columns[0]; 20 | 21 | size_t hist_size = BINS * 2 * sizeof(int); 22 | i32 input_count = num_rows(frame_col); 23 | 24 | u8* output_block = new_block_buffer_size(device_, hist_size, input_count); 25 | 26 | for (i32 i = 0; i < input_count; ++i) { 27 | cv::Mat img = frame_to_mat(frame_col[i].as_const_frame()); 28 | std::vector xy_planes; 29 | split( img, xy_planes ); 30 | 31 | cv::Mat deg, mag; 32 | 33 | cv::cartToPolar(xy_planes[0], xy_planes[1], mag, deg, true); 34 | 35 | float magRange[] = {0, 64.0}; 36 | const float* magHistRange = {magRange}; 37 | 38 | float degRange[] = {0, 360}; 39 | const float* degHistRange = {degRange}; 40 | 41 | u8* output_buf = output_block + i * hist_size; 42 | 43 | for (i32 j = 0; j < 2; ++j) { 44 | int channels[] = {0}; 45 | cv::Mat hist; 46 | cv::calcHist(j == 0 ? &mag : °, 1, channels, cv::Mat(), 47 | hist, 48 | 1, &BINS, 49 | j == 0 ? &magHistRange : & degHistRange); 50 | cv::Mat out(BINS, 1, CV_32SC1, output_buf + j * BINS * sizeof(int)); 51 | hist.convertTo(out, CV_32SC1); 52 | } 53 | 54 | insert_element(output_columns[0], output_buf, hist_size); 55 | } 56 | } 57 | 58 | private: 59 | DeviceHandle device_; 60 | }; 61 | 62 | REGISTER_OP(FlowHistogram).frame_input("flow").output("histogram"); 63 | 64 | REGISTER_KERNEL(FlowHistogram, FlowHistogramKernelCPU) 65 | .device(DeviceType::CPU) 66 | .batch() 67 | .num_devices(1); 68 | } 69 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/cpp_ops/imgproc.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | 4 | message ImgProcArgs { 5 | int32 width = 1; 6 | int32 height = 2; 7 | } 8 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/face_landmark_detection.py: -------------------------------------------------------------------------------- 1 | from .prelude import Pipeline 2 | from scannerpy import FrameType, DeviceType 3 | from scannerpy.stdlib import readers 4 | import scannerpy 5 | import pickle 6 | import numpy as np 7 | 8 | # NOTE: This pipeline will *NOT* work out of the box at the moment 9 | # This pipeline relies on torch > 0.4, but other parts of scannertools rely 10 | # on torch = 0.3.1 (clothing detection in particular) 11 | # Until then, you'll have to manually change the install version 12 | 13 | 14 | @scannerpy.register_python_op(device_sets=[[DeviceType.CPU, 0], [DeviceType.GPU, 1]]) 15 | class DetectFaceLandmarks(scannerpy.Kernel): 16 | def __init__(self, config): 17 | import face_alignment 18 | 19 | self.config = config 20 | cpu_only = True 21 | for handle in config.devices: 22 | if handle.type == DeviceType.GPU.value: 23 | cpu_only = False 24 | break 25 | 26 | self.fa = face_alignment.FaceAlignment( 27 | face_alignment.LandmarksType._2D, 28 | device=('cpu' if cpu_only else 'cuda'), 29 | flip_input=False) 30 | 31 | def execute(self, frame: FrameType, bboxes: bytes) -> bytes: 32 | [h, w] = frame.shape[:2] 33 | bboxes = readers.bboxes(bboxes, self.config.protobufs) 34 | if len(bboxes) == 0: 35 | return pickle.dumps([]) 36 | 37 | # This returns a numpy array of size (68, 2) for every bbox) 38 | predictions = self.fa.get_landmarks_from_image( 39 | frame, 40 | detected_faces=[(bbox.x1 * w, bbox.y1 * h, bbox.x2 * w, bbox.y2 * h) 41 | for bbox in bboxes]) 42 | 43 | predictions = [ 44 | np.array([[width / w, height / h] for [width, height] in prediction]) 45 | for prediction in predictions 46 | ] 47 | 48 | return pickle.dumps(predictions) 49 | 50 | 51 | class FaceLandmarkDetectionPipeline(Pipeline): 52 | job_suffix = 'face_landmarks' 53 | parser_fn = lambda _: lambda buf, _: pickle.loads(buf) 54 | additional_sources = ['bboxes'] 55 | run_opts = {'pipeline_instances_per_node': 1} 56 | 57 | def build_pipeline(self): 58 | return { 59 | 'face_landmarks': 60 | self._db.ops.DetectFaceLandmarks( 61 | frame=self._sources['frame_sampled'].op, 62 | bboxes=self._sources['bboxes'].op, 63 | device=self._device) 64 | } 65 | 66 | 67 | detect_face_landmarks = FaceLandmarkDetectionPipeline.make_runner() 68 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/hairstyle_detection.py: -------------------------------------------------------------------------------- 1 | from .prelude import Pipeline, try_import 2 | from scannerpy import Kernel, FrameType, DeviceType 3 | from scannerpy.stdlib.util import download_temp_file 4 | from scannerpy.stdlib import readers 5 | from scannerpy.stdlib.torch import TorchKernel 6 | import scannerpy 7 | import pickle 8 | import sys 9 | import os 10 | import numpy as np 11 | from typing import Sequence 12 | 13 | MODEL_URL = 'https://storage.googleapis.com/esper/models/clothing/model_hairstyle.tar' 14 | MODEL_DEF_URL = 'https://raw.githubusercontent.com/Haotianz94/video-analysis/master/streetstyle-classifier/classifier/model_hairstyle.py' 15 | 16 | 17 | ATTRIBUTES = [ 18 | { 19 | 'key': 'Hair color 3', 20 | 'values': {'black' : 0, 'white': 1, 'blond': 2}, 21 | }, 22 | { 23 | 'key': 'Hair color 5', 24 | 'values': {'black' : 0, 'white': 1, 'blond' : 2, 'brown' : 3, 'gray' : 4}, 25 | }, 26 | { 27 | 'key': 'Hair length', 28 | 'values': {'long' : 0, 'medium' : 1, 'short' : 2, 'bald' : 3} 29 | } 30 | ] # yapf: disable 31 | 32 | 33 | class HairStyle: 34 | def __init__(self, predictions): 35 | self._predictions = predictions 36 | 37 | def to_dict(self): 38 | return { 39 | attribute['key']: {v: k 40 | for k, v in attribute['values'].items()}[prediction] 41 | for prediction, attribute in zip(self._predictions, ATTRIBUTES) 42 | } 43 | 44 | def __str__(self): 45 | pieces = [] 46 | for prediction, attribute in zip(self._predictions, ATTRIBUTES): 47 | reverse_map = {v: k for k, v in attribute['values'].items()} 48 | pieces.append('{}: {}'.format(attribute['key'], reverse_map[prediction])) 49 | return '\n'.join(pieces) 50 | 51 | 52 | BATCH_SIZE = 2 53 | 54 | 55 | @scannerpy.register_python_op( 56 | device_sets=[[DeviceType.CPU, 0], [DeviceType.GPU, 1]], batch=BATCH_SIZE) 57 | class DetectHairStyle(TorchKernel): 58 | def __init__(self, config): 59 | from torchvision import transforms 60 | TorchKernel.__init__(self, config) 61 | 62 | self.transform = transforms.Compose([ 63 | transforms.Resize((299, 299)), 64 | transforms.ToTensor(), 65 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 66 | ]) 67 | 68 | def execute(self, frame: Sequence[FrameType], bboxes: Sequence[bytes]) -> Sequence[bytes]: 69 | from PIL import Image 70 | from torch.autograd import Variable 71 | import torch 72 | 73 | H, W = frame[0].shape[:2] 74 | 75 | counts = [] 76 | images = [] 77 | for i, (fram, bbs) in enumerate(zip(frame, bboxes)): 78 | bbs = readers.bboxes(bbs, self.config.protobufs) 79 | counts.append((counts[i - 1][0] + counts[i - 1][1] if i > 0 else 0, len(bbs))) 80 | if len(bboxes) == 0: 81 | raise Exception("No bounding boxes") 82 | 83 | for i, bbox in enumerate(bbs): 84 | x1 = int(bbox.x1 * W) 85 | y1 = int(bbox.y1 * H) 86 | x2 = int(bbox.x2 * W) 87 | y2 = int(bbox.y2 * H) 88 | w = max(y2 - y1, x2 - x1) * 3 // 4 89 | cx, cy = (x1 + x2) // 2, (y1 + y2) // 2 90 | x1 = cx - w if cx - w > 0 else 0 91 | x2 = cx + w if cx + w < W else W 92 | y1 = cy - w if cy - w > 0 else 0 93 | y2 = cy + w if cy + w < H else H 94 | cropped = fram[y1:y2, x1:x2, :] 95 | images.append(cropped) 96 | 97 | all_scores = [] 98 | for i in range(0, len(images), BATCH_SIZE): 99 | tensor = self.images_to_tensor( 100 | [self.transform(Image.fromarray(img)) for img in images[i:i + BATCH_SIZE]]) 101 | var = Variable(tensor if self.cpu_only else tensor.cuda(), requires_grad=False) 102 | all_scores.append(self.model(var)) 103 | 104 | scores = [ 105 | torch.cat([scores[i] for scores in all_scores], dim=0) 106 | for i in range(len(all_scores[0])) 107 | ] 108 | 109 | all_att = [] 110 | for k in range(len(frame)): 111 | (idx, n) = counts[k] 112 | predicted_attributes = np.zeros((n, len(scores)), dtype=np.int32) 113 | for i, attrib_score in enumerate(scores): 114 | _, predicted = torch.max(attrib_score[idx:idx + n, :], 1) 115 | predicted_attributes[:, i] = predicted.cpu().data.numpy().astype(np.int32) 116 | all_att.append(pickle.dumps(predicted_attributes)) 117 | 118 | return all_att 119 | 120 | 121 | def parse_hairstyle(s, _proto): 122 | predictions = pickle.loads(s) 123 | return [HairStyle(predictions[i, :]) for i in range(len(predictions))] 124 | 125 | 126 | class HairStyleDetectionPipeline(Pipeline): 127 | job_suffix = 'hairstyle' 128 | parser_fn = lambda _: parse_hairstyle 129 | additional_sources = ['bboxes'] 130 | 131 | def fetch_resources(self): 132 | try_import('torch', __name__) 133 | try_import('torchvision', __name__) 134 | 135 | self._model_path = download_temp_file(MODEL_URL) 136 | self._model_def_path = download_temp_file(MODEL_DEF_URL) 137 | 138 | def build_pipeline(self, adjust_bboxes=True): 139 | return { 140 | 'hairstyle': 141 | self._db.ops.DetectHairStyle( 142 | frame=self._sources['frame_sampled'].op, 143 | bboxes=self._sources['bboxes'].op, 144 | model_path=self._model_path, 145 | model_def_path=self._model_def_path, 146 | model_key='best_model', 147 | adjust_bboxes=adjust_bboxes, 148 | device=self._device) 149 | } 150 | 151 | 152 | detect_hairstyle = HairStyleDetectionPipeline.make_runner() 153 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/histograms.py: -------------------------------------------------------------------------------- 1 | from .prelude import * 2 | import numpy as np 3 | import os 4 | 5 | 6 | class HistogramPipeline(Pipeline): 7 | job_suffix = 'hist' 8 | parser_fn = lambda _: readers.histograms 9 | 10 | def build_pipeline(self, batch=1): 11 | return { 12 | 'histogram': 13 | self._db.ops.Histogram( 14 | frame=self._sources['frame'].op, device=self._device, batch=batch) 15 | } 16 | 17 | 18 | compute_histograms = HistogramPipeline.make_runner() 19 | 20 | 21 | class HSVHistogramPipeline(Pipeline): 22 | job_suffix = 'hsv_hist' 23 | parser_fn = lambda _: readers.histograms 24 | 25 | def fetch_resources(self): 26 | cwd = os.path.dirname(os.path.abspath(__file__)) 27 | 28 | self._db.load_op( 29 | os.path.join(cwd, 'cpp_ops/build/libimgproc_op.so'), 30 | os.path.join(cwd, 'cpp_ops/build/imgproc_pb2.py')) 31 | 32 | def build_pipeline(self, batch=1): 33 | hsv_frames = self._db.ops.ConvertToHSVCPP(frame=self._sources['frame'].op) 34 | 35 | return { 36 | 'histogram': self._db.ops.Histogram(frame=hsv_frames, device=self._device, batch=batch) 37 | } 38 | 39 | 40 | compute_hsv_histograms = HSVHistogramPipeline.make_runner() 41 | 42 | 43 | def flow_hist_reader(buf, protobufs): 44 | if buf is None: 45 | return None 46 | return np.split(np.frombuffer(buf, dtype=np.dtype(np.int32)), 2) 47 | 48 | 49 | class OpticalFlowHistogramPipeline(Pipeline): 50 | """ 51 | Computes histograms of optical flow on a video. 52 | """ 53 | job_suffix = 'flow_hist' 54 | parser_fn = lambda _: flow_hist_reader 55 | 56 | def fetch_resources(self): 57 | cwd = os.path.dirname(os.path.abspath(__file__)) 58 | 59 | self._db.load_op( 60 | os.path.join(cwd, 'cpp_ops/build/libimgproc_op.so'), 61 | os.path.join(cwd, 'cpp_ops/build/imgproc_pb2.py')) 62 | 63 | def build_pipeline(self): 64 | small_frames = self._db.ops.Resize( 65 | frame=self._sources['frame_sampled'].op, 66 | device=DeviceType.GPU if self._db.has_gpu() else DeviceType.CPU, 67 | width=426, 68 | height=240) 69 | 70 | flow = self._db.ops.OpticalFlow( 71 | frame=small_frames, 72 | device=DeviceType.GPU if self._db.has_gpu() else DeviceType.CPU) 73 | return { 74 | 'flow_hist': 75 | self._db.ops.FlowHistogram( 76 | flow=flow, 77 | device=DeviceType.CPU) 78 | } 79 | 80 | 81 | compute_flow_histograms = OpticalFlowHistogramPipeline.make_runner() 82 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/imgproc.py: -------------------------------------------------------------------------------- 1 | from .prelude import Pipeline 2 | from scannerpy import FrameType 3 | import scannerpy 4 | import cv2 5 | import pickle 6 | import numpy as np 7 | import struct 8 | import os 9 | 10 | 11 | @scannerpy.register_python_op(name='Brightness') 12 | def brightness(config, frame: FrameType) -> bytes: 13 | frame = cv2.cvtColor(frame, cv2.COLOR_RGB2YUV) 14 | 15 | # Calculate the mean value of the intensity channel 16 | brightness = np.mean(frame, axis=(0, 1))[0] 17 | return pickle.dumps(brightness) 18 | 19 | 20 | @scannerpy.register_python_op(name='Contrast') 21 | def contrast(config, frame: FrameType) -> bytes: 22 | frame = cv2.cvtColor(frame, cv2.COLOR_RGB2YUV) 23 | 24 | (h, w, c) = frame.shape 25 | intensities = frame.reshape((h * w * c))[::3] 26 | 27 | # Calculate the average intensity 28 | average_intensity = np.mean(intensities) 29 | contrast = np.sqrt(np.mean((intensities - average_intensity)**2)) 30 | return pickle.dumps(contrast) 31 | 32 | 33 | @scannerpy.register_python_op(name='Sharpness') 34 | def sharpness(config, frame: FrameType) -> bytes: 35 | sharpness = cv2.Laplacian(frame, cv2.CV_64F).var() 36 | return pickle.dumps(sharpness) 37 | 38 | 39 | @scannerpy.register_python_op(name='ConvertToHSV') 40 | def convert_to_hsv(config, frame: FrameType) -> FrameType: 41 | return cv2.cvtColor(frame, cv2.COLOR_RGB2HSV) 42 | 43 | 44 | @scannerpy.register_python_op(name='SharpnessBBox') 45 | def sharpness_bbox(config, frame: FrameType, bboxes: bytes) -> bytes: 46 | bboxes = readers.bboxes(bboxes.self.config.protobufs) 47 | 48 | results = [] 49 | for bbox in bboxes: 50 | img = frame[int(bbox.y1):int(bbox.y2), int(bbox.x1):int(bbox.x2), :] 51 | img = cv2.resize(img, (200, 200)) 52 | results.append(cv2.Laplacian(img, cv2.CV_64F).var()) 53 | 54 | return pickle.dumps(results) 55 | 56 | 57 | class BrightnessPipeline(Pipeline): 58 | job_suffix = 'brightness' 59 | parser_fn = lambda _: lambda buf, _: pickle.loads(buf) 60 | run_opts = {'pipeline_instances_per_node': 1} 61 | 62 | def build_pipeline(self): 63 | return {'brightness': self._db.ops.Brightness(frame=self._sources['frame_sampled'].op)} 64 | 65 | 66 | class BrightnessCPPPipeline(Pipeline): 67 | job_suffix = 'brightness_cpp' 68 | parser_fn = lambda _: lambda buf, _: struct.unpack('f', buf)[0] 69 | 70 | def fetch_resources(self): 71 | cwd = os.path.dirname(os.path.abspath(__file__)) 72 | 73 | self._db.load_op( 74 | os.path.join(cwd, 'cpp_ops/build/libimgproc_op.so'), 75 | os.path.join(cwd, 'cpp_ops/build/imgproc_pb2.py')) 76 | 77 | def build_pipeline(self): 78 | return {'brightness': self._db.ops.BrightnessCPP(frame=self._sources['frame_sampled'].op)} 79 | 80 | 81 | class ContrastPipeline(Pipeline): 82 | job_suffix = 'contrast' 83 | parser_fn = lambda _: lambda buf, _: pickle.loads(buf) 84 | run_opts = {'pipeline_instances_per_node': 1} 85 | 86 | def build_pipeline(self): 87 | return {'contrast': self._db.ops.Contrast(frame=self._sources['frame_sampled'].op)} 88 | 89 | 90 | class ContrastCPPPipeline(Pipeline): 91 | job_suffix = 'contrast_cpp' 92 | parser_fn = lambda _: lambda buf, _: struct.unpack('f', buf)[0] 93 | 94 | def fetch_resources(self): 95 | cwd = os.path.dirname(os.path.abspath(__file__)) 96 | 97 | self._db.load_op( 98 | os.path.join(cwd, 'cpp_ops/build/libimgproc_op.so'), 99 | os.path.join(cwd, 'cpp_ops/build/imgproc_pb2.py')) 100 | 101 | def build_pipeline(self): 102 | return {'contrast': self._db.ops.ContrastCPP(frame=self._sources['frame_sampled'].op)} 103 | 104 | 105 | class SharpnessPipeline(Pipeline): 106 | job_suffix = 'sharpness' 107 | parser_fn = lambda _: lambda buf, _: pickle.loads(buf) 108 | run_opts = {'pipeline_instances_per_node': 1} 109 | 110 | def build_pipeline(self): 111 | return {'sharpness': self._db.ops.Sharpness(frame=self._sources['frame_sampled'].op)} 112 | 113 | 114 | class SharpnessCPPPipeline(Pipeline): 115 | job_suffix = 'sharpness_cpp' 116 | parser_fn = lambda _: lambda buf, _: struct.unpack('f', buf)[0] 117 | 118 | def fetch_resources(self): 119 | cwd = os.path.dirname(os.path.abspath(__file__)) 120 | 121 | self._db.load_op( 122 | os.path.join(cwd, 'cpp_ops/build/libimgproc_op.so'), 123 | os.path.join(cwd, 'cpp_ops/build/imgproc_pb2.py')) 124 | 125 | def build_pipeline(self): 126 | return {'sharpness': self._db.ops.SharpnessCPP(frame=self._sources['frame_sampled'].op)} 127 | 128 | 129 | class SharpnessBBoxPipeline(Pipeline): 130 | job_suffix = 'sharpness_bbox' 131 | parser_fn = lambda _: lambda buf, _: pickle.loads(buf) 132 | additional_sources = ['bboxes'] 133 | run_opts = {'pipeline_instances_per_node': 1} 134 | 135 | def build_pipeline(self): 136 | return { 137 | 'sharpness_bbox': 138 | self._db.ops.SharpnessBBox( 139 | frame=self._sources['frame_sampled'].op, bboxes=self._sources['bboxes'].op) 140 | } 141 | 142 | 143 | class SharpnessBBoxCPPPipeline(Pipeline): 144 | job_suffix = 'sharpness_bbox_cpp' 145 | parser_fn = lambda _: lambda buf, _: struct.unpack('{}f'.format(int(len(buf) / 4)), buf) 146 | 147 | def fetch_resources(self): 148 | cwd = os.path.dirname(os.path.abspath(__file__)) 149 | 150 | self._db.load_op( 151 | os.path.join(cwd, 'cpp_ops/build/libimgproc_op.so'), 152 | os.path.join(cwd, 'cpp_ops/build/imgproc_pb2.py')) 153 | 154 | def build_pipeline(self): 155 | return { 156 | 'sharpness_bbox': self._db.ops.SharpnessBBoxCPP(frame=self._sources['frame_sampled'].op) 157 | } 158 | 159 | 160 | compute_brightness = BrightnessPipeline.make_runner() 161 | compute_brightness_cpp = BrightnessCPPPipeline.make_runner() 162 | compute_contrast = ContrastPipeline.make_runner() 163 | compute_contrast_cpp = ContrastCPPPipeline.make_runner() 164 | compute_sharpness = SharpnessPipeline.make_runner() 165 | compute_sharpness_cpp = SharpnessCPPPipeline.make_runner() 166 | compute_sharpness_bbox = SharpnessBBoxPipeline.make_runner() 167 | compute_sharpness_bbox_cpp = SharpnessBBoxCPPPipeline.make_runner() 168 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/optical_flow.py: -------------------------------------------------------------------------------- 1 | from .prelude import * 2 | import os 3 | import pickle 4 | 5 | SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) 6 | 7 | 8 | class OpticalFlowPipeline(Pipeline): 9 | """ 10 | Computes optical flow on a video. 11 | 12 | Unlike other functions, flow fields aren't materialized into memory as they're simply 13 | too large. 14 | """ 15 | 16 | job_suffix = 'flow' 17 | parser_fn = lambda _: lambda x: x 18 | 19 | def build_pipeline(self): 20 | return { 21 | 'flow': 22 | self._db.ops.OpticalFlow(frame=self._sources['frame_sampled'].op, device=self._device) 23 | } 24 | 25 | 26 | compute_flow = OpticalFlowPipeline.make_runner() 27 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/pipeline3.py: -------------------------------------------------------------------------------- 1 | import scannerpy 2 | 3 | def frame_gather_graph(db, video, frame_indices): 4 | frame = db.sources.FrameColumn(video=video) 5 | frame_gathered = db.streams.Gather(frame, indices=frame_indices) 6 | return frame_gathered 7 | 8 | def default_graph(kernel, sink=None): 9 | if not sink: 10 | sink = db.sinks.Column(columnsd={'column': output}, table_name=table_name) 11 | def graph(db, video, frame_indices, table_name): 12 | frame_gathered = frame_gather_graph(db, video, frame_indices) 13 | output = kernel(db, frame_gathered) 14 | sink = db.sinks.Column(columns={'column': output}, table_name=table_name) 15 | return sink 16 | return graph 17 | 18 | 19 | histogram_graph = default_graph(lambda db, frame: db.ops.Histogram(frame=frame)) 20 | 21 | from .prelude import sample_video 22 | with sample_video(delete=False) as video: 23 | db = scannerpy.Database() 24 | print(histogram_graph(db, video=[video], frame_indices=[[0]], table_name=['test'])) 25 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/test_audio.py: -------------------------------------------------------------------------------- 1 | from .prelude import * 2 | import numpy as np 3 | from scannerpy import FrameType 4 | import pickle 5 | 6 | 7 | @scannerpy.register_python_op(name='AverageVolume') 8 | def average_volume(config, audio: FrameType) -> bytes: 9 | return pickle.dumps(np.average(audio)) 10 | 11 | 12 | class AverageVolumePipeline(Pipeline): 13 | job_suffix = 'avgvolume' 14 | base_sources = ['audio'] 15 | run_opts = {'pipeline_instances_per_node': 1} 16 | parser_fn = lambda _: lambda buf, _: pickle.loads(buf) 17 | 18 | def build_pipeline(self): 19 | return {'avgvolume': self._db.ops.AverageVolume(audio=self._sources['audio'].op)} 20 | 21 | def build_sink(self): 22 | return BoundOp( 23 | op=self._db.sinks.Column(columns=self._output_ops), 24 | args=[ 25 | '{}_{}'.format(arg['path'], self.job_suffix) for arg in self._sources['audio'].args 26 | ]) 27 | 28 | 29 | compute_average_volume = AverageVolumePipeline.make_runner() 30 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/test_caption.py: -------------------------------------------------------------------------------- 1 | from .prelude import * 2 | import numpy as np 3 | from scannerpy import FrameType 4 | import pickle 5 | 6 | import json 7 | 8 | 9 | @scannerpy.register_python_op('PrintCap') 10 | def print_cap(config, captions: bytes) -> bytes: 11 | print(json.loads(captions.decode('utf-8'))) 12 | return b' ' 13 | 14 | 15 | class PrintCaptionsPipeline(Pipeline): 16 | job_suffix = 'printcap' 17 | base_sources = ['captions'] 18 | run_opts = {'pipeline_instances_per_node': 1} 19 | parser_fn = lambda _: lambda buf, _: () 20 | 21 | def build_pipeline(self): 22 | return {'cap': self._db.ops.PrintCap(captions=self._sources['captions'].op)} 23 | 24 | def build_sink(self): 25 | return BoundOp( 26 | op=self._db.sinks.Column(columns=self._output_ops), 27 | args=[ 28 | '{}_{}'.format(arg['path'], self.job_suffix) 29 | for arg in self._sources['captions'].args 30 | ]) 31 | 32 | 33 | print_captions = PrintCaptionsPipeline.make_runner() 34 | -------------------------------------------------------------------------------- /scannertools/scannertools/old/video.py: -------------------------------------------------------------------------------- 1 | from .prelude import * 2 | import os 3 | 4 | 5 | class Audio: 6 | """ 7 | Reference to an audio file on disk. 8 | """ 9 | 10 | def __init__(self, audio_path): 11 | self._path = audio_path 12 | 13 | def extract(self, path=None, ext='.wav', segment=None): 14 | return ffmpeg_extract( 15 | input_path=self.path(), output_path=path, output_ext=ext, segment=segment) 16 | 17 | def path(self): 18 | return self._path 19 | 20 | 21 | class Video: 22 | """ 23 | Reference to a video file on disk. 24 | 25 | Currently only supports mp4. 26 | """ 27 | 28 | def __init__(self, path, scanner_name=None): 29 | """ 30 | Args: 31 | path (str): Path to video file 32 | """ 33 | 34 | self._path = path 35 | self._decoder_handle = None 36 | self._scanner_name = scanner_name 37 | 38 | # Lazily load decoder 39 | def _decoder(self): 40 | if self._decoder_handle is None: 41 | try: 42 | video_file = storehouse.RandomReadFile(get_storage(), self._path.encode('ascii')) 43 | except UserWarning: 44 | raise Exception('Path to video `{}` does not exist.'.format(self._path)) 45 | self._decoder_handle = hwang.Decoder(video_file) 46 | return self._decoder_handle 47 | 48 | def path(self): 49 | """ 50 | Returns: 51 | str: Video file path. 52 | """ 53 | return self._path 54 | 55 | def scanner_name(self): 56 | """ 57 | Returns: 58 | str: Name of the video file in the Scanner database. 59 | """ 60 | 61 | return self._scanner_name or os.path.basename(self.path()) 62 | 63 | def width(self): 64 | """ 65 | Returns: 66 | int: Width in pixels of the video. 67 | """ 68 | return self._decoder().video_index.frame_width() 69 | 70 | def height(self): 71 | """ 72 | Returns: 73 | int: Height in pixels of the video. 74 | """ 75 | return self._decoder().video_index.frame_height() 76 | 77 | def fps(self): 78 | """ 79 | Returns: 80 | float: Frames per seconds of the video. 81 | """ 82 | return self._decoder().video_index.fps() 83 | 84 | def num_frames(self): 85 | """ 86 | Returns: 87 | int: Number of frames in the video. 88 | """ 89 | return self._decoder().video_index.frames() 90 | 91 | def duration(self): 92 | """ 93 | Returns: 94 | int: Length of the video in seconds. 95 | """ 96 | return self._decoder().video_index.duration() 97 | 98 | def frame(self, number=None, time=None): 99 | """ 100 | Extract a single frame from the video into memory. 101 | 102 | Exactly one of number or time should be specified. 103 | 104 | Args: 105 | number (int, optional): The index of the frame to access. 106 | time (float, optional): The time in seconds of the frame to access. 107 | 108 | Returns: 109 | np.array: (h x w x 3) np.uint8 image. 110 | """ 111 | 112 | if time is not None: 113 | return self.frames(times=[time])[0] 114 | else: 115 | return self.frames(numbers=[number])[0] 116 | 117 | def frames(self, numbers=None, times=None): 118 | """ 119 | Extract multiple frames from the video into memory. 120 | 121 | Args: 122 | numbers (List[int], optional): The indices of the frames to access. 123 | times (List[float], optional): The times in seconds of the frames to access. 124 | 125 | Returns: 126 | List[np.array]: List of (h x w x 3) np.uint8 images. 127 | """ 128 | if times is not None: 129 | numbers = [int(n * self.fps()) for n in times] 130 | 131 | # Hwang expects: a) no duplicates, and b) sorted frame indices 132 | to_fetch = sorted(list(set(numbers))) 133 | frames = self._decoder().retrieve(to_fetch) 134 | idx_map = {n: i for i, n in enumerate(to_fetch)} 135 | 136 | return [frames[idx_map[n]].copy() for n in numbers] 137 | 138 | def audio(self): 139 | """ 140 | Extract the audio from the video. 141 | 142 | Returns: 143 | Audio: Reference to the audio file. 144 | """ 145 | audio_path = ffmpeg_extract(input_path=self.path(), output_ext='.wav') 146 | return Audio(audio_path) 147 | 148 | def extract(self, path=None, ext='.mp4', segment=None): 149 | """ 150 | Extract an mp4 out of the video. 151 | 152 | Args: 153 | path (str, optional): Path to write the video. 154 | ext (str, optional): Video extension to write 155 | segment (Tuple(int, int), optional): Start/end in seconds 156 | 157 | Returns: 158 | str: Path to the created video 159 | """ 160 | 161 | return ffmpeg_extract( 162 | input_path=self.path(), output_path=path, output_ext=ext, segment=segment) 163 | 164 | def montage(self, frames, rows=None, cols=None): 165 | """ 166 | Create a tiled montage of frames in the video. 167 | 168 | Args: 169 | frames (List[int]): List of frame indices. 170 | rows (List[int], optional): Number of rows in the montage. 171 | cols (List[int], optional): Number of columns in the montage. 172 | 173 | Returns: 174 | np.array: Image of the montage. 175 | """ 176 | 177 | frames = self.frames(frames) 178 | return tile(frames, rows=rows, cols=cols) 179 | -------------------------------------------------------------------------------- /scannertools/scannertools/shot_detection.py: -------------------------------------------------------------------------------- 1 | from scipy.spatial import distance 2 | import numpy as np 3 | from typing import Sequence, Any 4 | import scannerpy 5 | from scannerpy.types import Histogram 6 | 7 | WINDOW_SIZE = 500 8 | BOUNDARY_BATCH = 10000000 9 | 10 | 11 | @scannerpy.register_python_op(name='ShotBoundaries', batch=BOUNDARY_BATCH) 12 | def shot_boundaries(config, histograms: Sequence[Histogram]) -> Sequence[Any]: 13 | # Compute the mean difference between each pair of adjacent frames 14 | diffs = np.array([ 15 | np.mean([distance.chebyshev(histograms[i - 1][j], histograms[i][j]) for j in range(3)]) 16 | for i in range(1, len(histograms)) 17 | ]) 18 | diffs = np.insert(diffs, 0, 0) 19 | n = len(diffs) 20 | 21 | # Do simple outlier detection to find boundaries between shots 22 | boundaries = [] 23 | for i in range(1, n): 24 | window = diffs[max(i - WINDOW_SIZE, 0):min(i + WINDOW_SIZE, n)] 25 | if diffs[i] - np.mean(window) > 2.5 * np.std(window): 26 | boundaries.append(i) 27 | 28 | return [boundaries] + [None for _ in range(len(histograms) - 1)] 29 | -------------------------------------------------------------------------------- /scannertools/scannertools/storage/__init__.py: -------------------------------------------------------------------------------- 1 | from scannertools_infra import _register_module 2 | import os 3 | _register_module(os.path.join(__file__, '..'), "scannertools_storage") 4 | -------------------------------------------------------------------------------- /scannertools/scannertools/storage/audio.py: -------------------------------------------------------------------------------- 1 | from scannerpy.storage import StorageBackend, StoredStream 2 | 3 | 4 | class AudioStorage(StorageBackend): 5 | """StorageBackend for stream of elements from a compressed audio file. 6 | 7 | Currently input-only.""" 8 | 9 | def source(self, sc, streams): 10 | return sc.sources.Audio( 11 | frame_size=[s._frame_size for s in streams], 12 | path=[s._path for s in streams]) 13 | 14 | 15 | class AudioStream(StoredStream): 16 | """Stream of elements from a compressed audio file.""" 17 | 18 | def __init__(self, path, frame_size, storage=None): 19 | """ 20 | Parameters 21 | ---------- 22 | path: str 23 | Path on filesystem to audio file. 24 | 25 | frame_size: float 26 | Size (in seconds) of each element, e.g. a 2s frame size with a a 44.1 kHz generates 27 | stream elements of 88.2k samples per element. 28 | 29 | storage: AudioStorage 30 | """ 31 | 32 | if storage is None: 33 | self._storage = AudioStorage() 34 | else: 35 | self._storage = storage 36 | 37 | self._path = path 38 | self._frame_size = frame_size 39 | 40 | def storage(self): 41 | return self._storage 42 | -------------------------------------------------------------------------------- /scannertools/scannertools/storage/caption.py: -------------------------------------------------------------------------------- 1 | from scannerpy.storage import StorageBackend, StoredStream 2 | 3 | class CaptionStorage(StorageBackend): 4 | """StorageBackend for caption streams.""" 5 | 6 | def source(self, sc, streams): 7 | return sc.sources.Captions( 8 | window_size=[s._window_size for s in streams], 9 | path=[s._path for s in streams], 10 | max_time=[s._max_time for s in streams]) 11 | 12 | 13 | class CaptionStream(StoredStream): 14 | """Stream of captions out of a caption file. 15 | 16 | In order for the number of stream elements to be predictable (e.g. to zip a caption stream 17 | with an audio stream for transcript alignment), we represent caption streams as uniform time 18 | intervals. You provide a stream duration (e.g. 10 minutes) and a window size (e.g. 5 seconds) 19 | and the stream contains 10 minutes / 5 seconds number of elements, where each element contains 20 | all of the text for captions that overlap with that window. 21 | """ 22 | 23 | def __init__(self, path, window_size, max_time, storage=None): 24 | """ 25 | Parameters 26 | ---------- 27 | path: str 28 | Path on the filesystem to the caption file. 29 | 30 | window_size: float 31 | Size of window in time (seconds) for each element. See class description. 32 | 33 | max_time: float 34 | Total time for the entire stream. See class description. 35 | 36 | storage: CaptionStorage 37 | """ 38 | 39 | if storage is None: 40 | self._storage = CaptionStorage() 41 | else: 42 | self._storage = storage 43 | 44 | self._path = path 45 | self._max_time = max_time 46 | self._window_size = window_size 47 | 48 | def storage(self): 49 | return self._storage 50 | -------------------------------------------------------------------------------- /scannertools/scannertools/storage/files.py: -------------------------------------------------------------------------------- 1 | from scannerpy.storage import StorageBackend, StoredStream 2 | from typing import List 3 | import os 4 | 5 | 6 | class FilesStorage(StorageBackend): 7 | """Storage of streams where each element is its own file.""" 8 | 9 | def __init__(self, storage_type: str = "posix", bucket: str = None, region: str = None, endpoint: str = None): 10 | """ 11 | Parameters 12 | ---------- 13 | storage_type: str 14 | Kind of filesystem the files are on. Either "posix" or "gcs" supported. 15 | 16 | bucket: str 17 | If filesystem is gcs, name of bucket 18 | 19 | region: str 20 | If filesytem is gcs, region name of bucket, e.g. us-west1 21 | 22 | endpoint: str 23 | If filesystem is gcs, URL of storage endpoint 24 | """ 25 | 26 | self._storage_type = storage_type 27 | self._bucket = bucket 28 | self._region = region 29 | self._endpoint = endpoint 30 | 31 | def source(self, sc, streams): 32 | return sc.sources.Files( 33 | storage_type=self._storage_type, 34 | bucket=self._bucket, 35 | region=self._region, 36 | endpoint=self._endpoint, 37 | paths=[s._paths for s in streams]) 38 | 39 | def sink(self, sc, op, streams): 40 | return sc.sinks.Files( 41 | input=op, 42 | storage_type=self._storage_type, 43 | bucket=self._bucket, 44 | region=self._region, 45 | endpoint=self._endpoint, 46 | paths=[s._paths for s in streams]) 47 | 48 | def delete(self, sc, streams): 49 | # TODO 50 | pass 51 | 52 | 53 | class FilesStream(StoredStream): 54 | """Stream where each element is a file.""" 55 | 56 | def __init__(self, paths: List[str], storage: FilesStorage = None): 57 | """ 58 | Parameters 59 | ---------- 60 | paths: List[str] 61 | List of paths to the files in the stream. 62 | 63 | storage: FilesStorage 64 | """ 65 | if storage is None: 66 | self._storage = FilesStorage() 67 | else: 68 | self._storage = storage 69 | 70 | self._paths = paths 71 | 72 | def load_bytes(self, rows=None): 73 | paths = self._paths 74 | if rows is not None: 75 | paths = [paths[i] for i in rows] 76 | 77 | for path in paths: 78 | yield open(path, 'rb').read() 79 | 80 | def storage(self): 81 | return self._storage 82 | 83 | def committed(self): 84 | # TODO 85 | return all(os.path.isfile(p) for p in self._paths) 86 | 87 | def exists(self): 88 | # TODO 89 | return any(os.path.isfile(p) for p in self._paths) 90 | 91 | def type(self): 92 | return None 93 | -------------------------------------------------------------------------------- /scannertools/scannertools/storage/python.py: -------------------------------------------------------------------------------- 1 | from scannerpy.storage import StorageBackend, StoredStream 2 | import pickle 3 | from typing import List, Any 4 | 5 | 6 | class PythonStorage(StorageBackend): 7 | """StorageBackend for a stream of elements directly from the current Python process. 8 | 9 | Only supports input, not output. 10 | """ 11 | 12 | def source(self, sc, streams): 13 | return sc.sources.Python(data=[pickle.dumps(stream._data) for stream in streams]) 14 | 15 | 16 | class PythonStream(StoredStream): 17 | """Stream of elements directly in the current Python process.""" 18 | 19 | def __init__(self, data: List[Any]): 20 | """ 21 | Parameters 22 | ---------- 23 | data: List[Any] 24 | Arbitrary data to stream. Must be pickleable. 25 | """ 26 | self._data = data 27 | 28 | def storage(self): 29 | return PythonStorage() 30 | -------------------------------------------------------------------------------- /scannertools/scannertools/tensorflow.py: -------------------------------------------------------------------------------- 1 | from scannerpy.kernel import Kernel 2 | from scannerpy import DeviceType 3 | 4 | 5 | class TensorFlowKernel(Kernel): 6 | def __init__(self, config): 7 | import tensorflow as tf 8 | 9 | # If this is a CPU kernel, tell TF that it should not use 10 | # any GPUs for its graph operations 11 | cpu_only = True 12 | visible_device_list = [] 13 | tf_config = tf.ConfigProto() 14 | for handle in config.devices: 15 | if handle.type == DeviceType.GPU.value: 16 | visible_device_list.append(str(handle.id)) 17 | tf_config.gpu_options.allow_growth = True 18 | cpu_only = False 19 | if cpu_only: 20 | tf_config.device_count['GPU'] = 0 21 | else: 22 | tf_config.gpu_options.visible_device_list = ','.join(visible_device_list) 23 | # TODO: wrap this in "with device" 24 | self.config = config 25 | self.tf_config = tf_config 26 | 27 | def close(self): 28 | self.sess.close() 29 | 30 | def setup_with_resources(self): 31 | import tensorflow as tf 32 | 33 | self.graph = self.build_graph() 34 | self.sess = tf.Session(config=self.tf_config, graph=self.graph) 35 | self.sess.as_default() 36 | 37 | def build_graph(self): 38 | raise NotImplementedError 39 | 40 | def execute(self): 41 | raise NotImplementedError 42 | -------------------------------------------------------------------------------- /scannertools/scannertools/torch.py: -------------------------------------------------------------------------------- 1 | from scannerpy.kernel import Kernel 2 | from scannerpy import DeviceType 3 | import sys 4 | import numpy as np 5 | 6 | class TorchKernel(Kernel): 7 | def __init__(self, config): 8 | import torch 9 | 10 | self.config = config 11 | 12 | self.cpu_only = True 13 | visible_device_list = [] 14 | for handle in config.devices: 15 | if int(handle.type) == DeviceType.GPU.value: 16 | visible_device_list.append(handle.id) 17 | self.cpu_only = False 18 | 19 | self.model = self.build_model() 20 | 21 | if not self.cpu_only: 22 | print('Using GPU: {}'.format(visible_device_list[0])) 23 | torch.cuda.set_device(visible_device_list[0]) 24 | self.model = self.model.cuda() 25 | else: 26 | print('Using CPU') 27 | 28 | # Not sure if this is necessary? Haotian had it in his code 29 | self.model.eval() 30 | 31 | def images_to_tensor(self, images): 32 | import torch 33 | return torch.from_numpy(np.concatenate(np.expand_dims(images, axis=0), axis=0)) 34 | 35 | def build_model(self): 36 | import torch 37 | 38 | sys.path.insert(0, os.path.split(self.config.args['model_def_path'])[0]) 39 | kwargs = { 40 | 'map_location': lambda storage, location: storage 41 | } if self.cpu_only else {} 42 | return torch.load(self.config.args['model_path'], 43 | **kwargs)[self.config.args['model_key']] 44 | 45 | def close(self): 46 | del self.model 47 | 48 | def execute(self): 49 | raise NotImplementedError 50 | -------------------------------------------------------------------------------- /scannertools/scannertools/tracker.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import scannerpy 4 | import scannerpy.stdlib.bboxes 5 | import scannerpy.stdlib.readers 6 | import scannerpy.stdlib.writers 7 | 8 | from scannerpy import Client, Job, DeviceType, FrameType 9 | from scannerpy.stdlib import pipelines 10 | from typing import Sequence 11 | 12 | @scannerpy.register_python_op(bounded_state=5) 13 | class TrackObjects(scannerpy.Kernel): 14 | def __init__(self, config): 15 | self.config = config 16 | 17 | self.last_merge = [] 18 | self.trackers = [] 19 | self.prev_bboxes = [] 20 | 21 | def reset(self): 22 | self.last_merge = [] 23 | self.trackers = [] 24 | self.prev_bboxes = [] 25 | 26 | def execute(self, frame: FrameType, bboxes: bytes) -> bytes: 27 | # If we have new input boxes, track them 28 | if bboxes: 29 | bboxes = scannerpy.stdlib.readers.bboxes(bboxes, self.config.protobufs) 30 | # Create new trackers for each bbox 31 | for b in bboxes: 32 | # If this input box is the same as a tracked box from a previous 33 | # frame, then we don't need to start a new tracker 34 | is_same = False 35 | for i, prev_bbox in enumerate(self.prev_bboxes): 36 | if scannerpy.stdlib.bboxes.iou(prev_bbox, b) > 0.25: 37 | # We found a match, so ignore this box 38 | is_same = True 39 | self.last_merge[i] = 0 40 | break 41 | if is_same: 42 | continue 43 | 44 | t = cv2.TrackerMIL_create() 45 | t.init(frame, (b.x1, b.y1, b.x2 - b.x1, b.y2 - b.y1)) 46 | self.trackers.append(t) 47 | self.last_merge.append(0) 48 | 49 | out_bboxes = [] 50 | new_trackers = [] 51 | new_last_merge = [] 52 | for i, t in enumerate(self.trackers): 53 | self.last_merge[i] += 1 54 | if self.last_merge[i] > 10: 55 | continue 56 | 57 | ok, newbox = t.update(frame) 58 | if ok: 59 | # Track was successful, so keep the tracker around 60 | new_trackers.append(t) 61 | new_last_merge.append(self.last_merge[i]) 62 | 63 | # Convert from opencv format to protobuf for serialization 64 | newbox_proto = self.config.protobufs.BoundingBox() 65 | newbox_proto.x1 = newbox[0] 66 | newbox_proto.y1 = newbox[1] 67 | newbox_proto.x2 = newbox[0] + newbox[2] 68 | newbox_proto.y2 = newbox[1] + newbox[3] 69 | out_bboxes.append(newbox_proto) 70 | else: 71 | # Tracker failed, so do nothing 72 | pass 73 | 74 | print('num trackers', len(new_trackers)) 75 | self.trackers = new_trackers 76 | self.prev_bboxes = out_bboxes 77 | self.last_merge = new_last_merge 78 | 79 | return scannerpy.stdlib.writers.bboxes(out_bboxes, 80 | self.config.protobufs) 81 | -------------------------------------------------------------------------------- /scannertools/scannertools/types.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import struct 3 | 4 | from scannerpy.stdlib.poses import Pose 5 | 6 | 7 | 8 | def poses(buf, protobufs): 9 | if len(buf) == 4: 10 | return [] 11 | 12 | kp_size = ( 13 | Pose.POSE_KEYPOINTS + Pose.FACE_KEYPOINTS + Pose.HAND_KEYPOINTS * 2 14 | ) * 3 + Pose.POSE_SCORES 15 | poses = [] 16 | all_kp = np.frombuffer(buf, dtype=np.float32) 17 | for j in range(0, len(all_kp), kp_size): 18 | pose = Pose.from_buffer(all_kp[j:(j + kp_size)].tobytes()) 19 | poses.append(pose) 20 | return poses 21 | 22 | 23 | def histograms(buf, protobufs): 24 | # bufs[0] is None when element is null 25 | if buf is None: 26 | return None 27 | return np.split(np.frombuffer(buf, dtype=np.dtype(np.int32)), 3) 28 | 29 | 30 | def frame_info(buf, protobufs): 31 | info = protobufs.FrameInfo() 32 | info.ParseFromString(buf) 33 | return info 34 | 35 | 36 | def flow(bufs, protobufs): 37 | if bufs[0] is None: 38 | return None 39 | output = np.frombuffer(bufs, dtype=np.dtype(np.float32)) 40 | info = frame_info(bufs[1], db) 41 | return output.reshape((info.height, info.width, 2)) 42 | 43 | 44 | def array(ty, size=None): 45 | def parser(buf, protobufs): 46 | if buf == b'\0': 47 | return None 48 | else: 49 | npbuf = np.frombuffer(buf, dtype=np.dtype(ty)) 50 | if size is not None: 51 | return np.split(npbuf, npbuf.shape[0] / size) 52 | else: 53 | return npbuf 54 | 55 | return parser 56 | 57 | 58 | def image(buf): 59 | import cv2 60 | return cv2.imdecode( 61 | np.frombuffer(buf, dtype=np.dtype(np.uint8)), cv2.IMREAD_COLOR) 62 | 63 | 64 | 65 | import struct 66 | 67 | 68 | def bboxes(buf, protobufs): 69 | s = struct.pack('=Q', len(buf)) 70 | for bbox in buf: 71 | bs = bbox.SerializeToString() 72 | s += struct.pack('=Q', len(bs)) 73 | s += bs 74 | return s 75 | 76 | 77 | def poses(poses, protobufs): 78 | if len(poses) == 0: 79 | return b'\0' 80 | else: 81 | return b''.join([pose.keypoints.tobytes() for pose in poses]) 82 | -------------------------------------------------------------------------------- /scannertools/scannertools/vis.py: -------------------------------------------------------------------------------- 1 | import scannerpy as sp 2 | from scannerpy import FrameType 3 | from scannerpy.types import BboxList 4 | import numpy as np 5 | import cv2 6 | 7 | 8 | @sp.register_python_op(name='DrawFlow') 9 | def draw_flow(config, frame: FrameType, flow: FrameType) -> FrameType: 10 | flow_vis = np.repeat(np.expand_dims(np.average(flow, axis=2), 2), 3, axis=2) 11 | return np.hstack((frame, (np.clip(flow_vis / np.max(flow_vis), None, 1.0) * 255).astype( 12 | np.uint8))) 13 | 14 | 15 | @sp.register_python_op(name='DrawBboxes') 16 | def draw_bboxes(config, frame: FrameType, bboxes: BboxList) -> FrameType: 17 | [h, w] = frame.shape[:2] 18 | for bbox in bboxes: 19 | cv2.rectangle( 20 | frame, 21 | (int(bbox.x1 * w), int(bbox.y1 * h)), 22 | (int(bbox.x2 * w), int(bbox.y2 * h)), 23 | (255, 0, 0)) 24 | return frame 25 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.2.0 FATAL_ERROR) 2 | 3 | execute_process( 4 | OUTPUT_VARIABLE SCANNER_CMAKE_PATH 5 | COMMAND python3 -c "import scannerpy.build_flags as b; b.print_cmake()") 6 | include(${SCANNER_CMAKE_PATH}) 7 | 8 | set(SOURCES) 9 | add_subdirectory(imgproc) 10 | add_subdirectory(storage) 11 | add_subdirectory(misc) 12 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | set(SOURCES 2 | blur_kernel_cpu.cpp 3 | histogram_kernel_cpu.cpp 4 | montage_kernel_cpu.cpp 5 | image_decoder_kernel_cpu.cpp 6 | resize_kernel.cpp 7 | convert_color_kernel.cpp 8 | optical_flow_kernel_cpu.cpp) 9 | 10 | # Temporarily disable CUDA imgproc operations since opencv4 broke the build 11 | if (BUILD_CUDA) 12 | list(APPEND SOURCES 13 | histogram_kernel_gpu.cpp 14 | montage_kernel_gpu.cpp 15 | optical_flow_kernel_gpu.cpp) 16 | #image_decoder_kernel_gpu.cpp 17 | endif() 18 | 19 | build_op( 20 | LIB_NAME scannertools_imgproc 21 | CPP_SRCS ${SOURCES} 22 | PROTO_SRC scannertools_imgproc.proto 23 | BUILD_CUDA ${BUILD_CUDA}) 24 | 25 | set(LIBRARIES) 26 | 27 | list(APPEND OPENCV_COMPONENTS core highgui imgproc optflow) 28 | 29 | if (BUILD_CUDA) 30 | list(APPEND OPENCV_COMPONENTS cudafeatures2d cudacodec cudaoptflow cudawarping) 31 | endif() 32 | 33 | find_package(OpenCV REQUIRED COMPONENTS "${OPENCV_COMPONENTS}") 34 | list(APPEND LIBRARIES "${OpenCV_LIBRARIES}") 35 | 36 | target_link_libraries(scannertools_imgproc PUBLIC ${LIBRARIES}) 37 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/blur_kernel_cpu.cpp: -------------------------------------------------------------------------------- 1 | /* Copyright 2016 Carnegie Mellon University 2 | * 3 | * Licensed under the Apache License, Version 2.0 (the "License"); 4 | * you may not use this file except in compliance with the License. 5 | * You may obtain a copy of the License at 6 | * 7 | * http://www.apache.org/licenses/LICENSE-2.0 8 | * 9 | * Unless required by applicable law or agreed to in writing, software 10 | * distributed under the License is distributed on an "AS IS" BASIS, 11 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | * See the License for the specific language governing permissions and 13 | * limitations under the License. 14 | */ 15 | 16 | #include "scanner/api/kernel.h" 17 | #include "scanner/api/op.h" 18 | #include "scanner/util/memory.h" 19 | #include "scannertools_imgproc.pb.h" 20 | 21 | #include 22 | 23 | namespace scanner { 24 | 25 | class BlurKernel : public Kernel, public VideoKernel { 26 | public: 27 | BlurKernel(const KernelConfig& config) : Kernel(config) { 28 | BlurArgs args; 29 | bool parsed = args.ParseFromArray(config.args.data(), config.args.size()); 30 | if (!parsed || config.args.size() == 0) { 31 | RESULT_ERROR(&valid_, "Could not parse BlurArgs"); 32 | return; 33 | } 34 | 35 | kernel_size_ = args.kernel_size(); 36 | sigma_ = args.sigma(); 37 | 38 | filter_left_ = std::ceil(kernel_size_ / 2.0) - 1; 39 | filter_right_ = kernel_size_ / 2; 40 | 41 | valid_.set_success(true); 42 | } 43 | 44 | void validate(Result* result) override { result->CopyFrom(valid_); } 45 | 46 | void new_frame_info() { 47 | frame_width_ = frame_info_.width(); 48 | frame_height_ = frame_info_.height(); 49 | } 50 | 51 | void execute(const Elements& input_columns, 52 | Elements& output_columns) override { 53 | auto& frame_col = input_columns[0]; 54 | check_frame(CPU_DEVICE, frame_col); 55 | 56 | i32 width = frame_width_; 57 | i32 height = frame_height_; 58 | size_t frame_size = width * height * 3 * sizeof(u8); 59 | FrameInfo info = frame_col.as_const_frame()->as_frame_info(); 60 | Frame* output_frame = new_frame(CPU_DEVICE, info); 61 | 62 | const u8* frame_buffer = frame_col.as_const_frame()->data; 63 | u8* blurred_buffer = output_frame->data; 64 | for (i32 y = filter_left_; y < height - filter_right_; ++y) { 65 | for (i32 x = filter_left_; x < width - filter_right_; ++x) { 66 | for (i32 c = 0; c < 3; ++c) { 67 | u32 value = 0; 68 | for (i32 ry = -filter_left_; ry < filter_right_ + 1; ++ry) { 69 | for (i32 rx = -filter_left_; rx < filter_right_ + 1; ++rx) { 70 | value += frame_buffer[(y + ry) * width * 3 + (x + rx) * 3 + c]; 71 | } 72 | } 73 | blurred_buffer[y * width * 3 + x * 3 + c] = 74 | value / ((filter_right_ + filter_left_ + 1) * 75 | (filter_right_ + filter_left_ + 1)); 76 | } 77 | } 78 | } 79 | insert_frame(output_columns[0], output_frame); 80 | } 81 | 82 | private: 83 | i32 kernel_size_; 84 | i32 filter_left_; 85 | i32 filter_right_; 86 | f64 sigma_; 87 | 88 | i32 frame_width_; 89 | i32 frame_height_; 90 | Result valid_; 91 | }; 92 | 93 | REGISTER_OP(Blur).frame_input("frame").frame_output("frame").protobuf_name( 94 | "BlurArgs"); 95 | 96 | REGISTER_KERNEL(Blur, BlurKernel).device(DeviceType::CPU).num_devices(1); 97 | } 98 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/frame_difference_kernel_cpu.cpp: -------------------------------------------------------------------------------- 1 | /* Copyright 2017 Carnegie Mellon University 2 | * 3 | * Licensed under the Apache License, Version 2.0 (the "License"); 4 | * you may not use this file except in compliance with the License. 5 | * You may obtain a copy of the License at 6 | * 7 | * http://www.apache.org/licenses/LICENSE-2.0 8 | * 9 | * Unless required by applicable law or agreed to in writing, software 10 | * distributed under the License is distributed on an "AS IS" BASIS, 11 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | * See the License for the specific language governing permissions and 13 | * limitations under the License. 14 | */ 15 | 16 | #include "scanner/api/kernel.h" 17 | #include "scanner/api/op.h" 18 | #include "scanner/util/memory.h" 19 | #include "scannertools_imgproc.pb.h" 20 | 21 | #include 22 | 23 | namespace scanner { 24 | 25 | class FrameDifferenceKernel : public StenciledKernel { 26 | public: 27 | BlurKernel(const KernelConfig& config) : StenciledKernel(config) { 28 | valid_.set_success(true); 29 | } 30 | 31 | void validate(Result* result) override { 32 | result->CopyFrom(valid_); 33 | } 34 | 35 | void new_frame_info() { 36 | frame_width_ = frame_info_.width(); 37 | frame_height_ = frame_info_.height(); 38 | } 39 | 40 | void execute(const StenciledElements& input_columns, 41 | Elements& output_columns) override { 42 | auto& frame_col = input_columns[0]; 43 | check_frame(CPU_DEVICE, frame_col); 44 | 45 | FrameInfo info = frame_col.as_const_frame()->as_frame_info(); 46 | i32 width = info.width(); 47 | i32 height = info.height(); 48 | i32 channels = info.channels(); 49 | size_t frame_size = width * height * channels * sizeof(u8); 50 | 51 | const u8* secondary_frame_buffer = frame_col[0].as_const_frame()->data; 52 | const u8* primary_frame_buffer = frame_col[1].as_const_frame()->data; 53 | 54 | Frame* output_frame = new_frame(CPU_DEVICE, info); 55 | u8* output_buffer = output_frame->data; 56 | for (i32 y = 0; y < height; ++y) { 57 | for (i32 x = 0; x < width; ++x) { 58 | for (i32 c = 0; c < channels; ++c) { 59 | i64 offset = y * width * channels + width * channels + c; 60 | output_buffer[offset] = 61 | primary_frame_buffer[offset] - secondary_frame_buffer[offset] 62 | } 63 | } 64 | } 65 | insert_frame(output_columns[0], output_frame); 66 | } 67 | 68 | private: 69 | i32 frame_width_; 70 | i32 frame_height_; 71 | Result valid_; 72 | }; 73 | 74 | REGISTER_OP(FrameDifference).frame_input("frame").frame_output("frame"); 75 | 76 | REGISTER_KERNEL(FrameDifference, FrameDifferenceKernel) 77 | .device(DeviceType::CPU) 78 | .batch(); 79 | .stencil({-1, 0}); 80 | .num_devices(1); 81 | } 82 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/histogram_kernel_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/memory.h" 4 | #include "scanner/util/opencv.h" 5 | 6 | namespace scanner { 7 | namespace { 8 | const i32 BINS = 16; 9 | } 10 | 11 | class HistogramKernelCPU : public BatchedKernel { 12 | public: 13 | HistogramKernelCPU(const KernelConfig& config) 14 | : BatchedKernel(config), device_(config.devices[0]) {} 15 | 16 | void execute(const BatchedElements& input_columns, 17 | BatchedElements& output_columns) override { 18 | auto& frame_col = input_columns[0]; 19 | 20 | size_t hist_size = BINS * 3 * sizeof(int); 21 | i32 input_count = num_rows(frame_col); 22 | 23 | u8* output_block = new_block_buffer_size(device_, hist_size, input_count); 24 | 25 | for (i32 i = 0; i < input_count; ++i) { 26 | cv::Mat img = frame_to_mat(frame_col[i].as_const_frame()); 27 | 28 | float range[] = {0, 256}; 29 | const float* histRange = {range}; 30 | 31 | u8* output_buf = output_block + i * hist_size; 32 | 33 | for (i32 j = 0; j < 3; ++j) { 34 | int channels[] = {j}; 35 | cv::Mat hist; 36 | cv::calcHist(&img, 1, channels, cv::Mat(), 37 | hist, 38 | 1, &BINS, 39 | &histRange); 40 | cv::Mat out(BINS, 1, CV_32SC1, output_buf + j * BINS * sizeof(int)); 41 | hist.convertTo(out, CV_32SC1); 42 | } 43 | 44 | insert_element(output_columns[0], output_buf, hist_size); 45 | } 46 | } 47 | 48 | private: 49 | DeviceHandle device_; 50 | }; 51 | 52 | REGISTER_OP(Histogram).frame_input("frame").output("histogram", ColumnType::Bytes, "Histogram"); 53 | 54 | REGISTER_KERNEL(Histogram, HistogramKernelCPU) 55 | .device(DeviceType::CPU) 56 | .batch() 57 | .num_devices(1); 58 | } 59 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/histogram_kernel_gpu.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/cuda.h" 4 | #include "scanner/util/memory.h" 5 | #include "scanner/util/opencv.h" 6 | 7 | namespace scanner { 8 | namespace { 9 | const i32 BINS = 16; 10 | } 11 | 12 | class HistogramKernelGPU : public BatchedKernel, public VideoKernel { 13 | public: 14 | HistogramKernelGPU(const KernelConfig& config) 15 | : BatchedKernel(config), 16 | device_(config.devices[0]), 17 | num_cuda_streams_(32), 18 | streams_(num_cuda_streams_) {} 19 | 20 | void new_frame_info() override { 21 | set_device(); 22 | streams_.resize(0); 23 | streams_.resize(num_cuda_streams_); 24 | planes_.clear(); 25 | for (i32 i = 0; i < 3; ++i) { 26 | planes_.push_back( 27 | cvc::GpuMat(frame_info_.width(), frame_info_.height(), CV_8UC1)); 28 | } 29 | } 30 | 31 | void execute(const BatchedElements& input_columns, 32 | BatchedElements& output_columns) override { 33 | auto& frame_col = input_columns[0]; 34 | 35 | set_device(); 36 | check_frame(device_, frame_col[0]); 37 | 38 | size_t hist_size = BINS * 3 * sizeof(float); 39 | i32 input_count = num_rows(frame_col); 40 | u8* output_block = 41 | new_block_buffer(device_, hist_size * input_count, input_count); 42 | 43 | for (i32 i = 0; i < input_count; ++i) { 44 | i32 sid = i % num_cuda_streams_; 45 | cv::cuda::Stream& s = streams_[sid]; 46 | 47 | // TODO(wcrichto): implement correctly w/ streams 48 | cvc::GpuMat img = frame_to_gpu_mat(frame_col[i].as_const_frame()); 49 | cvc::split(img, planes_); 50 | 51 | u8* output_buf = output_block + i * hist_size; 52 | cvc::GpuMat out_mat(1, BINS * 3, CV_32S, output_buf); 53 | 54 | for (i32 j = 0; j < 3; ++j) { 55 | cvc::histEven(planes_[j], out_mat(cv::Rect(j * BINS, 0, BINS, 1)), BINS, 56 | 0, 256); 57 | } 58 | 59 | insert_element(output_columns[0], output_buf, hist_size); 60 | } 61 | 62 | for (cv::cuda::Stream& s : streams_) { 63 | s.waitForCompletion(); 64 | } 65 | } 66 | 67 | void set_device() { 68 | CUDA_PROTECT({ CU_CHECK(cudaSetDevice(device_.id)); }); 69 | cvc::setDevice(device_.id); 70 | } 71 | 72 | private: 73 | DeviceHandle device_; 74 | i32 num_cuda_streams_; 75 | std::vector streams_; 76 | std::vector planes_; 77 | }; 78 | 79 | REGISTER_KERNEL(Histogram, HistogramKernelGPU) 80 | .device(DeviceType::GPU) 81 | .batch() 82 | .num_devices(1); 83 | } 84 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/image_decoder_kernel_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/memory.h" 4 | #include "scanner/util/opencv.h" 5 | #include "scannertools_imgproc.pb.h" 6 | #include "scanner/util/thread_pool.h" 7 | 8 | namespace scanner { 9 | 10 | class ImageDecoderKernelCPU : public BatchedKernel { 11 | public: 12 | ImageDecoderKernelCPU(const KernelConfig& config) : BatchedKernel(config), pool_(32) {} 13 | 14 | void execute(const BatchedElements& input_columns, 15 | BatchedElements& output_columns) override { 16 | i32 input_count = num_rows(input_columns[0]); 17 | 18 | auto decode = [&](i32 i) { 19 | std::vector input_buf( 20 | input_columns[0][i].buffer, 21 | input_columns[0][i].buffer + input_columns[0][i].size); 22 | cv::Mat img = cv::imdecode(input_buf, cv::IMREAD_UNCHANGED); 23 | if (img.channels() == 4) { 24 | cv::cvtColor(img, img, cv::COLOR_BGRA2RGBA); 25 | } else if (img.channels() == 3) { 26 | cv::cvtColor(img, img, cv::COLOR_BGR2RGB); 27 | } 28 | LOG_IF(FATAL, img.empty() || !img.data) << "Failed to decode image"; 29 | return img; 30 | }; 31 | 32 | std::vector> futures; 33 | for (i32 i = 0; i < input_count; ++i) { 34 | futures.push_back(pool_.enqueue(decode, i)); 35 | } 36 | 37 | std::vector images; 38 | for (i32 i = 0; i < input_count; ++i) { 39 | images.push_back(futures[i].get()); 40 | } 41 | 42 | std::vector output_frames = new_frames(CPU_DEVICE, mat_to_frame_info(images[0]), input_count); 43 | for (i32 i = 0; i < input_count; ++i) { 44 | size_t size = images[i].total() * images[i].elemSize(); 45 | std::memcpy(output_frames[i]->data, images[i].data, size); 46 | insert_frame(output_columns[0], output_frames[i]); 47 | } 48 | } 49 | 50 | ThreadPool pool_; 51 | }; 52 | 53 | REGISTER_OP(ImageDecoder).input("img").frame_output("frame"); 54 | 55 | REGISTER_KERNEL(ImageDecoder, ImageDecoderKernelCPU) 56 | .device(DeviceType::CPU) 57 | .batch() 58 | .num_devices(1); 59 | } 60 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/image_decoder_kernel_gpu.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/memory.h" 4 | #include "scanner/util/opencv.h" 5 | #include "scanner/util/cuda.h" 6 | #include "scannertools_imgproc.pb.h" 7 | 8 | namespace scanner { 9 | 10 | namespace codec = cv::cudacodec; 11 | 12 | class ImageSource : public codec::RawVideoSource { 13 | public: 14 | ImageSource(const BatchedElements& input_columns, const cv::Mat& img) 15 | : input_columns_(input_columns), img_(img) {} 16 | 17 | bool getNextPacket(unsigned char** data, int* size, bool* endOfFile) override { 18 | const Element& element = input_columns_[0][i_]; 19 | *data = element.buffer; 20 | *size = element.size; 21 | *endOfFile = false; 22 | // Theoretically we should be able to set endOfFile to true at the last 23 | // frame, but the OpenCV VideoReader appears to return false on a valid 24 | // nextFrame request if I do this, so instead I just keep feeding packets 25 | // until the loader thread dies. 26 | i_ = (i_ + 1) % input_columns_[0].size(); 27 | return true; 28 | } 29 | 30 | codec::FormatInfo format() const override { 31 | codec::FormatInfo format_info; 32 | format_info.codec = codec::Codec::JPEG; 33 | format_info.chromaFormat = codec::ChromaFormat::YUV420; 34 | format_info.width = img_.cols; 35 | format_info.height = img_.rows; 36 | return format_info; 37 | } 38 | 39 | private: 40 | int i_ = 0; 41 | const cv::Mat& img_; 42 | const BatchedElements& input_columns_; 43 | }; 44 | 45 | class ImageDecoderKernelGPU : public Kernel { 46 | public: 47 | ImageDecoderKernelGPU(const Kernel::Config& config) 48 | : Kernel(config), device_(config.devices[0]) { 49 | if (!args_.ParseFromArray(config.args.data(), config.args.size())) { 50 | LOG(FATAL) << "Failed to parse args"; 51 | } 52 | } 53 | 54 | void execute(const BatchedElements& input_columns, 55 | BatchedElements& output_columns) override { 56 | i32 input_count = num_rows(input_columns[0]); 57 | 58 | set_device(); 59 | 60 | // Assumes all images are the same size 61 | size_t sz = input_columns[0][0].size; 62 | u8* cpu_buf = new_buffer(CPU_DEVICE, sz); 63 | memcpy_buffer(cpu_buf, CPU_DEVICE, input_columns[0][0].buffer, 64 | device_, sz); 65 | std::vector input_buf(cpu_buf, cpu_buf + sz); 66 | cv::Mat img = cv::imdecode(input_buf, cv::IMREAD_UNCHANGED); 67 | FrameInfo frame_info(img.rows, img.cols, 3, FrameType::U8); 68 | delete_buffer(CPU_DEVICE, cpu_buf); 69 | 70 | ImageDecoderArgs_ImageType image_type = args_.image_type(); 71 | std::vector frames = new_frames(device_, frame_info, input_count); 72 | 73 | cv::Ptr src = 74 | cv::Ptr(new ImageSource(input_columns, img)); 75 | cv::Ptr d_reader = codec::createVideoReader(src); 76 | 77 | for (i32 i = 0; i < input_count; ++i) { 78 | if (image_type == ImageDecoderArgs_ImageType_JPEG) { 79 | cvc::GpuMat gpu_mat = frame_to_gpu_mat(frames[i]); 80 | if (!d_reader->nextFrame(gpu_mat)) { 81 | LOG(FATAL) << "Failed to decode image"; 82 | } 83 | cvc::cvtColor(gpu_mat, gpu_mat, cv::COLOR_BGR2RGB); 84 | insert_frame(output_columns[0], frames[i]); 85 | } else if (image_type == ImageDecoderArgs_ImageType_ANY) { 86 | LOG(FATAL) << "Not yet supported"; 87 | } else { 88 | LOG(FATAL) << "Invalid image type"; 89 | } 90 | } 91 | } 92 | 93 | void set_device() { 94 | CU_CHECK(cudaSetDevice(device_.id)); 95 | cvc::setDevice(device_.id); 96 | } 97 | 98 | private: 99 | ImageDecoderArgs args_; 100 | DeviceHandle device_; 101 | }; 102 | 103 | REGISTER_KERNEL(ImageDecoder, ImageDecoderKernelGPU) 104 | .device(DeviceType::GPU) 105 | .num_devices(1); 106 | } 107 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/montage_kernel_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/memory.h" 4 | #include "scanner/util/opencv.h" 5 | #include "scannertools_imgproc.pb.h" 6 | 7 | namespace scanner { 8 | 9 | class MontageKernelCPU : public BatchedKernel, public VideoKernel { 10 | public: 11 | MontageKernelCPU(const KernelConfig& config) 12 | : BatchedKernel(config), 13 | device_(config.devices[0]), 14 | frames_seen_(0), 15 | montage_width_(0), 16 | montage_buffer_(nullptr) { 17 | valid_.set_success(true); 18 | if (!args_.ParseFromArray(config.args.data(), config.args.size())) { 19 | RESULT_ERROR(&valid_, "MontageKernel could not parse protobuf args"); 20 | return; 21 | } 22 | 23 | num_frames_ = args_.num_frames(); 24 | target_width_ = args_.target_width(); 25 | frames_per_row_ = args_.frames_per_row(); 26 | } 27 | 28 | ~MontageKernelCPU() { 29 | if (montage_buffer_ != nullptr) { 30 | delete_buffer(device_, montage_buffer_); 31 | } 32 | } 33 | 34 | void reset() { 35 | if (montage_width_ != 0) { 36 | if (montage_buffer_ != nullptr) { 37 | delete_buffer(device_, montage_buffer_); 38 | } 39 | montage_buffer_ = 40 | new_buffer(device_, montage_width_ * montage_height_ * 3); 41 | montage_image_ = 42 | cv::Mat(montage_height_, montage_width_, CV_8UC3, montage_buffer_); 43 | montage_image_.setTo(0); 44 | frames_seen_ = 0; 45 | } 46 | } 47 | 48 | void new_frame_info() override { 49 | frame_width_ = frame_info_.width(); 50 | frame_height_ = frame_info_.height(); 51 | 52 | target_height_ = (target_width_ / (1.0 * frame_width_) * frame_height_); 53 | 54 | montage_width_ = frames_per_row_ * target_width_; 55 | montage_height_ = 56 | std::ceil(num_frames_ / (1.0 * frames_per_row_)) * target_height_; 57 | reset(); 58 | } 59 | 60 | void execute(const BatchedElements& input_columns, 61 | BatchedElements& output_columns) override { 62 | auto& frame_col = input_columns[0]; 63 | check_frame(device_, frame_col[0]); 64 | 65 | assert(montage_buffer_ != nullptr); 66 | i32 input_count = num_rows(frame_col); 67 | for (i32 i = 0; i < input_count; ++i) { 68 | cv::Mat img = frame_to_mat(frame_col[i].as_const_frame()); 69 | i64 x = frames_seen_ % frames_per_row_; 70 | i64 y = frames_seen_ / frames_per_row_; 71 | cv::Mat montage_subimg = 72 | montage_image_(cv::Rect(target_width_ * x, target_height_ * y, 73 | target_width_, target_height_)); 74 | cv::resize(img, montage_subimg, cv::Size(target_width_, target_height_)); 75 | 76 | frames_seen_++; 77 | if (frames_seen_ == num_frames_) { 78 | assert(montage_buffer_ != nullptr); 79 | FrameInfo info(montage_height_, montage_width_, 3, FrameType::U8); 80 | insert_frame(output_columns[0], new Frame(info, montage_buffer_)); 81 | montage_image_ = cv::Mat(); 82 | montage_buffer_ = nullptr; 83 | } else { 84 | FrameInfo info(montage_height_, montage_width_, 3, FrameType::U8); 85 | insert_frame(output_columns[0], new_frame(device_, info)); 86 | } 87 | } 88 | } 89 | 90 | private: 91 | Result valid_; 92 | DeviceHandle device_; 93 | MontageArgs args_; 94 | i64 num_frames_; 95 | i32 frame_width_; 96 | i32 frame_height_; 97 | i32 target_width_; 98 | i32 target_height_; 99 | i32 frames_per_row_; 100 | 101 | i64 montage_width_; 102 | i64 montage_height_; 103 | 104 | u8* montage_buffer_; 105 | cv::Mat montage_image_; 106 | i64 frames_seen_; 107 | }; 108 | 109 | REGISTER_OP(Montage) 110 | .frame_input("frame") 111 | .frame_output("montage") 112 | .unbounded_state() 113 | .protobuf_name("MontageArgs"); 114 | 115 | REGISTER_KERNEL(Montage, MontageKernelCPU) 116 | .device(DeviceType::CPU) 117 | .num_devices(1); 118 | } 119 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/montage_kernel_gpu.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/cuda.h" 4 | #include "scanner/util/memory.h" 5 | #include "scanner/util/opencv.h" 6 | #include "scannertools_imgproc.pb.h" 7 | 8 | namespace scanner { 9 | 10 | class MontageKernelGPU : public BatchedKernel, public VideoKernel { 11 | public: 12 | MontageKernelGPU(const KernelConfig& config) 13 | : BatchedKernel(config), 14 | device_(config.devices[0]), 15 | frames_seen_(0), 16 | montage_width_(0), 17 | montage_buffer_(nullptr) { 18 | valid_.set_success(true); 19 | if (!args_.ParseFromArray(config.args.data(), config.args.size())) { 20 | RESULT_ERROR(&valid_, "MontageKernel could not parse protobuf args"); 21 | return; 22 | } 23 | 24 | num_frames_ = args_.num_frames(); 25 | target_width_ = args_.target_width(); 26 | frames_per_row_ = args_.frames_per_row(); 27 | } 28 | 29 | ~MontageKernelGPU() { 30 | if (montage_buffer_ != nullptr) { 31 | delete_buffer(device_, montage_buffer_); 32 | } 33 | } 34 | 35 | void reset() { 36 | set_device(); 37 | if (montage_width_ != 0) { 38 | if (montage_buffer_ != nullptr) { 39 | delete_buffer(device_, montage_buffer_); 40 | } 41 | montage_buffer_ = 42 | new_buffer(device_, montage_width_ * montage_height_ * 3); 43 | montage_image_ = cvc::GpuMat(montage_height_, montage_width_, CV_8UC3, 44 | montage_buffer_); 45 | montage_image_.setTo(0); 46 | frames_seen_ = 0; 47 | } 48 | } 49 | 50 | void new_frame_info() override { 51 | set_device(); 52 | frame_width_ = frame_info_.width(); 53 | frame_height_ = frame_info_.height(); 54 | 55 | target_height_ = (target_width_ / (1.0 * frame_width_) * frame_height_); 56 | 57 | montage_width_ = frames_per_row_ * target_width_; 58 | montage_height_ = 59 | std::ceil(num_frames_ / (1.0 * frames_per_row_)) * target_height_; 60 | reset(); 61 | } 62 | 63 | void execute(const BatchedElements& input_columns, 64 | BatchedElements& output_columns) override { 65 | auto& frame_col = input_columns[0]; 66 | check_frame(device_, frame_col[0]); 67 | 68 | set_device(); 69 | 70 | assert(montage_buffer_ != nullptr); 71 | i32 input_count = num_rows(frame_col); 72 | for (i32 i = 0; i < input_count; ++i) { 73 | cvc::GpuMat img = frame_to_gpu_mat(frame_col[i].as_const_frame()); 74 | i64 x = frames_seen_ % frames_per_row_; 75 | i64 y = frames_seen_ / frames_per_row_; 76 | cvc::GpuMat montage_subimg = 77 | montage_image_(cv::Rect(target_width_ * x, target_height_ * y, 78 | target_width_, target_height_)); 79 | cvc::resize(img, montage_subimg, cv::Size(target_width_, target_height_)); 80 | 81 | frames_seen_++; 82 | if (frames_seen_ == num_frames_) { 83 | assert(montage_buffer_ != nullptr); 84 | FrameInfo info(montage_height_, montage_width_, 3, FrameType::U8); 85 | insert_frame(output_columns[0], new Frame(info, montage_buffer_)); 86 | montage_image_ = cvc::GpuMat(); 87 | montage_buffer_ = nullptr; 88 | } else { 89 | FrameInfo info(montage_height_, montage_width_, 3, FrameType::U8); 90 | insert_frame(output_columns[0], new_frame(device_, info)); 91 | } 92 | } 93 | } 94 | 95 | void set_device() { 96 | CUDA_PROTECT({ CU_CHECK(cudaSetDevice(device_.id)); }); 97 | cvc::setDevice(device_.id); 98 | } 99 | 100 | private: 101 | Result valid_; 102 | DeviceHandle device_; 103 | MontageArgs args_; 104 | i64 num_frames_; 105 | i32 frame_width_; 106 | i32 frame_height_; 107 | i32 target_width_; 108 | i32 target_height_; 109 | i32 frames_per_row_; 110 | 111 | i64 montage_width_; 112 | i64 montage_height_; 113 | 114 | u8* montage_buffer_; 115 | cvc::GpuMat montage_image_; 116 | i64 frames_seen_; 117 | }; 118 | 119 | REGISTER_KERNEL(Montage, MontageKernelGPU) 120 | .device(DeviceType::GPU) 121 | .batch() 122 | .num_devices(1); 123 | } 124 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/optical_flow_kernel_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/memory.h" 4 | #include "scanner/util/opencv.h" 5 | 6 | #include 7 | 8 | namespace scanner { 9 | 10 | class OpticalFlowKernelCPU : public StenciledKernel, public VideoKernel { 11 | public: 12 | OpticalFlowKernelCPU(const KernelConfig& config) 13 | : StenciledKernel(config), 14 | device_(config.devices[0]) { 15 | flow_finder_ = 16 | cv::FarnebackOpticalFlow::create(3, 0.5, false, 15, 3, 5, 1.2, 0); 17 | } 18 | 19 | void new_frame_info() override { 20 | grayscale_.resize(0); 21 | for (i32 i = 0; i < 2; ++i) { 22 | grayscale_.emplace_back(frame_info_.height(), frame_info_.width(), 23 | CV_8UC1); 24 | } 25 | } 26 | 27 | void execute(const StenciledElements& input_columns, 28 | Elements& output_columns) override { 29 | auto& frame_col = input_columns[0]; 30 | check_frame(device_, frame_col[0]); 31 | 32 | FrameInfo out_frame_info(frame_info_.height(), frame_info_.width(), 2, 33 | FrameType::F32); 34 | Frame* output_frame = new_frame(device_, out_frame_info); 35 | 36 | cv::Mat input0 = frame_to_mat(frame_col[0].as_const_frame()); 37 | cv::Mat input1 = frame_to_mat(frame_col[1].as_const_frame()); 38 | cv::cvtColor(input0, grayscale_[0], cv::COLOR_BGR2GRAY); 39 | cv::cvtColor(input1, grayscale_[1], cv::COLOR_BGR2GRAY); 40 | cv::Mat flow = frame_to_mat(output_frame); 41 | flow_finder_->calc(grayscale_[0], grayscale_[1], flow); 42 | insert_frame(output_columns[0], output_frame); 43 | } 44 | 45 | private: 46 | DeviceHandle device_; 47 | cv::Ptr flow_finder_; 48 | std::vector grayscale_; 49 | }; 50 | 51 | REGISTER_OP(OpticalFlow) 52 | .frame_input("frame") 53 | .frame_output("flow") 54 | .stencil({0, 1}); 55 | 56 | REGISTER_KERNEL(OpticalFlow, OpticalFlowKernelCPU) 57 | .device(DeviceType::CPU) 58 | .num_devices(1); 59 | } 60 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/optical_flow_kernel_gpu.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/cuda.h" 4 | #include "scanner/util/cycle_timer.h" 5 | #include "scanner/util/memory.h" 6 | #include "scanner/util/opencv.h" 7 | 8 | #include 9 | 10 | namespace scanner { 11 | 12 | class OpticalFlowKernelGPU : public StenciledBatchedKernel, public VideoKernel { 13 | public: 14 | OpticalFlowKernelGPU(const KernelConfig& config) 15 | : StenciledBatchedKernel(config), 16 | device_(config.devices[0]), 17 | num_cuda_streams_(8) { 18 | set_device(); 19 | cv::cuda::setBufferPoolUsage(true); 20 | cv::cuda::setBufferPoolConfig(device_.id, 50 * 1024 * 1024, 5); 21 | streams_.resize(num_cuda_streams_); 22 | for (i32 i = 0; i < num_cuda_streams_; ++i) { 23 | flow_finders_.push_back( 24 | cvc::FarnebackOpticalFlow::create(3, 0.5, false, 15, 3, 5, 1.2, 0)); 25 | } 26 | } 27 | 28 | ~OpticalFlowKernelGPU() { 29 | set_device(); 30 | flow_finders_.clear(); 31 | streams_.clear(); 32 | cv::cuda::setBufferPoolConfig(device_.id, 0, 0); 33 | cv::cuda::setBufferPoolUsage(false); 34 | } 35 | 36 | void new_frame_info() override { 37 | set_device(); 38 | } 39 | 40 | void reset() override { 41 | set_device(); 42 | initial_frame_ = cvc::GpuMat(); 43 | } 44 | 45 | void execute(const StenciledBatchedElements& input_columns, 46 | BatchedElements& output_columns) override { 47 | set_device(); 48 | 49 | auto& frame_col = input_columns[0]; 50 | check_frame(device_, frame_col[0][0]); 51 | 52 | i32 input_count = (i32)frame_col.size(); 53 | std::vector input_frames; 54 | for (i32 i = 0; i < input_count; ++i) { 55 | input_frames.push_back(frame_col[i][0].as_const_frame()); 56 | } 57 | input_frames.push_back(frame_col.back()[1].as_const_frame()); 58 | 59 | grayscale_.resize(input_count + 1); 60 | 61 | FrameInfo out_frame_info(frame_info_.height(), frame_info_.width(), 2, 62 | FrameType::F32); 63 | std::vector output_frames = 64 | new_frames(device_, out_frame_info, input_count); 65 | 66 | for (i32 i = 0; i < input_count + 1; ++i) { 67 | i32 sidx = i % num_cuda_streams_; 68 | streams_[sidx].waitForCompletion(); 69 | cvc::GpuMat input = frame_to_gpu_mat(input_frames[i]); 70 | cvc::cvtColor(input, grayscale_[i], cv::COLOR_BGR2GRAY, 0, streams_[sidx]); 71 | } 72 | for (auto& s : streams_) { 73 | s.waitForCompletion(); 74 | } 75 | 76 | for (i32 i = 1; i < input_count + 1; ++i) { 77 | i32 sidx = i % num_cuda_streams_; 78 | 79 | i32 curr_idx = i; 80 | i32 prev_idx = (i - 1); 81 | 82 | cvc::GpuMat& input0 = grayscale_[curr_idx]; 83 | cvc::GpuMat& input1 = grayscale_[prev_idx]; 84 | 85 | //streams_[sidx].waitForCompletion(); 86 | cvc::GpuMat output_mat = frame_to_gpu_mat(output_frames[i - 1]); 87 | flow_finders_[0]->calc(input0, input1, output_mat); 88 | insert_frame(output_columns[0], output_frames[i - 1]); 89 | } 90 | for (auto& s : streams_) { 91 | s.waitForCompletion(); 92 | } 93 | } 94 | 95 | private: 96 | void set_device() { 97 | CU_CHECK(cudaSetDevice(device_.id)); 98 | cvc::setDevice(device_.id); 99 | } 100 | 101 | DeviceHandle device_; 102 | std::vector> flow_finders_; 103 | cvc::GpuMat initial_frame_; 104 | std::vector grayscale_; 105 | i32 num_cuda_streams_; 106 | std::vector streams_; 107 | }; 108 | 109 | REGISTER_KERNEL(OpticalFlow, OpticalFlowKernelGPU) 110 | .device(DeviceType::GPU) 111 | .batch() 112 | .num_devices(1); 113 | } 114 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/resize_kernel.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/cuda.h" 4 | #include "scanner/util/memory.h" 5 | #include "scanner/util/opencv.h" 6 | #include "scannertools_imgproc.pb.h" 7 | 8 | namespace scanner { 9 | namespace { 10 | const std::map INTERP_TYPES = { 11 | {u8"INTER_NEAREST", cv::INTER_NEAREST}, 12 | {u8"INTER_LINEAR", cv::INTER_LINEAR}, 13 | {u8"INTER_CUBIC", cv::INTER_CUBIC}, 14 | {u8"INTER_AREA", cv::INTER_AREA}, 15 | {u8"INTER_LANCZOS4", cv::INTER_LANCZOS4}, 16 | {u8"INTER_MAX", cv::INTER_MAX}, 17 | {u8"WARP_FILL_OUTLIERS", cv::WARP_FILL_OUTLIERS}, 18 | {u8"WARP_INVERSE_MAP", cv::WARP_INVERSE_MAP}, 19 | }; 20 | } 21 | 22 | class ResizeKernel : public BatchedKernel { 23 | public: 24 | ResizeKernel(const KernelConfig& config) 25 | : BatchedKernel(config), device_(config.devices[0]) { 26 | } 27 | 28 | void new_stream(const std::vector& args) override { 29 | args_.ParseFromArray(args.data(), args.size()); 30 | 31 | interp_type_ = cv::INTER_LINEAR; 32 | if (INTERP_TYPES.count(args_.interpolation()) > 0) { 33 | interp_type_ = INTERP_TYPES.at(args_.interpolation()); 34 | } 35 | } 36 | 37 | void execute(const BatchedElements& input_columns, 38 | BatchedElements& output_columns) override { 39 | auto& frame_col = input_columns[0]; 40 | set_device(); 41 | 42 | const Frame* frame = frame_col[0].as_const_frame(); 43 | 44 | i32 target_width = args_.width(); 45 | i32 target_height = args_.height(); 46 | if (args_.preserve_aspect()) { 47 | if (target_width == 0) { 48 | target_width = 49 | frame->width() * target_height / frame->height(); 50 | } else { 51 | target_height = 52 | frame->height() * target_width / frame->width(); 53 | } 54 | } 55 | if (args_.min()) { 56 | if (frame->width() <= target_width && 57 | frame->height() <= target_height) { 58 | target_width = frame->width(); 59 | target_height = frame->height(); 60 | } 61 | } 62 | 63 | i32 input_count = num_rows(frame_col); 64 | FrameInfo info(target_height, target_width, frame->channels(), frame->type); 65 | std::vector output_frames = new_frames(device_, info, input_count); 66 | 67 | for (i32 i = 0; i < input_count; ++i) { 68 | if (device_.type == DeviceType::CPU) { 69 | cv::Mat img = frame_to_mat(frame_col[i].as_const_frame()); 70 | cv::Mat out_mat = frame_to_mat(output_frames[i]); 71 | cv::resize(img, out_mat, cv::Size(target_width, target_height), 72 | 0, 0, 73 | interp_type_); 74 | } else { 75 | CUDA_PROTECT({ 76 | cvc::GpuMat img = frame_to_gpu_mat(frame_col[i].as_const_frame()); 77 | cvc::GpuMat out_mat = frame_to_gpu_mat(output_frames[i]); 78 | cvc::resize(img, out_mat, cv::Size(target_width, target_height), 79 | 0, 0, 80 | interp_type_); 81 | }); 82 | } 83 | insert_frame(output_columns[0], output_frames[i]); 84 | } 85 | } 86 | 87 | void set_device() { 88 | if (device_.type == DeviceType::GPU) { 89 | CUDA_PROTECT({ 90 | CU_CHECK(cudaSetDevice(device_.id)); 91 | cvc::setDevice(device_.id); 92 | }); 93 | } 94 | } 95 | 96 | private: 97 | DeviceHandle device_; 98 | ResizeArgs args_; 99 | int interp_type_; 100 | }; 101 | 102 | REGISTER_OP(Resize).frame_input("frame").frame_output("frame").stream_protobuf_name( 103 | "ResizeArgs"); 104 | 105 | REGISTER_KERNEL(Resize, ResizeKernel).device(DeviceType::CPU).batch().num_devices(1); 106 | 107 | #ifdef HAVE_CUDA 108 | REGISTER_KERNEL(Resize, ResizeKernel).device(DeviceType::GPU).batch().num_devices(1); 109 | #endif 110 | } 111 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/imgproc/scannertools_imgproc.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | message BlurArgs { 4 | int32 kernel_size = 1; 5 | float sigma = 2; 6 | } 7 | 8 | message MontageArgs { 9 | int64 num_frames = 1; 10 | int32 target_width = 4; 11 | int32 frames_per_row = 6; 12 | } 13 | 14 | 15 | enum ExtractorType { 16 | SIFT = 0; 17 | SURF = 1; 18 | } 19 | 20 | message FeatureExtractorArgs { 21 | ExtractorType feature_type = 1; 22 | } 23 | 24 | message Keypoint { 25 | float x = 1; 26 | float y = 2; 27 | } 28 | 29 | message ConvertColorArgs { 30 | string conversion = 1; 31 | } 32 | 33 | message ResizeArgs { 34 | int32 width = 1; 35 | int32 height = 2; 36 | bool min = 3; 37 | bool preserve_aspect = 4; 38 | string interpolation = 5; 39 | } 40 | 41 | message ImageDecoderArgs { 42 | enum ImageType { 43 | PNG = 0; 44 | JPEG = 1; 45 | ANY = 2; 46 | } 47 | 48 | ImageType image_type = 1; 49 | } 50 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/misc/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | set(SOURCES 2 | info_from_frame_kernel.cpp 3 | pass_kernel.cpp 4 | discard_kernel.cpp) 5 | 6 | build_op( 7 | LIB_NAME scannertools_misc 8 | CPP_SRCS ${SOURCES}) 9 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/misc/discard_kernel.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/memory.h" 4 | 5 | namespace scanner { 6 | 7 | class DiscardKernel : public BatchedKernel { 8 | public: 9 | DiscardKernel(const KernelConfig& config) 10 | : BatchedKernel(config), 11 | device_(config.devices[0]) {} 12 | 13 | void execute(const BatchedElements& input_columns, 14 | BatchedElements& output_columns) override { 15 | i32 input_count = (i32)num_rows(input_columns[0]); 16 | u8* output_block = new_block_buffer(device_, input_count, input_count); 17 | for (i32 i = 0; i < input_count; ++i) { 18 | insert_element(output_columns[0], output_block + i, 1); 19 | } 20 | } 21 | 22 | private: 23 | DeviceHandle device_; 24 | }; 25 | 26 | REGISTER_OP(Discard).input("ignore").output("dummy"); 27 | 28 | REGISTER_OP(DiscardFrame).frame_input("ignore").output("dummy"); 29 | 30 | REGISTER_KERNEL(Discard, DiscardKernel).device(DeviceType::CPU).batch().num_devices(1); 31 | 32 | REGISTER_KERNEL(Discard, DiscardKernel).device(DeviceType::GPU).batch().num_devices(1); 33 | 34 | REGISTER_KERNEL(DiscardFrame, DiscardKernel) 35 | .device(DeviceType::CPU) 36 | .batch() 37 | .num_devices(1); 38 | 39 | REGISTER_KERNEL(DiscardFrame, DiscardKernel) 40 | .device(DeviceType::GPU) 41 | .batch() 42 | .num_devices(1); 43 | } 44 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/misc/info_from_frame_kernel.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/memory.h" 4 | 5 | namespace scanner { 6 | 7 | class InfoFromFrameKernel : public BatchedKernel { 8 | public: 9 | InfoFromFrameKernel(const KernelConfig& config) 10 | : BatchedKernel(config), 11 | device_(config.devices[0]) {} 12 | 13 | void execute(const BatchedElements& input_columns, 14 | BatchedElements& output_columns) override { 15 | i32 input_count = (i32)num_rows(input_columns[0]); 16 | u8* output_block = 17 | new_block_buffer(device_, sizeof(FrameInfo) * input_count, input_count); 18 | for (i32 i = 0; i < input_count; ++i) { 19 | const Frame* frame = input_columns[0][i].as_const_frame(); 20 | 21 | u8* buffer = output_block + i * sizeof(FrameInfo); 22 | FrameInfo* info = reinterpret_cast(buffer); 23 | FrameInfo info_cpu = frame->as_frame_info(); 24 | memcpy_buffer((u8*) info, device_, 25 | (u8*) &info_cpu, CPU_DEVICE, 26 | sizeof(FrameInfo)); 27 | insert_element(output_columns[0], buffer, sizeof(FrameInfo)); 28 | } 29 | } 30 | 31 | private: 32 | DeviceHandle device_; 33 | }; 34 | 35 | REGISTER_OP(InfoFromFrame).frame_input("frame").output("frame_info"); 36 | 37 | REGISTER_KERNEL(InfoFromFrame, InfoFromFrameKernel) 38 | .device(DeviceType::CPU) 39 | .num_devices(1); 40 | 41 | REGISTER_KERNEL(InfoFromFrame, InfoFromFrameKernel) 42 | .device(DeviceType::GPU) 43 | .num_devices(1); 44 | } 45 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/misc/pass_kernel.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/memory.h" 4 | 5 | namespace scanner { 6 | 7 | class PassKernel : public BatchedKernel { 8 | public: 9 | PassKernel(const KernelConfig& config) 10 | : BatchedKernel(config), 11 | device_(config.devices[0]) {} 12 | 13 | void execute(const BatchedElements& input_columns, 14 | BatchedElements& output_columns) override { 15 | for (auto& element : input_columns[0]) { 16 | add_buffer_ref(device_, element.buffer); 17 | } 18 | output_columns[0] = input_columns[0]; 19 | } 20 | 21 | private: 22 | DeviceHandle device_; 23 | }; 24 | 25 | REGISTER_OP(Pass).input("input").output("output"); 26 | 27 | REGISTER_KERNEL(Pass, PassKernel) 28 | .device(DeviceType::CPU) 29 | .batch() 30 | .num_devices(1); 31 | 32 | REGISTER_KERNEL(Pass, PassKernel) 33 | .device(DeviceType::GPU) 34 | .batch() 35 | .num_devices(1); 36 | 37 | } 38 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/storage/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | set(SOURCES 2 | packed_file_source.cpp 3 | python_source.cpp 4 | files_source.cpp 5 | files_sink.cpp 6 | captions_source.cpp 7 | audio_source.cpp) 8 | 9 | build_op( 10 | LIB_NAME scannertools_storage 11 | CPP_SRCS ${SOURCES} 12 | PROTO_SRC scannertools_storage.proto) 13 | 14 | set(PYBIND11_PYTHON_VERSION 3) 15 | find_package(pybind11 REQUIRED) 16 | target_include_directories(scannertools_storage PUBLIC "${PYTHON_INCLUDE_DIRS}") 17 | 18 | target_compile_definitions(scannertools_storage PUBLIC -DHAVE_FFMPEG) 19 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/storage/files_sink.cpp: -------------------------------------------------------------------------------- 1 | /* Copyright 2018 Carnegie Mellon University 2 | * 3 | * Licensed under the Apache License, Version 2.0 (the "License"); 4 | * you may not use this file except in compliance with the License. 5 | * You may obtain a copy of the License at 6 | * 7 | * http://www.apache.org/licenses/LICENSE-2.0 8 | * 9 | * Unless required by applicable law or agreed to in writing, software 10 | * distributed under the License is distributed on an "AS IS" BASIS, 11 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | * See the License for the specific language governing permissions and 13 | * limitations under the License. 14 | */ 15 | 16 | #include "scanner/api/sink.h" 17 | #include "scanner/util/storehouse.h" 18 | #include "scannertools_storage.pb.h" 19 | 20 | #include "storehouse/storage_backend.h" 21 | 22 | #include 23 | #include 24 | 25 | using storehouse::StorageBackend; 26 | using storehouse::StorageConfig; 27 | using storehouse::StoreResult; 28 | using storehouse::WriteFile; 29 | 30 | namespace scanner { 31 | 32 | class FilesSink : public Sink { 33 | public: 34 | FilesSink(const SinkConfig& config) : 35 | Sink(config) { 36 | FilesSinkArgs args; 37 | if (config.args.size() == 0) { 38 | // Sane defaults 39 | args.set_storage_type("posix"); 40 | } else { 41 | bool parsed = args.ParseFromArray(config.args.data(), config.args.size()); 42 | if (!parsed) { 43 | RESULT_ERROR(&valid_, "Could not parse ColumnSinkArgs"); 44 | return; 45 | } 46 | } 47 | 48 | // Setup storagebackend using config arguments 49 | std::map storage_args; 50 | storage_args["bucket"] = args.bucket(); 51 | storage_args["region"] = args.region(); 52 | storage_args["endpoint"] = args.endpoint(); 53 | StorageConfig* sc_config = StorageConfig::make_config(args.storage_type(), storage_args); 54 | if (sc_config == nullptr) { 55 | LOG(FATAL) << "Invalid storage config"; 56 | } 57 | storage_.reset(storehouse::StorageBackend::make_from_config(sc_config)); 58 | assert(storage_.get()); 59 | } 60 | 61 | void new_stream(const std::vector& args) override { 62 | paths_.clear(); 63 | 64 | FilesSinkStreamArgs sargs; 65 | if (args.size() != 0) { 66 | bool parsed = sargs.ParseFromArray(args.data(), args.size()); 67 | if (!parsed) { 68 | RESULT_ERROR(&valid_, "Could not parse FilesSinkStreamArgs"); 69 | return; 70 | } 71 | paths_ = 72 | std::vector(sargs.paths().begin(), sargs.paths().end()); 73 | } 74 | } 75 | 76 | void write(const BatchedElements& input_columns) override { 77 | // Write the data 78 | auto write_start = now(); 79 | for (size_t i = 0; i < input_columns[0].size(); ++i) { 80 | u64 offset = input_columns[0][i].index; 81 | assert(offset < paths_.size()); 82 | std::unique_ptr file; 83 | BACKOFF_FAIL( 84 | make_unique_write_file( 85 | storage_.get(), paths_.at(offset), file), 86 | "while trying to make write file for " + paths_.at(offset)); 87 | 88 | s_write(file.get(), input_columns[0][i].buffer, input_columns[0][i].size); 89 | } 90 | profiler_->add_interval("files_sink:write", write_start, now()); 91 | } 92 | 93 | private: 94 | Result valid_; 95 | std::unique_ptr 96 | storage_; // Setup a distinct storage backend for each IO thread 97 | std::vector paths_; 98 | }; 99 | 100 | REGISTER_SINK(Files, FilesSink) 101 | .input("input") 102 | .per_element_output() 103 | .protobuf_name("FilesSinkArgs") 104 | .stream_protobuf_name("FilesSinkStreamArgs"); 105 | } // namespace scanner 106 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/storage/python_source.cpp: -------------------------------------------------------------------------------- 1 | /* Copyright 2018 Carnegie Mellon University 2 | * 3 | * Licensed under the Apache License, Version 2.0 (the "License"); 4 | * you may not use this file except in compliance with the License. 5 | * You may obtain a copy of the License at 6 | * 7 | * http://www.apache.org/licenses/LICENSE-2.0 8 | * 9 | * Unless required by applicable law or agreed to in writing, software 10 | * distributed under the License is distributed on an "AS IS" BASIS, 11 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | * See the License for the specific language governing permissions and 13 | * limitations under the License. 14 | */ 15 | 16 | #include "scanner/api/source.h" 17 | #include "scanner/api/enumerator.h" 18 | #include "scanner/util/storehouse.h" 19 | #include "scannertools_storage.pb.h" 20 | 21 | #include 22 | #include 23 | 24 | #include 25 | #include 26 | #include 27 | #include 28 | 29 | namespace scanner { 30 | 31 | namespace py = pybind11; 32 | using namespace pybind11::literals; 33 | 34 | class PythonEnumerator : public Enumerator { 35 | public: 36 | PythonEnumerator(const EnumeratorConfig& config) 37 | : Enumerator(config) { 38 | PythonEnumeratorArgs args; 39 | bool parsed = args.ParseFromArray(config.args.data(), config.args.size()); 40 | if (!parsed) { 41 | RESULT_ERROR(&valid_, "Could not parse PythonEnumeratorArgs"); 42 | return; 43 | } 44 | 45 | { 46 | // HACK(apoms): to fix this issue: https://github.com/pybind/pybind11/issues/1364 47 | pybind11::get_shared_data(""); 48 | } 49 | // Unpickle arguments and repickle per array element 50 | try { 51 | py::module main = py::module::import("__main__"); 52 | py::object scope = main.attr("__dict__"); 53 | 54 | py::module pickle = py::module::import("pickle"); 55 | 56 | // Pass pickled data into var 57 | py::bytes pickled_data(reinterpret_cast(args.data().data()), 58 | args.data().size()); 59 | 60 | std::vector unpickled_data = 61 | pickle.attr("loads")(pickled_data).cast>(); 62 | 63 | for (const auto& obj : unpickled_data) { 64 | // If this is a python 'bytes' object already, we don't need to pickle it 65 | if (py::isinstance(obj)) { 66 | pickled_data_.push_back(obj.cast()); 67 | } else { 68 | pickled_data_.push_back(pickle.attr("dumps")(obj).cast()); 69 | } 70 | } 71 | } catch (py::error_already_set& e) { 72 | LOG(FATAL) << e.what(); 73 | } catch (py::cast_error& e) { 74 | LOG(FATAL) << e.what(); 75 | } 76 | } 77 | 78 | i64 total_elements() override { 79 | return pickled_data_.size(); 80 | } 81 | 82 | ElementArgs element_args_at(i64 element_idx) override { 83 | PythonElementArgs args; 84 | args.set_data(pickled_data_.at(element_idx)); 85 | size_t size = args.ByteSizeLong(); 86 | 87 | ElementArgs element_args; 88 | element_args.args.resize(size); 89 | args.SerializeToArray(element_args.args.data(), size); 90 | element_args.row_id = element_idx; 91 | 92 | return element_args; 93 | } 94 | 95 | private: 96 | Result valid_; 97 | std::vector pickled_data_; 98 | }; 99 | 100 | class PythonSource : public Source { 101 | public: 102 | PythonSource(const SourceConfig& config) : 103 | Source(config) { 104 | PythonSourceArgs args; 105 | } 106 | 107 | void read(const std::vector& element_args, 108 | std::vector& output_columns) override { 109 | // Deserialize all ElementArgs 110 | std::vector elements; 111 | size_t total_size = 0; 112 | std::vector sizes; 113 | for (size_t i = 0; i < element_args.size(); ++i) { 114 | PythonElementArgs a; 115 | bool parsed = a.ParseFromArray(element_args[i].args.data(), 116 | element_args[i].args.size()); 117 | assert(parsed); 118 | LOG_IF(FATAL, !parsed) << "Could not parse element args in FilesSource"; 119 | 120 | elements.push_back(a.data()); 121 | 122 | total_size += a.data().size(); 123 | sizes.push_back(a.data().size()); 124 | } 125 | 126 | // Allocate a buffer for all the data 127 | u8* block_buffer = new_block_buffer(CPU_DEVICE, total_size, elements.size()); 128 | 129 | u64 offset = 0; 130 | for (size_t i = 0; i < element_args.size(); ++i) { 131 | u8* dest_buffer = block_buffer + offset; 132 | 133 | std::memcpy(dest_buffer, elements[i].data(), sizes[i]); 134 | 135 | insert_element(output_columns[0], dest_buffer, sizes[i]); 136 | 137 | offset += sizes[i]; 138 | } 139 | } 140 | 141 | private: 142 | Result valid_; 143 | }; 144 | 145 | REGISTER_ENUMERATOR(Python, PythonEnumerator) 146 | .protobuf_name("PythonEnumeratorArgs"); 147 | 148 | REGISTER_SOURCE(Python, PythonSource) 149 | .output("output") 150 | .protobuf_name("PythonSourceArgs"); 151 | 152 | } // namespace scanner 153 | -------------------------------------------------------------------------------- /scannertools/scannertools_cpp/storage/scannertools_storage.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | message PackedFileEnumeratorArgs { 4 | // For creating storehouse adapter 5 | string storage_type = 1; 6 | string bucket = 2; 7 | string region = 3; 8 | string endpoint = 4; 9 | // Path to the file to read 10 | string path = 5; 11 | } 12 | 13 | message PackedFileSourceArgs { 14 | // For creating storehouse adapter 15 | string storage_type = 1; 16 | string bucket = 2; 17 | string region = 3; 18 | string endpoint = 4; 19 | } 20 | 21 | message PackedFileElementArgs { 22 | // Path to read 23 | string path = 1; 24 | // Offset and size to read 25 | uint64 offset = 2; 26 | uint64 size = 3; 27 | } 28 | 29 | message PythonEnumeratorArgs { 30 | // Pickled list of arguments 31 | bytes data = 1; 32 | } 33 | 34 | message PythonSourceArgs { 35 | } 36 | 37 | message PythonElementArgs { 38 | // Pickled argument 39 | bytes data = 1; 40 | } 41 | 42 | 43 | message FilesEnumeratorArgs { 44 | // For creating storehouse adapter 45 | string storage_type = 1; 46 | string bucket = 2; 47 | string region = 3; 48 | string endpoint = 4; 49 | // Path to the files to read 50 | repeated string paths = 5; 51 | } 52 | 53 | message FilesSourceArgs { 54 | // For creating storehouse adapter 55 | string storage_type = 1; 56 | string bucket = 2; 57 | string region = 3; 58 | string endpoint = 4; 59 | } 60 | 61 | 62 | message FilesElementArgs { 63 | // Path to read 64 | string path = 1; 65 | } 66 | 67 | message FilesSinkArgs { 68 | // For creating storehouse adapter 69 | string storage_type = 1; 70 | string bucket = 2; 71 | string region = 3; 72 | string endpoint = 4; 73 | } 74 | 75 | message FilesSinkStreamArgs { 76 | // Path to read 77 | repeated string paths = 1; 78 | } 79 | 80 | message AudioSourceArgs { 81 | } 82 | 83 | message AudioEnumeratorArgs { 84 | string path = 1; 85 | double frame_size = 2; // seconds 86 | double duration = 3; // seconds 87 | } 88 | 89 | message AudioElementArgs { 90 | string path = 1; 91 | double frame_size = 2; // seconds 92 | } 93 | 94 | message CaptionsSourceArgs { 95 | } 96 | 97 | message CaptionsEnumeratorArgs { 98 | string path = 1; 99 | double window_size = 2; // seconds 100 | double max_time = 3; // seconds 101 | } 102 | 103 | message CaptionsElementArgs { 104 | string path = 1; 105 | double window_size = 2; // seconds 106 | } 107 | -------------------------------------------------------------------------------- /scannertools/setup.cfg: -------------------------------------------------------------------------------- 1 | [aliases] 2 | test=pytest 3 | 4 | [tool:pytest] 5 | addopts = -vvs 6 | testpaths = tests -------------------------------------------------------------------------------- /scannertools/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from scannertools_infra import CMakeExtension, CMakeBuild, CudaInstallCommand, CudaDevelopCommand 3 | 4 | if __name__ == "__main__": 5 | setup(name='scannertools', 6 | version='0.2.15', 7 | description='Video analytics toolkit', 8 | url='http://github.com/scanner-research/scannertools', 9 | author='Will Crichton', 10 | author_email='wcrichto@cs.stanford.edu', 11 | license='Apache 2.0', 12 | packages=['scannertools'], 13 | install_requires=['torch>=1.1.0', 'scipy>=1.3.0', 'scikit-learn>=0.21.0'], 14 | setup_requires=['pytest-runner', 'scannertools_infra'], 15 | tests_require=['pytest', 'requests'], 16 | cmdclass=dict(build_ext=CMakeBuild, 17 | install=CudaInstallCommand, 18 | develop=CudaDevelopCommand), 19 | ext_modules=[CMakeExtension('scannertools', 'scannertools_cpp')], 20 | zip_safe=False) 21 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe/__init__.py: -------------------------------------------------------------------------------- 1 | from scannertools_infra import _register_module 2 | 3 | _register_module(__file__, "scannertools_caffe") 4 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.2.0 FATAL_ERROR) 2 | 3 | execute_process( 4 | OUTPUT_VARIABLE SCANNER_CMAKE_PATH 5 | COMMAND python3 -c "import scannerpy.build_flags as b; b.print_cmake()") 6 | include(${SCANNER_CMAKE_PATH}) 7 | 8 | set(SOURCES 9 | caffe_kernel.cpp 10 | caffe_kernel_cpu.cpp 11 | # Removed until halide '_halide_downgrade_buffer_t' symbol issue resolved 12 | # caffe_input_kernel.cpp 13 | # caffe_input_kernel_cpu.cpp 14 | facenet_input_kernel_cpu.cpp 15 | facenet_kernel.cpp 16 | facenet_output_kernel_cpu.cpp 17 | yolo_output_kernel_cpu.cpp 18 | faster_rcnn_kernel.cpp 19 | faster_rcnn_output_kernel_cpu.cpp 20 | openpose_kernel.cpp) 21 | 22 | if (BUILD_CUDA) 23 | list(APPEND SOURCES 24 | caffe_kernel_gpu.cpp 25 | # Removed until halide '_halide_downgrade_buffer_t' symbol issue resolved 26 | #caffe_input_kernel_gpu.cpp 27 | # Removed until OpenCV4 GPU issue is resolved 28 | #facenet_input_kernel_gpu.cpp 29 | ) 30 | endif() 31 | 32 | build_op( 33 | LIB_NAME scannertools_caffe 34 | CPP_SRCS ${SOURCES} 35 | PROTO_SRC scannertools_caffe.proto 36 | BUILD_CUDA ${BUILD_CUDA}) 37 | 38 | set(LIBRARIES) 39 | set(INCLUDES) 40 | set(DEFS) 41 | 42 | list(APPEND CMAKE_MODULE_PATH "${CMAKE_SOURCE_DIR}/cmake") 43 | 44 | #halide_library(caffe_input_transformer_cpu 45 | # SRCS caffe_input_transformer_cpu.cpp) 46 | #add_dependencies(caffe caffe_input_transformer_cpu) 47 | #include_directories(${CMAKE_BINARY_DIR}/genfiles/caffe_input_transformer_cpu) 48 | if (BUILD_CUDA) 49 | # halide_library(caffe_input_transformer_gpu 50 | # SRCS caffe_input_transformer_gpu.cpp 51 | # HALIDE_TARGET cuda) 52 | # include_directories(${CMAKE_BINARY_DIR}/genfiles/caffe_input_transformer_gpu) 53 | # add_dependencies(caffe caffe_input_transformer_gpu) 54 | else() 55 | # For Caffe 56 | list(APPEND DEFS -DCPU_ONLY) 57 | 58 | endif() 59 | 60 | set(OPENCV_COMPONENTS core highgui imgproc video videoio) 61 | if (BUILD_CUDA) 62 | list(APPEND OPENCV_COMPONENTS cudaarithm cudawarping cudaimgproc) 63 | endif() 64 | find_package(OpenCV REQUIRED COMPONENTS "${OPENCV_COMPONENTS}") 65 | list(APPEND INCLUDES "${OpenCV_INCLUDE_DIRS}") 66 | list(APPEND LIBRARIES "${OpenCV_LIBRARIES}") 67 | list(APPEND DEFS -DUSE_OPENCV) 68 | 69 | find_package(Caffe REQUIRED) 70 | list(APPEND INCLUDES "${CAFFE_INCLUDE_DIRS}") 71 | list(APPEND LIBRARIES "${CAFFE_LIBRARIES}") 72 | list(APPEND DEFS -DHAVE_CAFFE) 73 | 74 | set(PYBIND11_PYTHON_VERSION 3) 75 | find_package(pybind11 REQUIRED) 76 | list(APPEND INCLUDES "${PYTHON_INCLUDE_DIRS}") 77 | 78 | find_package(OpenPose REQUIRED) 79 | list(APPEND INCLUDES "${OPENPOSE_INCLUDE_DIRS}") 80 | list(APPEND LIBRARIES "${OPENPOSE_LIBRARIES}") 81 | 82 | target_link_libraries(scannertools_caffe PUBLIC ${LIBRARIES}) 83 | target_include_directories(scannertools_caffe PUBLIC ${INCLUDES}) 84 | target_compile_definitions(scannertools_caffe PUBLIC ${DEFS}) 85 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/caffe_input_kernel.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "scanner/api/kernel.h" 4 | #include "scanner/api/op.h" 5 | #include "scanner/util/cuda.h" 6 | #include "scanner/util/opencv.h" 7 | #include "scannertools_caffe.pb.h" 8 | 9 | #ifdef HAVE_CUDA 10 | #include "caffe_input_transformer_gpu.h" 11 | #endif 12 | #include "caffe_input_transformer_cpu.h" 13 | 14 | namespace scanner { 15 | 16 | class CaffeInputKernel : public BatchedKernel, public VideoKernel { 17 | public: 18 | CaffeInputKernel(const KernelConfig& config); 19 | ~CaffeInputKernel(); 20 | 21 | void new_frame_info() override; 22 | 23 | void execute(const BatchedElements& input_columns, 24 | BatchedElements& output_columns) override; 25 | 26 | void set_device(); 27 | 28 | virtual void extra_inputs(const BatchedElements& input_columns, 29 | BatchedElements& output_columns) {} 30 | 31 | protected: 32 | void set_halide_buf(halide_buffer_t& halide_buf, u8* buf, size_t size); 33 | void unset_halide_buf(halide_buffer_t& halide_buf); 34 | void transform_halide(const u8* input_buffer, u8* output_buffer); 35 | void transform_caffe(u8* input_buffer, u8* output_buffer); 36 | 37 | DeviceHandle device_; 38 | proto::CaffeInputArgs args_; 39 | i32 net_input_width_; 40 | i32 net_input_height_; 41 | #ifdef HAVE_CUDA 42 | CUcontext context_; 43 | #endif 44 | }; 45 | } 46 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/caffe_input_kernel_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include "caffe_input_kernel.h" 2 | 3 | namespace scanner { 4 | 5 | REGISTER_OP(CaffeInput) 6 | .frame_input("frame") 7 | .frame_output("caffe_frame") 8 | .protobuf_name("CaffeInputArgs"); 9 | 10 | REGISTER_KERNEL(CaffeInput, CaffeInputKernel) 11 | .device(DeviceType::CPU) 12 | .batch() 13 | .num_devices(1); 14 | } 15 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/caffe_input_kernel_gpu.cpp: -------------------------------------------------------------------------------- 1 | #include "caffe_input_kernel.h" 2 | 3 | namespace scanner { 4 | 5 | REGISTER_KERNEL(CaffeInput, CaffeInputKernel) 6 | .device(DeviceType::GPU) 7 | .batch() 8 | .num_devices(1); 9 | } 10 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/caffe_input_transformer_base.h: -------------------------------------------------------------------------------- 1 | #include "Halide.h" 2 | 3 | #define HAS_AUTOSCHEDULER 4 | 5 | using namespace Halide; 6 | using namespace Halide::Internal; 7 | 8 | // Resize code taken from 9 | // https://github.com/halide/Halide/blob/master/apps/resize/resize.cpp 10 | Expr kernel_box(Expr x) { 11 | Expr xx = abs(x); 12 | return select(xx <= 0.5f, 1.0f, 0.0f); 13 | } 14 | 15 | Expr sinc(Expr x) { return sin(float(M_PI) * x) / x; } 16 | 17 | Expr kernel_lanczos(Expr x) { 18 | Expr value = sinc(x) * sinc(x / 3); 19 | value = select(x == 0.0f, 1.0f, value); // Take care of singularity at zero 20 | value = select(x > 3 || x < -3, 0.0f, value); // Clamp to zero out of bounds 21 | return value; 22 | } 23 | 24 | struct KernelInfo { 25 | const char* name; 26 | float size; 27 | Expr (*kernel)(Expr); 28 | }; 29 | 30 | static KernelInfo kernelInfo[] = {{"box", 0.5f, kernel_box}, 31 | // { "linear", 1.0f, kernel_linear }, 32 | // { "cubic", 2.0f, kernel_cubic }, 33 | {"lanczos", 3.0f, kernel_lanczos}}; 34 | 35 | class CaffeInputTransformer : public Halide::Generator { 36 | public: 37 | ImageParam input{UInt(8), 3, "input"}; 38 | Param input_width{"input_width"}, input_height{"input_height"}; 39 | Param target_width{"target_width"}, target_height{"target_height"}; 40 | Param normalize{"normalize"}; 41 | Param mean_r{"mean_r"}, mean_g{"mean_g"}, mean_b{"mean_b"}; 42 | 43 | Func build() { 44 | Var x("x"), y("y"), c("c"), k("k"); 45 | 46 | Func clamped = BoundaryConditions::repeat_edge(input); 47 | 48 | Func scaled("scaled"); 49 | scaled(x, y, c) = cast(clamped(x, y, c)); 50 | 51 | Expr scaleX = target_width / cast(input_width); 52 | Expr scaleY = target_height / cast(input_height); 53 | 54 | const KernelInfo& info = kernelInfo[0]; 55 | Expr kernelSizeX = info.size / scaleX; 56 | Expr kernelSizeY = info.size / scaleY; 57 | 58 | Expr sourcex = (x + 0.5f) / scaleX; 59 | Expr sourcey = (y + 0.5f) / scaleY; 60 | 61 | Func kernelx("kernelx"), kernely("kernely"); 62 | Expr beginx = cast(sourcex - kernelSizeX + 0.5f); 63 | Expr beginy = cast(sourcey - kernelSizeY + 0.5f); 64 | RDom domx(0, 2.0f * kernelSizeX + 1, "domx"); 65 | RDom domy(0, 2.0f * kernelSizeY + 1, "domy"); 66 | { 67 | Func kx, ky; 68 | kx(x, k) = info.kernel((k + beginx - sourcex) * scaleX); 69 | ky(y, k) = info.kernel((k + beginy - sourcey) * scaleY); 70 | kernelx(x, k) = kx(x, k) / sum(kx(x, domx)); 71 | kernely(y, k) = ky(y, k) / sum(ky(y, domy)); 72 | } 73 | 74 | Func resized_x("resized_x"); 75 | Func resized_y("resized_y"); 76 | resized_x(x, y, c) = sum(kernelx(x, domx) * scaled(domx + beginx, y, c)); 77 | resized_y(x, y, c) = sum(kernely(y, domy) * resized_x(x, domy + beginy, c)); 78 | 79 | Func resized_final("resized_final"); 80 | resized_final(x, y, c) = clamp(resized_y(x, y, c), 0.0f, 255.0f); 81 | 82 | Func mean_subtract("mean_subtract"); 83 | mean_subtract(x, y, c) = 84 | resized_final(x, y, c) - 85 | select(c == 0, mean_r, select(c == 1, mean_g, mean_b)); 86 | 87 | Func rescaled("rescaled"); 88 | rescaled(x, y, c) = 89 | mean_subtract(x, y, 2 - c) / select(normalize, 255.0f, 1.0f); 90 | rescaled.bound(c, 0, 3); 91 | 92 | input.dim(0).set_stride(3).dim(2).set_stride(1); 93 | 94 | Target target = Halide::get_target_from_environment(); 95 | #ifdef HALIDE_USE_GPU 96 | target.set_feature(Target::CUDA); 97 | target.set_feature(Target::CUDACapability50); 98 | resized_x.compute_root().reorder(c, x, y).unroll(c).gpu_tile(x, y, 8, 8); 99 | rescaled.reorder(c, x, y).unroll(c).gpu_tile(x, y, 8, 8); 100 | #else 101 | // Pipeline p(rescaled); 102 | // p.auto_schedule(target); 103 | #endif 104 | 105 | return rescaled; 106 | } 107 | }; 108 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/caffe_input_transformer_cpu.cpp: -------------------------------------------------------------------------------- 1 | #define HALIDE_USE_CPU 2 | #include "caffe_input_transformer_base.h" 3 | 4 | HALIDE_REGISTER_GENERATOR(CaffeInputTransformer, caffe_input_transformer_cpu); 5 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/caffe_input_transformer_gpu.cpp: -------------------------------------------------------------------------------- 1 | #define HALIDE_USE_GPU 2 | #include "caffe_input_transformer_base.h" 3 | HALIDE_REGISTER_GENERATOR(CaffeInputTransformer, caffe_input_transformer_gpu); 4 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/caffe_kernel.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "scanner/api/kernel.h" 4 | #include "scanner/api/op.h" 5 | #include "scanner/util/cuda.h" 6 | #include "scanner/util/memory.h" 7 | #include "scannertools_caffe.pb.h" 8 | 9 | #include "caffe/blob.hpp" 10 | #include "caffe/common.hpp" 11 | #include "caffe/net.hpp" 12 | #include "caffe/proto/caffe.pb.h" 13 | #include "caffe/util/db.hpp" 14 | #include "caffe/util/io.hpp" 15 | 16 | namespace scanner { 17 | 18 | using CustomNetConfiguration = void (*)(const FrameInfo& frame_info, 19 | caffe::Net* net); 20 | 21 | class CaffeKernel : public BatchedKernel, public VideoKernel { 22 | public: 23 | CaffeKernel(const KernelConfig& config); 24 | void validate(proto::Result* result) override; 25 | void new_frame_info() override; 26 | void execute(const BatchedElements& input_columns, 27 | BatchedElements& output_columns) override; 28 | void set_device(); 29 | 30 | virtual void net_config() {} 31 | 32 | protected: 33 | proto::Result valid_; 34 | DeviceHandle device_; 35 | proto::CaffeArgs args_; 36 | std::unique_ptr> net_; 37 | CustomNetConfiguration net_config_; 38 | }; 39 | 40 | proto::NetDescriptor descriptor_from_net_file(const std::string& path); 41 | } 42 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/caffe_kernel_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include "caffe_kernel.h" 2 | 3 | namespace scanner { 4 | 5 | REGISTER_OP(Caffe) 6 | .frame_input("caffe_frame") 7 | .frame_output("caffe_output") 8 | .protobuf_name("CaffeArgs"); 9 | 10 | REGISTER_KERNEL(Caffe, CaffeKernel) 11 | .device(DeviceType::CPU) 12 | .batch() 13 | .num_devices(Kernel::UnlimitedDevices); 14 | } 15 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/caffe_kernel_gpu.cpp: -------------------------------------------------------------------------------- 1 | #include "caffe_kernel.h" 2 | 3 | namespace scanner { 4 | 5 | REGISTER_KERNEL(Caffe, CaffeKernel) 6 | .device(DeviceType::GPU) 7 | .batch() 8 | .num_devices(1); 9 | } 10 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/cmake/FindCaffe.cmake: -------------------------------------------------------------------------------- 1 | # - Try to find Caffe 2 | # 3 | # The following variables are optionally searched for defaults 4 | # CAFFE_ROOT_DIR: Base directory where all Caffe components are found 5 | # 6 | # The following are set after configuration is done: 7 | # CAFFE_FOUND 8 | # CAFFE_INCLUDE_DIRS 9 | # CAFFE_LIBRARIES 10 | # CAFFE_LIBRARY_DIRS 11 | 12 | include(FindPackageHandleStandardArgs) 13 | 14 | set(CAFFE_ROOT_DIR "" CACHE PATH "Folder contains Caffe") 15 | 16 | if (NOT "$ENV{Caffe_DIR}" STREQUAL "") 17 | set(CAFFE_ROOT_DIR $ENV{Caffe_DIR}) 18 | endif() 19 | 20 | # We are testing only a couple of files in the include directories 21 | if(WIN32) 22 | find_path(CAFFE_INCLUDE_DIR caffe/caffe.hpp 23 | PATHS ${CAFFE_ROOT_DIR}/src/windows) 24 | else() 25 | find_path(CAFFE_INCLUDE_DIR caffe/caffe.hpp 26 | PATHS ${CAFFE_ROOT_DIR}/include) 27 | endif() 28 | 29 | find_library(CAFFE_LIBRARY caffe PATHS ${CAFFE_ROOT_DIR}/lib) 30 | 31 | find_package_handle_standard_args(CAFFE DEFAULT_MSG 32 | CAFFE_INCLUDE_DIR CAFFE_LIBRARY) 33 | 34 | if(CAFFE_FOUND) 35 | set(CAFFE_INCLUDE_DIRS ${CAFFE_INCLUDE_DIR}) 36 | set(CAFFE_LIBRARIES ${CAFFE_LIBRARY}) 37 | endif() 38 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/cmake/FindOpenPose.cmake: -------------------------------------------------------------------------------- 1 | # - Try to find OpenPose 2 | # 3 | # The following variables are optionally searched for defaults 4 | # OPENPOSE_ROOT_DIR: Base directory where all Caffe components are found 5 | # 6 | # The following are set after configuration is done: 7 | # OPENPOSE_FOUND 8 | # OPENPOSE_INCLUDE_DIRS 9 | # OPENPOSE_LIBRARIES 10 | # OPENPOSE_LIBRARY_DIRS 11 | 12 | include(FindPackageHandleStandardArgs) 13 | 14 | set(OPENPOSE_ROOT_DIR "" CACHE PATH "Folder contains OpenPose") 15 | 16 | if (NOT "$ENV{OpenPose_DIR}" STREQUAL "") 17 | set(OPENPOSE_ROOT_DIR $ENV{OpenPose_DIR}) 18 | endif() 19 | 20 | # We are testing only a couple of files in the include directories 21 | if(WIN32) 22 | find_path(OPENPOSE_INCLUDE_DIR openpose/headers.hpp 23 | PATHS ${OPENPOSE_ROOT_DIR}/src/windows) 24 | else() 25 | find_path(OPENPOSE_INCLUDE_DIR openpose/headers.hpp 26 | PATHS ${OPENPOSE_ROOT_DIR}/include) 27 | endif() 28 | 29 | find_library(OPENPOSE_LIBRARY openpose PATHS ${OPENPOSE_ROOT_DIR}/lib) 30 | 31 | find_package_handle_standard_args(OPENPOSE DEFAULT_MSG 32 | OPENPOSE_INCLUDE_DIR OPENPOSE_LIBRARY) 33 | 34 | if(OPENPOSE_FOUND) 35 | set(OPENPOSE_INCLUDE_DIRS ${OPENPOSE_INCLUDE_DIR}) 36 | set(OPENPOSE_LIBRARIES ${OPENPOSE_LIBRARY}) 37 | endif() 38 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/cpm2_kernel.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/op.h" 2 | #include "caffe_kernel.h" 3 | 4 | #include "caffe/cpm/layers/imresize_layer.hpp" 5 | 6 | namespace scanner { 7 | 8 | class CPM2Kernel : public CaffeKernel { 9 | public: 10 | CPM2Kernel(const KernelConfig& config) 11 | : CaffeKernel(get_caffe_config(config)) {} 12 | 13 | void net_config() override { 14 | int net_input_width = frame_info_.shape[2]; 15 | int net_input_height = frame_info_.shape[1]; 16 | 17 | caffe::ImResizeLayer* resize_layer = 18 | (caffe::ImResizeLayer*)net_->layer_by_name("resize").get(); 19 | 20 | resize_layer->SetStartScale(1); 21 | resize_layer->SetScaleGap(0.1); 22 | resize_layer->setTargetDimenions(net_input_width, net_input_height); 23 | 24 | const boost::shared_ptr> input_blob{ 25 | net_->blob_by_name("image")}; 26 | input_blob->Reshape({input_blob->shape(0), input_blob->shape(1), 27 | net_input_height, net_input_width}); 28 | } 29 | 30 | KernelConfig get_caffe_config(const KernelConfig& config) { 31 | proto::CPM2Args args; 32 | args.ParseFromArray(config.args.data(), config.args.size()); 33 | scale_ = args.scale(); 34 | 35 | KernelConfig new_config(config); 36 | std::string caffe_string; 37 | args.caffe_args().SerializeToString(&caffe_string); 38 | new_config.args = std::vector(caffe_string.begin(), caffe_string.end()); 39 | return new_config; 40 | } 41 | 42 | private: 43 | f32 scale_; 44 | }; 45 | 46 | REGISTER_OP(CPM2) 47 | .frame_input("cpm2_input") 48 | .frame_output("cpm2_resized_map") 49 | .frame_output("cpm2_joints"); 50 | 51 | REGISTER_KERNEL(CPM2, CPM2Kernel).device(DeviceType::CPU).num_devices(1).batch(); 52 | REGISTER_KERNEL(CPM2, CPM2Kernel).device(DeviceType::GPU).num_devices(1).batch(); 53 | } 54 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/facenet_input_kernel_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/util/memory.h" 4 | #include "scanner/util/opencv.h" 5 | #include "scannertools_caffe.pb.h" 6 | 7 | namespace scanner { 8 | 9 | class FacenetInputKernelCPU : public BatchedKernel, public VideoKernel { 10 | public: 11 | FacenetInputKernelCPU(const KernelConfig& config) 12 | : BatchedKernel(config), 13 | device_(config.devices[0]) 14 | { 15 | proto::FacenetArgs args; 16 | args.ParseFromArray(config.args.data(), config.args.size()); 17 | args_.CopyFrom(args.caffe_args()); 18 | scale_ = args.scale(); 19 | } 20 | 21 | void new_frame_info() override { 22 | net_input_width_ = std::floor(frame_info_.width() * scale_); 23 | net_input_height_ = std::floor(frame_info_.height() * scale_); 24 | if (net_input_width_ % 8 != 0) { 25 | net_input_width_ += 8 - (net_input_width_ % 8); 26 | }; 27 | if (net_input_height_ % 8 != 0) { 28 | net_input_height_ += 8 - (net_input_height_ % 8); 29 | } 30 | 31 | mean_mat_ = 32 | cv::Mat(net_input_height_, net_input_width_, CV_32FC3, 33 | cv::Scalar(args_.net_descriptor().mean_colors(0), 34 | args_.net_descriptor().mean_colors(1), 35 | args_.net_descriptor().mean_colors(2))); 36 | 37 | frame_input_.clear(); 38 | resized_input_.clear(); 39 | float_input_.clear(); 40 | flipped_planes_.clear(); 41 | normalized_input_.clear(); 42 | input_planes_.clear(); 43 | planar_input_.clear(); 44 | flipped_planes_.clear(); 45 | for (size_t i = 0; i < 1; ++i) { 46 | frame_input_.push_back( 47 | cv::Mat(frame_info_.height(), frame_info_.width(), CV_8UC3)); 48 | resized_input_.push_back( 49 | cv::Mat(net_input_height_, net_input_width_, CV_8UC3)); 50 | float_input_.push_back( 51 | cv::Mat(net_input_height_, net_input_width_, CV_32FC3)); 52 | normalized_input_.push_back( 53 | cv::Mat(net_input_height_, net_input_width_, CV_32FC3)); 54 | std::vector planes; 55 | std::vector flipped_planes; 56 | for (i32 i = 0; i < 3; ++i) { 57 | planes.push_back( 58 | cv::Mat(net_input_height_, net_input_width_, CV_32FC1)); 59 | flipped_planes.push_back( 60 | cv::Mat(net_input_width_, net_input_height_, CV_32FC1)); 61 | } 62 | input_planes_.push_back(planes); 63 | flipped_planes_.push_back(flipped_planes); 64 | planar_input_.push_back( 65 | cv::Mat(net_input_width_ * 3, net_input_height_, CV_32FC1)); 66 | } 67 | } 68 | 69 | void execute(const BatchedElements& input_columns, 70 | BatchedElements& output_columns) override { 71 | auto& frame_col = input_columns[0]; 72 | check_frame(device_, frame_col[0]); 73 | 74 | i32 input_count = (i32)frame_col.size(); 75 | FrameInfo net_input_info(3, net_input_width_, net_input_height_, 76 | FrameType::F32); 77 | i32 net_input_size = net_input_info.size(); 78 | std::vector output_frames = 79 | new_frames(device_, net_input_info, input_count); 80 | 81 | for (i32 i = 0; i < input_count; ++i) { 82 | Frame* output_frame = output_frames[i]; 83 | f32* net_input = (f32*)output_frame->data; 84 | 85 | i32 sid = 0; 86 | 87 | // Convert input frame to gpu mat 88 | frame_input_[sid] = frame_to_mat(frame_col[i].as_const_frame()); 89 | 90 | cv::resize(frame_input_[sid], resized_input_[sid], 91 | cv::Size(net_input_width_, net_input_height_), 0, 0, 92 | cv::INTER_LINEAR); 93 | resized_input_[sid].convertTo(float_input_[sid], CV_32FC3); 94 | cv::subtract(float_input_[sid], mean_mat_, normalized_input_[sid], 95 | cv::noArray(), -1); 96 | // Changed from interleaved RGB to planar RGB 97 | cv::split(normalized_input_[sid], input_planes_[sid]); 98 | cv::transpose(input_planes_[sid][0], flipped_planes_[sid][0]); 99 | cv::transpose(input_planes_[sid][1], flipped_planes_[sid][1]); 100 | cv::transpose(input_planes_[sid][2], flipped_planes_[sid][2]); 101 | auto& plane1 = flipped_planes_[sid][0]; 102 | auto& plane2 = flipped_planes_[sid][1]; 103 | auto& plane3 = flipped_planes_[sid][2]; 104 | auto& planar_input = planar_input_[sid]; 105 | plane1.copyTo(planar_input(cv::Rect( 106 | 0, net_input_width_ * 0, net_input_height_, net_input_width_))); 107 | plane2.copyTo(planar_input(cv::Rect( 108 | 0, net_input_width_ * 1, net_input_height_, net_input_width_))); 109 | plane3.copyTo(planar_input(cv::Rect( 110 | 0, net_input_width_ * 2, net_input_height_, net_input_width_))); 111 | assert(planar_input.cols == net_input_height_); 112 | for (int j = 0; j < net_input_width_ * 3; ++j) { 113 | memcpy(net_input + j * net_input_height_, 114 | planar_input.data + j * planar_input.step, 115 | net_input_height_ * sizeof(float)); 116 | } 117 | insert_frame(output_columns[0], output_frame); 118 | } 119 | } 120 | 121 | private: 122 | DeviceHandle device_; 123 | proto::CaffeArgs args_; 124 | f32 scale_; 125 | i32 net_input_width_; 126 | i32 net_input_height_; 127 | 128 | cv::Mat mean_mat_; 129 | std::vector frame_input_; 130 | std::vector resized_input_; 131 | std::vector float_input_; 132 | std::vector normalized_input_; 133 | std::vector> input_planes_; 134 | std::vector> flipped_planes_; 135 | std::vector planar_input_; 136 | }; 137 | 138 | REGISTER_OP(FacenetInput) 139 | .frame_input("frame") 140 | .frame_output("facenet_input") 141 | .protobuf_name("FacenetArgs"); 142 | 143 | REGISTER_KERNEL(FacenetInput, FacenetInputKernelCPU) 144 | .device(DeviceType::CPU) 145 | .num_devices(1); 146 | } 147 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/facenet_kernel.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/op.h" 2 | #include "caffe_kernel.h" 3 | 4 | namespace scanner { 5 | 6 | class FacenetKernel : public CaffeKernel { 7 | public: 8 | FacenetKernel(const KernelConfig& config) 9 | : CaffeKernel(get_caffe_config(config)) {} 10 | 11 | void net_config() override { 12 | int net_input_width = frame_info_.shape[1]; 13 | int net_input_height = frame_info_.shape[2]; 14 | 15 | const boost::shared_ptr> input_blob{ 16 | net_->blob_by_name("data")}; 17 | input_blob->Reshape({input_blob->shape(0), input_blob->shape(1), 18 | net_input_width, net_input_height}); 19 | } 20 | 21 | KernelConfig get_caffe_config(const KernelConfig& config) { 22 | proto::FacenetArgs args; 23 | args.ParseFromArray(config.args.data(), config.args.size()); 24 | scale_ = args.scale(); 25 | 26 | KernelConfig new_config(config); 27 | std::string caffe_string; 28 | args.caffe_args().SerializeToString(&caffe_string); 29 | new_config.args = std::vector(caffe_string.begin(), caffe_string.end()); 30 | return new_config; 31 | } 32 | 33 | private: 34 | f32 scale_; 35 | }; 36 | 37 | REGISTER_OP(Facenet) 38 | .frame_input("facenet_input") 39 | .frame_output("facenet_output") 40 | .protobuf_name("FacenetArgs"); 41 | 42 | REGISTER_KERNEL(Facenet, FacenetKernel) 43 | .device(DeviceType::CPU) 44 | .num_devices(Kernel::UnlimitedDevices); 45 | 46 | REGISTER_KERNEL(Facenet, FacenetKernel).device(DeviceType::GPU).num_devices(1); 47 | } 48 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/faster_rcnn_kernel.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/op.h" 2 | #include "caffe_kernel.h" 3 | 4 | namespace scanner { 5 | 6 | class FasterRCNNKernel : public CaffeKernel { 7 | public: 8 | FasterRCNNKernel(const KernelConfig& config) : CaffeKernel(config) {} 9 | 10 | void net_config() override { 11 | boost::shared_ptr> blob = net_->blob_by_name("im_info"); 12 | f32 buf[3] = {(f32)frame_info_.shape[2], (f32)frame_info_.shape[1], 1.0}; 13 | f32* blob_data = device_.type == DeviceType::GPU ? blob->mutable_gpu_data() 14 | : blob->mutable_cpu_data(); 15 | memcpy_buffer((u8*)blob_data, device_, (u8*)buf, CPU_DEVICE, 16 | 3 * sizeof(f32)); 17 | } 18 | }; 19 | 20 | REGISTER_OP(FasterRCNN) 21 | .frame_input("caffe_input") 22 | .frame_output("cls_prob") 23 | .frame_output("rois") 24 | .frame_output("fc7") 25 | .protobuf_name("CaffeArgs"); 26 | 27 | REGISTER_KERNEL(FasterRCNN, FasterRCNNKernel) 28 | .device(DeviceType::CPU) 29 | .num_devices(1); 30 | REGISTER_KERNEL(FasterRCNN, FasterRCNNKernel) 31 | .device(DeviceType::GPU) 32 | .num_devices(1); 33 | } 34 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/faster_rcnn_output_kernel_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/api/kernel.h" 2 | #include "scanner/api/op.h" 3 | #include "scanner/types.pb.h" 4 | #include "scanner/util/bbox.h" 5 | #include "scanner/util/opencv.h" 6 | #include "scanner/util/serialize.h" 7 | #include "scannertools_caffe.pb.h" 8 | 9 | namespace scanner { 10 | 11 | #define CLASSES 81 12 | #define SCORE_THRESHOLD 0.7 13 | #define BOX_SIZE 5 14 | #define FEATURES 4096 15 | 16 | class FasterRCNNOutputKernel : public BatchedKernel { 17 | public: 18 | FasterRCNNOutputKernel(const KernelConfig& config) : BatchedKernel(config) {} 19 | 20 | void execute(const BatchedElements& input_columns, 21 | BatchedElements& output_columns) override { 22 | assert(input_columns.size() == 3); 23 | 24 | i32 input_count = num_rows(input_columns[0]); 25 | i32 cls_prob_idx = 0; 26 | i32 rois_idx = 1; 27 | i32 fc7_idx = 2; 28 | const Elements &cls_prob = input_columns[cls_prob_idx], 29 | &rois = input_columns[rois_idx], 30 | &fc7 = input_columns[fc7_idx]; 31 | 32 | for (i32 i = 0; i < input_count; ++i) { 33 | const Frame* cls_p = cls_prob[i].as_const_frame(); 34 | const Frame* roi = rois[i].as_const_frame(); 35 | const Frame* fc = fc7[i].as_const_frame(); 36 | 37 | i32 proposal_count = roi->size() / (BOX_SIZE * sizeof(f32)); 38 | assert(roi->size() == BOX_SIZE * sizeof(f32) * proposal_count); 39 | assert(cls_p->size() == CLASSES * sizeof(f32) * proposal_count); 40 | std::vector bboxes; 41 | for (i32 j = 0; j < proposal_count; ++j) { 42 | f32* ro = (f32*)(roi->data + (j * BOX_SIZE * sizeof(f32))); 43 | f32 x1 = ro[1], y1 = ro[2], x2 = ro[3], y2 = ro[4]; 44 | 45 | BoundingBox bbox; 46 | bbox.set_x1(x1); 47 | bbox.set_y1(y1); 48 | bbox.set_x2(x2); 49 | bbox.set_y2(y2); 50 | 51 | f32 max_score = std::numeric_limits::min(); 52 | i32 max_cls = 0; 53 | // Start at cls = 1 to skip background 54 | for (i32 cls = 1; cls < CLASSES; ++cls) { 55 | f32* scores = 56 | (f32*)(cls_p->data + (j * CLASSES * sizeof(f32))); 57 | f32 score = scores[cls]; 58 | if (score > max_score) { 59 | max_score = score; 60 | max_cls = cls; 61 | } 62 | } 63 | 64 | if (max_score > SCORE_THRESHOLD) { 65 | assert(max_cls != 0); 66 | bbox.set_score(max_score); 67 | bbox.set_track_id(j); 68 | bbox.set_label(max_cls); 69 | bboxes.push_back(bbox); 70 | } 71 | } 72 | 73 | std::vector best_bboxes; 74 | best_bboxes = best_nms(bboxes, 0.3); 75 | 76 | { 77 | size_t size; 78 | u8* buffer; 79 | serialize_bbox_vector(best_bboxes, buffer, size); 80 | insert_element(output_columns[0], buffer, size); 81 | } 82 | 83 | if (best_bboxes.size() == 0) { 84 | u8* buffer = new_buffer(CPU_DEVICE, 1); 85 | insert_element(output_columns[1], buffer, 1); 86 | } else { 87 | { 88 | size_t size = 89 | std::max(best_bboxes.size() * FEATURES * sizeof(f32), (size_t)1); 90 | u8* buffer = new_buffer(CPU_DEVICE, size); 91 | for (i32 k = 0; k < best_bboxes.size(); ++k) { 92 | i32 j = best_bboxes[k].track_id(); 93 | f32* fvec = (f32*)(fc->data + (j * FEATURES * sizeof(f32))); 94 | std::memcpy(buffer + (k * FEATURES * sizeof(f32)), fvec, 95 | FEATURES * sizeof(f32)); 96 | } 97 | insert_element(output_columns[1], buffer, size); 98 | } 99 | } 100 | } 101 | } 102 | }; 103 | 104 | REGISTER_OP(FasterRCNNOutput) 105 | .frame_input("cls_prob") 106 | .frame_input("rois") 107 | .frame_input("fc7") 108 | .output("bboxes") 109 | .output("features"); 110 | 111 | REGISTER_KERNEL(FasterRCNNOutput, FasterRCNNOutputKernel) 112 | .device(DeviceType::CPU) 113 | .num_devices(Kernel::UnlimitedDevices); 114 | } 115 | -------------------------------------------------------------------------------- /scannertools_caffe/scannertools_caffe_cpp/scannertools_caffe.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package scanner.proto; 4 | 5 | message NetDescriptor { 6 | string model_path = 1; 7 | string model_weights_path = 2; 8 | 9 | repeated string input_layer_names = 3; 10 | repeated string output_layer_names = 4; 11 | 12 | int32 input_width = 5; 13 | int32 input_height = 6; 14 | 15 | repeated float mean_colors = 7; 16 | // or 17 | repeated float mean_image = 8; 18 | int32 mean_width = 9; 19 | int32 mean_height = 10; 20 | 21 | bool normalize = 11; 22 | bool preserve_aspect_ratio = 12; 23 | bool transpose = 13; 24 | int32 pad_mod = 14; 25 | bool uses_python = 15; 26 | } 27 | 28 | message CaffeInputArgs { 29 | NetDescriptor net_descriptor = 1; 30 | int32 batch_size = 2; 31 | } 32 | 33 | message CaffeArgs { 34 | NetDescriptor net_descriptor = 1; 35 | int32 batch_size = 2; 36 | } 37 | 38 | message FacenetArgs { 39 | CaffeArgs caffe_args = 1; 40 | string templates_path = 2; 41 | float scale = 3; 42 | float threshold = 4; 43 | } 44 | 45 | message CPM2Args { 46 | CaffeArgs caffe_args = 1; 47 | float scale = 2; 48 | } 49 | 50 | message OpenPoseArgs { 51 | string model_directory = 1; 52 | 53 | int32 pose_num_scales = 2; 54 | float pose_scale_gap = 3; 55 | 56 | bool compute_hands = 4; 57 | int32 hand_num_scales = 5; 58 | float hand_scale_gap = 6; 59 | 60 | bool compute_face = 7; 61 | } -------------------------------------------------------------------------------- /scannertools_caffe/setup.cfg: -------------------------------------------------------------------------------- 1 | [aliases] 2 | test=pytest 3 | 4 | [tool:pytest] 5 | addopts = -vvs 6 | testpaths = tests -------------------------------------------------------------------------------- /scannertools_caffe/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from scannertools_infra import CMakeExtension, CMakeBuild, CudaInstallCommand, CudaDevelopCommand 3 | 4 | if __name__ == "__main__": 5 | setup(name='scannertools_caffe', 6 | version='0.2.15', 7 | description='Video analytics toolkit', 8 | url='http://github.com/scanner-research/scannertools', 9 | author='Will Crichton', 10 | author_email='wcrichto@cs.stanford.edu', 11 | license='Apache 2.0', 12 | packages=['scannertools_caffe'], 13 | install_requires=[], 14 | setup_requires=['pytest-runner', 'scannertools_infra'], 15 | tests_require=['pytest', 'scannertools_infra'], 16 | cmdclass=dict(build_ext=CMakeBuild, 17 | install=CudaInstallCommand, 18 | develop=CudaDevelopCommand), 19 | ext_modules=[CMakeExtension('scannertools_caffe', 'scannertools_caffe_cpp')], 20 | zip_safe=False) 21 | -------------------------------------------------------------------------------- /scannertools_caffe/tests/test_all.py: -------------------------------------------------------------------------------- 1 | from scannertools_infra.tests import sc, needs_gpu 2 | from scannerpy.storage import NamedVideoStream, NamedStream 3 | from scannerpy import PerfParams, DeviceType 4 | import scannertools_caffe.pose_detection 5 | 6 | @needs_gpu() 7 | def test_pose(sc): 8 | vid = [NamedVideoStream(sc, 'test1')] 9 | frame = sc.io.Input(vid) 10 | frame_sample = sc.streams.Gather(frame, [list(range(0, 1000, 100))]) 11 | pose = sc.ops.OpenPose( 12 | frame=frame_sample, 13 | device=DeviceType.GPU, 14 | pose_num_scales=6, 15 | pose_scale_gap=0.16, 16 | compute_hands=True, 17 | hand_num_scales=6, 18 | hand_scale_gap=0.16, 19 | compute_face=True, 20 | batch=5 21 | ) 22 | output = NamedStream(sc, 'test1-pose') 23 | output_op = sc.io.Output(pose, [output]) 24 | 25 | sc.run(output_op, PerfParams.estimate()) 26 | -------------------------------------------------------------------------------- /scannertools_infra/scannertools_infra/__init__.py: -------------------------------------------------------------------------------- 1 | from scannerpy.op import register_module 2 | from setuptools import Extension 3 | from setuptools.command.build_ext import build_ext 4 | from setuptools.command.install import install 5 | from setuptools.command.develop import develop 6 | from scannerpy import protobufs 7 | from multiprocessing import cpu_count 8 | import os 9 | import subprocess as sp 10 | import sys 11 | 12 | 13 | build_cuda = None 14 | 15 | 16 | class CudaOptMixin(object): 17 | user_options = install.user_options + [ 18 | ('build-cuda=', None, 'Path to Cuda install, e.g. /usr/local/cuda') 19 | ] 20 | 21 | def initialize_options(self): 22 | super().initialize_options() 23 | self.build_cuda = None 24 | 25 | def run(self): 26 | global build_cuda 27 | build_cuda = self.build_cuda 28 | super().run() 29 | 30 | 31 | class CudaInstallCommand(CudaOptMixin, install): 32 | user_options = getattr(install, 'user_options', []) + CudaOptMixin.user_options 33 | 34 | 35 | class CudaDevelopCommand(CudaOptMixin, develop): 36 | user_options = getattr(develop, 'user_options', []) + CudaOptMixin.user_options 37 | 38 | 39 | # See link below for explanation of CMake extensions 40 | # http://www.benjack.io/2018/02/02/python-cpp-revisited.html 41 | class CMakeExtension(Extension): 42 | def __init__(self, name, sourcedir=''): 43 | Extension.__init__(self, name, sources=[]) 44 | self.sourcedir = os.path.abspath(sourcedir) 45 | 46 | 47 | class CMakeBuild(build_ext): 48 | def run(self): 49 | try: 50 | out = sp.check_output(['cmake', '--version']) 51 | except OSError: 52 | raise RuntimeError( 53 | "CMake must be installed to build the following extensions: " + 54 | ", ".join(e.name for e in self.extensions)) 55 | 56 | for ext in self.extensions: 57 | self.build_extension(ext) 58 | 59 | def build_extension(self, ext): 60 | output_dir = os.path.join(ext.sourcedir, '..', ext.name, 'build') 61 | cfg = 'Debug' if self.debug else 'RelWithDebugInfo' 62 | build_args = ['--config', cfg, '--', '-j{}'.format(cpu_count())] 63 | cmake_args = [ 64 | '-DCMAKE_INSTALL_PREFIX=' + output_dir, 65 | '-DCMAKE_BUILD_TYPE=' + cfg 66 | ] 67 | 68 | if build_cuda is not None: 69 | cmake_args.extend([ 70 | '-DBUILD_CUDA=ON', 71 | '-DCUDA_TOOLKIT_ROOT_DIR=' + build_cuda 72 | ]) 73 | 74 | env = os.environ.copy() 75 | env['CXXFLAGS'] = '{} -DVERSION_INFO=\\"{}\\"'.format( 76 | env.get('CXXFLAGS', ''), 77 | self.distribution.get_version()) 78 | if not os.path.exists(self.build_temp): 79 | os.makedirs(self.build_temp) 80 | sp.check_call( 81 | ['cmake', ext.sourcedir] + cmake_args, 82 | cwd=self.build_temp, env=env) 83 | sp.check_call( 84 | ['cmake', '--build', '.'] + build_args, 85 | cwd=self.build_temp) 86 | sp.check_call(['make', 'install'], cwd=self.build_temp) 87 | print() # Add an empty line for cleaner output 88 | 89 | 90 | def _register_module(path, module): 91 | cwd = os.path.dirname(os.path.abspath(path)) 92 | so_path = os.path.join(cwd, 'build/lib{}.so'.format(module)) 93 | proto_path = os.path.join(cwd, 'build/{}_pb2.py'.format(module)) 94 | 95 | if not os.path.isfile(so_path): 96 | raise Exception("Error: missing so file {}".format(so_path)) 97 | 98 | register_module(so_path, proto_path if os.path.isfile(proto_path) else None) 99 | if os.path.isfile(proto_path): 100 | protobufs.add_module(proto_path) 101 | -------------------------------------------------------------------------------- /scannertools_infra/scannertools_infra/tests.py: -------------------------------------------------------------------------------- 1 | from scannerpy import Config, Client 2 | import pytest 3 | import tempfile 4 | import socket 5 | import requests 6 | import toml 7 | from subprocess import check_call as run 8 | import subprocess 9 | import GPUtil 10 | 11 | def needs_gpu(): 12 | try: 13 | return pytest.mark.skipif(len(GPUtil.getGPUs()) == 0, reason='need GPU to run') 14 | except Exception: 15 | return pytest.mark.skipif(True, reason='need GPU to run') 16 | 17 | def make_config(master_port=None, worker_port=None, path=None): 18 | cfg = Config.default_config() 19 | cfg['network']['master'] = 'localhost' 20 | cfg['storage']['db_path'] = tempfile.mkdtemp() 21 | if master_port is not None: 22 | cfg['network']['master_port'] = master_port 23 | if worker_port is not None: 24 | cfg['network']['worker_port'] = worker_port 25 | 26 | if path is not None: 27 | with open(path, 'w') as f: 28 | cfg_path = path 29 | f.write(toml.dumps(cfg)) 30 | else: 31 | with tempfile.NamedTemporaryFile(delete=False) as f: 32 | cfg_path = f.name 33 | f.write(bytes(toml.dumps(cfg), 'utf-8')) 34 | return (cfg_path, cfg) 35 | 36 | 37 | def download_videos(): 38 | # Download video from GCS 39 | url = "https://storage.googleapis.com/scanner-data/test/short_video.mp4" 40 | with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as f: 41 | host = socket.gethostname() 42 | # HACK: special proxy case for Ocean cluster 43 | if host in ['ocean', 'crissy', 'pismo', 'stinson']: 44 | resp = requests.get( 45 | url, 46 | stream=True, 47 | proxies={'https': 'http://proxy.pdl.cmu.edu:3128/'}) 48 | else: 49 | resp = requests.get(url, stream=True) 50 | assert resp.ok 51 | for block in resp.iter_content(1024): 52 | f.write(block) 53 | vid1_path = f.name 54 | 55 | # Make a second one shorter than the first 56 | with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as f: 57 | vid2_path = f.name 58 | run([ 59 | 'ffmpeg', '-y', '-i', vid1_path, '-ss', '00:00:00', '-t', '00:00:10', 60 | '-c:v', 'libx264', '-strict', '-2', vid2_path 61 | ]) 62 | 63 | return (vid1_path, vid2_path) 64 | 65 | @pytest.fixture(scope="module") 66 | def sc(): 67 | # Create new config 68 | (cfg_path, cfg) = make_config() 69 | 70 | # Setup and ingest video 71 | with Client(config_path=cfg_path, debug=True) as sc: 72 | (vid1_path, vid2_path) = download_videos() 73 | 74 | sc.ingest_videos([('test1', vid1_path), ('test2', vid2_path)]) 75 | 76 | sc.ingest_videos( 77 | [('test1_inplace', vid1_path), ('test2_inplace', vid2_path)], 78 | inplace=True) 79 | 80 | yield sc 81 | 82 | # Tear down 83 | run([ 84 | 'rm', '-rf', cfg['storage']['db_path'], cfg_path, vid1_path, 85 | vid2_path 86 | ]) 87 | -------------------------------------------------------------------------------- /scannertools_infra/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | if __name__ == "__main__": 4 | setup(name='scannertools_infra', 5 | version='0.1.0', 6 | description='Scannertools infrastructure', 7 | url='http://github.com/scanner-research/scannertools', 8 | author='Will Crichton', 9 | author_email='wcrichto@cs.stanford.edu', 10 | license='Apache 2.0', 11 | packages=['scannertools_infra'], 12 | install_requires=['pytest', 'toml', 'requests'], 13 | zip_safe=False) 14 | -------------------------------------------------------------------------------- /scannertools_sql/scannertools_sql/__init__.py: -------------------------------------------------------------------------------- 1 | from scannertools_infra import _register_module 2 | 3 | _register_module(__file__, "scannertools_sql") 4 | -------------------------------------------------------------------------------- /scannertools_sql/scannertools_sql/storage.py: -------------------------------------------------------------------------------- 1 | from scannerpy.storage import StorageBackend, StoredStream 2 | from scannerpy.protobufs import protobufs 3 | 4 | 5 | class SQLStorage(StorageBackend): 6 | """StorageBackend backend for streams from a SQL database. 7 | 8 | Currently only supports Postgres.""" 9 | 10 | def __init__(self, config, job_table): 11 | """ 12 | Parameters 13 | ---------- 14 | config: protobufs.SQLConfig 15 | Database connection parameters 16 | 17 | job_table: str 18 | Name of table in the database to track completed jobs 19 | """ 20 | 21 | self._config = config 22 | self._job_table = job_table 23 | 24 | def source(self, sc, streams): 25 | num_elements = [s._num_elements for s in streams] \ 26 | if streams[0]._num_elements is not None else None 27 | return sc.sources.SQL( 28 | query=streams[0]._query, 29 | config=self._config, 30 | enum_config=[self._config for _ in streams], 31 | enum_query=[s._query for s in streams], 32 | filter=[s._filter for s in streams], 33 | num_elements=num_elements) 34 | 35 | def sink(self, sc, op, streams): 36 | return sc.sinks.SQL( 37 | input=op, 38 | config=self._config, 39 | table=streams[0]._table, 40 | job_table=self._job_table, 41 | job_name=[s._job_name for s in streams], 42 | insert=streams[0]._insert) 43 | 44 | def delete(self, sc, streams): 45 | # TODO 46 | pass 47 | 48 | 49 | class SQLInputStream(StoredStream): 50 | """Stream of elements from a SQL database used as input.""" 51 | 52 | def __init__(self, query, filter, storage, num_elements=None): 53 | """ 54 | Parameters 55 | ---------- 56 | query: protobufs.SQLQuery 57 | Query that generates a table 58 | 59 | filter: str 60 | Filter on the query that picks the rows/elements only in this stream 61 | 62 | storage: SQLStorage 63 | 64 | num_elements: int 65 | Number of elements in this stream. Optional optimization to avoid Scanner having to count. 66 | """ 67 | 68 | assert isinstance(storage, SQLStorage) 69 | self._query = query 70 | self._storage = storage 71 | self._filter = filter 72 | self._num_elements = num_elements 73 | 74 | def storage(self): 75 | return self._storage 76 | 77 | 78 | class SQLOutputStream(StoredStream): 79 | """Stream of elements into a SQL database used as output.""" 80 | 81 | def __init__(self, table, job_name, storage, insert=True): 82 | """ 83 | Parameters 84 | ---------- 85 | table: str 86 | Name of table to stream into. 87 | 88 | job_name: str 89 | Name of job to insert into the job table. 90 | 91 | storage: SQLStorage 92 | 93 | insert: bool 94 | Whether to insert new rows or update existing rows. 95 | """ 96 | 97 | assert isinstance(storage, SQLStorage) 98 | self._storage = storage 99 | self._table = table 100 | self._job_name = job_name 101 | self._insert = insert 102 | 103 | def storage(self): 104 | return self._storage 105 | 106 | def exists(self): 107 | # TODO 108 | return False 109 | 110 | def committed(self): 111 | # TODO 112 | return False 113 | -------------------------------------------------------------------------------- /scannertools_sql/scannertools_sql_cpp/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.2.0 FATAL_ERROR) 2 | 3 | execute_process( 4 | OUTPUT_VARIABLE SCANNER_CMAKE_PATH 5 | COMMAND python3 -c "import scannerpy.build_flags as b; b.print_cmake()") 6 | include(${SCANNER_CMAKE_PATH}) 7 | 8 | set(SOURCES sql.cpp sql_sink.cpp sql_source.cpp) 9 | 10 | build_op( 11 | LIB_NAME scannertools_sql 12 | CPP_SRCS ${SOURCES} 13 | PROTO_SRC scannertools_sql.proto) 14 | 15 | set(LIBRARIES) 16 | 17 | find_package(PkgConfig REQUIRED) 18 | pkg_check_modules(libpqxx REQUIRED libpqxx) 19 | target_include_directories(scannertools_sql PUBLIC ${libpqxx_INCLUDE_DIRS}) 20 | target_link_libraries(scannertools_sql PUBLIC "${libpqxx_LIBRARY_DIRS}/libpqxx.a" pq) 21 | -------------------------------------------------------------------------------- /scannertools_sql/scannertools_sql_cpp/scannertools_sql.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | message SQLConfig { 4 | string hostaddr = 1; 5 | int32 port = 2; 6 | string dbname = 3; 7 | string user = 4; 8 | string password = 5; 9 | string adapter = 6; 10 | } 11 | 12 | message SQLQuery { 13 | string fields = 1; 14 | string table = 2; 15 | string id = 4; 16 | string group = 5; 17 | } 18 | 19 | message SQLEnumeratorArgs { 20 | SQLConfig enum_config = 1; 21 | SQLQuery enum_query = 2; 22 | string filter = 3; 23 | int32 num_elements = 4; 24 | } 25 | 26 | message SQLSourceArgs { 27 | SQLConfig config = 1; 28 | SQLQuery query = 2; 29 | } 30 | 31 | message SQLElementArgs { 32 | string filter = 1; 33 | } 34 | 35 | message SQLSinkArgs { 36 | SQLConfig config = 1; 37 | string job_table = 2; 38 | string table = 3; 39 | bool insert = 4; 40 | bool ignore_conflicts = 5; 41 | } 42 | 43 | message SQLSinkStreamArgs { 44 | string job_name = 1; 45 | } 46 | -------------------------------------------------------------------------------- /scannertools_sql/scannertools_sql_cpp/sql.cpp: -------------------------------------------------------------------------------- 1 | #include "scanner/util/common.h" 2 | #include "scanner/util/tinyformat.h" 3 | #include "sql.h" 4 | 5 | namespace scanner { 6 | std::unique_ptr sql_connect(SQLConfig sql_config) { 7 | LOG_IF(FATAL, sql_config.adapter() != "postgres") 8 | << "Requested adapter " << sql_config.adapter() 9 | << " does not exist. Only \"postgres\" is supported right now."; 10 | try { 11 | return std::make_unique(tfm::format( 12 | "hostaddr=%s port=%d dbname=%s user=%s password=%s", 13 | sql_config.hostaddr(), sql_config.port(), sql_config.dbname(), 14 | sql_config.user(), sql_config.password())); 15 | } catch (pqxx::pqxx_exception& e) { 16 | LOG(FATAL) << e.base().what(); 17 | } 18 | } 19 | 20 | } 21 | -------------------------------------------------------------------------------- /scannertools_sql/scannertools_sql_cpp/sql.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "scannertools_sql.pb.h" 4 | #include 5 | #include 6 | 7 | namespace scanner { 8 | std::unique_ptr sql_connect(SQLConfig sql_config); 9 | } 10 | -------------------------------------------------------------------------------- /scannertools_sql/scannertools_sql_cpp/sql_sink.cpp: -------------------------------------------------------------------------------- 1 | /* Copyright 2018 Carnegie Mellon University 2 | * 3 | * Licensed under the Apache License, Version 2.0 (the "License"); 4 | * you may not use this file except in compliance with the License. 5 | * You may obtain a copy of the License at 6 | * 7 | * http://www.apache.org/licenses/LICENSE-2.0 8 | * 9 | * Unless required by applicable law or agreed to in writing, software 10 | * distributed under the License is distributed on an "AS IS" BASIS, 11 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | * See the License for the specific language governing permissions and 13 | * limitations under the License. 14 | */ 15 | 16 | #include "scanner/api/sink.h" 17 | #include "scanner/util/json.hpp" 18 | #include "scanner/util/tinyformat.h" 19 | #include "sql.h" 20 | 21 | #include 22 | #include 23 | #include 24 | 25 | using nlohmann::json; 26 | 27 | namespace scanner { 28 | 29 | std::string join(const std::vector& v, const std::string& c) { 30 | std::stringstream ss; 31 | for(size_t i = 0; i < v.size(); ++i) { 32 | ss << v[i]; 33 | if(i != v.size() - 1) { ss << c; } 34 | } 35 | return ss.str(); 36 | } 37 | 38 | class SQLSink : public Sink { 39 | public: 40 | SQLSink(const SinkConfig& config) : Sink(config) { 41 | bool parsed = args_.ParseFromArray(config.args.data(), config.args.size()); 42 | if (!parsed) { 43 | RESULT_ERROR(&valid_, "Could not parse SQLSinkArgs"); 44 | return; 45 | } 46 | 47 | } 48 | 49 | void new_stream(const std::vector& args) { 50 | bool parsed = sargs_.ParseFromArray(args.data(), args.size()); 51 | if (!parsed) { 52 | RESULT_ERROR(&valid_, "Could not parse SQLSinkStreamArgs"); 53 | return; 54 | } 55 | 56 | if (args_.job_table() != "" && sargs_.job_name() == "") { 57 | RESULT_ERROR(&valid_, "job_table was specified in SQLSinkArgs but job_name was not provided in SQLSinkStreamArgs"); 58 | return; 59 | } 60 | } 61 | 62 | void finished() override { 63 | // If the user provided a job table, then save the stream's job name to the table as completed 64 | std::string job_table = args_.job_table(); 65 | if (job_table != "") { 66 | std::unique_ptr conn = sql_connect(args_.config()); 67 | pqxx::work txn{*conn}; 68 | txn.exec(tfm::format("INSERT INTO %s (name) VALUES ('%s')", job_table, sargs_.job_name())); 69 | txn.commit(); 70 | } 71 | } 72 | 73 | void write(const BatchedElements& input_columns) override { 74 | std::unique_ptr conn = sql_connect(args_.config()); 75 | pqxx::work txn{*conn}; 76 | 77 | // Each element should be a list of JSON objects where each object is one row of the database. 78 | for (size_t i = 0; i < input_columns[0].size(); ++i) { 79 | auto& element = input_columns[0][i]; 80 | json jrows = json::parse(std::string((char*)element.buffer, element.size)); 81 | 82 | for (json& jrow : jrows) { 83 | // For a given JSON object, collect all of the key/value pairs and serialize them 84 | // into a SQL-injectable format 85 | std::map updates; 86 | i64 id = -1; 87 | for (json::iterator it = jrow.begin(); it != jrow.end(); ++it) { 88 | auto k = it.key(); 89 | auto v = it.value(); 90 | if (k == "id") { 91 | id = v; 92 | } else { 93 | // Have to special case strings since SQL queries expect single 94 | // quotes, while json to string formatter uses double quotes 95 | if (v.is_string()) { 96 | updates[k] = tfm::format("'%s'", v.get()); 97 | } else { 98 | updates[k] = tfm::format("%s", v); 99 | } 100 | } 101 | } 102 | 103 | // Construct the SQL string for either inserting or updating the database 104 | std::string update_str; 105 | if (args_.insert()) { 106 | std::vector column_list; 107 | std::vector value_list; 108 | for (auto it = updates.begin(); it != updates.end(); it++) { 109 | column_list.push_back(it->first); 110 | value_list.push_back(it->second); 111 | } 112 | 113 | update_str = 114 | tfm::format("INSERT INTO %s (%s) VALUES (%s) %s", args_.table(), 115 | join(column_list, ", "), join(value_list, ", "), 116 | args_.ignore_conflicts() ? "ON CONFLICT DO NOTHING" : ""); 117 | } else { 118 | std::vector update_list; 119 | for (auto it = updates.begin(); it != updates.end(); it++) { 120 | update_list.push_back(tfm::format("%s = %s", it->first, it->second)); 121 | } 122 | 123 | LOG_IF(FATAL, id == -1) << "SQLSink updates must have an `id` field set to know which row to update."; 124 | 125 | update_str = tfm::format("UPDATE %s SET %s WHERE id = %d", args_.table(), join(update_list, ", "), id); 126 | } 127 | 128 | txn.exec(update_str); 129 | } 130 | } 131 | 132 | // Only commit once all the changes have been made 133 | txn.commit(); 134 | } 135 | 136 | private: 137 | SQLSinkArgs args_; 138 | SQLSinkStreamArgs sargs_; 139 | Result valid_; 140 | }; 141 | 142 | REGISTER_SINK(SQL, SQLSink) 143 | .input("input") 144 | .per_element_output() 145 | .protobuf_name("SQLSinkArgs") 146 | .stream_protobuf_name("SQLSinkStreamArgs"); 147 | } // namespace scanner 148 | -------------------------------------------------------------------------------- /scannertools_sql/setup.cfg: -------------------------------------------------------------------------------- 1 | [aliases] 2 | test=pytest 3 | 4 | [tool:pytest] 5 | addopts = -vvs 6 | testpaths = tests -------------------------------------------------------------------------------- /scannertools_sql/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from scannertools_infra import CMakeExtension, CMakeBuild, CudaInstallCommand, CudaDevelopCommand 3 | 4 | if __name__ == "__main__": 5 | setup(name='scannertools_sql', 6 | version='0.2.15', 7 | description='Video analytics toolkit', 8 | url='http://github.com/scanner-research/scannertools', 9 | author='Will Crichton', 10 | author_email='wcrichto@cs.stanford.edu', 11 | license='Apache 2.0', 12 | packages=['scannertools_sql'], 13 | install_requires=[], 14 | setup_requires=['pytest-runner', 'scannertools_infra'], 15 | tests_require=[ 16 | 'pytest', 'psycopg2-binary == 2.7.6.1', 'testing.postgresql == 1.3.0', 17 | 'scannertools_infra' 18 | ], 19 | cmdclass=dict(build_ext=CMakeBuild, 20 | install=CudaInstallCommand, 21 | develop=CudaDevelopCommand), 22 | ext_modules=[CMakeExtension('scannertools_sql', 'scannertools_sql_cpp')], 23 | zip_safe=False) 24 | -------------------------------------------------------------------------------- /scannertools_sql/tests/test_all.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from testing.postgresql import Postgresql 3 | import psycopg2 4 | import scannerpy 5 | import json 6 | from scannertools_infra.tests import sc 7 | from scannerpy import protobufs, PerfParams 8 | import scannertools_sql 9 | from scannertools_sql.storage import SQLStorage, SQLInputStream, SQLOutputStream 10 | 11 | @pytest.fixture(scope='module') 12 | def sql_sc(sc): 13 | with Postgresql() as postgresql: 14 | conn = psycopg2.connect(**postgresql.dsn()) 15 | cur = conn.cursor() 16 | 17 | cur.execute( 18 | 'CREATE TABLE test (id serial PRIMARY KEY, a integer, b integer, c text, d varchar(255), e boolean, f float, grp integer)' 19 | ) 20 | cur.execute( 21 | "INSERT INTO test (a, b, c, d, e, f, grp) VALUES (10, 0, 'hi', 'hello', true, 2.0, 0)" 22 | ) 23 | cur.execute( 24 | "INSERT INTO test (a, b, c, d, e, f, grp) VALUES (20, 0, 'hi', 'hello', true, 2.0, 0)" 25 | ) 26 | cur.execute( 27 | "INSERT INTO test (a, b, c, d, e, f, grp) VALUES (30, 0, 'hi', 'hello', true, 2.0, 1)" 28 | ) 29 | cur.execute('CREATE TABLE jobs (id serial PRIMARY KEY, name text)') 30 | cur.execute( 31 | 'CREATE TABLE test2 (id serial PRIMARY KEY, b integer, s text)') 32 | conn.commit() 33 | 34 | sql_params = postgresql.dsn() 35 | sql_config = protobufs.SQLConfig( 36 | hostaddr=sql_params['host'], 37 | port=sql_params['port'], 38 | dbname=sql_params['database'], 39 | user=sql_params['user'], 40 | adapter='postgres') 41 | 42 | yield sc, SQLStorage(config=sql_config, job_table='jobs'), cur 43 | 44 | cur.close() 45 | conn.close() 46 | 47 | 48 | @scannerpy.register_python_op(name='AddOne') 49 | def add_one(config, row: bytes) -> bytes: 50 | row = json.loads(row.decode('utf-8')) 51 | return json.dumps([{'id': r['id'], 'b': r['a'] + 1} for r in row]) 52 | 53 | 54 | def test_sql(sql_sc): 55 | (sc, storage, cur) = sql_sc 56 | 57 | cur.execute('SELECT COUNT(*) FROM test'); 58 | n, = cur.fetchone() 59 | 60 | row = sc.io.Input([SQLInputStream( 61 | query=protobufs.SQLQuery( 62 | fields='test.id as id, test.a, test.c, test.d, test.e, test.f', 63 | table='test', 64 | id='test.id', 65 | group='test.id'), 66 | filter='true', 67 | storage=storage, 68 | num_elements=n)]) 69 | row2 = sc.ops.AddOne(row=row) 70 | output_op = sc.io.Output(row2, [SQLOutputStream( 71 | table='test', 72 | storage=storage, 73 | job_name='foobar', 74 | insert=False)]) 75 | sc.run(output_op, PerfParams.estimate()) 76 | 77 | cur.execute('SELECT b FROM test') 78 | assert cur.fetchone()[0] == 11 79 | 80 | cur.execute('SELECT name FROM jobs') 81 | assert cur.fetchone()[0] == 'foobar' 82 | 83 | @scannerpy.register_python_op(name='AddAll') 84 | def add_all(config, row: bytes) -> bytes: 85 | row = json.loads(row.decode('utf-8')) 86 | total = sum([r['a'] for r in row]) 87 | return json.dumps([{'id': r['id'], 'b': total} for r in row]) 88 | 89 | 90 | def test_sql_grouped(sql_sc): 91 | (sc, storage, cur) = sql_sc 92 | 93 | row = sc.io.Input([SQLInputStream( 94 | storage=storage, 95 | query=protobufs.SQLQuery( 96 | fields='test.id as id, test.a', 97 | table='test', 98 | id='test.id', 99 | group='test.grp'), 100 | filter='true')]) 101 | row2 = sc.ops.AddAll(row=row) 102 | output_op = sc.io.Output( 103 | row2, [SQLOutputStream(storage=storage, table='test', job_name='test', insert=False)]) 104 | sc.run(output_op, PerfParams.estimate()) 105 | 106 | cur.execute('SELECT b FROM test') 107 | assert cur.fetchone()[0] == 30 108 | 109 | 110 | @scannerpy.register_python_op(name='SQLInsertTest') 111 | def sql_insert_test(config, row: bytes) -> bytes: 112 | row = json.loads(row.decode('utf-8')) 113 | return json.dumps([{'s': 'hello world', 'b': r['a'] + 1} for r in row]) 114 | 115 | 116 | def test_sql_insert(sql_sc): 117 | (sc, storage, cur) = sql_sc 118 | 119 | row = sc.io.Input([SQLInputStream( 120 | storage=storage, 121 | query=protobufs.SQLQuery( 122 | fields='test.id as id, test.a', 123 | table='test', 124 | id='test.id', 125 | group='test.grp'), 126 | filter='true')]) 127 | row2 = sc.ops.SQLInsertTest(row=row) 128 | output_op = sc.io.Output( 129 | row2, [SQLOutputStream( 130 | storage=storage, table='test2', job_name='test', insert=True)]) 131 | sc.run(output_op, PerfParams.estimate()) 132 | 133 | cur.execute('SELECT s FROM test2') 134 | assert cur.fetchone()[0] == "hello world" 135 | -------------------------------------------------------------------------------- /scripts/clean-all.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | 4 | for subdir in scannertools*/ ; do 5 | pushd $subdir 6 | rm -rf build *.egg-info ${subdir}build 7 | popd 8 | done 9 | -------------------------------------------------------------------------------- /scripts/install-all.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | 4 | pushd scannertools_infra 5 | pip3 install -e . 6 | popd 7 | 8 | for subdir in scannertools*/ ; do 9 | if [ "$subdir" = "scannertools_infra/" ]; then 10 | continue 11 | fi 12 | 13 | pushd $subdir 14 | if [ "$tag2" = "cpu" ]; then 15 | CUDA_OPT= 16 | INSTALL_TAG=cpu 17 | else 18 | CUDA_OPT=--install-option="--build-cuda=/usr/local/cuda" 19 | INSTALL_TAG=gpu 20 | fi 21 | 22 | # Install dependencies first before package 23 | python3 setup.py egg_info 24 | requires_path="${subdir%/}.egg-info/requires.txt" 25 | if [ -f "${requires_path}" ]; then 26 | pip3 install -r ${requires_path} 27 | fi 28 | 29 | pip3 install ${CUDA_OPT} -e. 30 | popd 31 | done 32 | -------------------------------------------------------------------------------- /scripts/test-all.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | 4 | for subdir in scannertools*/ ; do 5 | if [ "$subdir" = "scannertools_infra" ]; then 6 | continue 7 | fi 8 | 9 | pushd $subdir 10 | python3 setup.py test 11 | popd 12 | done 13 | -------------------------------------------------------------------------------- /scripts/travis-build.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | 4 | yes | docker login -u="$DOCKER_USER" -p="$DOCKER_PASS" 5 | 6 | docker build \ 7 | -t ${DOCKER_REPO}:${TAG}-latest \ 8 | --build-arg tag=${TAG} \ 9 | --build-arg tag2=${TAG} \ 10 | . 11 | 12 | if [ "${TAG}" = "cpu" ]; 13 | then 14 | docker run ${DOCKER_REPO}:${TAG}-latest bash \ 15 | -c "adduser --disabled-password --gecos \"\" user && chown -R user /opt/scannertools && su -c \"cd /opt/scannertools && ./scripts/test-all.sh\" user" 16 | fi 17 | 18 | 19 | if [[ "$TRAVIS_BRANCH" = "master" ]]; then 20 | TRAVIS_TAG="latest" 21 | fi 22 | 23 | docker push ${DOCKER_REPO}:${TAG}-${TRAVIS_TAG} 24 | -------------------------------------------------------------------------------- /scripts/travis-publish.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | 4 | exit 0 5 | 6 | # Exit if this is a PR or not the CPU build 7 | if ! [ "$TRAVIS_PULL_REQUEST" = "false" -a "$TAG" = "cpu"]; then 8 | exit 0 9 | fi 10 | 11 | # Push to Pypi if there is a tag on this build 12 | if [ -n "$TRAVIS_TAG" ]; 13 | then 14 | docker run -v $(pwd):/opt/scannertools $DOCKER_REPO:cpu-latest bash -c " 15 | cd /opt/scannertools 16 | pip3 install twine 17 | python3 setup.py bdist_wheel 18 | twine upload -u 'wcrichto' -p '${PYPI_PASS}' dist/* 19 | " 20 | fi 21 | 22 | # Exit if it's not the master branch 23 | if ! ["$TRAVIS_BRANCH" = "master" ]; then 24 | exit 0 25 | fi 26 | 27 | pip3 install travis-sphinx 28 | 29 | docker run -v $(pwd):/opt/scannertools $DOCKER_REPO:cpu-latest bash -c " 30 | cd /opt/scannertools 31 | 32 | export LC_ALL=C.UTF-8 33 | export LANG=C.UTF-8 34 | pip3 install travis-sphinx sphinx-nameko-theme 35 | 36 | sphinx-apidoc -f -o doc/source scannertools 37 | travis-sphinx build -s doc -n 38 | " 39 | 40 | travis-sphinx deploy 41 | --------------------------------------------------------------------------------