app.py

import streamlit as st
import os
import glob
import base64
import json
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from langchain_openai import ChatOpenAI
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import chain
from PIL import Image as PILImage
from io import BytesIO

# Streamlit title
st.title("Vehicle Information Extraction from Images")

# Prompt user for OpenAI API key
openai_api_key = st.text_input("Enter your OpenAI API Key:", type="password")

# Set the OpenAI API key if provided
if openai_api_key:
    os.environ["OPENAI_API_KEY"] = openai_api_key

# Vehicle class (same as in the original code)
class Vehicle(BaseModel):
    Type: str = Field(..., examples=["Car", "Truck", "Motorcycle", 'Bus', 'Van'], description="The type of the vehicle.")
    License: str = Field(..., description="The license plate number of the vehicle.")
    Make: str = Field(..., examples=["Toyota", "Honda", "Ford", "Suzuki"], description="The Make of the vehicle.")
    Model: str = Field(..., examples=["Corolla", "Civic", "F-150"], description="The Model of the vehicle.")
    Color: str = Field(..., example=["Red", "Blue", "Black", "White"], description="Return the color of the vehicle.")

# Parser for vehicle details
parser = JsonOutputParser(pydantic_object=Vehicle)
instructions = parser.get_format_instructions()

# Image encoding function (for base64 encoding)
def image_encoding(inputs):
    """Load and convert image to base64 encoding"""
    with open(inputs["image_path"], "rb") as image_file:
        image_base64 = base64.b64encode(image_file.read()).decode("utf-8")
    return {"image": image_base64}

# Image display in grid (for multiple images)
def display_image_grid(image_paths, rows=2, cols=3, figsize=(10, 7)):
    fig = plt.figure(figsize=figsize)
    max_images = rows * cols
    image_paths = image_paths[:max_images]

    for idx, path in enumerate(image_paths):
        ax = fig.add_subplot(rows, cols, idx + 1)
        img = mpimg.imread(path)
        ax.imshow(img)
        ax.axis('off')
        filename = path.split('/')[-1]
        ax.set_title(filename)

    plt.tight_layout()
    st.pyplot(fig)

# Create the prompt for the AI model
@chain
def prompt(inputs):
    prompt = [
        SystemMessage(content="""You are an AI assistant whose job is to inspect an image and provide the desired information from the image. If the desired field is not clear or not well detected, return None for this field. Do not try to guess."""), 
        HumanMessage(
            content=[{"type": "text", "text": "Examine the main vehicle type, license plate number, make, model and color."},
                     {"type": "text", "text": instructions},
                     {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{inputs['image']}", "detail": "low"}}]
        )
    ]
    return prompt

# Invoke the model for extracting vehicle details
@chain
def MLLM_response(inputs):
    model: ChatOpenAI = ChatOpenAI(model="gpt-4o-2024-08-06", temperature=0.0, max_tokens=1024)
    output = model.invoke(inputs)
    return output.content

# The complete pipeline for extracting vehicle details
pipeline = image_encoding | prompt | MLLM_response | parser

# Streamlit Interface for uploading images and showing results
st.header("Upload a Vehicle Image for Information Extraction")

uploaded_image = st.file_uploader("Choose a JPEG image", type="jpeg")

if uploaded_image is not None:
    # Display the uploaded image
    image = PILImage.open(uploaded_image)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Convert the uploaded image to base64
    image_path = "/tmp/uploaded_image.jpeg"
    with open(image_path, "wb") as f:
        f.write(uploaded_image.getbuffer())

    # Process the image through the pipeline
    output = pipeline.invoke({"image_path": image_path})

    # Show the results in a user-friendly format
    st.subheader("Extracted Vehicle Information")
    st.json(output)

    # Optionally, display more vehicle images from the folder
    img_dir = "/content/images"
    image_paths = glob.glob(os.path.join(img_dir, "*.jpeg"))
    display_image_grid(image_paths)

# You can also allow users to upload and process a batch of images
st.sidebar.header("Batch Image Upload")

batch_images = st.sidebar.file_uploader("Upload Images", type="jpeg", accept_multiple_files=True)

if batch_images:
    batch_input = [{"image_path": f"/tmp/{file.name}"} for file in batch_images]
    for file in batch_images:
        with open(f"/tmp/{file.name}", "wb") as f:
            f.write(file.getbuffer())
    
    # Process the batch and display the results in a DataFrame
    batch_output = pipeline.batch(batch_input)
    df = pd.DataFrame(batch_output)
    st.dataframe(df)

    # Show images in a grid
    image_paths = [f"/tmp/{file.name}" for file in batch_images]
    display_image_grid(image_paths)