Generate alt tags for images (using llava)

img2alt.py

import os
import requests
from bs4 import BeautifulSoup
import argparse
import base64
import json

def download_image(url, save_path):
    response = requests.get(url)
    if response.status_code == 200:
        with open(save_path, 'wb') as f:
            f.write(response.content)
        return True
    return False

def get_alt_text(image_path, llava_url):
    with open(image_path, "rb") as image_file:
        image_base64 = base64.b64encode(image_file.read()).decode('utf-8')

    payload = {
        "image": image_base64,
        "prompt": "Describe this image in a concise manner suitable for an alt text."
    }

    response = requests.post(llava_url, json=payload)
    if response.status_code == 200:
        return response.json()['response']
    else:
        return None

def process_html(html_path, output_path, llava_url):
    with open(html_path, 'r') as f:
        soup = BeautifulSoup(f, 'html.parser')

    img_tags = soup.find_all('img')
    for img in img_tags:
        src = img.get('src')
        if src and src.startswith(('http://', 'https://')):
            filename = os.path.basename(src)
            save_path = os.path.join('images', filename)

            if not os.path.exists('images'):
                os.makedirs('images')

            if download_image(src, save_path):
                alt_text = get_alt_text(save_path, llava_url)
                if alt_text:
                    img['alt'] = alt_text
                    print(f"Added alt text to {src}: {alt_text}")
                else:
                    print(f"Failed to generate alt text for {src}")
            else:
                print(f"Failed to download image: {src}")

    with open(output_path, 'w') as f:
        f.write(str(soup))

def main():
    parser = argparse.ArgumentParser(description="Process HTML and add alt text to images using LLaVA")
    parser.add_argument("html_path", help="Path to the input HTML file")
    parser.add_argument("output_path", help="Path to save the output HTML file")
    parser.add_argument("llava_url", help="URL of the LLaVA server")
    args = parser.parse_args()

    process_html(args.html_path, args.output_path, args.llava_url)

if __name__ == "__main__":
    main()