Create winrate.py

winrate.py

@@ -0,0 +1,147 @@
+import os
+import pandas as pd
+from collections import defaultdict
+
+def calculate_avg_comet(root_dir):
+    results = []
+    
+    for category in os.listdir(root_dir):
+        category_path = os.path.join(root_dir, category)
+        if not os.path.isdir(category_path):
+            continue
+            
+        for subcategory in os.listdir(category_path):
+            subcategory_path = os.path.join(category_path, subcategory)
+            if not os.path.isdir(subcategory_path):
+                continue
+                
+            report_path = os.path.join(subcategory_path, 'report', 'full_image_comet.csv')
+            if os.path.exists(report_path):
+                try:
+                    df = pd.read_csv(report_path)
+                    avg_score = df['avg_comet_sentence'].mean()
+                    results.append({
+                        'category': category,
+                        'subcategory': subcategory,
+                        'avg_comet': avg_score
+                    })
+                except Exception as e:
+                    print(f"Error processing {report_path}: {e}")
+    
+    result_df = pd.DataFrame(results)
+    result_df.to_csv('comet_avg_scores.csv', index=False)
+    return result_df
+
+def compare_win_rate(dir1, dir2):
+    data1 = defaultdict(dict)
+    data2 = defaultdict(dict)
+    
+    # 读取第一个目录的数据
+    for category in os.listdir(dir1):
+        category_path = os.path.join(dir1, category)
+        if not os.path.isdir(category_path):
+            continue
+            
+        for subcategory in os.listdir(category_path):
+            subcategory_path = os.path.join(category_path, subcategory)
+            if not os.path.isdir(subcategory_path):
+                continue
+                
+            report_path = os.path.join(subcategory_path, 'report', 'full_image_comet.csv')
+            if os.path.exists(report_path):
+                try:
+                    df = pd.read_csv(report_path)
+                    for _, row in df.iterrows():
+                        data1[(category, subcategory)][row['file_name']] = row['avg_comet_sentence']
+                except Exception as e:
+                    print(f"Error processing {report_path}: {e}")
+    
+    # 读取第二个目录的数据
+    for category in os.listdir(dir2):
+        category_path = os.path.join(dir2, category)
+        if not os.path.isdir(category_path):
+            continue
+            
+        for subcategory in os.listdir(category_path):
+            subcategory_path = os.path.join(category_path, subcategory)
+            if not os.path.isdir(subcategory_path):
+                continue
+                
+            report_path = os.path.join(subcategory_path, 'report', 'full_image_comet.csv')
+            if os.path.exists(report_path):
+                try:
+                    df = pd.read_csv(report_path)
+                    for _, row in df.iterrows():
+                        data2[(category, subcategory)][row['file_name']] = row['avg_comet_sentence']
+                except Exception as e:
+                    print(f"Error processing {report_path}: {e}")
+    
+    # 计算win rate
+    win_rate_results = []
+    category_stats = defaultdict(lambda: {'wins': 0, 'total': 0})
+    
+    for key in set(data1.keys()).union(set(data2.keys())):
+        category, subcategory = key
+        scores1 = data1.get(key, {})
+        scores2 = data2.get(key, {})
+        
+        common_files = set(scores1.keys()).intersection(set(scores2.keys()))
+        if not common_files:
+            continue
+            
+        wins = 0
+        total = 0
+        
+        for file in common_files:
+            total += 1
+            if scores1[file] > scores2[file]:
+                wins += 1
+        
+        win_rate = wins / total if total > 0 else 0
+        win_rate_results.append({
+            'category': category,
+            'subcategory': subcategory,
+            'win_rate': win_rate,
+            'wins': wins,
+            'total_comparisons': total
+        })
+        
+        # 更新大类统计
+        category_stats[category]['wins'] += wins
+        category_stats[category]['total'] += total
+    
+    # 添加大类统计结果
+    for category, stats in category_stats.items():
+        cat_win_rate = stats['wins'] / stats['total'] if stats['total'] > 0 else 0
+        win_rate_results.append({
+            'category': category,
+            'subcategory': 'ALL',
+            'win_rate': cat_win_rate,
+            'wins': stats['wins'],
+            'total_comparisons': stats['total']
+        })
+    
+    # 创建DataFrame并保存
+    win_rate_df = pd.DataFrame(win_rate_results)
+    win_rate_df.to_csv('win_rate_comparison.csv', index=False)
+    
+    # 打印统计数据
+    print("\n=== 大类统计结果 ===")
+    for category, stats in category_stats.items():
+        win_rate = stats['wins'] / stats['total'] if stats['total'] > 0 else 0
+        print(f"大类: {category}")
+        print(f"总比较次数: {stats['total']}")
+        print(f"获胜次数: {stats['wins']}")
+        print(f"Win Rate: {win_rate:.2%}")
+        print("-" * 30)
+    
+    return win_rate_df
+
+if __name__ == "__main__":
+    print("Processing task 1...")
+    avg_scores = calculate_avg_comet('organized_data_1')
+    print("Average comet scores saved to comet_avg_scores.csv")
+    
+    print("\nProcessing task 2...")
+    win_rate_df = compare_win_rate('organized_data_1', 'organized_data_2')
+    print("\nWin rate comparison saved to win_rate_comparison.csv")