TinyLLaVA_Factory/tinyllava/utils/message.py

'''
@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