Update app.py

app.py

@@ -6,7 +6,7 @@ import torch
 from transformers import AutoProcessor, Pix2StructForConditionalGeneration
 from PIL import Image
 
-# 配置日志
+# 配置日志格式
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(levelname)s - %(message)s',
@@ -17,26 +17,59 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+def print_section(title, char='='):
+    """打印格式化的章节标题"""
+    print(f"\n{char * 50}")
+    print(f"{title.center(50)}")
+    print(f"{char * 50}\n")
+
+def print_table(data):
+    """格式化打印表格数据"""
+    if not data:
+        print("No data available")
+        return
+        
+    # 计算每列的最大宽度
+    col_widths = []
+    for i in range(len(data[0])):
+        col_width = max(len(str(row[i])) for row in data)
+        col_widths.append(col_width)
+    
+    # 打印表头
+    header = data[0]
+    header_str = " | ".join(str(header[i]).ljust(col_widths[i]) for i in range(len(header)))
+    print(header_str)
+    print("-" * len(header_str))
+    
+    # 打印数据行
+    for row in data[1:]:
+        row_str = " | ".join(str(row[i]).ljust(col_widths[i]) for i in range(len(row)))
+        print(row_str)
+
 class ChartAnalyzer:
     def __init__(self):
         try:
-            logger.info("Initializing model and processor...")
+            print_section("初始化模型")
+            print("正在加载模型和处理器...")
             self.model = Pix2StructForConditionalGeneration.from_pretrained("google/deplot")
             self.processor = AutoProcessor.from_pretrained("google/deplot")
-            logger.info("Model and processor initialized successfully")
+            print("✓ 模型加载完成")
         except Exception as e:
+            print("✗ 模型加载失败")
             logger.error(f"Error initializing model: {str(e)}")
             raise
 
     def process_image(self, image_path, prompt=None):
         """处理图片并生成数据表格"""
         try:
+            print_section("图片处理", char='-')
+            
             # 验证文件存在
             if not os.path.exists(image_path):
-                raise FileNotFoundError(f"Image file not found: {image_path}")
+                raise FileNotFoundError(f"找不到图片文件: {image_path}")
 
             # 打开并处理图片
-            logger.info(f"Processing image: {image_path}")
+            print(f"正在处理图片: {image_path}")
             image = Image.open(image_path)
             
             # 准备输入
@@ -50,8 +83,8 @@ class ChartAnalyzer:
             )
 
             # 生成预测
-            logger.info("Generating predictions...")
-            with torch.no_grad():  # 提高性能并减少内存使用
+            print("\n正在生成数据分析...")
+            with torch.no_grad():
                 predictions = self.model.generate(
                     **inputs,
                     max_new_tokens=512,
@@ -77,31 +110,36 @@ class ChartAnalyzer:
                 for row in result_array:
                     file.write(" | ".join(row) + "\n")
             
-            logger.info(f"Results saved to {output_file}")
+            print(f"\n✓ 结果已保存至: {output_file}")
             return result_array
 
         except Exception as e:
+            print("\n✗ 处理失败")
             logger.error(f"Error processing image: {str(e)}")
             raise
 
 def main():
     try:
+        print_section("图表数据提取系统", char='*')
+        
         # 创建分析器实例
         analyzer = ChartAnalyzer()
         
-        # 指定图片路径（在Space中使用上传的图片路径）
+        # 指定图片路径
         image_path = '05e57f1c9acff69f1eb6fa72d4805d0.jpg'
         
         # 处理图片
         results = analyzer.process_image(image_path)
         
         # 打印结果
-        print("\nAnalysis Results:")
-        for row in results:
-            print(" | ".join(row))
+        print_section("分析结果")
+        print_table(results)
+        
+        print_section("处理完成", char='*')
             
     except Exception as e:
         logger.error(f"Application error: {str(e)}")
+        print("\n✗ 程序执行出错，请查看日志获取详细信息")
         raise
 
 if __name__ == "__main__":