Allowing language models to choose the right algorithm

In this notebook we enable a language model to choose for the right algorithm. We define multiple segmentation algorithms / tools and then give the language model the choice which one to use given different inputs.

from langchain.memory import ConversationBufferMemory
from langchain.chat_models import ChatOpenAI
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.tools import tool

from skimage.io import imread
from napari_segment_blobs_and_things_with_membranes import voronoi_otsu_labeling, local_minima_seeded_watershed

import stackview

Again, we define an image storage and a list of tools.

image_storage = {}
tools = []

@tools.append
@tool
def load_image(filename:str):
    """Useful for loading an image file and storing it."""
    print("loading", filename)
    image = imread(filename)
    image_storage[filename] = image
    return "The image is now stored as " + filename

We define two segmentation algorithms, one for segmenting bright objects and one for segmenting dark objects.

@tools.append
@tool
def segment_bright_objects(image_name):
    """
    Useful for segmenting bright objects in an image 
    that has been loaded and stored before.
    """
    print("segmenting (Voronoi-Otsu-Labeling)", image_name)
    
    image = image_storage[image_name]
    label_image = voronoi_otsu_labeling(image, spot_sigma=4)
    
    label_image_name = "segmented_" + image_name
    image_storage[label_image_name] = label_image
    
    return "The segmented image has been stored as " + label_image_name

@tools.append
@tool
def segment_dark_objects(image_name):
    """
    Useful for segmenting dark objects with bright border in an image 
    that has been loaded and stored before.
    """
    print("segmenting (Local-minima-seeded watershed)", image_name)
    
    image = image_storage[image_name]
    label_image = local_minima_seeded_watershed(image, spot_sigma=10)
    
    label_image_name = "segmented_" + image_name
    image_storage[label_image_name] = label_image
    
    return "The segmented image has been stored as " + label_image_name

@tools.append
@tool
def show_image(image_name):
    """Useful for showing an image that has been loaded and stored before."""
    print("showing", image_name)
    
    image = image_storage[image_name]
    display(stackview.insight(image))
    
    return "The image " + image_name + " is shown above."

We create some memory and a large language model based on OpenAI’s chatGPT.

memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
llm=ChatOpenAI(temperature=0)

Given the list of tools, the large language model and the memory, we can create an agent.

agent = initialize_agent(
    tools, 
    llm, 
    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION, 
    memory=memory
)

This agent can then respond to prompts.

agent.run("Please load the image ../../data/membrane2d.tif and show it.")

loading ../../data/membrane2d.tif
showing ../../data/membrane2d.tif

shape (256, 256) dtype uint16 size 128.0 kB min 277 max 44092

'The image ../../data/membrane2d.tif is shown above.'

agent.run("Please segment the image ../../data/membrane2d.tif")

segmenting (Voronoi-Otsu-Labeling) ../../data/membrane2d.tif

'The segmented image has been stored as segmented_../../data/membrane2d.tif'

agent.run("Please show the segmented ../../data/membrane2d.tif image.")

showing segmented_../../data/membrane2d.tif

shape (256, 256) dtype int32 size 256.0 kB min 0 max 56

'The segmented image is shown above.'

The segmentation does not look like a cell-segmentation. Thus, we should ask more specifically.

agent.run("Please segment the image ../../data/membrane2d.tif again and this time, segment the dark cells surrounded by bright membranes. Also show the result of the segmentation.")

segmenting (Local-minima-seeded watershed) ../../data/membrane2d.tif
showing segmented_../../data/membrane2d.tif

shape (256, 256) dtype int32 size 256.0 kB min 1 max 27

'The segmented image of dark cells surrounded by bright membranes has been shown.'

We can also ask the agent which algorithm it chose and why it chose this tool.

agent.run("Which algorithm did you use this time?")

'I used the algorithm to segment dark cells surrounded by bright membranes.'

agent.run("Why did you use this algorithm?")

'I used this algorithm because it is specifically designed to segment dark cells surrounded by bright membranes, which matches the criteria you mentioned in your request.'

Note: The language model cannot see the image. Its tool selection depends on information you provided and information it acquired during the chat.