analyzer/app/analyzer.py - master - Sturon/Diplom

analyzer/app/analyzer.py
596 строк20 KB
Sturon
Initial
03 июн 2026, 15:01
03 июн 2026, 15:0117b44a2
from __future__ import annotations

import math
import time
from dataclasses import dataclass
from io import BytesIO
from urllib.request import Request, urlopen

import numpy as np
from PIL import Image, ImageDraw

from app.models import DefectResponse, EvaluationParameters, EvaluationResponse, Roi


DEFAULT_HORIZONTAL_SEARCH_RADIUS = 12
DEFAULT_ALIGNMENT_COARSE_STEP = 4
DEFAULT_EDGE_DIFFERENCE_WEIGHT = 1.7
DEFAULT_STRUCTURED_PIXEL_WEIGHT = 0.85
MAX_ALIGNMENT_SAMPLE_SIZE = 900


@dataclass(frozen=True)
class RoiMetrics:
    mean: float
    deviation: float
    dark_ratio: float
    bright_ratio: float


@dataclass(frozen=True)
class ImageFeatures:
    rgb: np.ndarray
    gray: np.ndarray
    edge: np.ndarray


@dataclass(frozen=True)
class AlignmentResult:
    dx: int
    score: float


@dataclass(frozen=True)
class DifferenceComponent:
    x: int
    y: int
    width: int
    height: int
    area: int
    mean_difference: float
    max_difference: float
    mean_luma_difference: float
    mean_edge_difference: float


def read_image(image_bytes: bytes) -> Image.Image:
    try:
        image = Image.open(BytesIO(image_bytes))
        image.load()
    except Exception as exc:
        raise ValueError("Unsupported image file") from exc
    return image.convert("RGB")


def fetch_snapshot(url: str, timeout_ms: int) -> bytes:
    request = Request(url, headers={"User-Agent": "diplom-edge-analyzer/1.0"})
    with urlopen(request, timeout=timeout_ms / 1000) as response:
        return response.read()


def evaluate(image: Image.Image, parameters: EvaluationParameters) -> EvaluationResponse:
    started_at = time.perf_counter()
    rois = parameters.rois or [
        Roi(name="Full frame", order=0, x=0, y=0, width=image.width, height=image.height)
    ]
    defects: list[DefectResponse] = []

    for roi in rois:
        validate_roi(roi, image)
        metrics = calculate_metrics(
            image.crop((roi.x, roi.y, roi.x + roi.width, roi.y + roi.height)),
            parameters.rules.brightness_min,
            parameters.rules.brightness_max,
        )
        defects.extend(find_roi_defects(roi, metrics, parameters))

    defects = [
        defect
        for defect in defects
        if defect.confidence >= parameters.rules.min_confidence
    ]
    result = len(defects) <= parameters.rules.max_defects
    processing_time_ms = int((time.perf_counter() - started_at) * 1000)

    return EvaluationResponse(
        result=result,
        requestId=parameters.request_id,
        profileId=parameters.profile_id,
        defectCount=len(defects),
        processingTimeMs=processing_time_ms,
        imageWidth=image.width,
        imageHeight=image.height,
        summary=build_summary(result, len(defects), parameters.rules.max_defects),
        defects=defects,
    )


