File size: 3,903 Bytes
87337b1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
"""Utility functions for BedrockV2V extension."""
from io import BytesIO
from PIL import Image
from typing import List, Tuple, Any
def is_punctuation(char: str) -> bool:
"""Check if a character is a punctuation mark."""
return char in [",", "，", ".", "。", "?", "？", "!", "！"]
def parse_sentence(sentence: str, content: str) -> Tuple[str, str, bool]:
"""Parse a sentence and return the complete sentence, remaining content, and completion status."""
remain = ""
found_punc = False
for char in content:
if not found_punc:
sentence += char
else:
remain += char
if not found_punc and is_punctuation(char):
found_punc = True
return sentence, remain, found_punc
def rgb2base64jpeg(rgb_data: bytes, width: int, height: int) -> bytes:
"""Convert RGB data to JPEG format."""
# Convert the RGB image to a PIL Image
pil_image = Image.frombytes("RGBA", (width, height), bytes(rgb_data))
pil_image = pil_image.convert("RGB")
# Resize the image while maintaining its aspect ratio
pil_image = resize_image_keep_aspect(pil_image, 640)
# Save the image to a BytesIO object in JPEG format
buffered = BytesIO()
pil_image.save(buffered, format="JPEG")
return buffered.getvalue()
def resize_image_keep_aspect(image: Image.Image, max_size: int = 512) -> Image.Image:
"""Resize an image while maintaining its aspect ratio."""
width, height = image.size
if width <= max_size and height <= max_size:
return image
aspect_ratio = width / height
if width > height:
new_width = max_size
new_height = int(max_size / aspect_ratio)
else:
new_height = max_size
new_width = int(max_size * aspect_ratio)
return image.resize((new_width, new_height))
def filter_images(image_array: List[Any], max_images: int = 10) -> List[Any]:
"""Filter images to maintain a maximum count while preserving temporal distribution."""
if len(image_array) <= max_images:
return image_array
result = []
skip = len(image_array) // max_images
for i in range(0, len(image_array), skip):
result.append(image_array[i])
if len(result) == max_images:
break
return result
# merge images into one image with a grid layout
def merge_images(image_array: List[Any], max_images: int = 4, width: int = 512) -> bytes:
"""Merge multiple images into one image."""
if len(image_array) == 0:
return b""
if len(image_array) > max_images:
# Filter images to maintain a maximum count while preserving temporal distribution
image_array = filter_images(image_array, max_images)
total_images = len(image_array)
# Calculate the number of rows and columns for the grid
rows = int((total_images - 1) / 2) + 1
cols = 2 if total_images > 1 else 1
# Calculate the size of each image in the grid
image_width = width // cols
image_height = image_width
# Create a new image to store the grid
grid = Image.new("RGB", (width, image_height * rows))
# Paste each image into the grid
for i, image in enumerate(image_array):
row = i // cols
col = i % cols
image = Image.open(BytesIO(image))
image = resize_image_keep_aspect(image, image_width)
grid.paste(image, (col * image_width, row * image_height))
# Save the grid to a BytesIO object in JPEG format
buffered = BytesIO()
grid.save(buffered, format="JPEG")
return buffered.getvalue()
def get_greeting_text(language: str) -> str:
"""Get appropriate greeting text based on language."""
greetings = {
"zh-CN": "你好。",
"ja-JP": "こんにちは",
"ko-KR": "안녕하세요",
"en-US": "Hi, there."
}
return greetings.get(language, "Hi, there.")
|