Shapiro/part2/python/chapter_2.py - main - pentryyy/ComputerVision

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


def rgb_to_gray(img: np.ndarray) -> np.ndarray:
    return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)


def gamma_correction(img: np.ndarray, gamma: float) -> np.ndarray:
    inv = 1.0 / gamma
    table = np.array([((i / 255.0) ** inv) * 255 for i in range(256)]).astype("uint8")
    return cv2.LUT(img, table)


def add_gaussian_noise(img: np.ndarray, mean: float = 0.0, sigma: float = 10.0) -> np.ndarray:
    gauss = np.random.normal(mean, sigma, img.shape).astype(np.float32)
    noisy = img.astype(np.float32) + gauss
    noisy = np.clip(noisy, 0, 255).astype(np.uint8)
    return noisy


def add_salt_and_pepper(img: np.ndarray, amount: float = 0.005) -> np.ndarray:
    out = img.copy()
    total = img.shape[0] * img.shape[1]
    num_salt = np.ceil(amount * total * 0.5).astype(int)
    num_pepper = np.ceil(amount * total * 0.5).astype(int)
    coords = [np.random.randint(0, i - 1, num_salt) for i in img.shape[:2]]
    out[coords[0], coords[1]] = 255
    coords = [np.random.randint(0, i - 1, num_pepper) for i in img.shape[:2]]
    out[coords[0], coords[1]] = 0
    return out


def chromatic_aberration(img: np.ndarray, shift: int = 2) -> np.ndarray:
    b, g, r = cv2.split(img)
    rows, cols = b.shape
    M = np.float32([[1, 0, shift], [0, 1, 0]])
    b2 = cv2.warpAffine(b, M, (cols, rows), borderMode=cv2.BORDER_REFLECT)
    M2 = np.float32([[1, 0, -shift], [0, 1, 0]])
    r2 = cv2.warpAffine(r, M2, (cols, rows), borderMode=cv2.BORDER_REFLECT)
    return cv2.merge([b2, g, r2])


def bloom(img: np.ndarray, thresh: int = 200, ksize: int = 21) -> np.ndarray:
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    mask = gray >= thresh
    bright = img.copy()
    bright[~mask] = 0
    blur = cv2.GaussianBlur(bright, (ksize, ksize), 0)
    out = cv2.addWeighted(img, 1.0, blur, 0.6, 0)
    return out


def radial_distortion(img: np.ndarray, k1: float = -1e-4, k2: float = 0.0) -> np.ndarray:
    h, w = img.shape[:2]
    x_c = w / 2
    y_c = h / 2
    map_x = np.zeros((h, w), dtype=np.float32)
    map_y = np.zeros((h, w), dtype=np.float32)
    for y in range(h):
        for x in range(w):
            x_n = (x - x_c) / x_c
            y_n = (y - y_c) / y_c
            r2 = x_n * x_n + y_n * y_n
            factor = 1 + k1 * r2 + k2 * r2 * r2
            x_dist = x_c + x_n * factor * x_c
            y_dist = y_c + y_n * factor * y_c
            map_x[y, x] = x_dist
            map_y[y, x] = y_dist
    return cv2.remap(img, map_x, map_y, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)


def affine_transform(img: np.ndarray, angle: float = 5.0, scale: float = 1.0, tx: float = 0.0, ty: float = 0.0) -> np.ndarray:
    h, w = img.shape[:2]
    M = cv2.getRotationMatrix2D((w/2, h/2), angle, scale)
    M[0, 2] += tx
    M[1, 2] += ty
    return cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)


def downsample_upsample(img: np.ndarray, factor: int = 4) -> np.ndarray:
    h, w = img.shape[:2]
    small = cv2.resize(img, (w // factor, h // factor), interpolation=cv2.INTER_AREA)
    back = cv2.resize(small, (w, h), interpolation=cv2.INTER_LINEAR)
    return back


def vignetting(img: np.ndarray, strength: float = 0.5) -> np.ndarray:
    h, w = img.shape[:2]
    X = np.linspace(-1, 1, w)
    Y = np.linspace(-1, 1, h)
    xv, yv = np.meshgrid(X, Y)
    mask = 1 - strength * (xv*xv + yv*yv)
    mask = np.clip(mask, 0, 1)
    if img.ndim == 3:
        mask = mask[:, :, None]
    out = (img.astype(np.float32) * mask).astype(np.uint8)
    return out


def simulate_sensor_effects(img: np.ndarray):
    out = img.copy()
    out = vignetting(out, strength=0.4)
    out = add_gaussian_noise(out, sigma=8.0)
    out = add_salt_and_pepper(out, amount=0.002)
    out = chromatic_aberration(out, shift=1)
    out = bloom(out, thresh=220, ksize=15)
    out = radial_distortion(out, k1=-5e-5)
    out = downsample_upsample(out, factor=3)
    return out


if __name__ == "__main__":
    import sys
    import os
    path = sys.argv[1]
    img = cv2.imread(path, cv2.IMREAD_COLOR)
    if img is None:
        raise FileNotFoundError(path)
    out = simulate_sensor_effects(img)
    base = os.path.splitext(os.path.basename(path))[0]
    cv2.imwrite(f"{base}_sensor_sim_2.png", out)
    gray = rgb_to_gray(img)
    cv2.imwrite(f"{base}_gray_2.png", gray)
    gamma = gamma_correction(img, 2.2)
    cv2.imwrite(f"{base}_gamma_2.png", gamma)
    print("Done")