CV-Agent / tool_utils /object_extractor.py
Samarth991's picture
adding CV agent file
0e78cbf
from pydantic import BaseModel , Field
from langchain_core.prompts import PromptTemplate
class objects_identified(BaseModel):
objects : str = Field(...,description="Generate a list of objects identified from the given description of the image")
def objectextractor_prompt():
template = """
You are an AI assistant provided with a context below. The context is a description of an image. Your task is to identify the objects within the image.
The objects must be living beings or physical items or things that one can view, feel, and touch. these objects must be nouns and not verbs or adjectives or words describing an object like
'large', 'beautiful' etc .
Only provide the name of the object and not the description of the object.Refer the example input and output for better understanding.
If the conntext mention a boy , a girl , women , girls they will come under the category of "People". so the object for such classes will be people .
Example Input: Context: "A park filled with men and women , a large oak tree standing in the center, a dog running near a bench, and a bicycle leaning against a nearby fence."
Example Output:"People" , "Dog" , "Bench" , "Bicycle" ,"Fence"
Context: {context}
"""
prompt = PromptTemplate(template=template,input_variables=["context"])
return prompt
def create_object_extraction_chain(llm):
object_extraction_chain = objectextractor_prompt() | llm.with_structured_output(objects_identified)
return object_extraction_chain