Raspberry Pi camera capture code

capture.py

#!/usr/bin/env python3
import io
import time
import picamera
import cv2
import numpy as np
from fractions import Fraction


def capture(camera, ss, start):
    camera.awb_mode = 'off'
    camera.awb_gains = (Fraction(3, 1), Fraction(3, 2))
    camera.iso = 800
    camera.exposure_mode = 'off'
    camera.framerate = Fraction(1, 200)
    camera.shutter_speed = int(round(ss))
    frame = np.zeros((3040, 4032, 3), dtype=np.float32)

    tic = time.time()
    for i in range(10):
        camera.resolution = (4032, 3040)
        output = np.empty((3040, 4032, 3), dtype=np.uint8)
        camera.capture(output, 'rgb')
        image = output[:, :, ::-1]
        frame = frame + image.astype(np.float32)
        toc = time.time()
        print('time', toc - tic)
        if toc - start > 110:
            break
    frame = undistort(frame)
    frame *= 255
    cv2.imwrite('output_{}.jpg'.format(int(time.time())), frame)


def undistort(src):
    height, width, _ = src.shape
    zr = 1 - 9e-4
    zb = 1 - 1e-4

    dst = src
    r = cv2.warpAffine(
        dst[:, :, 2],
        np.matrix([[zr, 0, (1 - zr) * width / 2],
                   [0, zr, (1 - zr) * height / 2]]), (width, height))
    dst[:, :, 2] = r

    dst[:, :, 1] *= 0.8
    b = cv2.warpAffine(
        dst[:, :, 0],
        np.matrix([[zb, 0, (1 - zb) * width / 2],
                   [0, zb, (1 - zb) * height / 2]]), (width, height))
    dst[:, :, 0] = 0.683 * b

    distCoeff = np.zeros((4, 1), np.float64)
    distCoeff[0, 0] = -1.2e-6
    distCoeff[1, 0] = 0
    distCoeff[2, 0] = 0
    distCoeff[3, 0] = 0

    cam = np.eye(3, dtype=np.float32)

    cam[0, 2] = width / 2.0  # define center x
    cam[1, 2] = height / 2.0  # define center y
    cam[0, 0] = 10.  # define focal length x
    cam[1, 1] = 10.  # define focal length y

    dst = cv2.undistort(src, cam, distCoeff)
    dst = dst / np.percentile(dst, 97)
    bdst = cv2.GaussianBlur(dst, (0, 0), 1.5)
    dst = dst * 1.5 - 0.5 * bdst
    dst = np.clip(dst, 0, 1)
    dst = np.sqrt(dst)
    return dst


last_ss = 10000


def autoexposure(camera):
    global last_ss
    ss = last_ss
    camera.resolution = (1280, 720)
    for attempt in range(100):
        ssi = int(round(ss))
        print('setting shutter speed', ssi)
        camera.shutter_speed = ssi
        time.sleep(0.5)
        print(camera.shutter_speed)
        print(camera.exposure_speed)
        #output = np.empty((4056, 3040, 3), dtype=np.uint8)
        output = np.empty((720, 1280, 3), dtype=np.uint8)
        camera.capture(output, 'rgb')
        image = output[:, :, ::-1]
        high_r = np.percentile(image[:, :, 2], 95)
        high_g = np.percentile(image[:, :, 1], 95)
        high_b = np.percentile(image[:, :, 0], 95)
        high = max(high_r, high_g, high_b)
        #cv2.imwrite('qwer{:0>4d}.png'.format(attempt), image)
        if 150 < high < 200:
            ss *= 190 / (high + 1)
            last_ss = ss
            return ss
        if high <= 150:
            ss *= min(190 / (high + 1), 10)
            if ss > 2e7:
                last_ss = 2e7
                return 2e7
        elif high >= 200:
            ss /= 2


start = time.time()
with picamera.PiCamera(sensor_mode=3) as camera:
    picamera.PiCamera.CAPTURE_TIMEOUT = 600
    camera.resolution = camera.MAX_RESOLUTION
    camera.awb_mode = 'off'
    camera.awb_gains = (Fraction(3, 1), Fraction(3, 2))

    camera.iso = 800
    camera.exposure_mode = 'off'
    camera.image_denoise = False
    camera.framerate = Fraction(1, 200)
    ss = autoexposure(camera)
    capture(camera, ss, start)