app.py

import numpy as np
import matplotlib.pyplot as plt
import gradio as gr
import cv2
import scipy

class GMM:
    def __init__(self, k, max_iter=5, variances = None):
        self.k ,self.max_iter = k, max_iter # sets the number of clusters and the maximum number of iterations
        if variances is None:
            self.variances = np.ones(self.k)
        else:
            self.variances = variances # sets the variances for each cluster
    
    def initialize(self, X):
        
        rows, cols = X.shape
        self.cols = cols
        self.rows = rows
        indexes = np.arange(rows)
        np.random.shuffle(indexes)
        self.mu = X[indexes[:self.k]] # randomly set up the means
        self.sigma = np.array([np.eye(cols) * self.variances[i] for i in range(self.k)])# set up the covariance matrices
        # print("Initial means:")
        # print(self.mu)
        # print("Initial covariances:")
        # print(self.sigma)
        
        
    def e_step(self, X):
        # E-Step: 
        self.weights = self.predict_proba(X) # compute the probability of each data point belonging to each cluster
        self.phi = np.mean(self.weights, axis=0)# calculate the fraction of points belonging to each cluster
        
    def m_step(self, X, fix_variance = False,perc_thresh = 80):
        # M-Step: 
        for class_idx in range(self.k):
            w = self.weights[:, class_idx].copy() # weight assigned to each data point for the class
            if w.sum() == 0:
                    self.sigma[class_idx] = np.identity(self.sigma[class_idx].shape[0]) * 0.00001
                    self.mu[class_idx] = (self.mu[class_idx] * 0) 
                    continue
            probs = w/(w.sum()) # probability of each data point belonging to the class
            lower_percentile = np.percentile(probs, perc_thresh)
            self.weights[:, class_idx+1:][(probs >= lower_percentile)] = 0 # set the weights of the data points assigned to the class to 0
            self.mu[class_idx] = (w/w.sum()) @ X    # mean of the data points assigned to the class
            if not fix_variance:
                    self.sigma[class_idx] = (((w/w.sum()) * (X - self.mu[class_idx]).T)@ (X - self.mu[class_idx]))  # covariance matrix of the data points assigned to the class
                    vv = self.sigma[class_idx].copy()
                    vv_det = abs(np.linalg.det(vv))
                    vv_side = (vv_det/3.14)**(1/vv.shape[0])
                    self.sigma[class_idx] = np.diag([vv_side]*vv.shape[0])
    
    def fit(self, X, fix_variance = False, perc_thresh = 80):
        self.initialize(X)  # sets the means and covariance matrices
        for _ in range(self.max_iter):
            self.e_step(X)      # E-step
            self.m_step(X, fix_variance, perc_thresh)      # M-step
    
    def predict_proba(self, X):
        Px_Pk = np.array([scipy.stats.multivariate_normal.pdf(X, mean=self.mu[i], cov=self.sigma[i]) for i in range(self.k)]).T  # calculate the probability of each data point belonging to each cluster
        return Px_Pk / Px_Pk.sum(axis=1)[:, np.newaxis]  # return the probability of each data point belonging to each cluster

    def predict(self, X):
        return self.predict_proba(X).argmax(axis=1) # return the cluster to which each data point belongs


