analyze_laion_data.py

"""
Walk through a LAION dataset and analyze it.
"""

import os
import json
import concurrent.futures
from PIL import Image

def get_aspect_ratio(image_path):
    try:
        with Image.open(image_path) as img:
            width, height = img.size
            return image_path, width / height
    except Exception as e:
        os.remove(image_path)
        return None

def analyze_images(directory_path):
    aspect_ratios = {}
    good_count = 0
    bad_count = 0
    image_files = [os.path.join(directory_path, filename) 
                   for filename in os.listdir(directory_path) 
                   if filename.endswith(".jpg") or filename.endswith(".png")]
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=64) as executor:
        futures = {executor.submit(get_aspect_ratio, image_path): image_path for image_path in image_files}
        for future in concurrent.futures.as_completed(futures):
            image_path = futures[future]
            try:
                aspect_ratio = future.result()
                if aspect_ratio is not None:
                    image_path, aspect_ratio = aspect_ratio
                    aspect_ratio = round(aspect_ratio, 2)  # round to 2 decimal places
                    if aspect_ratio not in aspect_ratios:
                        aspect_ratios[aspect_ratio] = []
                    aspect_ratios[aspect_ratio].append(image_path)
                    good_count += 1
                else:
                    bad_count += 1
            except Exception as e:
                pass
    print(f'Good images: {good_count}, Bad images: {bad_count}')
    return aspect_ratios

def write_to_json(data, filename):
    with open(filename, 'w') as outfile:
        json.dump(data, outfile)

if __name__ == "__main__":
    image_directory = "/notebooks/datasets/laion-high-resolution/downloaded_images"
    output_file = "aspect_ratios.json"
    aspect_ratios = analyze_images(image_directory)
    write_to_json(aspect_ratios, output_file)