Process plate images for stopmotion video

plates.py

#!/usr/bin/env python3

from tqdm import tqdm
import glob
import os
import sys

import cv2
import matplotlib.pyplot as plt
import numpy as np


def load_image(src):
    img = cv2.imread(src)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img


def show_image(img, title=""):
    if img.ndim == 2:
        plt.imshow(img, cmap="gray")
    else:
        plt.imshow(img)
    plt.title(title)
    plt.show()


def unsharp_mask(img, blur_size=(9, 9), imgWeight=1.5, gaussianWeight=-0.5):
    # gaussian = cv2.GaussianBlur(img, (5,5), 0)
    gaussian = cv2.GaussianBlur(img, blur_size, 0)
    return cv2.addWeighted(img, imgWeight, gaussian, gaussianWeight, 0)


def get_images(path):
    files = glob.glob(os.path.join(path, "*.jpg"))
    # files.sort(key=os.path.getatime)
    files.sort()
    return files


def detect_plate_region(img, debug=False):
    print("Detecting plate region...")
    scale = 0.2

    print(f"Downsampling factor = {0.2}")
    img_resized = cv2.resize(img, None, fx=scale, fy=scale)
    canvas = img_resized.copy()

    gray = cv2.cvtColor(img_resized, cv2.COLOR_RGB2GRAY)
    gray_blurred = cv2.blur(gray, (3, 3))

    thresh = cv2.threshold(
        gray_blurred, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU
    )[1]

    kernel = np.ones((5, 5))
    thresh = cv2.dilate(thresh, kernel, 7)
    # thresh = cv2.erode(thresh, kernel, 1)
    if debug:
        show_image(thresh, "thresh")

    h, w = img_resized.shape[:2]
    img_r = h // 2
    circles = cv2.HoughCircles(
        thresh,
        cv2.HOUGH_GRADIENT,
        1,
        20,
        param1=50,
        param2=30,
        minRadius=int(0.40 * img_r),
        maxRadius=int(1.5 * img_r),
    )
    # circles = cv2.HoughCircles(gray, cv2.HOUGH_GRADIENT, 1.2, 100)
    center_thresh = 10
    img_cx, img_cy = w // 2, h // 2
    circles = np.squeeze(circles)
    circles = list(filter(lambda x: 0.8 * img_r <= x[-1] <= 1.1 * img_r, circles))
    circles = np.array(circles)
    print(f"Total circles found = {len(circles)}")
    if not len(circles):
        return []

    if debug:
        for c in circles:
            x, y, r = c.astype(int)
            # bool_x = x - center_thresh <= img_cx <= img_cx + center_thresh
            # bool_y = y - center_thresh <= img_cy <= img_cy + center_thresh
            # if not any([bool_x, bool_y]):
            #     continue
            cv2.circle(canvas, (x, y), r, (0, 255, 0), 4)
            cv2.rectangle(canvas, (x - 5, y - 5), (x + 5, y + 5), (0, 128, 255), -1)

    spans = [(int(x - r), int(x + r)) for (x, y, r) in circles]
    center = np.mean(circles[:, :2], axis=0).astype(int)
    cx, cy = center
    lx, rx = zip(*spans)
    lx, rx = min(lx), max(rx)
    if debug:
        cv2.rectangle(canvas, (cx - 5, cy - 5), (cx + 5, cy + 5), (255, 0, 0), 5)
        cv2.rectangle(canvas, (lx, 0), (rx, h), (255, 0, 0), 5)
        show_image(canvas, "scaled")

    canvas = img.copy()
    rescale = 1 / scale
    print(f"Re-scaling by {rescale}")
    cx, cy = int(cx * rescale), int(cy * rescale)
    lx, rx = int(lx * rescale), int(rx * rescale)
    h = int(h * rescale)
    print(f"Plate mean center = {cx}, {cy}")
    print(f"Plate resultant span = {lx}, {rx}")

    if debug:
        cv2.rectangle(canvas, (cx - 10, cy - 10), (cx + 10, cy + 10), (255, 0, 0), 5)
        cv2.rectangle(canvas, (lx, 0), (rx, h), (255, 0, 0), 5)
        show_image(canvas, "original")
    return (cx, cy), (lx, rx)


def eat_plates(sources, outdir="tmp/"):
    # bh, bw = 1080, 1920
    bw, bh = 3840, 2160
    bcx, bcy = bw // 2, bh // 2
    scale = 0.6
    for i, src in tqdm(enumerate(sources)):
        base_img = np.zeros((bh, bw, 3))
        img = load_image(src)
        (px, py), (lx, rx) = detect_plate_region(img)
        img = img[:, lx:rx]
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        img = unsharp_mask(img)
        img = unsharp_mask(img)
        # img = cv2.resize(img, None, fx=scale, fy=scale)
        h, w = img.shape[:2]
        cx, cy = w // 2, h // 2

        # translate to new image
        print(f"Translating to base image of {bw}x{bh}")
        tx, ty = (bcx - cx), (bcy - cy)
        mat = np.float32([[1, 0, tx], [0, 1, ty]])
        img = cv2.warpAffine(img, mat, (bw, bh), borderValue=(0, 0, 0))

        fname = os.path.join(outdir, "img_{0:0=2d}.png".format(i))
        cv2.imwrite(fname, img)


def main():
    path = "/tmp/data/images"
    sources = get_images(path)
    eat_plates(sources)


if __name__ == "__main__":
    main()