Python+OpenCVで流星群画像を仕分ける

import math
import os
import sys
import typing
# 3rd-party libraries
import cv2
import numpy

def load_gray(filepath: str) -> numpy.array:
    """
    グレイスケール画像読み込み
    :param str filepath: 入力ファイルパス
    :return: グレイスケール画像データ
    """
    img = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
    return img

def detect_lines(img: numpy.array) -> typing.List[numpy.array]:
    """
    直線検出
    :param numpy.array img: 入力画像（グレイスケール）
    :return: 検出直線リスト
    """
    ret, thr = cv2.threshold(img, 127, 255, cv2.ADAPTIVE_THRESH_MEAN_C)
    lines = cv2.HoughLinesP(thr, rho=1, theta=math.pi/180, threshold=200, minLineLength=20)
    return lines

result = []
for filepath in tqdm(image_list):
    filepath = str(filepath)
    img = load_gray(filepath)
    lines = detect_lines(img)
    if lines is not None:
        result.append((filepath, lines))

print("detected:{}/{}", len(result), len(image_list))

def detect_area(img: numpy.array, threshold: float = 0.0001) -> typing.List[numpy.array]:
    """
    閾値を超える面積を持つ輪郭の検出
    :param numpy.array img: 入力画像
    :param float threshold: 閾値（画像全体の何%を`(0, 1]`で指定）
    :return: 閾値を超えた面積の領域リスト
    """
    height, width = img.shape
    img_area = width * height
    ret, thr = cv2.threshold(img, 127, 255, cv2.ADAPTIVE_THRESH_MEAN_C)
    contours, hierarchy = cv2.findContours(thr, 1, 2)
    contours = [cnt for cnt in contours if (cv2.contourArea(cnt) / img_area) > threshold]
    return contours

T = typing.TypeVar("T")
def clamp(v: T, min_v: T, max_v: T) -> T:
    """
    clamp関数
    入力値を`[min_v, max_v]`の範囲に収める
    :param T v: 入力値
    :param T min_v: 最小値
    :param T max_v: 最大値
    :return: `[min_v, max_v]`内に収めた値
    """
    return min(max_v, max(v, min_v))

def fill_area(img: numpy.array, contours: typing.List[numpy.array], buffer_ratio: float = 0.01, color: float = None) -> numpy.array:
    """
    領域の外接矩形で塗りつぶす
    :param numpy.array img: 入力画像
    :param contours: 領域リスト
    :param float buffer_ratio: バッファ率 (e.g. 6000 * 0.01 => 60px)
    :param float color: 塗りつぶしの色（未指定の場合は入力画像の中央値で塗りつぶす）
    :return: 塗りつぶし後の画像
    """
    height, width = img.shape
    x_buffer = int(width * buffer_ratio)
    y_buffer = int(height * buffer_ratio)
    # detect fill color
    if color is None:
        color = numpy.median(img)
    # fill bounding rect
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        left = clamp(x - x_buffer, 0, width)
        top = clamp(y - y_buffer, 0, height)
        right = clamp(x + w + x_buffer, 0, width)
        bottom = clamp(y + h + y_buffer, 0, height)
        pts = numpy.asarray([
            [left, top],
            [left, bottom],
            [right, bottom],
            [right, top],
        ])
        img = cv2.fillPoly(img, pts=[pts], color=(color,))
    return img

def line_length(line: numpy.array) -> float:
    """
    直線の長さの算出
    `cv2.HoughLinesP()`の返り値は`[ [ [start_x, start_y, end_x, end_y] ], ... ]`形式になっている
    :param line: `[ [start_x, start_y, end_x, end_y] ]`であること
    :return: 直線の長さ
    """
    assert len(line) == 1
    sx, sy, ex, ey = line[0]
    dx = ex - sx
    dy = ey - sy
    return math.sqrt(dx * dx + dy * dy)

def detect_meteor(filepath: str, area_threshold: float = 0.0001, line_threshold: float = 100) -> typing.Optional[typing.Tuple[str, typing.List[numpy.array]]]:
    """
    流星の検出
    :param str filepath: 入力画像ファイルパス
    :param float area_threshold: 面積のある領域検知用閾値（`detect_area()`関数`threshold`参照）
    :param float line_threshold: 検出した直線を流星と判定する最小の長さ
    :return: (画像ファイルパス, 検出した直線) or None
    """
    img = load_gray(filepath)
    area_contours = detect_area(img, area_threshold)
    if area_contours:
        img = fill_area(img, area_contours)
    lines = detect_lines(img)
    if lines is not None:
        length = max([line_length(x) for x in lines])
        if length > line_threshold:
            return lines
    return None

result = []
for filepath in tqdm(image_list):
    filepath = str(filepath) # fqdm -> str
    lines = detect_meteor(filepath)
    if lines is not None:
        result.append((filepath, lines))

print("detected:{}/{}", len(result), len(image_list))