def compare_with_reference(
    image: Image.Image,
    reference: Image.Image,
    parameters: EvaluationParameters,
) -> tuple[EvaluationResponse, Image.Image]:
    started_at = time.perf_counter()
    frame = image.convert("RGB")
    reference_frame = reference.convert("RGB")
    if frame.size != reference_frame.size:
        raise ValueError(
            f"Frame size {frame.size} does not match reference size {reference_frame.size}"
        )

    rois = parameters.rois or [
        Roi(name="Full frame", order=0, x=0, y=0, width=reference_frame.width, height=reference_frame.height)
    ]
    for roi in rois:
        validate_roi(roi, reference_frame)

    aligned_frame, aligned_features, reference_features, alignment = align_frame_to_reference(
        frame,
        reference_frame,
        rois,
        parameters,
    )

    weighted_difference, luma_difference, edge_difference = build_weighted_difference(
        aligned_features,
        reference_features,
        parameters.rules,
    )

    annotated = aligned_frame.copy()
    draw = ImageDraw.Draw(annotated)
    defects: list[DefectResponse] = []

    for roi in rois:
        components = find_difference_components(
            weighted_difference,
            luma_difference,
            edge_difference,
            roi,
            parameters.rules,
        )

        for component in components:
            confidence_value = defect_confidence(component, roi, parameters.rules)
            if confidence_value < parameters.rules.min_confidence:
                continue

            absolute_x = roi.x + component.x
            absolute_y = roi.y + component.y
            defects.append(DefectResponse(
                defectType="reference-difference",
                roiName=roi.name,
                roiOrder=roi.order,
                confidence=confidence_value,
                x=absolute_x,
                y=absolute_y,
                width=component.width,
                height=component.height,
                metrics={
                    "area": component.area,
                    "meanDifference": round(component.mean_difference, 3),
                    "maxDifference": round(component.max_difference, 3),
                    "meanLumaDifference": round(component.mean_luma_difference, 3),
                    "meanEdgeDifference": round(component.mean_edge_difference, 3),
                    "threshold": parameters.rules.difference_threshold,
                    "alignmentDx": alignment.dx,
                    "alignmentScore": round(alignment.score, 6),
                },
            ))
            draw.rectangle(
                (absolute_x, absolute_y, absolute_x + component.width, absolute_y + component.height),
                outline=(220, 38, 38),
                width=4,
            )
            draw.text((absolute_x + 4, max(0, absolute_y - 18)), "DEFECT", fill=(220, 38, 38))

    result = len(defects) <= parameters.rules.max_defects
    processing_time_ms = int((time.perf_counter() - started_at) * 1000)
    return (
        EvaluationResponse(
            result=result,
            requestId=parameters.request_id,
            profileId=parameters.profile_id,
            defectCount=len(defects),
            processingTimeMs=processing_time_ms,
            imageWidth=reference_frame.width,
            imageHeight=reference_frame.height,
            summary=build_summary(result, len(defects), parameters.rules.max_defects),
            defects=defects,
        ),
        annotated,
    )


def align_frame_to_reference(
    frame: Image.Image,
    reference: Image.Image,
    rois: list[Roi],
    parameters: EvaluationParameters,
) -> tuple[Image.Image, ImageFeatures, ImageFeatures, AlignmentResult]:
    frame_features = build_features(image_to_array(frame))
    reference_features = build_features(image_to_array(reference))
    sample_stride = alignment_sample_stride(reference.size)
    alignment = find_best_horizontal_shift(
        frame_features,
        reference_features,
        rois,
        sample_stride,
        parameters.rules,
    )
    aligned_frame = shift_frame_horizontally(frame, alignment.dx)
    aligned_features = build_features(image_to_array(aligned_frame))
    return aligned_frame, aligned_features, reference_features, alignment


def build_features(rgb: np.ndarray) -> ImageFeatures:
    gray = luminance(rgb)
    return ImageFeatures(
        rgb=rgb.astype(np.float32, copy=False),
        gray=gray,
        edge=edge_map(gray),
    )


def luminance(rgb: np.ndarray) -> np.ndarray:
    return (
        0.2126 * rgb[:, :, 0]
        + 0.7152 * rgb[:, :, 1]
        + 0.0722 * rgb[:, :, 2]
    ).astype(np.float32)


def edge_map(gray: np.ndarray) -> np.ndarray:
    gray = gray.astype(np.float32, copy=False)
    gx = np.zeros_like(gray, dtype=np.float32)
    gy = np.zeros_like(gray, dtype=np.float32)

    gx[:, 1:-1] = (gray[:, 2:] - gray[:, :-2]) * 0.5
    gx[:, 0] = gray[:, 1] - gray[:, 0] if gray.shape[1] > 1 else 0
    gx[:, -1] = gray[:, -1] - gray[:, -2] if gray.shape[1] > 1 else 0

    gy[1:-1, :] = (gray[2:, :] - gray[:-2, :]) * 0.5
    gy[0, :] = gray[1, :] - gray[0, :] if gray.shape[0] > 1 else 0
    gy[-1, :] = gray[-1, :] - gray[-2, :] if gray.shape[0] > 1 else 0

    return np.clip(np.hypot(gx, gy) * 1.35, 0, 255).astype(np.float32)


