import argparse import os import random import shutil from pathlib import Path import yaml # Classes finais do projeto (devem coincidir com threat_kb.CLASSES). TARGET_CLASSES = [ "usuario", "servidor_aplicacao", "banco_de_dados", "api_gateway", "load_balancer", "firewall_waf", "fila_mensagens", "cache", "servico_externo", "identity_provider", "storage", "cdn", ] # Mapeamento heurístico: nome original (dataset AWS Icon Detector, Roboflow) # -> classe final do projeto. Ícones sem correspondência clara caem em # "servico_externo" (categoria genérica de componente externo/auxiliar). AWS_TO_TARGET = { "ACM": "firewall_waf", "ALB": "load_balancer", "AMI": "servidor_aplicacao", "API-Gateway": "api_gateway", "Active Directory Service": "identity_provider", "Airflow": "servico_externo", "Amplify": "servidor_aplicacao", "Analytics Services": "servico_externo", "AppFlow": "servico_externo", "Appsync": "api_gateway", "Athena": "banco_de_dados", "Aurora": "banco_de_dados", "Auto Scaling Group": "servidor_aplicacao", "Auto Scaling": "servidor_aplicacao", "Automated Tests": "servico_externo", "Availability Zone": "servico_externo", "Backup": "storage", "Build Environment": "servico_externo", "CDN": "cdn", "CUR": "servico_externo", "Call Metrics": "servico_externo", "Call Recordings": "storage", "Certificate Manager": "firewall_waf", "Client": "usuario", "Cloud Connector": "servico_externo", "Cloud Map": "servico_externo", "Cloud Search": "banco_de_dados", "Cloud Trail": "servico_externo", "Cloud Watch": "servico_externo", "CloudFormation Stack": "servico_externo", "CloudHSM": "identity_provider", "CloudWatch Alarm": "servico_externo", "Cloudfront": "cdn", "CodeBuild": "servico_externo", "CodeCommit": "servico_externo", "CodeDeploy": "servico_externo", "CodePipeline": "servico_externo", "Cognito": "identity_provider", "Comprehend": "servico_externo", "Config": "servico_externo", "Connect Contact Lens": "servico_externo", "Connect": "servico_externo", "Container": "servidor_aplicacao", "Control Tower": "servico_externo", "Customer Gateway": "firewall_waf", "DSI": "servico_externo", "Data Pipeline": "servico_externo", "DataSync": "storage", "Deploy Stage": "servico_externo", "Detective": "firewall_waf", "Direct Connect": "servico_externo", "Distribution": "cdn", "Docker Image": "servidor_aplicacao", "Dynamo DB": "banco_de_dados", "EBS": "storage", "EC2": "servidor_aplicacao", "EFS Mount Target": "storage", "EFS": "storage", "EKS": "servidor_aplicacao", "ELB": "load_balancer", "EMR": "servidor_aplicacao", "Edge Location": "cdn", "ElastiCache": "cache", "Elastic Container Registry": "storage", "Elastic Container Service": "servidor_aplicacao", "Elastic Search": "banco_de_dados", "Elemental MediaConvert": "servico_externo", "Elemental MediaPackage": "servico_externo", "Email": "servico_externo", "Endpoint": "api_gateway", "Event Bus": "fila_mensagens", "EventBridge": "fila_mensagens", "Experiment Duration": "servico_externo", "Experiments": "servico_externo", "Fargate": "servidor_aplicacao", "Fault Injection Simulator": "servico_externo", "Firewall Manager": "firewall_waf", "Flask": "servidor_aplicacao", "Flow logs": "servico_externo", "GameLift": "servidor_aplicacao", "Git": "servico_externo", "Github": "servico_externo", "Glacier": "storage", "Glue DataBrew": "servico_externo", "Glue": "servico_externo", "Grafana": "servico_externo", "GuardDuty": "firewall_waf", "IAM Role": "identity_provider", "IAM": "identity_provider", "IOT Core": "servico_externo", "Image Builder": "servidor_aplicacao", "Image": "storage", "Ingress": "load_balancer", "Inspector Agent": "firewall_waf", "Instances": "servidor_aplicacao", "Internet Gateway": "firewall_waf", "Internet": "servico_externo", "Jenkins": "servico_externo", "Key Management Service": "identity_provider", "Kibana": "servico_externo", "Kinesis Data Streams": "fila_mensagens", "Kubernetes": "servidor_aplicacao", "Lambda": "servidor_aplicacao", "Lex": "servico_externo", "MQ": "fila_mensagens", "Machine Learning": "servico_externo", "Macie": "firewall_waf", "Marketplace": "servico_externo", "Memcached": "cache", "Mobile Client": "usuario", "Mongo DB": "banco_de_dados", "MySQL": "banco_de_dados", "NAT Gateway": "firewall_waf", "Neptune": "banco_de_dados", "Network Adapter": "servico_externo", "Network Firewall": "firewall_waf", "Notebook": "servidor_aplicacao", "Order Controller": "servico_externo", "Organization Trail": "servico_externo", "Parameter Store": "storage", "Pinpoint": "servico_externo", "PostgreSQL": "banco_de_dados", "Private Link": "api_gateway", "Private Subnet": "servico_externo", "Prometheus": "servico_externo", "Public Subnet": "servico_externo", "Quarkus": "servidor_aplicacao", "Quicksight": "servico_externo", "RDS": "banco_de_dados", "React": "servico_externo", "Redis": "cache", "Redshift": "banco_de_dados", "Region": "servico_externo", "Rekognition": "servico_externo", "Results": "servico_externo", "Route 53": "cdn", "Route53": "cdn", "S3": "storage", "SAR": "servico_externo", "SDK": "servico_externo", "SES": "servico_externo", "SNS": "fila_mensagens", "SQS": "fila_mensagens", "SSM Agent": "servico_externo", "Sagemaker": "servidor_aplicacao", "Secret Manager": "identity_provider", "Security Group": "firewall_waf", "Security Hub": "firewall_waf", "Server": "servidor_aplicacao", "Service Catalog": "servico_externo", "Shield": "firewall_waf", "Sign-On": "identity_provider", "Slack": "servico_externo", "Snowball": "storage", "Stack": "servico_externo", "Step Function": "servico_externo", "Storage Gateway": "storage", "SwaggerHub": "servico_externo", "Systems Manager": "servico_externo", "TV": "servico_externo", "Table": "banco_de_dados", "Task Runner": "servico_externo", "Terraform": "servico_externo", "Text File": "storage", "Textract": "servico_externo", "Transcribe": "servico_externo", "Transfer Family": "storage", "Transit Gateway": "firewall_waf", "Translate": "servico_externo", "Trusted Advisor": "servico_externo", "Twilio": "servico_externo", "Users": "usuario", "VDA": "servico_externo", "VP Gateway": "firewall_waf", "VPC Router": "firewall_waf", "VPN Connection": "firewall_waf", "WAF": "firewall_waf", "Web Clients": "usuario", "Websites": "usuario", "X-Ray": "servico_externo", "aws": "servico_externo", "cache Worker": "cache", } def _bbox_from_polygon(coords): xs = coords[0::2] ys = coords[1::2] x0, x1 = min(xs), max(xs) y0, y1 = min(ys), max(ys) xc = (x0 + x1) / 2 yc = (y0 + y1) / 2 return xc, yc, x1 - x0, y1 - y0 def _convert_label_file(src_path, dest_path, source_names, class_lookup): out_lines = [] with open(src_path) as f: for raw_line in f: parts = raw_line.strip().split() if not parts: continue cls_id = int(parts[0]) coords = [float(v) for v in parts[1:]] if len(coords) == 4: xc, yc, w, h = coords else: xc, yc, w, h = _bbox_from_polygon(coords) src_name = source_names[cls_id] if cls_id < len(source_names) else None target_name = AWS_TO_TARGET.get(src_name, "servico_externo") target_id = class_lookup[target_name] out_lines.append(f"{target_id} {xc:.6f} {yc:.6f} {w:.6f} {h:.6f}") with open(dest_path, "w") as f: f.write("\n".join(out_lines) + ("\n" if out_lines else "")) def convert(source_dir: str, dest_dir: str, val_split: float, seed: int = 42): source_dir = Path(source_dir) dest_dir = Path(dest_dir) class_lookup = {name: i for i, name in enumerate(TARGET_CLASSES)} with open(source_dir / "data.yaml") as f: source_yaml = yaml.safe_load(f) source_names = source_yaml["names"] unmapped = sorted(set(source_names) - set(AWS_TO_TARGET)) if unmapped: print(f"[convert] Aviso: {len(unmapped)} classes sem mapeamento explícito, " f"usando 'servico_externo': {unmapped}") all_images = [] for split_dir in ["train", "test"]: img_dir = source_dir / split_dir / "images" if img_dir.exists(): for img_path in sorted(img_dir.iterdir()): if img_path.suffix.lower() in (".jpg", ".jpeg", ".png"): label_path = source_dir / split_dir / "labels" / (img_path.stem + ".txt") if label_path.exists(): all_images.append((img_path, label_path)) random.seed(seed) random.shuffle(all_images) n_val = max(1, int(len(all_images) * val_split)) val_set = all_images[:n_val] train_set = all_images[n_val:] for split_name, items in [("train", train_set), ("val", val_set)]: img_out = dest_dir / "images" / split_name lbl_out = dest_dir / "labels" / split_name img_out.mkdir(parents=True, exist_ok=True) lbl_out.mkdir(parents=True, exist_ok=True) for img_path, label_path in items: shutil.copy(img_path, img_out / img_path.name) _convert_label_file(label_path, lbl_out / (img_path.stem + ".txt"), source_names, class_lookup) print(f"[convert] {len(items)} imagens copiadas para '{split_name}'") data_yaml = { "path": str(dest_dir.resolve()), "train": "images/train", "val": "images/val", "names": {i: name for i, name in enumerate(TARGET_CLASSES)}, } yaml_path = dest_dir / f"{dest_dir.name}.yaml" with open(yaml_path, "w") as f: yaml.safe_dump(data_yaml, f, allow_unicode=True, sort_keys=False) print(f"[convert] {yaml_path} criado") def main(): parser = argparse.ArgumentParser(description="Converte dataset AWS Icon Detector (Roboflow) para as classes do projeto") parser.add_argument("--source", type=str, required=True, help="Pasta extraída do dataset Roboflow (contém train/, test/, data.yaml)") parser.add_argument("--dest", type=str, default="data_real_aws", help="Pasta de saída no layout esperado por finetune.py") parser.add_argument("--val-split", type=float, default=0.15) args = parser.parse_args() convert(args.source, args.dest, args.val_split) if __name__ == "__main__": main()