Spaces:

imperat0
/

strike

Sleeping

strike / src /dataset_generator.py

Projeto

73c5503 4 months ago

8.16 kB

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()

	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()