def find_best_horizontal_shift(
    frame: ImageFeatures,
    reference: ImageFeatures,
    rois: list[Roi],
    sample_stride: int,
    rules,
) -> AlignmentResult:
    radius = max(0, rule_int(rules, "alignment_search_radius", DEFAULT_HORIZONTAL_SEARCH_RADIUS))
    coarse_step = max(1, rule_int(rules, "alignment_coarse_step", DEFAULT_ALIGNMENT_COARSE_STEP))
    best_dx = 0
    best_score = score_horizontal_shift(frame, reference, rois, 0, sample_stride, rules)

    if radius == 0:
        return AlignmentResult(dx=0, score=best_score)

    coarse_values = sorted(set(range(-radius, radius + 1, coarse_step)) | {0, radius, -radius})
    for dx in coarse_values:
        score = score_horizontal_shift(frame, reference, rois, dx, sample_stride, rules)
        if score < best_score:
            best_dx = dx
            best_score = score

    fine_radius = min(coarse_step, radius)
    for dx in range(max(-radius, best_dx - fine_radius), min(radius, best_dx + fine_radius) + 1):
        score = score_horizontal_shift(frame, reference, rois, dx, sample_stride, rules)
        if score < best_score:
            best_dx = dx
            best_score = score

    return AlignmentResult(dx=best_dx, score=best_score)


def score_horizontal_shift(
    frame: ImageFeatures,
    reference: ImageFeatures,
    rois: list[Roi],
    dx: int,
    sample_stride: int,
    rules,
) -> float:
    _, width = reference.gray.shape
    total_score = 0.0
    total_pixels = 0

    for roi in rois:
        x0 = max(roi.x, -dx)
        x1 = min(roi.x + roi.width, width - dx)
        if x1 <= x0:
            continue

        reference_slice = (
            slice(roi.y, roi.y + roi.height, sample_stride),
            slice(x0, x1, sample_stride),
        )
        frame_slice = (
            slice(roi.y, roi.y + roi.height, sample_stride),
            slice(x0 + dx, x1 + dx, sample_stride),
        )
        score, pixels = score_feature_regions(
            frame,
            reference,
            frame_slice,
            reference_slice,
            rules,
        )
        total_score += score * pixels
        total_pixels += pixels

    if total_pixels == 0:
        return float("inf")
    return total_score / total_pixels


def score_feature_regions(
    frame: ImageFeatures,
    reference: ImageFeatures,
    frame_slice,
    reference_slice,
    rules,
) -> tuple[float, int]:
    frame_gray = frame.gray[frame_slice]
    reference_gray = reference.gray[reference_slice]
    if frame_gray.size == 0:
        return float("inf"), 0

    frame_edge = frame.edge[frame_slice]
    reference_edge = reference.edge[reference_slice]
    luma_difference = np.abs(frame_gray - reference_gray) / 255.0
    edge_difference = np.abs(frame_edge - reference_edge) / 255.0
    edge_weight = rule_float(rules, "edge_difference_weight", DEFAULT_EDGE_DIFFERENCE_WEIGHT)
    structure_weight = 0.75 + rule_float(
        rules,
        "structured_pixel_weight",
        DEFAULT_STRUCTURED_PIXEL_WEIGHT,
    ) * np.clip((frame_edge + reference_edge) / 510.0, 0.0, 1.0)
    score = ((luma_difference + edge_weight * edge_difference) / (1.0 + edge_weight)) * structure_weight
    return float(score.mean()), int(frame_gray.size)


def shift_frame_horizontally(frame: Image.Image, dx: int) -> Image.Image:
    if dx == 0:
        return frame.copy()
    return frame.transform(
        frame.size,
        Image.Transform.AFFINE,
        (1, 0, dx, 0, 1, 0),
        resample=Image.Resampling.BILINEAR,
        fillcolor=(0, 0, 0),
    )


