Shapiro/part1/python/chapter_1.py - main - Dev_Devil/ComputerVision

Shapiro/part1/python/chapter_1.py
115 строк4 KB
heydolono
shapiro 1_2 part
20 ноя 2025, 19:14
20 ноя 2025, 19:14c86186d
import cv2
import numpy as np


def count_holes(binary: np.ndarray) -> int:
    b = (binary > 0).astype(np.uint8)
    E = 0
    I = 0
    r, c = b.shape
    for i in range(r - 1):
        for j in range(c - 1):
            patch = b[i:i+2, j:j+2].ravel()
            ones = int(patch.sum())
            if ones == 3:
                if patch.tolist().count(1) == 3:
                    E += 1
            elif ones == 1:
                if patch.tolist().count(1) == 1:
                    I += 1
    return (E - I) // 4


def connected_components_stats(binary: np.ndarray):
    b = (binary > 0).astype(np.uint8)
    num_labels, labels = cv2.connectedComponents(b, connectivity=4)
    stats = []
    for lbl in range(1, num_labels):
        mask = labels == lbl
        area = int(mask.sum())
        ys, xs = np.nonzero(mask)
        cy = int(ys.mean()) if ys.size else 0
        cx = int(xs.mean()) if xs.size else 0
        stats.append({"label": lbl, "area": area, "centroid": (cx, cy)})
    return stats, labels


def local_contrast(image: np.ndarray, ksize: int = 3) -> np.ndarray:
    img = image.astype(np.float32)
    kernel = np.ones((ksize, ksize), dtype=np.float32) / (ksize * ksize)
    mean = cv2.filter2D(img, -1, kernel)
    sq_mean = cv2.filter2D(img*img, -1, kernel)
    std = np.sqrt(np.maximum(0, sq_mean - mean*mean))
    contrast = std
    cmin, cmax = contrast.min(), contrast.max()
    if cmax > cmin:
        contrast = ((contrast - cmin) / (cmax - cmin) * 255).astype(np.uint8)
    else:
        contrast = (contrast * 0).astype(np.uint8)
    return contrast


def subtract_images(img1: np.ndarray, img2: np.ndarray) -> np.ndarray:
    a = img1.astype(np.int16)
    b = img2.astype(np.int16)
    diff = a - b
    diff = np.clip(diff + 128, 0, 255).astype(np.uint8)
    return diff


def add_images(img1: np.ndarray, img2: np.ndarray) -> np.ndarray:
    s = img1.astype(np.int32) + img2.astype(np.int32)
    s = np.clip(s // 2, 0, 255).astype(np.uint8)
    return s


def extract_patch_values(image: np.ndarray, top: int, left: int, h: int = 8, w: int = 8) -> np.ndarray:
    return image[top:top+h, left:left+w].copy()


def binarize_otsu(gray: np.ndarray) -> np.ndarray:
    _, th = cv2.threshold(gray, 0, 255, cv2.THRESH_OTSU)
    return th


def demo_from_file(path: str):
    img = cv2.imread(path, cv2.IMREAD_COLOR)
    if img is None:
        raise FileNotFoundError(path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    contrast = local_contrast(gray, ksize=3)
    _, th = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY)
    holes = count_holes((th == 0).astype(np.uint8))
    stats, labels = connected_components_stats(th == 0)
    diff = subtract_images(gray, cv2.GaussianBlur(gray, (9, 9), 0))
    added = add_images(img, img)
    patch = extract_patch_values(gray, 10, 10, 8, 8)
    return {
        "gray": gray,
        "contrast": contrast,
        "binary": th,
        "holes": holes,
        "components": stats,
        "diff": diff,
        "added": added,
        "patch": patch,
        "labels": labels
    }


if __name__ == "__main__":
    import sys
    import os
    path = sys.argv[1]
    out = demo_from_file(path)
    base = os.path.splitext(os.path.basename(path))[0]
    cv2.imwrite(f"{base}_gray_1.png", out["gray"])
    cv2.imwrite(f"{base}_contrast_1.png", out["contrast"])
    cv2.imwrite(f"{base}_binar_1.png", out["binary"])
    cv2.imwrite(f"{base}_diff_1.png", out["diff"])
    cv2.imwrite(f"{base}_added_1.png", out["added"])
    labels_vis = (out["labels"] * (255 // (out["labels"].max() + 1))).astype(np.uint8)
    cv2.imwrite(f"{base}_labels_1.png", labels_vis)
    print("Holes (approx):", out["holes"])
    for s in out["components"][:10]:
        print(s)