table = wandb.Table(columns=["Image", "True Label", "Predicted Label"])

y_pred = []  
y_label = []  
all_images = []  

for img_batch, labels in test_data:
    batch_pred = tf.argmax(VGG_model(img_batch, training=False), axis=1).numpy()
    y_pred.extend(batch_pred)
    batch_labels = tf.argmax(labels, axis=1).numpy()
    y_label.extend(batch_labels)
    all_images.extend(img_batch.numpy())

y_pred = np.array(y_pred)
y_label = np.array(y_label)
all_images = np.array(all_images)


plt.figure(figsize=(10, 10))
for i in range(16):  
    ax = plt.subplot(4, 4, i + 1)
    img = (all_images[i] * 255).astype("uint8")  
    plt.imshow(img)
    plt.title(f"Actual: {class_names[y_label[i]]}\nPred: {class_names[y_pred[i]]}")
    plt.axis(False)

plt.tight_layout()
plt.show()


for i in range(16):
    table.add_data(
        wandb.Image(all_images[i]), 
        class_names[y_label[i]],     
        class_names[y_pred[i]]      
    )

run.log({"predictions_table": table})
run.finish()

        __
        __