Pandas style highlight diagonal and off diagonal elements

Question

I was attempting to style highlight both the main diagonal and opposite main diagonal elements in pandas dataframe.

I saw this link: Pandas style: How to highlight diagonal elements

This shows how to highlight the main diagonal but I was wondering how to highlight two diagonals with two different colors?

Here is my data frame:

import numpy as np
import pandas as pd


df = pd.DataFrame(data={'pred0': [10,   4],
                   'pred1': [0,   0],
            'total': [10,  4]},index=['true0','true1']
                   )

print(df)
       pred0  pred1  total
true0     10      0     10
true1      4      0      4

My attempt

# Credit: ALLOLZ
def highlight_diag(df):
    a = np.full(df.shape, '', dtype='<U24')
    np.fill_diagonal(a, 'background-color: yellow')
    return pd.DataFrame(a, index=df.index, columns=df.columns)

df.style.apply(highlight_diag, axis=None)

But this only hightlight one diagonal and does not highlight another diagonal. How to highlight both diagonals.

Required

          pred0          pred1         total
true0     10(green)      0(red)     10(no highlight)
true1      4(red)      0(green)      4(no highlight)

TIY.

Answer 1

This gives what you want

import numpy as np
import pandas as pd

def highlight_diags(data):
    attr1 = 'background-color: lightgreen'
    attr2 = 'background-color: salmon'

    df_style = data.replace(data, '')
    np.fill_diagonal(df_style.values, attr1)
    np.fill_diagonal(np.flipud(df_style), attr2) 
    return df_style

df = pd.DataFrame(data={'pred0': [10,   4],
                   'pred1': [0,   0],
            'total': [10,  4]},index=['true0','true1']
                   )

df.style.apply(highlight_diags,axis=None)