import os
import cv2
import numpy as np
from typing import Tuple
# Definicja kolorów i etykiet
CLASS_COLORS = {
0: {'color_rgb': (0, 0, 0), 'label': 'background'},
1: {'color_rgb': (255, 0, 0), 'label': 'chondroosseous border'},
2: {'color_rgb': (0, 255, 0), 'label': 'femoral head'},
3: {'color_rgb': (0, 0, 255), 'label': 'labrum'},
4: {'color_rgb': (255, 255, 0), 'label': 'cartilagineous roof'},
5: {'color_rgb': (0, 255, 255), 'label': 'bony roof'},
6: {'color_rgb': (159, 2, 250), 'label': 'bony rim'},
7: {'color_rgb': (255, 132, 0), 'label': 'lower limb'},
8: {'color_rgb': (255, 0, 255), 'label': 'baseline'},
9: {'color_rgb': (66, 135, 245), 'label': 'lower limb template'},
10: {'color_rgb': (255, 69, 0), 'label': 'lower limb - alg. v3'}
}
def visualize_masks(input_dir, output_dir):
# Tworzenie katalogu wyjściowego, jeśli nie istnieje
os.makedirs(output_dir, exist_ok=True)
# Iteracja po plikach PNG w katalogu wejściowym
for filename in os.listdir(input_dir):
if filename.endswith(".png"):
# Wczytywanie maski
mask_path = os.path.join(input_dir, filename)
mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
# Tworzenie pustego obrazu RGB o tych samych wymiarach co maska
color_image = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
# Zamiana etykiet na kolory RGB zgodnie z CLASS_COLORS
for class_id, color_info in CLASS_COLORS.items():
color_rgb = color_info['color_rgb']
color_image[mask == class_id] = color_rgb
# Znajdowanie punktów baseline
right_point, left_point = find_baseline(mask)
# Rysowanie baseline na obrazie, jeżeli punkty istnieją
if right_point[0] is not None and left_point[0] is not None:
cv2.line(color_image, right_point, left_point, (255, 255, 255), 1)
# Zapisanie obrazu kolorowego do katalogu wyjściowego
output_path = os.path.join(output_dir, filename)
cv2.imwrite(output_path, color_image)
print(f"Maska {filename} została zapisana w katalogu wyjściowym.")
# Przykładowe użycie
input_dir = "./Angles/dane/ta_jedna"
output_dir = "./Angles/dane/wizul"
visualize_masks(input_dir, output_dir)