CLASSES = [
'Triangle', 'Square', 'Pentagon', 'Hexagon',
'Heptagon', 'Circle', 'Ellipse', 'Star', 'Cross'
]
N_LADOS = {'Triangle': 3, 'Square': 4, 'Pentagon': 5, 'Hexagon': 6, 'Heptagon': 7}
def poligono_regular(cx, cy, r, n_lados, rot_graus):
ang0 = np.deg2rad(rot_graus - 90)
angs = ang0 + 2 * np.pi * np.arange(n_lados) / n_lados
return np.stack([cx + r * np.cos(angs), cy + r * np.sin(angs)], axis=1)
def poligono_estrela(cx, cy, r_externo, rot_graus, n_pontas=5):
r_interno = r_externo * 0.45
ang0 = np.deg2rad(rot_graus - 90)
angs = ang0 + np.pi * np.arange(2 * n_pontas) / n_pontas
raios = np.where(np.arange(2 * n_pontas) % 2 == 0, r_externo, r_interno)
return np.stack([cx + raios * np.cos(angs), cy + raios * np.sin(angs)], axis=1)
def poligono_cruz(cx, cy, r, rot_graus, espessura_rel=0.35):
w = r * espessura_rel
base = np.array([
(-w, -r), (w, -r), (w, -w), (r, -w), (r, w), (w, w),
(w, r), (-w, r), (-w, w), (-r, w), (-r, -w), (-w, -w),
])
theta = np.deg2rad(rot_graus)
R = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])
return base @ R.T + np.array([cx, cy])
def desenha_objeto(img, classe_idx, cx, cy, tamanho, rotacao, cor):
nome = CLASSES[classe_idx]
if nome in N_LADOS:
pts = poligono_regular(cx, cy, tamanho, N_LADOS[nome], rotacao)
cv2.fillPoly(img, [pts.astype(np.int32)], cor)
xs, ys = pts[:, 0], pts[:, 1]
elif nome == 'Star':
pts = poligono_estrela(cx, cy, tamanho, rotacao)
cv2.fillPoly(img, [pts.astype(np.int32)], cor)
xs, ys = pts[:, 0], pts[:, 1]
elif nome == 'Cross':
pts = poligono_cruz(cx, cy, tamanho, rotacao)
cv2.fillPoly(img, [pts.astype(np.int32)], cor)
xs, ys = pts[:, 0], pts[:, 1]
elif nome == 'Circle':
cv2.circle(img, (int(cx), int(cy)), int(tamanho), cor, -1)
xs, ys = np.array([cx - tamanho, cx + tamanho]), np.array([cy - tamanho, cy + tamanho])
else: # Ellipse
eixo = (int(tamanho), int(tamanho * 0.6))
cv2.ellipse(img, (int(cx), int(cy)), eixo, rotacao, 0, 360, cor, -1)
ang = np.deg2rad(rotacao)
dx = np.hypot(eixo[0] * np.cos(ang), eixo[1] * np.sin(ang))
dy = np.hypot(eixo[0] * np.sin(ang), eixo[1] * np.cos(ang))
xs, ys = np.array([cx - dx, cx + dx]), np.array([cy - dy, cy + dy])
return xs.min(), ys.min(), xs.max(), ys.max()
def adiciona_ruido_sal_pimenta(img, quantidade=0.05):
img_ruidosa = img.copy()
h, w, c = img_ruidosa.shape
num_ruido = int(quantidade * h * w)
# Sal (255, 255, 255)
coords_sal = [np.random.randint(0, i - 1, num_ruido) for i in (h, w)]
img_ruidosa[coords_sal[0], coords_sal[1]] = [255, 255, 255]
# Pimenta (0, 0, 0)
coords_pimenta = [np.random.randint(0, i - 1, num_ruido) for i in (h, w)]
img_ruidosa[coords_pimenta[0], coords_pimenta[1]] = [0, 0, 0]
return img_ruidosa
def gera_imagem_ruidosa(tam_img=160, n_objetos=(1, 3), taxa_ruido=0.01, rng=None):
rng = rng or random.Random()
img_limpa = np.full((tam_img, tam_img, 3), 255, dtype=np.uint8)
anotacoes = []
for _ in range(rng.randint(*n_objetos)):
classe_idx = rng.randrange(len(CLASSES))
tamanho = rng.randint(tam_img // 10, tam_img // 5)
cx = rng.randint(tamanho + 2, tam_img - tamanho - 2)
cy = rng.randint(tamanho + 2, tam_img - tamanho - 2)
rotacao = rng.uniform(0, 360)
cor = tuple(rng.sample(range(30, 226), 3))
x0, y0, x1, y1 = desenha_objeto(img_limpa, classe_idx, cx, cy, tamanho, rotacao, cor)
x0, y0 = max(x0, 0), max(y0, 0)
x1, y1 = min(x1, tam_img), min(y1, tam_img)
# Padrão YOLO: (classe, x_centro, y_centro, largura, altura) normalizados
xc, yc = (x0 + x1) / 2 / tam_img, (y0 + y1) / 2 / tam_img
w, h = (x1 - x0) / tam_img, (y1 - y0) / tam_img
anotacoes.append((classe_idx, xc, yc, w, h))
img_ruidosa = adiciona_ruido_sal_pimenta(img_limpa, quantidade=taxa_ruido)
return img_ruidosa, anotacoes