cv_work.py - master - MaxIs/maxlib

cv_work.py
166 строк9 KB
MaxIs
Некоторые полезные вещи для работы с контурами на изображениях.
04 дек 2025, 09:59
04 дек 2025, 09:59bcca3e1
"""
Некоторые полезные вещи по работе с изображениями, чего нет из коробки в cv2
"""

from collections import defaultdict
from typing import Tuple, List
import cv2
import numpy as np


class ContourWork:
    """Поможет найти все точки внутри контура и определить угол контура"""

    @staticmethod
    def get_all_indexes(cnt, xy: bool = False) -> Tuple[np.array, np.array]:
        """
        Получим индексы из всего контура не только границы, но и внутри все точки
        :param cnt: np.ndarray контур (строки, столбцы, 2)
        :param xy: bool признак как возвращать индексы, если xy=True, то (x, y) иначе (y, x)
        :return: Tuple[np.array[int], np.array[int]] индексы в зависимости от xy
        """
        cnt = cnt.flatten()
        d = defaultdict(list)
        for x, y in zip(cnt[0::2], cnt[1::2]):
            d[int(x)].append(int(y))
        d = dict(d)
        for k, v in d.items():
            d[k] = [x for x in range(min(v), max(v) + 1)]
        xs = []
        ys = []
        for x, v in d.items():
            xs.extend([x for _ in range(len(v))])
            ys.extend(v)
        if xy:
            d = (np.array(xs), np.array(ys))
        else:
            d = (np.array(ys), np.array(xs))
        return d

    @staticmethod
    def get_degrees(mom, return_degree: bool = True) -> float:
        """Посчитать наклон фигуры по моментам"""
        theta = .0
        nu20, nu11, nu02 = mom['nu20'], mom['nu11'], mom['nu02']
        num = 2 * nu11
        den = nu20 - nu02
        if num == 0 and den == 0:
            # print(f"форма близка к кругу, ориентация не определена")
            return theta
        elif num != 0 and den == 0:
            theta = 0.785398 if num > 0 else -0.785398
        else:
            theta = 0.5 * np.arctan(num / den)
        if return_degree:
            theta = np.degrees(theta)
        return theta


