File size: 1,838 Bytes
3cc53ab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """
Visualize sample images and masks from the CMP Facade dataset.
"""
import os
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from PIL import Image
import random
def visualize_samples(data_dir="data/cmp_facade", split="train", num_samples=5, save_path="outputs/dataset_visualization.png"):
os.makedirs(os.path.dirname(save_path), exist_ok=True)
split_dir = f"{data_dir}/{split}"
img_dir = f"{split_dir}/images"
mask_dir = f"{split_dir}/masks_binary"
img_files = sorted([f for f in os.listdir(img_dir) if f.endswith('.jpg')])[:20]
selected = random.sample(img_files, min(num_samples, len(img_files)))
fig, axes = plt.subplots(num_samples, 3, figsize=(12, 4 * num_samples))
if num_samples == 1:
axes = axes.reshape(1, -1)
for i, img_file in enumerate(selected):
img = Image.open(f"{img_dir}/{img_file}").convert('RGB')
mask = Image.open(f"{mask_dir}/{img_file.replace('.jpg', '.png')}").convert('L')
mask_arr = np.array(mask)
img_arr = np.array(img)
overlay = img_arr.copy()
overlay[mask_arr > 0] = [255, 0, 0]
blended = (img_arr * 0.6 + overlay * 0.4).astype(np.uint8)
axes[i, 0].imshow(img)
axes[i, 0].set_title(f"Image {img_file}")
axes[i, 0].axis('off')
axes[i, 1].imshow(mask, cmap='gray')
axes[i, 1].set_title("Facade Mask (class 1)")
axes[i, 1].axis('off')
axes[i, 2].imshow(blended)
axes[i, 2].set_title("Overlay")
axes[i, 2].axis('off')
plt.tight_layout()
plt.savefig(save_path, dpi=150)
plt.close()
print(f"Visualization saved to {save_path}")
if __name__ == "__main__":
visualize_samples("data/cmp_facade", "train", num_samples=5, save_path="outputs/dataset_visualization.png")
|