Add opencv-python to requirements

app.py

@@ -2,12 +2,10 @@ import gradio as gr
 import torch
 from transformers import GPT2Model, ViTModel, GPT2Tokenizer, ViTImageProcessor
 from captum.attr import IntegratedGradients
+import matplotlib.pyplot as plt
 from PIL import Image
 import numpy as np
 import cv2
-import matplotlib.pyplot as plt
-import io
-import base64
 
 # 定义多模态模型
 class MultiModalModel(torch.nn.Module):
@@ -39,57 +37,30 @@ def load_model():
     model.eval()
     return model
 
+# 转换张量为 PIL 图像
+def convert_tensor_to_pil(tensor_image):
+    if isinstance(tensor_image, torch.Tensor):
+        tensor_image = tensor_image.numpy()
+    image_np = np.transpose(tensor_image, (1, 2, 0))
+    if image_np.max() <= 1.0:
+        image_np = (image_np * 255).astype(np.uint8)
+    return Image.fromarray(image_np)
+
+# 自定义前向函数用于集成梯度
+def custom_forward(pixel_values, input_ids, attention_mask):
+    logits = model(input_ids=input_ids, attention_mask=attention_mask, pixel_values=pixel_values)
+    return logits
+
+# 初始化集成梯度
+integrated_gradients = IntegratedGradients(custom_forward)
+
 # 初始化模型和加载器
 model = load_model()
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 tokenizer.pad_token = tokenizer.eos_token
 feature_extractor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
 
-# 定义集成梯度解释性分析
-integrated_gradients = IntegratedGradients(model)
-
-# 转换为 PIL 图像
-def convert_to_pil(image_array):
-    if isinstance(image_array, torch.Tensor):
-        image_array = image_array.numpy()
-    image_array = np.transpose(image_array, (1, 2, 0))
-    if image_array.max() <= 1.0:
-        image_array = (image_array * 255).astype(np.uint8)
-    return Image.fromarray(image_array)
-
-# 可视化归因结果
-def visualize_attributions(attributions, pixel_values):
-    attribution_image = attributions.squeeze().cpu().numpy()
-    attribution_image = (attribution_image - attribution_image.min()) / (attribution_image.max() - attribution_image.min())
-    attribution_image = np.uint8(255 * attribution_image)
-    attribution_image_pil = convert_to_pil(attribution_image)
-
-    # 转换为灰度图并进行轮廓检测
-    attribution_gray = cv2.cvtColor(np.array(attribution_image_pil), cv2.COLOR_RGB2GRAY)
-    _, binary_mask = cv2.threshold(attribution_gray, 128, 255, cv2.THRESH_BINARY)
-    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-
-    contour_image = np.array(attribution_image_pil)
-    cv2.drawContours(contour_image, contours, -1, (255, 0, 0), 2)
-
-    # 创建可视化图像
-    fig, axes = plt.subplots(1, 2, figsize=(10, 5))
-    axes[0].imshow(convert_to_pil(pixel_values.squeeze(0).numpy()))
-    axes[0].axis('off')
-    axes[0].set_title("Original Image")
-    
-    axes[1].imshow(contour_image)
-    axes[1].axis('off')
-    axes[1].set_title("Attribution with Contours")
-
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png')
-    plt.close(fig)
-    buf.seek(0)
-    img_str = base64.b64encode(buf.getvalue()).decode('utf-8')
-    return f"data:image/png;base64,{img_str}"
-
-# 推理并进行归因分析
+# 定义推理和归因分析函数
 def predict(image, text):
     # 处理图像
     image = Image.fromarray(image)
@@ -104,8 +75,8 @@ def predict(image, text):
         padding="max_length"
     )
 
-    input_ids = inputs["input_ids"]
-    attention_mask = inputs["attention_mask"]
+    input_ids = inputs["input_ids"].long()
+    attention_mask = inputs["attention_mask"].long()
     pixel_values = image_features["pixel_values"]
 
     # 推理
@@ -114,18 +85,29 @@ def predict(image, text):
         prediction = torch.argmax(logits, dim=1).item()
         label = "yes" if prediction == 1 else "no"
 
-    # 集成梯度归因
+    # 归因分析
     attributions, _ = integrated_gradients.attribute(
         inputs=pixel_values,
         target=prediction,
         additional_forward_args=(input_ids, attention_mask),
         n_steps=1,
-        return_convergence_delta=False
+        return_convergence_delta=True
     )
 
-    attribution_visual = visualize_attributions(attributions, pixel_values)
-    
-    return label, attribution_visual
+    # 可视化归因结果
+    attribution_image = attributions.squeeze().cpu().numpy()
+    attribution_image = (attribution_image - attribution_image.min()) / (attribution_image.max() - attribution_image.min())
+    attribution_image = np.uint8(255 * attribution_image)
+    attribution_image_real = convert_tensor_to_pil(attribution_image)
+
+    # 轮廓检测
+    attribution_gray = cv2.cvtColor(np.array(attribution_image_real), cv2.COLOR_RGB2GRAY)
+    _, binary_mask = cv2.threshold(attribution_gray, 128, 255, cv2.THRESH_BINARY)
+    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    contour_image = np.array(attribution_image_real)
+    cv2.drawContours(contour_image, contours, -1, (255, 0, 0), 2)
+
+    return label, Image.fromarray(contour_image)
 
 # 创建 Gradio 界面
 iface = gr.Interface(