Initial commit

playground/2d_hist.py

@@ -0,0 +1,132 @@
+import json
+import os
+from PIL import Image
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+import numpy as np
+from multiprocessing import Pool
+import functools
+import argparse
+
+
+def load_data(json_path):
+    with open(json_path, "r") as f:
+        return json.load(f)
+
+
+def filter_data(data):
+    filtered_data = [item for item in data if "image" in item]
+    return filtered_data
+
+
+def calculate_image_dimension(image_path, images_folder):
+    full_path = os.path.join(images_folder, image_path)
+    try:
+        with Image.open(full_path) as img:
+            width, height = img.size
+            return width, height
+    except Exception as e:
+        print(f"Error opening {full_path}: {e}")
+        return None, None
+
+
+def calculate_image_dimensions_multiprocess(filtered_data, images_folder, num_processes=256):
+    image_paths = []
+    for item in filtered_data:
+        if isinstance(item["image"], list):
+            image_paths.extend(item["image"])
+        else:
+            image_paths.append(item["image"])
+
+    with Pool(num_processes) as p:
+        dimensions = list(
+            tqdm(
+                p.imap(functools.partial(calculate_image_dimension, images_folder=images_folder), image_paths),
+                total=len(image_paths),
+                desc="Calculating image dimensions",
+            )
+        )
+    widths, heights = zip(*[dim for dim in dimensions if dim[0] is not None])
+    return list(widths), list(heights)
+
+
+def tokenize(text):
+    return text.split()
+
+
+def calculate_tokenized_lengths(data):
+    lengths = []
+    for item in tqdm(data, desc="Tokenizing conversations"):
+        for conversation in item["conversations"]:
+            tokenized_value = tokenize(conversation["value"])
+            lengths.append(len(tokenized_value))
+    return lengths
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Process data for LLaVA_Next project.")
+    parser.add_argument(
+        "--json_path",
+        type=str,
+        help="Path to the JSON file containing data.",
+        default="/mnt/bn/vl-research/data/llava_instruct/real_vision_flan/llava_ofa_DEMON-FULL.json",
+    )
+    parser.add_argument(
+        "--images_folder",
+        type=str,
+        default="/mnt/bn/vl-research/data/llava_data",
+        help="Path to the folder containing images.",
+    )
+    args = parser.parse_args()
+
+    llava_instruct_name = os.path.basename(args.json_path).replace(".json", "")
+    images_folder = args.images_folder
+
+    data = load_data(args.json_path)
+    filtered_data = filter_data(data)
+
+    print(f"Total data items: {len(data)}, Filtered data items: {len(filtered_data)}")
+    widths, heights = calculate_image_dimensions_multiprocess(filtered_data, images_folder)
+    max_width, max_height = max(widths), max(heights)
+    print(f"Max width: {max_width}, Max height: {max_height}")
+
+    tokenized_lengths = calculate_tokenized_lengths(filtered_data)
+
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(24, 12))
+
+    # Plot 2D histogram
+    widths_bins = [min(widths), max(widths) + 1] if min(widths) == max(widths) else np.arange(min(widths), max(widths) + 100, 100)
+    heights_bins = [min(heights), max(heights) + 1] if min(heights) == max(heights) else np.arange(min(heights), max(heights) + 100, 100)
+
+    h, xedges, yedges, image = ax1.hist2d(widths, heights, bins=[widths_bins, heights_bins], cmap=plt.cm.jet, density=True)
+    fig.colorbar(image, ax=ax1)
+    ax1.set_xlabel("Width")
+    ax1.set_ylabel("Height")
+    ax1.set_title(
+        f"dist_{llava_instruct_name}_2d_w_h\nMax width: {max(widths)}, Max height: {max(heights)}",
+        fontsize=10,
+    )
+
+    # Plot histogram
+    hist, bin_edges = np.histogram(tokenized_lengths, bins=np.arange(0, max(tokenized_lengths) + 10, 10))
+    bins = np.arange(0, max(tokenized_lengths) + 10, 10)
+    ax2.bar(bin_edges[:-1], hist, width=7, edgecolor="black", log=True)
+
+    # Display every nth label on the x-axis
+    n = 8  # Adjust this value to control the number of labels displayed
+    ticks = bins[::n]
+    tick_labels = [int(tick) for tick in ticks]
+    ax2.set_xticks(ticks)
+    ax2.set_xticklabels(tick_labels, rotation=90, fontsize=8)
+
+    ax2.set_xlim(min(bin_edges), max(bin_edges))
+    ax2.set_xlabel("Tokenized Length")
+    ax2.set_ylabel("Count (log scale)")
+    ax2.set_title(f"dist_{llava_instruct_name}_tokenized_length", fontsize=8)
+
+    plt.tight_layout()
+    plt.savefig(f"./dist_{llava_instruct_name}_combined.png")
+
+
+if __name__ == "__main__":
+    main()