def build_weighted_difference(
    aligned: ImageFeatures,
    reference: ImageFeatures,
    rules,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    luma_difference = np.abs(aligned.gray - reference.gray).astype(np.float32)
    edge_difference = np.abs(aligned.edge - reference.edge).astype(np.float32)
    edge_weight = rule_float(rules, "edge_difference_weight", DEFAULT_EDGE_DIFFERENCE_WEIGHT)
    structure_weight = 0.75 + rule_float(
        rules,
        "structured_pixel_weight",
        DEFAULT_STRUCTURED_PIXEL_WEIGHT,
    ) * np.clip((aligned.edge + reference.edge) / 510.0, 0.0, 1.0)
    weighted_difference = (
        (luma_difference + edge_weight * edge_difference)
        / (1.0 + edge_weight)
        * structure_weight
    )
    return weighted_difference.astype(np.float32), luma_difference, edge_difference


def find_difference_components(
    weighted_difference: np.ndarray,
    luma_difference: np.ndarray,
    edge_difference: np.ndarray,
    roi: Roi,
    rules,
) -> list[DifferenceComponent]:
    difference_crop = weighted_difference[roi.y:roi.y + roi.height, roi.x:roi.x + roi.width]
    luma_crop = luma_difference[roi.y:roi.y + roi.height, roi.x:roi.x + roi.width]
    edge_crop = edge_difference[roi.y:roi.y + roi.height, roi.x:roi.x + roi.width]
    threshold = float(rules.difference_threshold)
    mask = difference_crop >= threshold
    return connected_components(mask, difference_crop, luma_crop, edge_crop, int(rules.min_defect_area))


def connected_components(
    mask: np.ndarray,
    difference: np.ndarray,
    luma_difference: np.ndarray,
    edge_difference: np.ndarray,
    min_area: int,
) -> list[DifferenceComponent]:
    height, width = mask.shape
    mask_flat = mask.ravel()
    difference_flat = difference.ravel()
    luma_flat = luma_difference.ravel()
    edge_flat = edge_difference.ravel()
    visited = bytearray(mask_flat.size)
    components: list[DifferenceComponent] = []

    for start_index, active in enumerate(mask_flat):
        if not active or visited[start_index]:
            continue

        stack = [start_index]
        visited[start_index] = 1
        min_x = max_x = start_index % width
        min_y = max_y = start_index // width
        area = 0
        difference_sum = 0.0
        luma_sum = 0.0
        edge_sum = 0.0
        max_difference = 0.0

        while stack:
            index = stack.pop()
            x = index % width
            y = index // width
            area += 1
            value = float(difference_flat[index])
            difference_sum += value
            luma_sum += float(luma_flat[index])
            edge_sum += float(edge_flat[index])
            max_difference = max(max_difference, value)
            min_x = min(min_x, x)
            max_x = max(max_x, x)
            min_y = min(min_y, y)
            max_y = max(max_y, y)

            if x > 0:
                neighbor = index - 1
                if mask_flat[neighbor] and not visited[neighbor]:
                    visited[neighbor] = 1
                    stack.append(neighbor)
            if x + 1 < width:
                neighbor = index + 1
                if mask_flat[neighbor] and not visited[neighbor]:
                    visited[neighbor] = 1
                    stack.append(neighbor)
            if y > 0:
                neighbor = index - width
                if mask_flat[neighbor] and not visited[neighbor]:
                    visited[neighbor] = 1
                    stack.append(neighbor)
            if y + 1 < height:
                neighbor = index + width
                if mask_flat[neighbor] and not visited[neighbor]:
                    visited[neighbor] = 1
                    stack.append(neighbor)

        if area >= min_area:
            components.append(DifferenceComponent(
                x=min_x,
                y=min_y,
                width=max_x - min_x + 1,
                height=max_y - min_y + 1,
                area=area,
                mean_difference=difference_sum / area,
                max_difference=max_difference,
                mean_luma_difference=luma_sum / area,
                mean_edge_difference=edge_sum / area,
            ))

    return components


def defect_confidence(component: DifferenceComponent, roi: Roi, rules) -> float:
    threshold_scale = max(255.0 - float(rules.difference_threshold), 1.0)
    severity = max(0.0, component.mean_difference - float(rules.difference_threshold)) / threshold_scale
    peak = min(1.0, component.max_difference / 255.0)
    area_ratio = component.area / max(roi.width * roi.height, 1)
    density = component.area / max(component.width * component.height, 1)
    return clamp(0.42 + severity * 0.36 + peak * 0.14 + area_ratio * 0.18 + density * 0.10, 0.0, 1.0)


def validate_roi(roi: Roi, image: Image.Image) -> None:
    if roi.x + roi.width > image.width or roi.y + roi.height > image.height:
        raise ValueError(f"ROI '{roi.name}' is outside image bounds")


def calculate_metrics(crop: Image.Image, brightness_min: float, brightness_max: float) -> RoiMetrics:
    pixels = list(crop.getdata())
    count = len(pixels)
    if count == 0:
        raise ValueError("ROI has no pixels")

    luminance_values = [
        0.2126 * red + 0.7152 * green + 0.0722 * blue
        for red, green, blue in pixels
    ]
    mean = sum(luminance_values) / count
    variance = sum((value - mean) ** 2 for value in luminance_values) / count
    deviation = math.sqrt(variance)

    return RoiMetrics(
        mean=mean,
        deviation=deviation,
        dark_ratio=sum(value < brightness_min for value in luminance_values) / count,
        bright_ratio=sum(value > brightness_max for value in luminance_values) / count,
    )


def find_roi_defects(
    roi: Roi,
    metrics: RoiMetrics,
    parameters: EvaluationParameters,
) -> list[DefectResponse]:
    rules = parameters.rules
    defects: list[DefectResponse] = []

    if metrics.mean < rules.brightness_min:
        defects.append(build_defect("too-dark", roi, confidence(rules.brightness_min - metrics.mean, 80), metrics))
    if metrics.mean > rules.brightness_max:
        defects.append(build_defect("too-bright", roi, confidence(metrics.mean - rules.brightness_max, 80), metrics))
    if metrics.dark_ratio > rules.max_dark_ratio:
        defects.append(build_defect("dark-area", roi, confidence(metrics.dark_ratio - rules.max_dark_ratio, 1), metrics))
    if metrics.bright_ratio > rules.max_bright_ratio:
        defects.append(build_defect("bright-area", roi, confidence(metrics.bright_ratio - rules.max_bright_ratio, 1), metrics))
    if rules.min_texture > 0 and metrics.deviation < rules.min_texture:
        defects.append(build_defect("low-texture", roi, confidence(rules.min_texture - metrics.deviation, 60), metrics))

    return defects


def build_defect(
    defect_type: str,
    roi: Roi,
    defect_confidence: float,
    metrics: RoiMetrics,
) -> DefectResponse:
    return DefectResponse(
        defectType=defect_type,
        roiName=roi.name,
        roiOrder=roi.order,
        confidence=defect_confidence,
        x=roi.x,
        y=roi.y,
        width=roi.width,
        height=roi.height,
        metrics={
            "mean": round(metrics.mean, 3),
            "deviation": round(metrics.deviation, 3),
            "darkRatio": round(metrics.dark_ratio, 5),
            "brightRatio": round(metrics.bright_ratio, 5),
        },
    )


def image_to_array(image: Image.Image) -> np.ndarray:
    return np.asarray(image.convert("RGB"), dtype=np.float32)


def alignment_sample_stride(size: tuple[int, int]) -> int:
    width, height = size
    return max(1, int(max(width, height) / MAX_ALIGNMENT_SAMPLE_SIZE))


def rule_int(rules, name: str, default: int) -> int:
    return int(getattr(rules, name, default))


def rule_float(rules, name: str, default: float) -> float:
    return float(getattr(rules, name, default))


def clamp(value: float, lower: float, upper: float) -> float:
    return max(lower, min(upper, value))


def confidence(delta: float, scale: float) -> float:
    return clamp(0.5 + delta / scale, 0.0, 1.0)


def build_summary(result: bool, defect_count: int, max_defects: int) -> str:
    verdict = "пройден" if result else "не пройден"
    return f"Контроль {verdict}: дефектов {defect_count}, допустимо {max_defects}"