AI Fabric Defect Detection System

Разработка AI для детекции дефектов ткани и текстиля

Инспекция тканей — обязательный процесс в текстильном производстве: каждый рулон ткани (50–100 м) проходит через раппортный контроль на скорости 20–60 м/мин. Ручная инспекция: 1 инспектор контролирует 20–30% производства и пропускает 15–25% дефектов из-за усталости. AI-система: 100% покрытие, скорость до 100 м/мин, стабильная точность. Датасеты: TILDA (Textile Inspection in the Loop with Deep Algorithms), AITEX Fabric Dataset.

Детектор дефектов ткани

import numpy as np
import cv2
import torch
from anomalib.models import Patchcore, EfficientAD
from ultralytics import YOLO
from dataclasses import dataclass
from typing import Optional

@dataclass
class FabricDefect:
    defect_type: str      # hole / stain / broken_thread / weaving_error / scratch / fold
    severity: str         # minor / major / critical
    bbox: list
    area_px2: int
    confidence: float
    location_pct: tuple   # (x%, y%) - относительное положение

class FabricDefectDetector:
    """
    Двухуровневая детекция дефектов ткани:
    Level 1: Anomaly detection (PatchCore) — обучается только на хорошей ткани
    Level 2: Defect classification (YOLO) — если нужна классификация по типу

    AITEX Fabric Dataset: 7 типов дефектов, 12 видов тканей.
    TILDA: производственные дефекты, 8 классов.
    """
    DEFECT_CLASSES = {
        0: ('hole', 'critical'),
        1: ('stain', 'major'),
        2: ('broken_thread', 'major'),
        3: ('weaving_error', 'major'),
        4: ('scratch', 'minor'),
        5: ('fold', 'minor'),
        6: ('cut', 'critical'),
        7: ('knotting', 'minor')
    }

    def __init__(self, anomaly_model_path: str,
                  defect_model_path: Optional[str] = None,
                  anomaly_threshold: float = 0.5,
                  device: str = 'cuda'):
        self.anomaly_model = Patchcore.load_from_checkpoint(anomaly_model_path)
        self.anomaly_model.eval()
        self.anomaly_threshold = anomaly_threshold

        self.defect_model = YOLO(defect_model_path) if defect_model_path else None
        self.device = device

        # Трансформация для tile-based инспекции
        from torchvision import transforms
        self.transform = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    def inspect_fabric_strip(self, strip: np.ndarray,
                               tile_size: int = 256,
                               overlap: int = 32) -> dict:
        """
        Инспекция полосы ткани (горизонтальный кадр с линейной камеры).
        Тайл-based обработка для скоростных линий.
        """
        h, w = strip.shape[:2]
        from PIL import Image

        anomaly_map = np.zeros((h, w), dtype=np.float32)
        count_map = np.zeros((h, w), dtype=np.float32)

        stride = tile_size - overlap

        # Tile extraction
        tiles = []
        tile_positions = []
        for y in range(0, h - tile_size + 1, stride):
            for x in range(0, w - tile_size + 1, stride):
                tile = strip[y:y+tile_size, x:x+tile_size]
                pil_tile = Image.fromarray(cv2.cvtColor(tile, cv2.COLOR_BGR2RGB))
                tensor = self.transform(pil_tile)
                tiles.append(tensor)
                tile_positions.append((x, y))

        if not tiles:
            return {'defects': [], 'anomaly_score': 0, 'pass': True}

        # Batch inference
        batch = torch.stack(tiles)
        with torch.no_grad():
            outputs = self.anomaly_model({'image': batch})
            scores = outputs['pred_score'].cpu().numpy()
            anomaly_maps = outputs.get('anomaly_map')

        # Сборка общей карты аномалий
        for i, (x, y) in enumerate(tile_positions):
            if anomaly_maps is not None:
                am = anomaly_maps[i].cpu().numpy()
                am_resized = cv2.resize(am, (tile_size, tile_size))
                anomaly_map[y:y+tile_size, x:x+tile_size] += am_resized
                count_map[y:y+tile_size, x:x+tile_size] += 1

        # Нормализация
        count_map = np.maximum(count_map, 1)
        anomaly_map /= count_map

        # Обнаружение дефектных зон
        defects = self._extract_defects(anomaly_map, strip, w, h)

        overall_score = float(np.max(scores))

        return {
            'defects': [d.__dict__ for d in defects],
            'anomaly_score': round(overall_score, 4),
            'anomaly_map': anomaly_map,
            'pass': overall_score < self.anomaly_threshold and len(defects) == 0
        }

    def _extract_defects(self, anomaly_map: np.ndarray,
                          original: np.ndarray,
                          w: int, h: int) -> list[FabricDefect]:
        """Извлечение bbox дефектов из карты аномалий"""
        defects = []

        if anomaly_map.max() < 0.3:
            return defects

        # Бинаризация карты аномалий
        norm_map = ((anomaly_map / anomaly_map.max()) * 255).astype(np.uint8)
        _, thresh = cv2.threshold(norm_map, int(self.anomaly_threshold * 255),
                                   255, cv2.THRESH_BINARY)

        # Морфологическая чистка
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        cleaned = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
        cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_OPEN, kernel)

        contours, _ = cv2.findContours(cleaned, cv2.RETR_EXTERNAL,
                                         cv2.CHAIN_APPROX_SIMPLE)

        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area < 100:  # слишком мелкие
                continue

            x, y, bw, bh = cv2.boundingRect(cnt)
            max_anomaly = float(anomaly_map[y:y+bh, x:x+bw].max())

            severity = ('critical' if max_anomaly > 0.85 else
                        'major' if max_anomaly > 0.65 else 'minor')

            # Попытка классификации дефекта
            defect_type = 'unknown'
            if self.defect_model:
                crop = original[y:y+bh, x:x+bw]
                if crop.size > 0:
                    results = self.defect_model(crop, conf=0.35, verbose=False)
                    if results[0].boxes and len(results[0].boxes):
                        cls_id = int(results[0].boxes.cls[0].item())
                        defect_type, severity = self.DEFECT_CLASSES.get(
                            cls_id, ('unknown', severity)
                        )

            defects.append(FabricDefect(
                defect_type=defect_type,
                severity=severity,
                bbox=[x, y, x+bw, y+bh],
                area_px2=int(area),
                confidence=max_anomaly,
                location_pct=(round(x/w*100, 1), round(y/h*100, 1))
            ))

        return defects