修改图像分类黑盒水印检测问题

watermark_verify/process/classification_pytorch_blackbox_process.py

@@ -51,7 +51,7 @@ class ModelWatermarkProcessor(BlackBoxWatermarkProcessDefine):
 
             # 提取预测结果
             for j, image_file in enumerate(batch_files):
-                predicted_class = np.argmax(outputs[0][j])  # 假设输出是每类的概率
+                predicted_class = np.argmax(outputs[j])  # 假设输出是每类的概率
                 total_predictions += 1
 
                 # 比较预测结果与目标分类

watermark_verify/process/classification_tensorflow_blackbox_process.py

@@ -51,7 +51,7 @@ class ModelWatermarkProcessor(BlackBoxWatermarkProcessDefine):
 
             # 提取预测结果
             for j, image_file in enumerate(batch_files):
-                predicted_class = np.argmax(outputs[0][j])  # 假设输出是每类的概率
+                predicted_class = np.argmax(outputs[j])  # 假设输出是每类的概率
                 total_predictions += 1
 
                 # 比较预测结果与目标分类