Depth of field from stereo

depth_blur.py

# Example:
#  depth_blur.py im3.png im4.png -16 32
#  where im3.png and im4.png:
#  http://vision.middlebury.edu/stereo/eval/newEval/tsukuba/im3.png
#  http://vision.middlebury.edu/stereo/eval/newEval/tsukuba/im4.png
# Usage:
#  move mouse over image to change depth of focus
#
# Requires pyopencl and OpenCV builded with Python support

import cv2
import sys
import numpy as np
import pyopencl as cl

blur_kernel = """
float gaussian_kernel(int dx, int dy, float sigma) {
    if (sigma == 0.0f) {
        return (dx == 0 && dy == 0) ? 1.0f : 0.0f;
    } else {
        return exp(-(dx*dx + dy*dy) / (2 * sigma * sigma)) / (sigma * sqrt(2 * M_PI));
    }
}

__kernel void gauss(__read_only  image2d_t img,
                    __read_only  image2d_t disp,
                    __write_only image2d_t res,
                                 int       disp0,
                                 int       focus_disp,
                                 float     sigma_per_disp
                    ) {
    int2 coord = (int2) (get_global_id(0), get_global_id(1));
    int2 img_size = (int2) (get_image_width(img), get_image_height(img));
    if (any(coord >= img_size)) {
        return;
    }

    const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_FILTER_NEAREST | CLK_ADDRESS_NONE;

    float sigma;

    int disp_cur = read_imagei(disp, smp, coord).x;
    if (disp_cur == focus_disp) {
        sigma = 0.0f;
    } else if (disp_cur < disp0) {
        sigma = 2.0f;
    } else {
        sigma = sigma_per_disp * abs(disp_cur - focus_disp);
    }

    float4 color = (float4) 0.0;
    float sum_weight = 0;

    int radius = min(15, max(2, (int) (sigma * 3)));
    for (int dy = -radius; dy <= radius; ++dy) {
        for (int dx = -radius; dx <= radius; ++dx) {
            int2 ncoord = coord + (int2) (dx, dy);
            if (any(ncoord < (int2) (0, 0)) || any(ncoord >= img_size)) {
                continue;
            }
            float weight = gaussian_kernel(dx, dy, sigma);
            color += read_imagef(img, smp, ncoord) * weight;
            sum_weight += weight;
        }
    }
    write_imagef(res, coord, color / sum_weight);
}
"""


class DepthBlurWindow:

    def __init__(self, im0, im1, disp0, ndisp):
        self.im0 = im0
        self.im1 = im1
        self.img = None
        self.disp0 = disp0
        self.ndisp = ndisp

        self.disp = None
        self.disp_no_holes = None
        self.focusedDisparity = self.disp0
        self.windowName = "Depth blur"

        self.cl_context = None
        self.cl_queue = None
        self.cl_program = None

        self.cl_img = None
        self.cl_disp = None
        self.cl_blur = None

    def prepare(self):
        window_size = 7
        sgm = cv2.StereoSGBM_create(minDisparity=self.disp0, numDisparities=self.ndisp, blockSize=window_size,
                                    P1=8 * 3 * window_size ** 2, P2=32 * 3 * window_size ** 2,
                                    mode=cv2.STEREO_SGBM_MODE_HH)

        self.disp = sgm.compute(im0, im1) // 16
        self.fill_holes()
        self.img = self.im0.copy()
        # self.img[self.disp < self.disp0] = 0
        cv2.namedWindow(self.windowName)
        cv2.setMouseCallback(self.windowName, self.update_focus)

        self._prepare_cl()

    def fill_holes(self):
        values = np.float32(self.disp).ravel()
        values[values < self.disp0] = np.nan
        mask = np.isnan(values)
        values[mask] = np.interp(np.flatnonzero(mask), np.flatnonzero(~mask), values[~mask])
        self.disp_no_holes = np.int16(values).reshape(self.disp.shape)

    def _prepare_cl(self):
        self.cl_context = cl.create_some_context(True)
        self.cl_queue = cl.CommandQueue(self.cl_context)
        self.cl_program = cl.Program(self.cl_context, blur_kernel).build()

        h, w, cn = self.img.shape
        assert cn == 3 or cn == 4
        if cn == 3:
            self.img = np.dstack([self.img, np.zeros((h, w), np.uint8)])
        self.cl_disp = cl.Image(self.cl_context, cl.mem_flags.READ_ONLY | cl.mem_flags.COPY_HOST_PTR,
                                cl.ImageFormat(cl.channel_order.R, cl.channel_type.SIGNED_INT16), (w, h), hostbuf=self.disp_no_holes)
        img_format = cl.ImageFormat(cl.channel_order.RGBA, cl.channel_type.UNORM_INT8)
        self.cl_img = cl.Image(self.cl_context, cl.mem_flags.READ_ONLY | cl.mem_flags.COPY_HOST_PTR, img_format, (w, h), hostbuf=self.img)
        self.cl_blur = cl.Image(self.cl_context, cl.mem_flags.WRITE_ONLY, img_format, (w, h))

    def update_focus(self, event, x, y, flags, param):
        if y < 0 or y >= self.disp.shape[0] or x < 0 or x >= self.disp.shape[1]:
            return
        if self.disp[y, x] >= self.disp0:
            self.focusedDisparity = self.disp[y, x]

    def blur_image(self):
        h, w = self.img.shape[:2]
        sigma = 10.0 / self.ndisp
        self.cl_program.gauss(self.cl_queue, (w, h), None,
                              self.cl_img, self.cl_disp, self.cl_blur,
                              np.int32(self.disp0), np.int32(self.focusedDisparity), np.float32(sigma))

        blur = np.zeros((h, w, 4), np.uint8)
        cl.enqueue_copy(self.cl_queue, blur, self.cl_blur, origin=(0, 0, 0), region=(w, h, 1)).wait()
        return blur

    def loop(self):
        # disp = self.disp#self.disp_no_holes # self.disp
        # disp_normalized = np.uint8(255.0 * (disp - disp.min()) / (disp.max() - disp.min()))
        while True:
            # cv2.imshow(self.windowName, disp_normalized)
            cv2.imshow(self.windowName, self.blur_image())
            cv2.waitKey(10)

if __name__ == '__main__':
    if len(sys.argv) != 5:
        print("Usage: depth_blur.py <image1> <image2> <disp0> <ndisp>")
    else:
        im0 = cv2.imread(sys.argv[1])
        im1 = cv2.imread(sys.argv[2])
        disp0 = int(sys.argv[3])
        ndisp = int(sys.argv[4])

        if im0 is None:
            print("Unable to read image: ", sys.argv[1])
        elif im1 is None:
            print("Unable to read image: ", sys.argv[2])
        else:
            window = DepthBlurWindow(im0, im1, disp0, ndisp)
            window.prepare()
            window.loop()