338 lines
10 KiB
Python
338 lines
10 KiB
Python
|
"""
|
||
|
|
ATM Anomaly Detection Model
|
||
|
|
|
||
|
|
YOLOv8-based anomaly detector for ATM images.
|
||
|
|
Detects: damage, graffiti, obstruction, skimming devices, out-of-service
|
||
|
|
"""
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
from ultralytics import YOLO
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import List, Dict, Tuple, Optional
|
||
|
|
import numpy as np
|
||
|
|
import cv2
|
||
|
|
|
||
|
|
class ATMAnomalyDetector:
|
||
|
|
"""
|
||
|
|
ATM Anomaly Detection using YOLOv8.
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
detector = ATMAnomalyDetector(model_path='models/best.pt')
|
||
|
|
results = detector.predict('path/to/image.jpg')
|
||
|
|
"""
|
||
|
|
|
||
|
|
# Anomaly class mapping
|
||
|
|
CLASS_NAMES = {
|
||
|
|
0: 'physical_damage',
|
||
|
|
1: 'vandalism',
|
||
|
|
2: 'graffiti',
|
||
|
|
3: 'dirt_debris',
|
||
|
|
4: 'obstruction',
|
||
|
|
5: 'skimming_device',
|
||
|
|
6: 'suspicious_attachment',
|
||
|
|
7: 'out_of_service',
|
||
|
|
8: 'screen_damage',
|
||
|
|
9: 'cash_jam',
|
||
|
|
10: 'receipt_jam',
|
||
|
|
11: 'lighting_failure',
|
||
|
|
12: 'camera_blind',
|
||
|
|
13: 'network_down'
|
||
|
|
}
|
||
|
|
|
||
|
|
SEVERITY_MAP = {
|
||
|
|
'skimming_device': 5,
|
||
|
|
'suspicious_attachment': 5,
|
||
|
|
'physical_damage': 4,
|
||
|
|
'vandalism': 4,
|
||
|
|
'camera_blind': 4,
|
||
|
|
'obstruction': 3,
|
||
|
|
'out_of_service': 3,
|
||
|
|
'screen_damage': 3,
|
||
|
|
'cash_jam': 3,
|
||
|
|
'lighting_failure': 3,
|
||
|
|
'network_down': 3,
|
||
|
|
'graffiti': 2,
|
||
|
|
'dirt_debris': 2,
|
||
|
|
'receipt_jam': 2
|
||
|
|
}
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
model_path: Optional[str] = None,
|
||
|
|
conf_threshold: float = 0.25,
|
||
|
|
iou_threshold: float = 0.45,
|
||
|
|
device: str = 'auto'
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Initialize detector.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
model_path: Path to trained YOLO model
|
||
|
|
conf_threshold: Confidence threshold for detections
|
||
|
|
iou_threshold: IoU threshold for NMS
|
||
|
|
device: 'cpu', 'cuda', or 'auto'
|
||
|
|
"""
|
||
|
|
self.conf_threshold = conf_threshold
|
||
|
|
self.iou_threshold = iou_threshold
|
||
|
|
|
||
|
|
# Set device
|
||
|
|
if device == 'auto':
|
||
|
|
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
|
|
else:
|
||
|
|
self.device = device
|
||
|
|
|
||
|
|
# Load model
|
||
|
|
if model_path and Path(model_path).exists():
|
||
|
|
self.model = YOLO(model_path)
|
||
|
|
else:
|
||
|
|
# Load pretrained COCO model as base
|
||
|
|
print("No trained model found. Loading YOLOv8n pretrained...")
|
||
|
|
self.model = YOLO('yolov8n.pt')
|
||
|
|
|
||
|
|
self.model.to(self.device)
|
||
|
|
|
||
|
|
def predict(
|
||
|
|
self,
|
||
|
|
image_path: str,
|
||
|
|
save: bool = False,
|
||
|
|
save_dir: Optional[str] = None
|
||
|
|
) -> List[Dict]:
|
||
|
|
"""
|
||
|
|
Run anomaly detection on image.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
image_path: Path to image file
|
||
|
|
save: Whether to save annotated image
|
||
|
|
save_dir: Directory to save annotations
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
List of detection dicts with keys:
|
||
|
|
- class_id: int
|
||
|
|
- class_name: str
|
||
|
|
- confidence: float
|
||
|
|
- bbox: [x1, y1, x2, y2]
|
||
|
|
- severity: int
|
||
|
|
"""
|
||
|
|
# Run inference
|
||
|
|
results = self.model(
|
||
|
|
image_path,
|
||
|
|
conf=self.conf_threshold,
|
||
|
|
iou=self.iou_threshold,
|
||
|
|
device=self.device,
|
||
|
|
verbose=False
|
||
|
|
)
|
||
|
|
|
||
|
|
detections = []
|
||
|
|
|
||
|
|
for result in results:
|
||
|
|
boxes = result.boxes
|
||
|
|
if boxes is None:
|
||
|
|
continue
|
||
|
|
|
||
|
|
for box in boxes:
|
||
|
|
class_id = int(box.cls.item())
|
||
|
|
confidence = float(box.conf.item())
|
||
|
|
bbox = box.xyxy[0].cpu().numpy().tolist()
|
||
|
|
|
||
|
|
class_name = self.CLASS_NAMES.get(class_id, 'unknown')
|
||
|
|
severity = self.SEVERITY_MAP.get(class_name, 1)
|
||
|
|
|
||
|
|
detection = {
|
||
|
|
'class_id': class_id,
|
||
|
|
'class_name': class_name,
|
||
|
|
'confidence': round(confidence, 4),
|
||
|
|
'bbox': [round(x, 2) for x in bbox],
|
||
|
|
'severity': severity,
|
||
|
|
'requires_action': severity >= 4
|
||
|
|
}
|
||
|
|
detections.append(detection)
|
||
|
|
|
||
|
|
# Sort by severity (highest first)
|
||
|
|
detections.sort(key=lambda x: x['severity'], reverse=True)
|
||
|
|
|
||
|
|
# Save annotated image if requested
|
||
|
|
if save and save_dir:
|
||
|
|
self._save_annotated(image_path, detections, save_dir)
|
||
|
|
|
||
|
|
return detections
|
||
|
|
|
||
|
|
def predict_batch(
|
||
|
|
self,
|
||
|
|
image_paths: List[str],
|
||
|
|
batch_size: int = 8
|
||
|
|
) -> List[List[Dict]]:
|
||
|
|
"""
|
||
|
|
Run detection on batch of images.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
image_paths: List of image paths
|
||
|
|
batch_size: Batch size for inference
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
List of detection lists
|
||
|
|
"""
|
||
|
|
all_detections = []
|
||
|
|
|
||
|
|
for i in range(0, len(image_paths), batch_size):
|
||
|
|
batch = image_paths[i:i + batch_size]
|
||
|
|
results = self.model(
|
||
|
|
batch,
|
||
|
|
conf=self.conf_threshold,
|
||
|
|
iou=self.iou_threshold,
|
||
|
|
device=self.device,
|
||
|
|
verbose=False
|
||
|
|
)
|
||
|
|
|
||
|
|
for result in results:
|
||
|
|
detections = []
|
||
|
|
boxes = result.boxes
|
||
|
|
|
||
|
|
if boxes is not None:
|
||
|
|
for box in boxes:
|
||
|
|
class_id = int(box.cls.item())
|
||
|
|
confidence = float(box.conf.item())
|
||
|
|
bbox = box.xyxy[0].cpu().numpy().tolist()
|
||
|
|
class_name = self.CLASS_NAMES.get(class_id, 'unknown')
|
||
|
|
|
||
|
|
detections.append({
|
||
|
|
'class_id': class_id,
|
||
|
|
'class_name': class_name,
|
||
|
|
'confidence': round(confidence, 4),
|
||
|
|
'bbox': [round(x, 2) for x in bbox],
|
||
|
|
'severity': self.SEVERITY_MAP.get(class_name, 1),
|
||
|
|
'requires_action': self.SEVERITY_MAP.get(class_name, 1) >= 4
|
||
|
|
})
|
||
|
|
|
||
|
|
detections.sort(key=lambda x: x['severity'], reverse=True)
|
||
|
|
all_detections.append(detections)
|
||
|
|
|
||
|
|
return all_detections
|
||
|
|
|
||
|
|
def _save_annotated(
|
||
|
|
self,
|
||
|
|
image_path: str,
|
||
|
|
detections: List[Dict],
|
||
|
|
save_dir: str
|
||
|
|
):
|
||
|
|
"""Save annotated image with bounding boxes."""
|
||
|
|
import os
|
||
|
|
os.makedirs(save_dir, exist_ok=True)
|
||
|
|
|
||
|
|
image = cv2.imread(image_path)
|
||
|
|
|
||
|
|
for det in detections:
|
||
|
|
x1, y1, x2, y2 = map(int, det['bbox'])
|
||
|
|
color = (0, 0, 255) if det['severity'] >= 4 else (0, 165, 255)
|
||
|
|
|
||
|
|
cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
|
||
|
|
|
||
|
|
label = f"{det['class_name']} {det['confidence']:.2f}"
|
||
|
|
cv2.putText(
|
||
|
|
image, label, (x1, y1 - 10),
|
||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2
|
||
|
|
)
|
||
|
|
|
||
|
|
filename = Path(image_path).name
|
||
|
|
save_path = os.path.join(save_dir, f"annotated_{filename}")
|
||
|
|
cv2.imwrite(save_path, image)
|
||
|
|
|
||
|
|
def train(
|
||
|
|
self,
|
||
|
|
data_yaml: str,
|
||
|
|
epochs: int = 100,
|
||
|
|
batch_size: int = 16,
|
||
|
|
img_size: int = 640,
|
||
|
|
output_dir: str = 'models/checkpoints'
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Train model on custom dataset.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
data_yaml: Path to data.yaml for YOLO
|
||
|
|
epochs: Number of training epochs
|
||
|
|
batch_size: Batch size
|
||
|
|
img_size: Input image size
|
||
|
|
output_dir: Where to save checkpoints
|
||
|
|
"""
|
||
|
|
self.model.train(
|
||
|
|
data=data_yaml,
|
||
|
|
epochs=epochs,
|
||
|
|
batch=batch_size,
|
||
|
|
imgsz=img_size,
|
||
|
|
project=output_dir,
|
||
|
|
name='atm_anomaly',
|
||
|
|
device=self.device
|
||
|
|
)
|
||
|
|
|
||
|
|
def export(
|
||
|
|
self,
|
||
|
|
format: str = 'onnx',
|
||
|
|
output_path: Optional[str] = None
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Export model to deployment format.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
format: 'onnx', 'torchscript', 'openvino', 'engine'
|
||
|
|
output_path: Where to save exported model
|
||
|
|
"""
|
||
|
|
self.model.export(format=format)
|
||
|
|
|
||
|
|
if output_path:
|
||
|
|
import shutil
|
||
|
|
default_path = f"models/checkpoints/atm_anomaly/weights/best.{format}"
|
||
|
|
if Path(default_path).exists():
|
||
|
|
shutil.copy2(default_path, output_path)
|
||
|
|
print(f"Exported to {output_path}")
|
||
|
|
|
||
|
|
|
||
|
|
def create_data_yaml(
|
||
|
|
train_dir: str,
|
||
|
|
val_dir: str,
|
||
|
|
test_dir: Optional[str] = None,
|
||
|
|
class_names: Optional[List[str]] = None,
|
||
|
|
output_path: str = 'data/data.yaml'
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Create YOLO data.yaml configuration file.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
train_dir: Path to train directory
|
||
|
|
val_dir: Path to validation directory
|
||
|
|
test_dir: Path to test directory (optional)
|
||
|
|
class_names: List of class names
|
||
|
|
output_path: Where to save yaml
|
||
|
|
"""
|
||
|
|
if class_names is None:
|
||
|
|
class_names = list(ATMAnomalyDetector.CLASS_NAMES.values())
|
||
|
|
|
||
|
|
import yaml
|
||
|
|
|
||
|
|
data = {
|
||
|
|
'path': str(Path(train_dir).parent),
|
||
|
|
'train': str(Path(train_dir).relative_to(Path(train_dir).parent)),
|
||
|
|
'val': str(Path(val_dir).relative_to(Path(val_dir).parent)),
|
||
|
|
'nc': len(class_names),
|
||
|
|
'names': class_names
|
||
|
|
}
|
||
|
|
|
||
|
|
if test_dir:
|
||
|
|
data['test'] = str(Path(test_dir).relative_to(Path(test_dir).parent))
|
||
|
|
|
||
|
|
with open(output_path, 'w') as f:
|
||
|
|
yaml.dump(data, f, default_flow_style=False)
|
||
|
|
|
||
|
|
print(f"Created {output_path}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
# Example usage
|
||
|
|
detector = ATMAnomalyDetector()
|
||
|
|
|
||
|
|
# Single image prediction
|
||
|
|
results = detector.predict('data/test/atm_001.jpg', save=True, save_dir='output')
|
||
|
|
print(f"Found {len(results)} anomalies")
|
||
|
|
for r in results:
|
||
|
|
print(f" - {r['class_name']}: {r['confidence']:.2f} (severity: {r['severity']})")
|