import os import random import yaml from dataclasses import dataclass from typing import List, Tuple from PIL import Image, ImageDraw, ImageFont from threat_kb import CLASSES random.seed(42) # Tamanhos de canvas variáveis — diagramas reais raramente são quadrados # pequenos; costumam ser retangulares e relativamente grandes. CANVAS_SIZES = [(640, 640), (960, 640), (1280, 800), (1024, 1024), (1280, 960)] # Faixa de tamanho de cada ícone (em pixels) — aproxima a proporção # ícone/canvas observada em diagramas reais de arquitetura em nuvem. ICON_SIZE_RANGE = (34, 64) CLASS_COLORS = { "usuario": (240, 200, 90), "servidor_aplicacao": (230, 120, 60), "banco_de_dados": (70, 130, 210), "api_gateway": (150, 90, 210), "load_balancer": (90, 180, 140), "firewall_waf": (210, 70, 70), "fila_mensagens": (200, 150, 60), "cache": (100, 190, 190), "servico_externo": (150, 150, 150), "identity_provider": (220, 90, 150), "storage": (90, 140, 90), "cdn": (140, 110, 200), } CONTAINER_LABELS = [ "VPC", "Subnet Publica", "Subnet Privada", "Zona A", "Zona B", "Zona C", "Resource Group", "Virtual Network", "Availability Zone", "Region", ] @dataclass class Box: cls_id: int x0: int y0: int x1: int y1: int def _load_font(size=11): try: return ImageFont.truetype("DejaVuSans-Bold.ttf", size) except Exception: return ImageFont.load_default() def _load_font_small(size=9): try: return ImageFont.truetype("DejaVuSans.ttf", size) except Exception: return ImageFont.load_default() def _rects_overlap(a, b, margin=4): ax0, ay0, ax1, ay1 = a bx0, by0, bx1, by1 = b return not (ax1 + margin < bx0 or bx1 + margin < ax0 or ay1 + margin < by0 or by1 + margin < ay0) def _draw_containers(draw: ImageDraw.ImageDraw, canvas_w, canvas_h, font_small): """Desenha 1-4 contêineres visuais (caixas pontilhadas tipo VPC/Subnet), sem gerar anotação — apenas contexto visual, igual diagramas reais.""" n_containers = random.randint(1, 4) used = [] for _ in range(n_containers): w = random.randint(int(canvas_w * 0.3), int(canvas_w * 0.7)) h = random.randint(int(canvas_h * 0.3), int(canvas_h * 0.7)) x0 = random.randint(0, max(1, canvas_w - w)) y0 = random.randint(0, max(1, canvas_h - h)) x1, y1 = x0 + w, y0 + h color = (random.randint(80, 160), random.randint(80, 160), random.randint(180, 230)) # borda pontilhada simulada com pequenos segmentos step = 8 for x in range(x0, x1, step * 2): draw.line([(x, y0), (min(x + step, x1), y0)], fill=color, width=1) draw.line([(x, y1), (min(x + step, x1), y1)], fill=color, width=1) for y in range(y0, y1, step * 2): draw.line([(x0, y), (x0, min(y + step, y1))], fill=color, width=1) draw.line([(x1, y), (x1, min(y + step, y1))], fill=color, width=1) label = random.choice(CONTAINER_LABELS) draw.text((x0 + 4, y0 + 2), label, fill=color, font=font_small) used.append((x0, y0, x1, y1)) return used def _draw_background_noise(draw, canvas_w, canvas_h, font_small): """Ruído leve de fundo: linhas de grade e textos soltos, para robustez.""" for _ in range(random.randint(0, 3)): y = random.randint(0, canvas_h) draw.line([(0, y), (canvas_w, y)], fill=(235, 235, 240), width=1) for _ in range(random.randint(0, 4)): x, y = random.randint(0, canvas_w - 60), random.randint(0, canvas_h - 12) draw.text((x, y), random.choice(["v1.2", "prod", "TODO", "#3", "id-042"]), fill=(210, 210, 215), font=font_small) def _place_icons(n_components: int, canvas_w: int, canvas_h: int): """Posiciona ícones pequenos evitando muita sobreposição, sem grade rígida.""" chosen_labels = [random.choice(CLASSES) for _ in range(n_components)] placed = [] # (x0,y0,x1,y1) layout = [] max_attempts = 60 for label in chosen_labels: w = random.randint(*ICON_SIZE_RANGE) h = random.randint(*ICON_SIZE_RANGE) for _ in range(max_attempts): x0 = random.randint(4, max(5, canvas_w - w - 4)) y0 = random.randint(4, max(5, canvas_h - h - 4)) candidate = (x0, y0, x0 + w, y0 + h) # permite alguma sobreposição leve (10% das vezes), como em # diagramas reais densos, mas evita sobreposição excessiva if not any(_rects_overlap(candidate, p) for p in placed) or random.random() < 0.1: placed.append(candidate) layout.append((label, x0, y0, x0 + w, y0 + h)) break else: # não achou espaço livre; posiciona mesmo assim (densidade alta) x0 = random.randint(4, max(5, canvas_w - w - 4)) y0 = random.randint(4, max(5, canvas_h - h - 4)) layout.append((label, x0, y0, x0 + w, y0 + h)) return layout def _draw_component(draw: ImageDraw.ImageDraw, label: str, x0, y0, x1, y1, font_small): color = CLASS_COLORS.get(label, (180, 180, 180)) draw.rectangle([x0, y0, x1, y1], fill=color, outline=(20, 20, 20), width=1) text = label.replace("_", " ") # trunca o texto se não couber no ícone (realista, sem "explodir" a caixa) max_chars = max(3, (x1 - x0) // 5) if len(text) > max_chars: text = text[:max_chars - 1] + "…" draw.text((x0 + 2, y1 + 1), text, fill=(40, 40, 40), font=font_small) def generate_image(idx: int, out_img_dir: str, out_lbl_dir: str, font_small): canvas_w, canvas_h = random.choice(CANVAS_SIZES) img = Image.new("RGB", (canvas_w, canvas_h), (247, 248, 250)) draw = ImageDraw.Draw(img) _draw_containers(draw, canvas_w, canvas_h, font_small) _draw_background_noise(draw, canvas_w, canvas_h, font_small) # densidade variável: diagramas simples (5-9) e diagramas densos (15-30) if random.random() < 0.5: n_components = random.randint(5, 9) else: n_components = random.randint(12, 28) layout = _place_icons(n_components, canvas_w, canvas_h) # conexões (setas) entre componentes consecutivos centers = [((x0 + x1) // 2, (y0 + y1) // 2) for _, x0, y0, x1, y1 in layout] for i in range(len(centers) - 1): if random.random() < 0.6: # nem tudo conectado, como em diagramas reais draw.line([centers[i], centers[i + 1]], fill=(130, 130, 130), width=1) boxes = [] for label, x0, y0, x1, y1 in layout: _draw_component(draw, label, x0, y0, x1, y1, font_small) boxes.append(Box(CLASSES.index(label), x0, y0, x1, y1)) img_path = os.path.join(out_img_dir, f"diagram_{idx:04d}.png") lbl_path = os.path.join(out_lbl_dir, f"diagram_{idx:04d}.txt") img.save(img_path) with open(lbl_path, "w") as f: for b in boxes: xc = ((b.x0 + b.x1) / 2) / canvas_w yc = ((b.y0 + b.y1) / 2) / canvas_h w = (b.x1 - b.x0) / canvas_w h = (b.y1 - b.y0) / canvas_h f.write(f"{b.cls_id} {xc:.6f} {yc:.6f} {w:.6f} {h:.6f}\n") def build_dataset(root="data", n_train=400, n_val=80): font_small = _load_font_small(10) for split, n in [("train", n_train), ("val", n_val)]: img_dir = os.path.join(root, "images", split) lbl_dir = os.path.join(root, "labels", split) os.makedirs(img_dir, exist_ok=True) os.makedirs(lbl_dir, exist_ok=True) for i in range(n): generate_image(i, img_dir, lbl_dir, font_small) print(f"[dataset_generator] {n} imagens geradas para '{split}' em {img_dir}") data_yaml = { "path": os.path.abspath(root), "train": "images/train", "val": "images/val", "names": {i: name for i, name in enumerate(CLASSES)}, } yaml_path = os.path.join(root, "data.yaml") with open(yaml_path, "w") as f: yaml.safe_dump(data_yaml, f, allow_unicode=True, sort_keys=False) print(f"[dataset_generator] data.yaml escrito em {yaml_path}") if __name__ == "__main__": build_dataset()