wehub-resource-sync
133 lines
4.4 KiB
Python
133 lines
4.4 KiB
Python
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()
|