File size: 2,419 Bytes
74b17e0 |
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 |
'''
@Description:
@Author: jiajunlong
@Date: 2024-06-19 19:30:17
@LastEditTime: 2024-06-19 19:32:47
@LastEditors: jiajunlong
'''
class Message:
def __init__(self, msg=None):
self._messages = msg if msg else []
self._images = []
self.skip_next = False
def add_message(self, question, answer=None):
quension_msg_dict = {'from': 'human'}
quension_msg_dict['value'] = question
answer_msg_dict = {'from': 'gpt'}
answer_msg_dict['value'] = answer
self._messages.append(quension_msg_dict)
self._messages.append(answer_msg_dict)
def add_image(self, image, index=0):
self._images.append((image, index))
@property
def images(self):
return self._images
@property
def messages(self):
return self._messages
def copy(self):
return Message(self._messages)
def to_gradio_chatbot(self):
ret = []
for i, msg in enumerate(self.messages):
if i % 2 == 0:
if len(self.images) != 0 and i == self.images[0][1]:
image = self.images[0][0]
import base64
from io import BytesIO
msg = msg['value']
max_hw, min_hw = max(image.size), min(image.size)
aspect_ratio = max_hw / min_hw
max_len, min_len = 800, 400
shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
longest_edge = int(shortest_edge * aspect_ratio)
W, H = image.size
if H > W:
H, W = longest_edge, shortest_edge
else:
H, W = shortest_edge, longest_edge
image = image.resize((W, H))
buffered = BytesIO()
image.save(buffered, format="JPEG")
img_b64_str = base64.b64encode(buffered.getvalue()).decode()
img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
msg = img_str + msg.replace('<image>', '').strip()
ret.append([msg, None])
else:
ret.append([msg['value'], None])
else:
ret[-1][-1] = msg['value']
return ret |