基于OpenCV的视频处理管道

$ git clone git://github.com/jagin/image-processing-pipeline.git$ cd image-processing-pipeline$ git checkout 7df1963247caa01b503980fe152138b88df6c526$ conda env create -f environment.yml$ conda activate pipeline

$ git pull$ git checkout 7df1963247caa01b503980fe152138b88df6c526$ conda env update -f environment.yml

import cv2from pipeline.pipeline import Pipeline
class CaptureVideo(Pipeline):    def __init__(self, src=0):        self.cap = cv2.VideoCapture(src)        if not self.cap.isOpened():            raise IOError(f"Cannot open video {src}")
        self.fps = int(self.cap.get(cv2.CAP_PROP_FPS))        self.frame_count = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        super(CaptureVideo, self).__init__()
    def generator(self):        image_idx = 0        while self.has_next():            ret, image = self.cap.read()            if not ret:                # no frames has been grabbed                break
            data = {                "image_id": f"{image_idx:05d}",                "image": image,            }
            if self.filter(data):                image_idx += 1                yield self.map(data)
    def cleanup(self):        # Closes video file or capturing device        self.cap.release()

data = {    "image_id": f"{image_idx:05d}",    "image": image,}

import cv2import numpy as np
class FaceDetector:    def __init__(self, prototxt, model, confidence=0.5):        self.confidence = confidence
        self.net = cv2.dnn.readNetFromCaffe(prototxt, model)
    def detect(self, images):        # convert images into blob        blob = self.preprocess(images)
        # pass the blob through the network and obtain the detections and predictions        self.net.setInput(blob)        detections = self.net.forward()        # Prepare storage for faces for every image in the batch        faces = dict(zip(range(len(images)), [[] for _ in range(len(images))]))
        # loop over the detections        for i in range(0, detections.shape[2]):            # extract the confidence (i.e., probability) associated with the prediction            confidence = detections[0, 0, i, 2]
            # filter out weak detections by ensuring the `confidence` is            # greater than the minimum confidence            if confidence < self.confidence:                continue
            # grab the image index            image_idx = int(detections[0, 0, i, 0])            # grab the image dimensions            (h, w) = images[image_idx].shape[:2]            # compute the (x, y)-coordinates of the bounding box for the object            box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
            # Add result            faces[image_idx].append((box, confidence))
        return faces
    def preprocess(self, images):        return cv2.dnn.blobFromImages(images, 1.0, (300, 300), (104.0, 177.0, 123.0))

import osimport cv2
from pipeline.libs.face_detector import FaceDetectorimport tests.config as config

class TestFaceDetector:    def test_face_detector(self):        prototxt = os.path.join(config.MODELS_FACE_DETECTOR_DIR, "deploy.prototxt.txt")        model = os.path.join(config.MODELS_FACE_DETECTOR_DIR, "res10_300x300_ssd_iter_140000.caffemodel")        detector = FaceDetector(prototxt, model)
        test_image = cv2.imread(os.path.join(config.ASSETS_IMAGES_DIR, "friends", "friends_01.jpg"))        faces = detector.detect([test_image])
        assert len(faces) == 1        assert len(faces[0])  # Should recognize some faces from friends_01.jpg

from pipeline.pipeline import Pipelinefrom pipeline.libs.face_detector import FaceDetector
class DetectFaces(Pipeline):    def __init__(self, prototxt, model, batch_size=1, confidence=0.5):        self.detector = FaceDetector(prototxt, model, confidence=confidence)        self.batch_size = batch_size
        super(DetectFaces, self).__init__()
    def generator(self):        batch = []        stop = False        while self.has_next() and not stop:            try:                # Buffer the pipeline stream                data = next(self.source)                batch.append(data)            except StopIteration:                stop = True
            # Check if there is anything in batch.            # Process it if the size match batch_size or there is the end of the input stream.            if len(batch) and (len(batch) == self.batch_size or stop):                # Prepare images batch                images = [data["image"] for data in batch]                # Detect faces on all images at once                faces = self.detector.detect(images)
                # Extract the faces and attache them to the proper image                for image_idx, image_faces in faces.items():                    batch[image_idx]["faces"] = image_faces
                # Yield all the data from buffer                for data in batch:                    if self.filter(data):                        yield self.map(data)
                batch = []

pipeline = (capture_video |            detect_faces |            save_faces |            annotate_image |            display_video |            save_video |            save_summary)

python process_video_pipeline.py -i assets/videos/faces.mp4 -p -d -ov faces.avi,M,];

$ python process_video_pipeline.py -i assets/videos/faces.mp4 -p--batch-size 1100%|███████████████████████████| 577/577 [00:11<00:00, 52.26it/s][INFO] Saving summary to output/summary.json...
$ python process_video_pipeline.py -i assets/videos/faces.mp4 -p--batch-size 4100%|███████████████████████████| 577/577 [00:09<00:00, 64.66it/s][INFO] Saving summary to output/summary.json...$ python process_video_pipeline.py -i assets/videos/faces.mp4 -p--batch-size 8100%|███████████████████████████| 577/577 [00:10<00:00, 56.04it/s][INFO] Saving summary to output/summary.json...