Agents-X/data-view
1
1import json2import os3from typing import Optional, Union4from PIL import Image5import base646from io import BytesIO7import gradio as gr8import markdown9import zipfile10import tempfile11from datetime import datetime12import re13 14def export_to_zip(images, conversations, format_type="original"):15 """16 Export images and conversation data to a ZIP file17 18 Args:19 images: List of extracted images20 conversations: Conversation JSON data21 format_type: Format type, "original" or "sharegpt"22 23 Returns:24 Path to the generated ZIP file25 """26 # Create a temporary directory27 temp_dir = tempfile.mkdtemp()28 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")29 zip_filename = os.path.join(temp_dir, f"export_{timestamp}.zip")30 31 # Create a ZIP file32 with zipfile.ZipFile(zip_filename, 'w') as zipf:33 # Save images34 for i, img in enumerate(images):35 img_path = os.path.join(temp_dir, f"image_{i}.png")36 img.save(img_path)37 zipf.write(img_path, f"images/image_{i}.png")38 os.remove(img_path) # Delete temporary image file39 40 # Save conversation data41 json_path = os.path.join(temp_dir, "conversations.json")42 with open(json_path, 'w', encoding='utf-8') as f:43 json.dump(conversations, f, ensure_ascii=False, indent=4)44 zipf.write(json_path, "conversations.json")45 os.remove(json_path) # Delete temporary JSON file46 47 return zip_filename48 49def base64_to_image(50 base64_str: str, 51 remove_prefix: bool = True, 52 convert_mode: Optional[str] = "RGB"53) -> Union[Image.Image, None]:54 """55 Convert a base64 encoded image string to a PIL Image object56 57 Args:58 base64_str: Base64 encoded image string (with or without data: prefix)59 remove_prefix: Whether to automatically remove the "data:image/..." prefix (default True)60 convert_mode: Convert to the specified mode (e.g., "RGB"/"RGBA", None means no conversion)61 62 Returns:63 PIL.Image.Image object, returns None if decoding fails64 """65 try:66 # 1. Handle Base64 prefix67 if remove_prefix and "," in base64_str:68 base64_str = base64_str.split(",")[1]69 70 # 2. Decode Base6471 image_data = base64.b64decode(base64_str)72 73 # 3. Convert to PIL Image74 image = Image.open(BytesIO(image_data))75 76 # 4. Optional mode conversion77 if convert_mode:78 image = image.convert(convert_mode)79 80 return image81 82 except (base64.binascii.Error, OSError, Exception) as e:83 print(f"Base64 decoding failed: {str(e)}")84 return None85 86def process_message_to_sharegpt_format(message):87 """88 Convert messages to ShareGPT format89 90 Args:91 message: Original message data92 93 Returns:94 Data in ShareGPT format95 """96 sharegpt_images = []97 sharegpt_conversation = []98 image_idx = 099 100 for i, message_item in enumerate(message):101 role = message_item['role']102 103 content_list = message_item['content']104 whole_content = ""105 for content_item in content_list:106 content_type = content_item['type']107 if content_type == "text":108 content_value = content_item['text']109 whole_content += content_value110 elif content_type == "image_url":111 content_value = content_item['image_url']['url']112 whole_content += "<image>"113 image = base64_to_image(content_value)114 if image:115 sharegpt_images.append(image)116 image_idx += 1117 118 if i == 0:119 sharegpt_conversation.append({"from": "human", "value": whole_content})120 continue121 122 if "<interpreter>" in whole_content:123 gpt_content, observation_content = whole_content.split("<interpreter>", -1)124 sharegpt_conversation.append({"from": "gpt", "value": gpt_content})125 sharegpt_conversation.append({"from": "observation", "value": "<interpreter>"+observation_content})126 elif i != 0:127 sharegpt_conversation.append({"from": "gpt", "value": whole_content})128 129 sharegpt_data_item = {130 "conversations": sharegpt_conversation,131 "images": sharegpt_images132 }133 134 return sharegpt_data_item135 136def extract_images_from_messages(messages):137 """138 Extract all images from messages139 140 Args:141 messages: Message JSON data142 143 Returns:144 Extracted image list and updated messages145 """146 images = []147 148 for message in messages:149 if 'content' in message and isinstance(message['content'], list):150 for content_item in message['content']:151 if content_item.get('type') == 'image_url':152 image_url = content_item.get('image_url', {}).get('url', '')153 if image_url.startswith('data:'):154 # Extract base64 image155 image = base64_to_image(image_url)156 if image:157 images.append(image)158 159 return images, messages160 161def process_message(file_path):162 try:163 # Read JSON file164 with open(file_path, "r", encoding="utf-8") as f:165 messages = json.load(f)166 167 # Extract images168 images, messages = extract_images_from_messages(messages)169 170 # Convert to ShareGPT format171 sharegpt_data = process_message_to_sharegpt_format(messages)172 173 # Create HTML output174 html_output = '<div style="color: black;">' # Add a wrapper div for all content, set text color black175 176 for message_item in messages:177 role = message_item['role']178 content = message_item['content']179 180 # Style based on role181 if role == "user" or role == "human":182 html_output += f'<div style="background-color: #f0f0f0; padding: 10px; margin: 10px 0; border-radius: 10px; color: black;"><strong>User:</strong><br>'183 elif role == "assistant":184 html_output += f'<div style="background-color: #e6f7ff; padding: 10px; margin: 10px 0; border-radius: 10px; color: black;"><strong>Assistant:</strong><br>'185 else:186 html_output += f'<div style="background-color: #f9f9f9; padding: 10px; margin: 10px 0; border-radius: 10px; color: black;"><strong>{role.capitalize()}:</strong><br>'187 188 # Handle content189 for content_item in content:190 content_type = content_item['type']191 192 if content_type == "text":193 # Convert Markdown text to HTML194 md_text = content_item['text']195 html_text = markdown.markdown(md_text, extensions=['fenced_code', 'codehilite'])196 html_output += f'<div style="color: black;">{html_text}</div>'197 198 elif content_type == "image_url":199 content_value = content_item['image_url']['url']200 # If base64 image201 if content_value.startswith("data:"):202 html_output += f'<img src="{content_value}" style="max-width: 100%; margin: 10px 0;">'203 else:204 html_output += f'<img src="{content_value}" style="max-width: 100%; margin: 10px 0;">'205 206 html_output += '</div>'207 208 html_output += '</div>' # Close outermost div209 return html_output, images, messages, sharegpt_data210 211 except Exception as e:212 return f"<div style='color: red;'>Error processing file: {str(e)}</div>", [], None, None213 214def upload_and_process(file):215 if file is None:216 return "Please upload a JSON file", [], None, None217 218 html_output, images, messages, sharegpt_data = process_message(file.name)219 return html_output, images, messages, sharegpt_data220 221def use_example():222 # Use example file223 example_path = "test_message_gpt.json"224 return process_message(example_path)225 226def handle_export_original(images, conversations):227 """Handle export request for original format"""228 if not images or conversations is None:229 return None230 231 zip_path = export_to_zip(images, conversations, "original")232 return zip_path233 234def handle_export_sharegpt(sharegpt_data):235 """Handle export request for ShareGPT format"""236 if sharegpt_data is None:237 return None238 239 images = sharegpt_data.get("images", [])240 conversations = sharegpt_data.get("conversations", [])241 242 if not images and not conversations:243 return None244 245 zip_path = export_to_zip(images, conversations, "sharegpt")246 return zip_path247 248# Ensure example file exists249def setup_example_file():250 # Here we need to create the example file because we don't have actual content251 # In a real application, you should place the original test_message_gpt.json file in the root directory252 example_path = "test_message_gpt.json"253 254 # Create a simple example if the file does not exist255 if not os.path.exists(example_path):256 example_messages = [257 {258 "role": "user",259 "content": [260 {261 "type": "text",262 "text": "Hello, please introduce yourself."263 }264 ]265 },266 {267 "role": "assistant",268 "content": [269 {270 "type": "text",271 "text": "Hello! I am an AI assistant. I can help answer questions, provide information, and have conversations. I am designed to assist users with a variety of tasks, from simple Q&A to more complex discussions.\n\nI can handle text information and also understand and describe images. Although I have some limitations, I will do my best to provide useful, accurate, and helpful responses.\n\nHow can I help you today?"272 }273 ]274 }275 ]276 277 with open(example_path, "w", encoding="utf-8") as f:278 json.dump(example_messages, f, ensure_ascii=False, indent=2)279 280# Set up the example file281setup_example_file()282 283# Create Gradio interface284with gr.Blocks(title="ChatGPT Conversation Visualizer", css="div.prose * {color: black !important;}") as demo:285 gr.Markdown("# ChatGPT Conversation Visualization Tool")286 gr.Markdown("Upload a JSON file containing ChatGPT conversation records or use the example file to view visualization results.")287 288 with gr.Row():289 file_input = gr.File(label="Upload JSON File", file_types=[".json"])290 291 with gr.Row():292 col1, col2 = gr.Column(), gr.Column()293 with col1:294 visualize_button = gr.Button("Visualize Uploaded Conversation")295 with col2:296 example_button = gr.Button("Use Example File")297 298 with gr.Row():299 output = gr.HTML(label="Conversation Content")300 301 # Add export buttons302 with gr.Row():303 with gr.Column():304 export_original_btn = gr.Button("Export Original Format")305 download_original_file = gr.File(label="Download Original Format ZIP")306 307 with gr.Column():308 export_sharegpt_btn = gr.Button("Export ShareGPT Format")309 download_sharegpt_file = gr.File(label="Download ShareGPT Format ZIP")310 311 # State variables to store current results312 current_images = gr.State([])313 current_json = gr.State(None)314 current_sharegpt = gr.State(None)315 316 visualize_button.click(317 fn=upload_and_process,318 inputs=[file_input],319 outputs=[output, current_images, current_json, current_sharegpt]320 )321 322 example_button.click(323 fn=use_example,324 inputs=[],325 outputs=[output, current_images, current_json, current_sharegpt]326 )327 328 export_original_btn.click(329 fn=handle_export_original,330 inputs=[current_images, current_json],331 outputs=[download_original_file]332 )333 334 export_sharegpt_btn.click(335 fn=handle_export_sharegpt,336 inputs=[current_sharegpt],337 outputs=[download_sharegpt_file]338 )339 340# Launch Gradio app341demo.launch()