Extract structured data from images

%pip install -qU pixeltable openai

import getpass
import os

if 'OPENAI_API_KEY' not in os.environ:
    os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key: ')

import pixeltable as pxt
from pixeltable.functions import openai

# Create a fresh directory
pxt.drop_dir('extraction_demo', force=True)
pxt.create_dir('extraction_demo')

t = pxt.create_table('extraction_demo/images', {'image': pxt.Image})

# Insert sample images
t.insert(
    [
        {
            'image': 'https://raw.githubusercontent.com/pixeltable/pixeltable/main/docs/resources/images/000000000036.jpg'
        },
        {
            'image': 'https://raw.githubusercontent.com/pixeltable/pixeltable/main/docs/resources/images/000000000090.jpg'
        },
    ]
)

# Add extraction column using openai.chat_completions with multimodal messages
PROMPT = """Analyze this image and extract the following as JSON:
- description: A brief description of the image
- objects: List of objects visible in the image
- dominant_colors: List of dominant colors
- scene_type: Type of scene (indoor, outdoor, etc.)"""

messages = [
    {
        'role': 'user',
        'content': [
            {'type': 'text', 'text': PROMPT},
            {'type': 'image_url', 'image_url': t.image},
        ],
    }
]

t.add_computed_column(
    data=openai.chat_completions(
        messages,
        model='gpt-4o-mini',
        model_kwargs={'response_format': {'type': 'json_object'}},
    )
)

# View extracted data
t.select(
    t.image, t.data, t.data['choices'][0]['message']['content']
).collect()

# You can also parse the JSON into individual columns if needed
import json


@pxt.udf
def parse_description(data: str) -> str:
    return json.loads(data).get('description', '')


t.select(
    t.image,
    description=parse_description(
        t.data['choices'][0]['message']['content']
    ),
).collect()