luisasousa/aws-architecture-vision-detector

AWS Architecture Vision Detector

Detector YOLOv8 (Ultralytics) que localiza componentes, trust boundaries e setas de fluxo de dados em diagramas de arquitetura AWS (desenhados no draw.io ou em imagens de referência reais), alimentando um pipeline de modelagem de ameaças STRIDE.

Código completo, scripts de treino/geração de dataset e a etapa de pós-processamento (containment de trust boundary, OCR, matching seta→componente) estão em: https://github.com/luisaoliveira1/posdiap_7iadt_techchallenge_5/tree/main/models/vision-detector

Classes (51)

15 arquétipos de arquitetura (compute, database, storage, api_gateway, load_balancer, external_actor, edge_network, messaging_eventing, identity_access, secrets_crypto, security_service, logging_monitoring, cicd_devops, analytics_ml, communication) + 28 serviços AWS específicos (ALB, EC2, Lambda, RDS, S3, DynamoDB, etc. — promovidos quando havia ≥20 instâncias reais anotadas) + boundary genérico + 6 sub-tipos de boundary (aws_cloud, vpc, region, availability_zone, public_subnet, private_subnet) + arrowhead (estrutural, marca onde uma seta toca o componente de destino).

Como usar

from ultralytics import YOLO

model = YOLO("best.pt")
results = model.predict("meu_diagrama.png", conf=0.25, imgsz=1280)

Para o pipeline completo (que também monta o ArchitectureDiagram — trust boundaries com containment resolvido, OCR de nomes/protocolos, matching seta→componente), use predict.py do repositório do GitHub linkado acima, apontando WEIGHTS para este best.pt.

Métricas

• mAP50=0.85, mAP50-95=0.80 (split de teste do dataset de treino, mistura real+sintético — sinal padrão do YOLO, mas mais otimista que a avaliação abaixo)
• Contra 5 diagramas reais anotados à mão, nunca usados em treino (fonte de verdade mais confiável — ver dataset aws-architecture-diagrams): boundary F1≈0.57, componente F1≈0.81 (acurácia de categoria≈0.79 dado localização correta), arrowhead F1≈0.54

Treino

YOLOv8s, transfer learning a partir dos pesos pré-treinados no COCO, 150 épocas (early stopping com patience=25), imgsz=640, batch=16, augmentation (hsv, rotação leve, flip horizontal apenas, mosaic, mixup). Ver train.py no repositório para os hiperparâmetros completos.

Licença

CC BY 4.0 — o dataset de treino inclui uma parcela derivada do aws-icon-detector (Roboflow Universe, CC BY 4.0), remapeada para arquétipos; o restante é sintético (gerado via draw.io) ou anotado à mão pelos autores.