acd23's picture
Upload folder using huggingface_hub
3cc53ab verified
"""
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")