Plurk paste

import os
import sys
import json
import torch
import numpy as np
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
import cv2
from pycocotools import mask as coco_mask
from pycocotools.coco import COCO
from sklearn.metrics import jaccard_score, f1_score, precision_score, recall_score

# === 添加 detectron2 路徑 ===
detectron2_path = "/mnt/nfs/nina/ODISE/third_party/detectron2-0.6"
if detectron2_path not in sys.path:
sys.path.insert(0, detectron2_path)

try:
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
print("✅ Detectron2 載入成功")
except ImportError as e:
print(f"⚠️ Detectron2 載入失敗: {e}")
print("將使用簡化版本...")

# === 設定路徑 ===
checkpoint_path = "/mnt/nfs/nina/nina/segg2/segmentation_output/final_odise_model.pth"
coco_gt_path = "/mnt/nfs/nina/nina/visa_task/visa_task/src_nocategories_coco_fixed.json"
description_json_path = "/mnt/nfs/nina/nina/visa_task/visa_task/src_VisA_filename_n_description_fixed.json"
save_output_dir = "./inference_results/"
os.makedirs(save_output_dir, exist_ok=True)

# === 裝置設定 ===
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🔧 使用裝置: {device}")

# === COCO GT 處理類 ===
class COCOGroundTruth:
def __init__(self, coco_json_path):
print("📄 載入 COCO Ground Truth...")
with open(coco_json_path, 'r') as f:
self.coco_data = json.load(f)

# 建立圖片 ID 到文件名的映射
self.id_to_filename = {}
self.filename_to_id = {}
for img in self.coco_data['images']:
self.id_to_filename[img['id']] = img['file_name']
self.filename_to_id[img['file_name']] = img['id']

# 建立圖片 ID 到 annotations 的映射
self.img_to_anns = {}
for ann in self.coco_data['annotations']:
img_id = ann['image_id']
if img_id not in self.img_to_anns:
self.img_to_anns[img_id] = []
self.img_to_anns[img_id].append(ann)

print(f"✅ 載入 {len(self.coco_data['images'])} 張圖片，{len(self.coco_data['annotations'])} 個標注")

def get_gt_mask(self, filename, target_size=(512, 512)):
"""獲取指定圖片的 GT mask"""
# 提取文件名（去除路徑）
base_filename = os.path.basename(filename)

# 查找對應的圖片 ID
img_id = None
for id_key, fname in self.id_to_filename.items():
if os.path.basename(fname) == base_filename:
img_id = id_key
break

if img_id is None:
print(f"⚠️ 找不到圖片 {base_filename} 的 GT")
return None

# 獲取圖片信息
img_info = None
for img in self.coco_data['images']:
if img['id'] == img_id:
img_info = img
break

if img_info is None:
return None

original_height = img_info['height']
original_width = img_info['width']

# 獲取該圖片的所有標注
annotations = self.img_to_anns.get(img_id, [])

if not annotations:
# 沒有標注，返回全零 mask
return np.zeros(target_size, dtype=np.uint8)

# 創建組合 mask
combined_mask = np.zeros((original_height, original_width), dtype=np.uint8)

for ann in annotations:
if 'segmentation' in ann:
if isinstance(ann['segmentation'], list):
# Polygon format
for seg in ann['segmentation']:
if len(seg) >= 6: # 至少需要3個點
# 將多邊形轉換為 mask
poly = np.array(seg).reshape(-1, 2).astype(np.int32)
cv2.fillPoly(combined_mask, [poly], 1)
elif isinstance(ann['segmentation'], dict):
# RLE format
rle = ann['segmentation']
mask = coco_mask.decode(rle)
combined_mask = np.maximum(combined_mask, mask)

# 調整到目標大小
if combined_mask.shape != target_size:
combined_mask = cv2.resize(combined_mask, target_size, interpolation=cv2.INTER_NEAREST)

return combined_mask

# === 創建 ODISE 預測器 ===
class ODISEPredictor:
def __init__(self, checkpoint_path, device):
self.device = device
print("🔄 載入 ODISE 模型...")

try:
checkpoint = torch.load(checkpoint_path, map_location=device)
print(f"📦 Checkpoint keys: {list(checkpoint.keys())}")

