import numpy as np

def calculate_class_weights(Y, n_classes, method="paszke", c=1.02):
    # CLASS PROBABILITIES - based on empirical observation of data
    ids, counts = np.unique(Y, return_counts=True)
    n_pixels = Y.size
    p_class = np.zeros(n_classes)
    p_class[ids] = counts/n_pixels

    # CLASS WEIGHTS
    if method == "paszke":
        weights = 1/np.log(c+p_class)
    elif method == "eigen":
        assert False, "TODO: Implement eigen method"
        # TODO: Implement eigen method
        # where class_freq is the number of pixels of class c divided by
        # the total number of pixels in images where c is actually present,
        # and median freq is the median of these frequencies.
    elif method in {"eigen2", "logeigen2"}:
        epsilon = 1e-8 # to prevent division by 0
        median = np.median(p_class)
        weights = median/(p_class+epsilon)
        if method == "logeigen2":
            weights = np.log(weights+1)
    else:
        assert False, "Incorrect choice for method"

    return weights

@@ -17,75 +17,1874 @@
 np%0A%0A
-def calculate_class_weights(Y, n_classes, method=%22paszke%22, c=1.02):
+# ==============================================================================%0A#                                                        CALCULATE_CLASS_WEIGHTS%0A# ==============================================================================%0Adef calculate_class_weights(Y, n_classes, method=%22paszke%22, c=1.02):%0A    %22%22%22 Given the training data labels Calculates the class weights.%0A%0A    Args:%0A        Y:      (numpy array) The training labels as class id integers.%0A                The shape does not matter, as long as each element represents%0A                a class id (ie, NOT one-hot-vectors).%0A        n_classes: (int) Number of possible classes.%0A        method: (str) The type of class weighting to use.%0A%0A                - %22paszke%22 = use the method from from Paszke et al 2016%0A                            %601/ln(c + class_probability)%60%0A                - %22eigen%22  = use the method from Eigen & Fergus 2014.%0A                             %60median_freq/class_freq%60%0A                             where %60class_freq%60 is based only on images that%0A                             actually contain that class.%0A                - %22eigen2%22 = Similar to %60eigen%60, except that class_freq is%0A                             based on the frequency of the class in the%0A                             entire dataset, not just images where it occurs.%0A                -%22logeigen2%22 = takes the log of %22eigen2%22 method, so that%0A                            incredibly rare classes do not completely overpower%0A                            other values.%0A        c:      (float) Coefficient to use, when using paszke method.%0A%0A    Returns:%0A        weights:    (numpy array) Array of shape %5Bn_classes%5D assigning a%0A                    weight value to each class.%0A%0A    References:%0A        Eigen & Fergus 2014: https://arxiv.org/abs/1411.4734%0A        Paszke et al 2016: https://arxiv.org/abs/1606.02147%0A    %22%22%22
 %0A   

Add docstring for calculate_class_weights()