class ContoursWork:
    """Поможет удалить контур и исправить иерархию без повторного поиска"""

    def __init__(self, img: np.ndarray):
        self.img = img
        self.gray = cv2.cvtColor(self.img, cv2.COLOR_RGB2GRAY)
        self.contours, self.hierarchy = cv2.findContours(self.gray, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

    def contour_delete(self, cnt) -> None:
        """Удаляем все связанные с контуром ссылки исхожу из того что контур уникален и 2 одинаковых контуров быть
        не может"""
        c = [idx for idx, x in enumerate(self.contours) if x.shape == cnt.shape and (self.contours[idx] == cnt).all()]
        if len(c) == 1:
            c = c[0]
            self.contours = tuple([x for idx, x in enumerate(self.contours) if idx != c])
        elif len(c) == 0:
            raise ValueError(f"cnt not found in self.contours")
        else:
            raise ValueError(f"cnt many found in self.contours")
        # удалим контур из иерархии
        self._del_idx_in_hierarchy(c)

    def _del_idx_in_hierarchy(self, idx_del) -> None:
        """
        Изменяем иерархию удаляем элемент с idx_del и изменяем все ссылки на него и на остальные элементы
        :param idx_del: индекс удаляемого элемента
        :return: None inplace удаляем индекс контура на месте, изменённая иерархия в self.hierarchy
        """
        first_child, parent_id = self.hierarchy[0, idx_del, 2:]
        # изменяем родителя в иерархии для детей удаляемого контура на "деда"
        if first_child != -1:
            childs_lvl = ([(first_child, self.hierarchy[0, first_child, 0])] +
                          ContoursWork.recursive_get_id(self.hierarchy, first_child, [], change=0))
            for ch in childs_lvl:
                ContoursWork.change_val(self.hierarchy, ch[0], parent_id, 3)

        # выставляем в уровне удаляемого контура значения по порядку без учёта удаляемого контура
        indexes = [i for i, v in enumerate(self.hierarchy[0, :, 3] == parent_id) if v and i != idx_del]
        lin = len(indexes)
        if lin > 0:
            for idx, index in enumerate(indexes):
                if 0 < idx < lin - 1:
                    # следующий
                    ContoursWork.change_val(self.hierarchy, index, indexes[idx + 1], 0)
                    # предыдущий
                    ContoursWork.change_val(self.hierarchy, index, indexes[idx - 1], 1)
                elif idx == lin - 1:
                    # следующий = -1
                    ContoursWork.change_val(self.hierarchy, index, -1, 0)
                    # предыдущий
                    if lin > 1:
                        ContoursWork.change_val(self.hierarchy, index, indexes[idx - 1], 1)
                else:
                    # предыдущий = -1
                    ContoursWork.change_val(self.hierarchy, index, -1, 1)
                    if lin > 1:
                        # следующий
                        ContoursWork.change_val(self.hierarchy, index, indexes[idx + 1], 0)
        # у всех контуров изменяем ссылки на ссылку -1 если они больше чем удаляемый контур
        for i in range(self.hierarchy.shape[2]):
            indexes = [idx for idx, x in enumerate(self.hierarchy[0, :, i] > idx_del) if x]
            self.hierarchy[0, indexes, i] -= 1

        self.hierarchy = np.delete(self.hierarchy, idx_del, 1)

    @staticmethod
    def change_val(hier, idx, new_val, change: int) -> None:
        """
        Изменяем значения в иерархии у определённого контура
        :param hier: np.ndarray иерархия размер (1, n, 4)
        :param idx: индекс в иерархии
        :param new_val: np.issubdtype(type(new_val), np.signedinteger) новое значение одно из типов np.signedinteger
        :param change: int выбор какое значение изменяем 0 - следующий, 1 - предыдущий, 2 - первый ребёнок, 3 - родитель
        :return: None изменяем на месте иерархию
        """
        if not isinstance(hier, np.ndarray):
            raise TypeError(f"hierarchy should be np.ndarray but got {type(hier)}")
        if hier.shape[::2] != (1, 4):
            raise ValueError(f"hierarchy shape not valid, size must be (1, n, 4)")
        ridx = range(hier.shape[1])
        if idx not in ridx:
            raise IndexError(f"idx out of range, current range in {ridx}")
        change_set = set(range(4))
        if not isinstance(change, int) or change not in change_set:
            raise IndexError(f"change {change} is not valid, maybe only one of {change_set}")
        if not np.issubdtype(type(new_val), np.signedinteger):
            raise TypeError(f"type new_val must be subset np.signedinteger")

        hier[0, idx, change] = new_val

    @staticmethod
    def recursive_get_id(hierarchy_tmp, idx_tmp: int, idxs: list,
                         change: int = 0) -> List[Tuple[np.signedinteger, np.signedinteger]] | []:
        """
        Рекурсивно вернуть список всех элементов ссылающихся на нужный
        :param hierarchy_tmp: np.ndarray иерархия
        :param idx_tmp: нужный нам контур
        :param idxs: список контуров связанных с искомым
        :param change: выбор какие связи ищем 0 - следующие контура, 1 - предыдущие контура, 2 - первые дети, 3 родители
        :return: List[Tuple[np.signedinteger, np.signedinteger]] список индексом и значением на кого ссылается далее
        """
        val = hierarchy_tmp[0][int(idx_tmp)][change]
        if int(val) == -1:
            idxs.append((idx_tmp, val))
            return idxs[1:] if len(idxs) > 1 else []
        else:
            idxs.append((idx_tmp, val))
            return ContoursWork.recursive_get_id(hierarchy_tmp, hierarchy_tmp[0][idx_tmp][change], idxs, change)