def create_plot_with_coverage(im, prompt_text,
                              height_length = 500, width_length = 1500, #in meters,
                              update_radii = False,
                              perc_thresh = 80
                              ):
    """
    Creates a plot with circles based on the given radii and calculates the percentage of points covered by the circles.

    Args:
    mask_img: The mask image as a NumPy array.
    radii: A list of radii for the circles.

    Returns:
    plot_array: A NumPy array representing the plot.
    coverage_percentage: The percentage of points covered by the circles.
    """
    height_length, width_length = float(height_length), float(width_length)
    # Parse the prompt text to get the radii
    radii = list(map(float, prompt_text.split(',')))
    img = im["background"][:,:,:3]
    # print("img shape: ", img.shape)
    mask_img = im["layers"][0][:,:,:3]
    
    # Determine the current dimensions of the mask image
    current_height, current_width = mask_img.shape[:2]

    # print("current_height, current_width: ", current_height, current_width)
    
    max_dim = max(current_height, current_width)
    if max_dim == current_height:
        px_to_len_ratio = current_height / height_length
        # print("height ratio: ", px_to_len_ratio)
    else:
        px_to_len_ratio = current_width / width_length
        # print("width ratio: ", px_to_len_ratio)
    
    # mask_img = cv2.resize(mask_img, (new_width, new_height), interpolation=cv2.INTER_NEAREST)
    data = np.argwhere(mask_img.sum(axis=2) != 0)
    if len(data) > 500:
        data  = data[np.random.choice(data.shape[0], 500, replace=False)]
    # data = data[~np.all(data == [0, 0], axis=1)]
    # sample_size = int(0.03 * len(data))
    # data = data[np.random.choice(data.shape[0], sample_size, replace=False)]
    # print("*"*100)
    # print(data)

    variances = [ [(r*px_to_len_ratio/3)**2, 
                   (r*px_to_len_ratio/3)**2] for r in radii ]  # Convert radii to variances
    
    def expectation_maximization(data=data, num_components=len(variances), 
                                 variances=variances, 
                                 max_iter = 100, 
                                 update_covariances= update_radii,
                                 perc_thresh = 80
                                 ):
        """
        """
        # Initialize the GMM with the given number of components and variances
        gmm = GMM(k=num_components, 
                  max_iter=max_iter,
                  variances=variances)

        # Fit the GMM to the data
        gmm.fit(data, fix_variance=not(update_covariances),
                perc_thresh=perc_thresh
                )

        # Extract the parameters
        means = gmm.mu
        covariances = gmm.sigma
        weights = gmm.weights
        responsibilities = gmm.predict_proba(data)

        return means, covariances, weights, responsibilities
    means, covariances, weights, responsibilities = expectation_maximization(data, 
                                                                          len(radii),
                                                                         variances = variances,
                                                                          max_iter = 1000,
                                                                         update_covariances = update_radii,
                                                                        perc_thresh = perc_thresh
                                                                         )
    # if update_radii:
    radii = [(covariances[i][0][0]**(.5))*(3/px_to_len_ratio) for i in range(len(radii))]
        # print("covariances: ", covariances)
    
    # Plot circles based on the radii
    # print("means: ", means)
    covered_points = set()
    # Calculate covered points
    for i, radius in enumerate(radii):
        for point in data:
                if np.linalg.norm(point - means[i]) <= (radius*px_to_len_ratio):
                    covered_points.add(tuple(point))

    # Calculate coverage percentage
    coverage_percentage = (len(covered_points) / len(data)) * 100
    
    # Overlay points, circles, and text directly on the image
    # print("mask_img max min:",mask_img.max(), mask_img.min())
    plot_array = (img * .6) + (mask_img * .4)
    plot_array = plot_array.astype(np.uint8)

    # # Plot the data points
    # data_sample = data[np.random.choice(data.shape[0], int(0.1 * len(data)), replace=False)]
    # for point in data_sample:
    #     cv2.circle(plot_array, (point[1], point[0]), 1, (152, 251, 152), -1)  # palegreen color

    # Plot circles and text based on the radii
    for i, radius in enumerate(radii):
        center = (int(means[i][1]), int(means[i][0]))
        thickness = int(np.ceil(max_dim * .005))
        cv2.circle(plot_array, center, int(radius * px_to_len_ratio), (0, 255, 255), thickness)  # white color
        cv2.circle(plot_array, center, thickness*2, (0, 255, 255), -1)  # white color
        coord_text = f"({means[i][1]:.1f},{means[i][0]:.1f})"
        radius_text = f"r={np.ceil(radius)}"
        text_color = (0, 255, 255)  # green color
        horizontal_margin = -50
        vertical_margin = 15
        cv2.putText(plot_array, coord_text, (center[0] + horizontal_margin, 
                                             center[1] + vertical_margin), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.4, 
                    text_color, 1, cv2.LINE_AA)  
        cv2.putText(plot_array, radius_text, (center[0] + horizontal_margin, 
                                             center[1] + (vertical_margin * 2)), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.4, 
                    text_color, 1, cv2.LINE_AA)

    return plot_array, coverage_percentage


with gr.Blocks() as demo:
  selected_pixels  = []
  with gr.Row():
    with gr.Column():
      im = gr.ImageEditor(
            type="numpy",
            # canvas_size = (500,800),
            # fixed_canvas = True,
            # height = 800,
            # crop_size="1:1",
        )   
      examples_CD = gr.Examples(
      examples=["./PEZ_map.png"],
      inputs=[im],
      )
    with gr.Column():
        output_img = gr.Image(label="Output")
        percentage_covered_txt = gr.Textbox(lines=1, label="Percentage Covered", value="0.0")

  with gr.Row():
    prompt_text = gr.Textbox(lines=1, value= "50,75,75,100,100,150",
                             label="radii of the clusters (m)")
    update_radii = gr.Checkbox(label="Update Radii", value=False)
    perc_thresh = gr.Slider(minimum=0, maximum=100, value=80, label="Percentage Threshold")
  with gr.Row():
    height_length = gr.Textbox(lines=1, value= "500",
                             label="Height Length (m)")
    width_length = gr.Textbox(lines=1, value= "500",
                             label="Width Length (m)")
    

  with gr.Row():
    submit = gr.Button("Submit")
  

  submit.click(create_plot_with_coverage, [im, prompt_text,
                                            height_length, width_length, update_radii,perc_thresh] , 
                                            outputs = [output_img,percentage_covered_txt]) 

demo.launch(
    # debug=True
)