bisnet0/riftshield-architecture-yolo

RiftShield — architecture diagram detector (YOLOv8n)

Object detector trained to read software architecture diagrams (UML class diagrams, network topology, cloud diagrams) and turn them into structured elements that an LLM agent can reason about. It is the visual front-end of RiftShield, a threat-modeling platform: the detected elements become the graph that a LangGraph agent walks to apply STRIDE, retrieve mitigations through RAG (FAISS) and simulate attacks.

• Model: YOLOv8n fine-tuned (yolov8n.pt base), Ultralytics
• Classes: 34 (UML shapes and relationships, network devices, cloud/hardware icons, plus code / diagram page-level classification)
• Epochs: 11 · Weights: weights/yolov8n-architecture-best.pt
• Training data: 7 public Roboflow projects merged into one 34-class set (10,041 annotated files) — see the companion dataset bisnet0/riftshield-architecture-diagrams

Metrics (best epoch, validation split)

metric	value
Precision (B)	0.7416
Recall (B)	0.5561
mAP50 (B)	0.5942
mAP50-95 (B)	0.4183
epoch	10

Full per-epoch history in training/results.csv; curves in training/.

Usage

from ultralytics import YOLO

model = YOLO("weights/yolov8n-architecture-best.pt")
results = model("diagram.png", conf=0.25)
for box in results[0].boxes:
    print(model.names[int(box.cls)], [round(v, 1) for v in box.xyxy[0].tolist()])

Limitations

• Trained on a merged set of small public datasets; class distribution is imbalanced, so the rare UML relationship classes score well below the page-level diagram/code classes.
• Best suited to white-background diagrams similar to the training sources; hand-drawn and dark-theme diagrams degrade noticeably.
• Detection only — semantic interpretation (STRIDE category, mitigation) happens in the LangGraph agent, not in this model.

Citation

Part of RiftShield (FIAP Software Security Hackathon 2026). Code: https://github.com/bisnet0/RiftShield