立体深度过滤器

1#!/usr/bin/env python3
2
3import argparse
4import random
5import cv2
6import depthai as dai
7import numpy as np
8
9FPS = 20
10
11def getRandomMedianFilterParams():
12    return random.choice(
13        [
14            dai.node.ImageFilters.MedianFilterParams.MEDIAN_OFF,
15            dai.node.ImageFilters.MedianFilterParams.KERNEL_3x3,
16            dai.node.ImageFilters.MedianFilterParams.KERNEL_5x5,
17        ]
18    )
19
20
21def getRandomTemporalFilterParams():
22    params = dai.node.ImageFilters.TemporalFilterParams()
23    params.enable = random.choice([True, False])
24    params.persistencyMode = random.choice(
25        [
26            dai.filters.params.TemporalFilter.PersistencyMode.PERSISTENCY_OFF,
27            dai.filters.params.TemporalFilter.PersistencyMode.VALID_2_IN_LAST_4,
28            dai.filters.params.TemporalFilter.PersistencyMode.VALID_1_IN_LAST_2,
29            dai.filters.params.TemporalFilter.PersistencyMode.VALID_1_IN_LAST_8,
30            dai.filters.params.TemporalFilter.PersistencyMode.PERSISTENCY_INDEFINITELY,
31        ]
32    )
33    params.alpha = random.uniform(0.3, 0.9)
34    return params
35
36
37def getRandomSpeckleFilterParams():
38    params = dai.node.ImageFilters.SpeckleFilterParams()
39    params.enable = random.choice([True, False])
40    return params
41
42
43def getRandomSpatialFilterParams():
44    params = dai.node.ImageFilters.SpatialFilterParams()
45    params.enable = random.choice([True, False])
46    return params
47
48
49def main(args: argparse.Namespace):
50    # Create pipeline
51    with dai.Pipeline() as pipeline:
52        monoLeft = pipeline.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_B)
53        monoRight = pipeline.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_C)
54        outLeft = monoLeft.requestOutput((640, 400), fps=FPS)
55        outRight = monoRight.requestOutput((640, 400), fps=FPS)
56
57        depth = pipeline.create(dai.node.StereoDepth)
58
59        filterPipeline = pipeline.create(dai.node.ImageFilters)
60        filterFactories = [
61            getRandomSpeckleFilterParams,
62            getRandomTemporalFilterParams,
63            getRandomSpatialFilterParams,
64            getRandomMedianFilterParams,
65        ]
66
67        filterPipeline.setRunOnHost(True)
68
69        depth.setLeftRightCheck(args.lr_check)
70        depth.setExtendedDisparity(args.extended_disparity)
71        depth.setSubpixel(args.subpixel)
72        depth.inputConfig.setBlocking(False)
73
74        # Linking
75        outLeft.link(depth.left)
76        outRight.link(depth.right)
77        depthQueue = depth.disparity.createOutputQueue()
78
79        filterPipeline.build(depth.disparity)
80
81        ## Create a new filter pipeline
82        filterPipeline.initialConfig.filterIndices = []
83        filterPipeline.initialConfig.filterParams = [
84            filterFactory() for filterFactory in filterFactories
85        ]
86
87        configInputQueue = filterPipeline.inputConfig.createInputQueue()
88        filterOutputQueue = filterPipeline.output.createOutputQueue()
89
90        pipeline.start()
91        import time
92
93        tSwitch = time.time()
94        while pipeline.isRunning():
95            inDisparity: dai.ImgFrame = (
96                depthQueue.get()
97            )  # blocking call, will wait until a new data has arrived
98            frame = inDisparity.getFrame()
99            filterFrame = filterOutputQueue.get()
100            filterFrame = (
101                filterFrame.getFrame() * (255 / depth.initialConfig.getMaxDisparity())
102            ).astype(np.uint8)
103            frame = (frame * (255 / depth.initialConfig.getMaxDisparity())).astype(
104                np.uint8
105            )
106            cv2.imshow("disparity", frame)
107            frame = cv2.applyColorMap(frame, cv2.COLORMAP_JET)
108            cv2.imshow("disparity_color", frame)
109            cv2.imshow(
110                "filtered_disparity_color",
111                cv2.applyColorMap(filterFrame, cv2.COLORMAP_JET),
112            )
113
114            ## Update filter pipeline
115            if time.time() - tSwitch > 1.0:
116                index = random.randint(0, len(filterFactories) - 1)
117                new_params = filterFactories[index]()
118                config = dai.ImageFiltersConfig().updateFilterAtIndex(
119                    index, new_params
120                )
121                configInputQueue.send(config)
122                tSwitch = time.time()
123                print(f"Filter at index {index} changed to {new_params}")
124
125            key = cv2.waitKey(1)
126            if key == ord("q"):
127                break
128
129
130if __name__ == "__main__":
131    parser = argparse.ArgumentParser()
132    parser.add_argument(
133        "--extended_disparity", action="store_true", help="Use extended disparity"
134    )
135    parser.add_argument("--subpixel", action="store_true", help="Use subpixel")
136    parser.add_argument("--lr_check", action="store_true", help="Use left-right check")
137    args = parser.parse_args()
138    main(args)