bitmapplots.py

import os
import numpy as np
from PIL import Image

import matplotlib.pyplot as plt


tool = 'carpet'
# Path to the root directory containing subfolders
root_dir = f'/home/marthijnvd/master/docs/proftpdafterdry/{tool}'

def find_fuzz_bitmap_files(root):
    fuzz_files = []
    for subdir, _, files in os.walk(root):
        for file in files:
            if file == 'fuzz_bitmap':
                fuzz_files.append(os.path.join(subdir, file))
    return fuzz_files

def plot_bitmap(file_path, output_dir):
    with open(file_path, 'rb') as f:
        data = f.read()
    # Try to make a square image, or as close as possible
    length = len(data)
    print(f"Processing {file_path} with length {length}")
    size = int(np.ceil(np.sqrt(length)))
    padded = np.frombuffer(data, dtype=np.uint8)
    if padded.size < size * size:
        padded = np.pad(padded, (0, size*size - padded.size), 'constant')
    img = padded.reshape((size, size))
    plt.imshow(img, cmap='gray', interpolation='nearest')
    plt.axis('off')
    out_name = os.path.basename(os.path.dirname(file_path)) + '_bitmap.png'
    plt.savefig(os.path.join(output_dir, out_name), bbox_inches='tight', pad_inches=0)
    plt.close()

def combine_bitmaps(input_dir, output_dir):

    # Find all generated bitmap pngs
    bitmap_files = []
    for f in os.listdir(input_dir):
        if f.endswith('.png'):
            bitmap_files.append(os.path.join(input_dir, f))
    if not bitmap_files:
        print("No bitmap images found to combine.")
        return

    # Open all images and convert to numpy arrays
    imgs = [np.array(Image.open(f).convert('L')) for f in bitmap_files]
    # Resize all to the smallest shape
    min_shape = np.min([img.shape for img in imgs], axis=0)
    imgs = [img[:min_shape[0], :min_shape[1]] for img in imgs]

    # Calculate overlap percentage (pixels nonzero in all images)
    stacked = np.stack([img > 0 for img in imgs], axis=0)
    overlap = np.all(stacked, axis=0)
    overlap_percentage = 100.0 * np.sum(overlap) / overlap.size
    print(f"Overlap percentage: {overlap_percentage:.2f}%")

    # Overlap percentage disregarding white pixels (255)
    non_white_mask = np.all([img < 255 for img in imgs], axis=0)
    if np.any(non_white_mask):
        overlap_nonwhite = np.all([img > 0 for img in imgs], axis=0) & non_white_mask
        overlap_nonwhite_percentage = 100.0 * np.sum(overlap_nonwhite) / np.sum(non_white_mask)
        print(f"Overlap percentage (non-white pixels only): {overlap_nonwhite_percentage:.2f}%")
    else:
        print("No non-white pixels found in all images for non-white overlap calculation.")

    # Stack into RGB channels (up to 3 images)
    combined = np.zeros((min_shape[0], min_shape[1], 3), dtype=np.uint8)
    for i in range(min(3, len(imgs))):
        combined[..., i] = imgs[i]

    # If more than 3 images, blend the rest into the RGB channels
    if len(imgs) > 3:
        for idx, img in enumerate(imgs[3:], start=3):
            channel = idx % 3
            combined[..., channel] = np.maximum(combined[..., channel], img)

    out_path = os.path.join(output_dir, 'combined_bitmap.png')
    Image.fromarray(combined).save(out_path)
    print(f"Combined bitmap saved to {out_path}")

def main():
    
    output_dir = f'/home/marthijnvd/master/docs/bitmapplots/{tool}'
    os.makedirs(output_dir, exist_ok=True)
    fuzz_files = find_fuzz_bitmap_files(root_dir)
    for file_path in fuzz_files:
        plot_bitmap(file_path, output_dir)
    print(f"Plotted {len(fuzz_files)} bitmaps to {output_dir}")
    new_output_dir = f'/home/marthijnvd/master/docs/bitmapplotscombined/{tool}'
    combine_bitmaps(output_dir, new_output_dir)

if __name__ == '__main__':
    main()