""" title: Coral TPU Preprocessor author: Luis Chacón description: Filter to preprocess a message with images through a Coral TPU to augment a non-vision model. author_url: https://github.com/luisgreen funding_url: https://github.com/luisgreen issues: https://github.com/luisgreen/openwebui-coral-preprocessor/issues/new requirements: tflite-runtime, numpy version: 0.4 changelog: v0.4 - Handle a proper response if there is an image but nothing could not be identified. - Enhaced response prompt. v0.3 - Added Counter so that the question can gather more context on cluster of things v0.2 - Improved query prompt v0.1 - Initial """ from PIL import Image from io import BytesIO import os import numpy as np import base64 from tflite_runtime.interpreter import Interpreter from tflite_runtime.interpreter import load_delegate from collections import Counter from pydantic import BaseModel, Field from typing import Optional from open_webui.utils.misc import ( get_last_user_message_item, get_last_user_message, get_content_from_message, ) import base64 import logging class Filter: class Valves(BaseModel): MODEL_FILE: str = Field( os.getenv( "MODEL_FILE", "tf2_ssd_mobilenet_v1_fpn_640x640_coco17_ptq_edgetpu.tflite", ), description="Coral TPU TFlite Model", ) LABEL_FILE: str = Field( os.getenv("LABEL_FILE", "coco_labels.txt"), description="Labels supplied with the model", ) CONFIDENCE_THRESHOLD: str = Field( os.getenv("CONFIDENCE_THRESHOLD", "0.35"), description="Confidence threshold for the inference", ) CONTEXT_WINDOW: int = Field( os.getenv("CONTEXT_WINDOW", 30000), description="The number of tokens to use in the context window", ) pass def __init__(self): self.valves = self.Valves() labels_path = f"/app/backend/data/coral/{self.valves.LABEL_FILE}" model_path = f"/app/backend/data/coral/{self.valves.MODEL_FILE}" # Initialize the interpreter self.interpreter = Interpreter( model_path, experimental_preserve_all_tensors=True, experimental_delegates=[ load_delegate("/app/backend/data/coral/libedgetpu.so.1.0") ], ) # Initialize the labels with open(labels_path, "r") as f: self.labels = {i: line.strip() for i, line in enumerate(f.readlines())} def analyze_results(self, boxes, classes, scores, num_boxes): objects_detected = list() for i in range(int(num_boxes)): score = scores[i] if float(score) >= float(self.valves.CONFIDENCE_THRESHOLD): objects_detected.append(self.labels[classes[i]]) return objects_detected def summarize_counts(self, objects): counts = Counter(objects) summary = [] for obj, count in counts.items(): if count > 1: summary.append(f"{count} {obj}s") else: summary.append(f"{count} {obj}") return ", ".join(summary) def assemble_prompt(self, message_prompt, found_objects): elements = self.summarize_counts(found_objects) return f""" Your task is to create an answer to the question, based on the context given here. This context is a serie of elements that were identified by a TPU previously, so dont make assumptions on the likelyhood of the elements. Pay attention to the given question, and elaborate an answer that correlates the identified elements with the question and in an accurate, clear and concise manner. Do your best to return a descriptive answer that would consider and explain what the question contain. If there are elements identified, try to infere, if possible, a very detailed scene where the elements can be present together and explain with details where all the elements can coexist. If you dont know, just skip this step. You are allowed to mention something like 'My Coral TPU identified the following elements <elements>' If the list of elements empty, you can say that 'My Coral TPU was not able to identify anything'. If you consider there is no other relevant information, you can stop responding inmediately. Otherwise if you find there is something interesting you can share based on the question, focus on it and answer. <CONDITIONS> - Don't include previous answers as context for this analysis. - Attached to this message there is one or more images, ignore them, just focus in the context. - Dont mention that the elements where separated by comma. - Don't mention that you are only a text based model, just use the context to answer. - Always mention which elements were given in the response. <RESPONSE> Explanation here if any Final answer here Question: {message_prompt} Elements: {elements} """ async def inlet( self, body: dict, __user__: Optional[dict] = None, __event_emitter__=None, ) -> dict: images_to_process = set() found_objects = list() exist_images = False messages = body.get("messages") message = get_last_user_message_item(messages) message_prompt = get_last_user_message(messages) if "images" in message: exist_images = True images_to_process = [image for image in message["images"]] for image in images_to_process: input_details = self.interpreter.get_input_details()[0] output_details = self.interpreter.get_output_details() self.interpreter.allocate_tensors() im = Image.open(BytesIO(base64.b64decode(image))).resize( (input_details["shape"][1], input_details["shape"][2]) ) input_data = np.expand_dims(im, axis=0) self.interpreter.set_tensor(input_details["index"], input_data) self.interpreter.invoke() scores = self.interpreter.get_tensor(output_details[0]["index"])[ 0 ] # Confidence of detected objects boxes = self.interpreter.get_tensor(output_details[1]["index"])[ 0 ] # Bounding box coordinates of detected objects num_boxes = self.interpreter.get_tensor(output_details[2]["index"])[ 0 ] # Number of boxed detected classes = self.interpreter.get_tensor(output_details[3]["index"])[ 0 ] # Class index of detected objects found_objects = found_objects + self.analyze_results( boxes, classes, scores, num_boxes ) # total_objects_found = len(found_objects) if exist_images: message_to_replace = body["messages"].pop() message_to_replace["content"] = self.assemble_prompt( message_prompt, found_objects ) body["messages"].append(message_to_replace) return body return body def outlet(self, body: dict, __user__: Optional[dict] = None) -> dict: return body