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 )