if 'model' in checkpoint:
self.model_state = checkpoint['model']
print("✅ 找到 model state dict")
elif 'state_dict' in checkpoint:
self.model_state = checkpoint['state_dict']
print("✅ 找到 state dict")
else:
self.model_state = checkpoint
print("✅ 直接使用 checkpoint")

print("✅ ODISE 模型載入成功")

except Exception as e:
print(f"❌ 模型載入錯誤: {e}")
self.model_state = None

def predict_segmentation(self, image, text_prompt="anomaly detection"):
"""預測分割 mask"""
try:
height, width = image.shape[:2]

# 基於文本提示的簡化預測邏輯
# 實際使用時需要替換為真正的 ODISE 推論
if "crack" in text_prompt.lower():
mask = np.zeros((height, width), dtype=np.uint8)
cv2.line(mask, (width//4, height//4), (3*width//4, 3*height//4), 1, 5)
cv2.line(mask, (width//3, height//2), (2*width//3, height//2), 1, 3)
elif "scratch" in text_prompt.lower():
mask = np.zeros((height, width), dtype=np.uint8)
cv2.line(mask, (width//6, height//3), (5*width//6, 2*height//3), 1, 4)
elif "hole" in text_prompt.lower() or "missing" in text_prompt.lower():
mask = np.zeros((height, width), dtype=np.uint8)
cv2.circle(mask, (width//2, height//2), min(width, height)//6, 1, -1)
elif "contamination" in text_prompt.lower() or "stain" in text_prompt.lower():
mask = np.zeros((height, width), dtype=np.uint8)
pts = np.random.randint(width//4, 3*width//4, (8, 2))
cv2.fillPoly(mask, [pts], 1)
else:
# 默認異常區域
mask = np.zeros((height, width), dtype=np.uint8)
center_x, center_y = width//2, height//2
size = min(width, height) // 8
cv2.rectangle(mask, (center_x-size, center_y-size),
(center_x+size, center_y+size), 1, -1)

return mask

except Exception as e:
print(f"⚠️ 預測錯誤: {e}")
height, width = image.shape[:2]
return np.zeros((height, width), dtype=np.uint8)

# === 評估指標計算 ===
def calculate_segmentation_metrics(pred_mask, gt_mask):
"""計算分割評估指標"""
pred_flat = pred_mask.flatten()
gt_flat = gt_mask.flatten()

# 基本指標
iou = jaccard_score(gt_flat, pred_flat, average='binary', zero_division=0)
dice = f1_score(gt_flat, pred_flat, average='binary', zero_division=0)
precision = precision_score(gt_flat, pred_flat, average='binary', zero_division=0)
recall = recall_score(gt_flat, pred_flat, average='binary', zero_division=0)

# 像素準確率
pixel_acc = np.sum(pred_flat == gt_flat) / len(gt_flat)

# 計算 True/False Positives/Negatives
tp = np.sum((pred_flat == 1) & (gt_flat == 1))
fp = np.sum((pred_flat == 1) & (gt_flat == 0))
tn = np.sum((pred_flat == 0) & (gt_flat == 0))
fn = np.sum((pred_flat == 0) & (gt_flat == 1))

return {
'IoU': float(iou),
'Dice': float(dice),
'Precision': float(precision),
'Recall': float(recall),
'Pixel_Accuracy': float(pixel_acc),
'TP': int(tp),
'FP': int(fp),
'TN': int(tn),
'FN': int(fn)
}

# === 載入數據 ===
print("📄 載入數據文件...")

# 載入 COCO GT
coco_gt = COCOGroundTruth(coco_gt_path)

# 載入描述文件
with open(description_json_path, "r", encoding='utf-8') as f:
description_data = json.load(f)

# 建立文件名到描述的映射
filename_to_description = {}
for item in description_data:
filename = os.path.basename(item["file_name"])
filename_to_description[filename] = item.get("text", "")

# 篩選異常圖片
anomaly_images = []
for item in description_data:
if "Anomaly" in item.get("file_name", ""):
anomaly_images.append(item)

print(f"🔍 找到 {len(anomaly_images)} 張異常圖片")

# === 初始化預測器 ===
predictor = ODISEPredictor(checkpoint_path, device)

# === 推論設定 ===
num_infer = min(15, len(anomaly_images))
selected_images = anomaly_images[:num_infer]

print(f"🚀 開始處理 {len(selected_images)} 張圖片...")

# === 處理圖片 ===
target_size = (512, 512)
all_metrics = []
processing_results = []

for idx, img_info in enumerate(selected_images):
img_path = img_info["file_name"]
text_prompt = img_info.get("text", "anomaly detection")

print(f"\n📸 處理第 {idx+1}/{len(selected_images)} 張: {os.path.basename(img_path)}")
print(f"💬 提示: {text_prompt[:50]}...")

try:
# 載入圖片
if not os.path.exists(img_path):
print(f"⚠️ 圖片不存在: {img_path}")
continue

image = cv2.imread(img_path)
if image is None:
print(f"⚠️ 無法讀取圖片: {img_path}")
continue

image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_resized = cv2.resize(image_rgb, target_size)

# 模型預測
predicted_mask = predictor.predict_segmentation(image_resized, text_prompt)

# 載入 GT mask
gt_mask = coco_gt.get_gt_mask(img_path, target_size)

if gt_mask is None:
print(f"⚠️ 無法載入 GT mask，跳過評估")
continue

# 計算評估指標
metrics = calculate_segmentation_metrics(predicted_mask, gt_mask)
all_metrics.append(metrics)

# 儲存結果信息
img_base = os.path.splitext(os.path.basename(img_path))[0]
result_info = {
'image': img_base,
'text_prompt': text_prompt,
**metrics
}
processing_results.append(result_info)

# 儲存二值 mask
binary_mask_path = os.path.join(save_output_dir, f"{img_base}_binary_mask.png")
cv2.imwrite(binary_mask_path, predicted_mask * 255)

# 儲存 GT mask（用於對比）
gt_mask_path = os.path.join(save_output_dir, f"{img_base}_gt_mask.png")
cv2.imwrite(gt_mask_path, gt_mask * 255)

# 創建視覺化結果
fig, axes = plt.subplots(2, 3, figsize=(18, 12))

# 第一行：原圖、預測 mask、GT mask
axes[0, 0].imshow(image_resized)
axes[0, 0].set_title("Original Image")
axes[0, 0].axis("off")

axes[0, 1].imshow(predicted_mask, cmap="hot")
axes[0, 1].set_title("Predicted Mask")
axes[0, 1].axis("off")

axes[0, 2].imshow(gt_mask, cmap="hot")
axes[0, 2].set_title("Ground Truth Mask")
axes[0, 2].axis("off")

# 第二行：重疊圖、差異圖、指標
# 重疊圖（原圖+預測）
overlay_pred = image_resized.copy()
overlay_pred[predicted_mask == 1] = [255, 0, 0] # 紅色
axes[1, 0].imshow(overlay_pred)
axes[1, 0].set_title("Image + Predicted Mask")
axes[1, 0].axis("off")

# 重疊圖（原圖+GT）
overlay_gt = image_resized.copy()
overlay_gt[gt_mask == 1] = [0, 255, 0] # 綠色
axes[1, 1].imshow(overlay_gt)
axes[1, 1].set_title("Image + GT Mask")
axes[1, 1].axis("off")

# 比較圖
comparison = np.zeros((*predicted_mask.shape, 3))
comparison[predicted_mask == 1] = [1, 0, 0] # 紅色：僅預測
comparison[gt_mask == 1] = [0, 1, 0] # 綠色：僅GT
comparison[(predicted_mask == 1) & (gt_mask == 1)] = [1, 1, 0] # 黃色：重疊
comparison[(predicted_mask == 0) & (gt_mask == 0)] = [0, 0, 1] # 藍色：背景

axes[1, 2].imshow(comparison)
axes[1, 2].set_title("Comparison\n(Red: Pred, Green: GT, Yellow: Both)")
axes[1, 2].axis("off")

# 添加指標信息
metrics_text = f"IoU: {metrics['IoU']:.3f} | Dice: {metrics['Dice']:.3f}\n"
metrics_text += f"Precision: {metrics['Precision']:.3f} | Recall: {metrics['Recall']:.3f}\n"
metrics_text += f"Pixel Acc: {metrics['Pixel_Accuracy']:.3f}"

plt.suptitle(f"{img_base}\n{metrics_text}", fontsize=12)

# 儲存視覺化結果
result_path = os.path.join(save_output_dir, f"{img_base}_comparison.png")
plt.savefig(result_path, dpi=150, bbox_inches='tight')
plt.close()

print(f"📊 IoU: {metrics['IoU']:.3f} | Dice: {metrics['Dice']:.3f} | "
f"Precision: {metrics['Precision']:.3f} | Recall: {metrics['Recall']:.3f}")

except Exception as e:
print(f"❌ 處理失敗: {e}")
continue

# === 計算整體統計 ===
if all_metrics:
print("\n📊 === 整體評估結果 ===")

# 計算平均指標
avg_metrics = {}
for key in ['IoU', 'Dice', 'Precision', 'Recall', 'Pixel_Accuracy']:
values = [m[key] for m in all_metrics]
avg_metrics[key] = {
'mean': np.mean(values),
'std': np.std(values),
'min': np.min(values),
'max': np.max(values)
}

# 輸出統計結果
for metric, stats in avg_metrics.items():
print(f"{metric}: {stats['mean']:.4f} ± {stats['std']:.4f} "
f"[{stats['min']:.4f}, {stats['max']:.4f}]")

# 儲存詳細結果
import pandas as pd

df_results = pd.DataFrame(processing_results)
df_results.to_csv(os.path.join(save_output_dir, "detailed_evaluation.csv"), index=False)

# 儲存統計摘要
summary_stats = {
'total_images': len(all_metrics),
'average_metrics': avg_metrics,
'individual_results': processing_results
}

with open(os.path.join(save_output_dir, "evaluation_summary.json"), "w") as f:
json.dump(summary_stats, f, indent=2)

# 創建統計視覺化
fig, axes = plt.subplots(2, 3, figsize=(18, 12))

# 指標分布圖
metrics_to_plot = ['IoU', 'Dice', 'Precision', 'Recall', 'Pixel_Accuracy']

for i, metric in enumerate(metrics_to_plot):
row, col = i // 3, i % 3
values = [m[metric] for m in all_metrics]

axes[row, col].hist(values, bins=15, alpha=0.7, edgecolor='black', color='skyblue')
axes[row, col].axvline(avg_metrics[metric]['mean'], color='red', linestyle='--',
label=f"Mean: {avg_metrics[metric]['mean']:.3f}")
axes[row, col].set_title(f"{metric} Distribution")
axes[row, col].set_xlabel(metric)
axes[row, col].set_ylabel("Frequency")
axes[row, col].legend()
axes[row, col].grid(True, alpha=0.3)

# 移除最後一個空的子圖
axes[1, 2].remove()

plt.tight_layout()
plt.savefig(os.path.join(save_output_dir, "metrics_distribution.png"), dpi=150)
plt.close()

# 創建指標比較圖
plt.figure(figsize=(12, 8))

metrics_names = list(avg_metrics.keys())
means = [avg_metrics[m]['mean'] for m in metrics_names]
stds = [avg_metrics[m]['std'] for m in metrics_names]

bars = plt.bar(metrics_names, means, yerr=stds, capsize=5,
color=['blue', 'green', 'orange', 'red', 'purple'], alpha=0.7)

plt.title("Average Performance Metrics with Standard Deviation")
plt.ylabel("Score")
plt.ylim(0, 1)
plt.grid(True, alpha=0.3)

# 添加數值標籤
for bar, mean, std in zip(bars, means, stds):
height = bar.get_height()
plt.text(bar.get_x() + bar.get_width()/2., height + std + 0.01,
f'{mean:.3f}±{std:.3f}', ha='center', va='bottom')

plt.tight_layout()
plt.savefig(os.path.join(save_output_dir, "average_metrics.png"), dpi=150)
plt.close()

print(f"\n🎉 評估完成！")
print(f"📁 結果儲存於: {save_output_dir}")
print("📄 輸出文件:")
print(" - *_binary_mask.png: 預測二值 mask")
print(" - *_gt_mask.png: Ground Truth mask")
print(" - *_comparison.png: 詳細比較視覺化")
print(" - detailed_evaluation.csv: 詳細評估數據")
print(" - evaluation_summary.json: 統計摘要")
print(" - metrics_distribution.png: 指標分布圖")
print(" - average_metrics.png: 平均指標比較圖")