huytofu92's picture
enhance VLM and fix bugs with step extraciton
f864d65
import cv2
from cv2 import dnn
import numpy as np
import pytesseract
import requests
import base64
import onnxruntime
import os
from io import BytesIO
from PIL import Image
from langchain_core.tools import tool as langchain_tool
from smolagents.tools import Tool, tool
def pre_processing(image: str, input_size=(416, 416))->tuple:
"""
Pre-process an image for YOLO model
Args:
image: The image in base64 format to process
input_size: The size to which the image should be resized
Returns:
tuple: (processed_image, original_shape)
"""
try:
# Decode base64 image
image_data = base64.b64decode(image)
np_image = np.frombuffer(image_data, np.uint8)
img = cv2.imdecode(np_image, cv2.IMREAD_COLOR)
if img is None:
raise ValueError("Failed to decode image")
# Store original shape for post-processing
original_shape = img.shape[:2] # (height, width)
# Ensure input_size is valid
if not isinstance(input_size, tuple) or len(input_size) != 2:
input_size = (416, 416)
# Resize and normalize the image
img = cv2.resize(img, input_size, interpolation=cv2.INTER_LINEAR)
if img is None:
raise ValueError("Failed to resize image")
# Ensure image is in BGR format (3 channels)
if len(img.shape) == 2: # If grayscale
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
elif img.shape[2] == 4: # If RGBA
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)
# Convert BGR to RGB and normalize
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # More reliable than array slicing
img = img.astype(np.float32) / 255.0 # Normalize to [0, 1]
# Convert to NCHW format (batch, channels, height, width)
img = np.transpose(img, (2, 0, 1)) # HWC to CHW
img = np.expand_dims(img, axis=0) # Add batch dimension
# Verify final shape
if img.shape != (1, 3, 416, 416):
print(f"Warning: Final shape is {img.shape}, expected (1, 3, 416, 416)")
img = np.reshape(img, (1, 3, 416, 416))
return img, original_shape
except Exception as e:
raise ValueError(f"Error in pre_processing: {str(e)}")
def post_processing(onnx_output, classes, original_shape, conf_threshold=0.5, nms_threshold=0.4)->list:
"""
Post-process the output of the YOLO model
Args:
onnx_output: The raw output from the ONNX model
classes: List of class names
original_shape: Original shape of the image
conf_threshold: Confidence threshold for filtering detections
nms_threshold: Non-max suppression threshold
Returns:
List of detected objects with labels, confidence, and bounding boxes
"""
class_ids = []
confidences = []
boxes = []
for detection in onnx_output[0]:
scores = detection[5:]
class_id = np.argmax(scores)
confidence = scores[class_id]
if confidence > conf_threshold:
center_x = int(detection[0] * original_shape[1])
center_y = int(detection[1] * original_shape[0])
w = int(detection[2] * original_shape[1])
h = int(detection[3] * original_shape[0])
x = int(center_x - w / 2)
y = int(center_y - h / 2)
boxes.append([x, y, w, h])
confidences.append(float(confidence))
class_ids.append(class_id)
# Apply non-max suppression
indices = dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)
detected_objects = []
for i in indices:
i = i[0]
box = boxes[i]
label = str(classes[class_ids[i]])
detected_objects.append((label, confidences[i], box))
return detected_objects
@tool
def extract_images_from_video(video_path: str) -> list:
"""
Extract images (frames) from a video
Args:
video_path: The path to the video file
Returns:
A list of images (frames) as numpy arrays
"""
cap = cv2.VideoCapture(video_path)
images = []
while cap.isOpened():
ret, image = cap.read()
if not ret:
break
images.append(image)
cap.release()
return images
@tool
def get_image_from_file_path(file_path: str)->str:
"""
Load an image from a file path and convert it to a base64 string
Args:
file_path: The path to the file
Returns:
The image as a base64 string
"""
try:
# Debug prints for original path
# print(f"Original file_path: {file_path}")
# print(f"Original path exists: {os.path.exists(file_path)}")
# if os.path.exists(file_path):
# print(f"Original path is file: {os.path.isfile(file_path)}")
# print(f"Original path permissions: {oct(os.stat(file_path).st_mode)[-3:]}")
# print(f"Original path absolute: {os.path.abspath(file_path)}")
# Try reading with cv2
img = cv2.imread(file_path)
if img is None:
raise FileNotFoundError(f"Could not read image at {file_path}")
# Use BytesIO to encode the image
with BytesIO() as buffer:
_, buffer_data = cv2.imencode('.jpg', img)
buffer.write(buffer_data.tobytes())
image = base64.b64encode(buffer.getvalue()).decode('utf-8')
except Exception as e:
print(f"First attempt failed: {str(e)}")
# Try with adjusted path
try:
current_file_path = os.path.abspath(__file__)
current_file_dir = os.path.dirname(current_file_path)
adjusted_path = os.path.join(current_file_dir, file_path)
# Debug prints for adjusted path
# print(f"Adjusted file_path: {adjusted_path}")
# print(f"Adjusted path exists: {os.path.exists(adjusted_path)}")
# if os.path.exists(adjusted_path):
# print(f"Adjusted path is file: {os.path.isfile(adjusted_path)}")
# print(f"Adjusted path permissions: {oct(os.stat(adjusted_path).st_mode)[-3:]}")
# print(f"Adjusted path absolute: {os.path.abspath(adjusted_path)}")
# Try reading with cv2
img = cv2.imread(adjusted_path)
if img is None:
raise FileNotFoundError(f"Could not read image at {adjusted_path}")
# Use BytesIO to encode the image
with BytesIO() as buffer:
_, buffer_data = cv2.imencode('.jpg', img)
buffer.write(buffer_data.tobytes())
image = base64.b64encode(buffer.getvalue()).decode('utf-8')
except Exception as e2:
print(f"Second attempt failed: {str(e2)}")
# List directory contents to help debug
try:
validation_dir = os.path.join(current_file_dir, "validation")
if os.path.exists(validation_dir):
print(f"Contents of validation directory: {os.listdir(validation_dir)}")
except Exception as e3:
print(f"Failed to list directory contents: {str(e3)}")
raise FileNotFoundError(f"Could not read image at {file_path} or {adjusted_path}")
return image
@tool
def get_video_from_file_path(file_path: str)->str:
"""
Load a video from a file path and convert it to a base64 string
Args:
file_path: The path to the file
Returns:
The video as a base64 string
"""
try:
# Use cv2 to read the video
cap = cv2.VideoCapture(file_path)
if not cap.isOpened():
raise FileNotFoundError(f"Could not read video at {file_path}")
# Get video properties
fps = cap.get(cv2.CAP_PROP_FPS)
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# Create a BytesIO buffer to store the images (frames)
images = []
while cap.isOpened():
ret, image = cap.read()
if not ret:
break
# Convert frame to jpg and store in memory
_, buffer = cv2.imencode('.jpg', image)
images.append(buffer.tobytes())
# Release the video capture
cap.release()
# Combine all images into a single buffer
with BytesIO() as buffer:
# Write each image to the buffer
for image_data in images:
buffer.write(image_data)
# Encode to base64
video_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
except Exception as e:
current_file_path = os.path.abspath(__file__)
current_file_dir = os.path.dirname(current_file_path)
file_path = os.path.join(current_file_dir, file_path)
# Try again with the new path
cap = cv2.VideoCapture(file_path)
if not cap.isOpened():
raise FileNotFoundError(f"Could not read video at {file_path}")
# Get video properties
fps = cap.get(cv2.CAP_PROP_FPS)
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# Create a BytesIO buffer to store the images (frames)
images = []
while cap.isOpened():
ret, image = cap.read()
if not ret:
break
# Convert image to jpg and store in memory
_, buffer = cv2.imencode('.jpg', image)
images.append(buffer.tobytes())
# Release the video capture
cap.release()
# Combine all images into a single buffer
with BytesIO() as buffer:
# Write each image to the buffer
for image_data in images:
buffer.write(image_data)
# Encode to base64
video_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
return video_base64
@tool
def image_processing(image: str, brightness: float = 1.0, contrast: float = 1.0)->str:
"""
Process an image
Args:
image: The image in base64 format to process
brightness: The brightness of the image on scale of 0-10
contrast: The contrast of the image on scale of 0-10
Returns:
The processed image
"""
image_data = base64.b64decode(image)
np_image = np.frombuffer(image_data, np.uint8)
img = cv2.imdecode(np_image, cv2.IMREAD_COLOR)
# Adjust brightness and contrast
img = cv2.convertScaleAbs(img, alpha=contrast, beta=brightness)
_, buffer = cv2.imencode('.jpg', img)
processed_image = base64.b64encode(buffer).decode('utf-8')
return processed_image
onnx_path = "vlm_assets/yolov3-8.onnx"
class ObjectDetectionTool(Tool):
name = "object_detection"
description = """
Detect objects in a list of images.
Input Requirements:
- Input must be a list of images, where each image is a base64-encoded string
- Each base64 string must be properly padded (length must be a multiple of 4)
- Images will be resized to 416x416 pixels during processing
- Images should be in RGB or BGR format (3 channels)
- Supported image formats: JPG, PNG
Processing:
- Images are automatically resized to 416x416
- Images are normalized to [0,1] range
- Model expects input shape: [1, 3, 416, 416] (batch, channels, height, width)
Output:
- Returns a list of detected objects for each image
- Each detection includes: (label, confidence, bounding_box)
- Bounding boxes are in format: [x, y, width, height]
- Confidence threshold: 0.5
- NMS threshold: 0.4
Example input format:
["base64_encoded_image1", "base64_encoded_image2"]
Example output format:
[
[("person", 0.95, [100, 200, 50, 100]), ("car", 0.88, [300, 400, 80, 60])], # detections for image1
[("dog", 0.92, [150, 250, 40, 80])] # detections for image2
]
"""
inputs = {
"images": {
"type": "any",
"description": "List of base64-encoded images. Each image must be a valid base64 string with proper padding (length multiple of 4). Images will be resized to 416x416."
}
}
output_type = "any"
def setup(self):
try:
# Load ONNX model
self.onnx_path = onnx_path
self.onnx_model = onnxruntime.InferenceSession(self.onnx_path)
# Get model input details
self.input_name = self.onnx_model.get_inputs()[0].name
self.input_shape = self.onnx_model.get_inputs()[0].shape
print(f"Model input shape: {self.input_shape}")
# Load class labels
self.classes = [
'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat',
'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book',
'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
except Exception as e:
raise RuntimeError(f"Error in setup: {str(e)}")
def forward(self, images: any)->any:
try:
if not isinstance(images, list):
images = [images] # Convert single image to list
detected_objects = []
for image in images:
try:
# Preprocess the image
img, original_shape = pre_processing(image)
# Verify input shape and convert to NCHW if needed
if len(img.shape) != 4: # Should be NCHW
raise ValueError(f"Invalid input shape: {img.shape}, expected NCHW format")
if img.shape[1] != 3: # Should have 3 channels
# If channels are last, transpose to NCHW
if img.shape[3] == 3:
img = np.transpose(img, (0, 3, 1, 2))
else:
raise ValueError(f"Invalid number of channels: {img.shape[1]}, expected 3")
# Verify final shape
if img.shape != (1, 3, 416, 416):
print(f"Warning: Reshaping input from {img.shape} to (1, 3, 416, 416)")
img = np.reshape(img, (1, 3, 416, 416))
# Run inference
onnx_input = {self.input_name: img}
onnx_output = self.onnx_model.run(None, onnx_input)
# Handle shape mismatch by transposing if needed
if len(onnx_output[0].shape) == 4: # If in NCHW format
if onnx_output[0].shape[1] == 255: # If channels first
onnx_output = [onnx_output[0].transpose(0, 2, 3, 1)] # Convert to NHWC
# Post-process the output
objects = post_processing(onnx_output, self.classes, original_shape)
detected_objects.append(objects)
except Exception as e:
print(f"Error processing image: {str(e)}")
detected_objects.append([]) # Add empty list for failed image
return detected_objects
except Exception as e:
raise RuntimeError(f"Error in forward pass: {str(e)}")
class OCRTool(Tool):
description = """
Scan an image for text using OCR (Optical Character Recognition).
Input Requirements:
- Input must be a list of images, where each image is a base64-encoded string
- Each base64 string must be properly padded (length must be a multiple of 4)
- Images should be in RGB or BGR format (3 channels)
- Supported image formats: JPG, PNG
- For best results:
* Text should be clear and well-lit
* Image should have good contrast
* Text should be properly oriented
* Avoid blurry or distorted images
Processing:
- Uses Tesseract OCR engine
- Automatically handles text orientation
- Supports multiple languages (default: English)
- Processes each image independently
Output:
- Returns a list of text strings, one for each input image
- Empty string is returned if no text is detected
- Text is returned in the order it appears in the image
- Line breaks are preserved in the output
Example input format:
["base64_encoded_image1", "base64_encoded_image2"]
Example output format:
[
"This is text from image 1\nSecond line of text", # text from image1
"Text from image 2" # text from image2
]
"""
name = "ocr_scan"
inputs = {
"images": {
"type": "any",
"description": "List of base64-encoded images. Each image must be a valid base64 string with proper padding (length multiple of 4). Images should be clear and well-lit for best OCR results."
}
}
output_type = "any"
def forward(self, images: any)->any:
scanned_text = []
for image in images:
image_data = base64.b64decode(image)
img = Image.open(BytesIO(image_data))
scanned_text.append(pytesseract.image_to_string(img))
return scanned_text
ocr_scan_tool = OCRTool()
object_detection_tool = ObjectDetectionTool()
#Test 3