pacomesimon/CameraCoverage
0
1import numpy as np
2import matplotlib.pyplot as plt
3import gradio as gr
4import cv2
5import scipy
6
7class GMM:
8 def __init__(self, k, max_iter=5, variances = None):
9 self.k ,self.max_iter = k, max_iter # sets the number of clusters and the maximum number of iterations
10 if variances is None:
11 self.variances = np.ones(self.k)
12 else:
13 self.variances = variances # sets the variances for each cluster
14
15 def initialize(self, X):
16
17 rows, cols = X.shape
18 self.cols = cols
19 self.rows = rows
20 indexes = np.arange(rows)
21 np.random.shuffle(indexes)
22 self.mu = X[indexes[:self.k]] # randomly set up the means
23 self.sigma = np.array([np.eye(cols) * self.variances[i] for i in range(self.k)])# set up the covariance matrices
24 # print("Initial means:")
25 # print(self.mu)
26 # print("Initial covariances:")
27 # print(self.sigma)
28
29
30 def e_step(self, X):
31 # E-Step:
32 self.weights = self.predict_proba(X) # compute the probability of each data point belonging to each cluster
33 self.phi = np.mean(self.weights, axis=0)# calculate the fraction of points belonging to each cluster
34
35 def m_step(self, X, fix_variance = False,perc_thresh = 80):
36 # M-Step:
37 for class_idx in range(self.k):
38 w = self.weights[:, class_idx].copy() # weight assigned to each data point for the class
39 if w.sum() == 0:
40 self.sigma[class_idx] = np.identity(self.sigma[class_idx].shape[0]) * 0.00001
41 self.mu[class_idx] = (self.mu[class_idx] * 0)
42 continue
43 probs = w/(w.sum()) # probability of each data point belonging to the class
44 lower_percentile = np.percentile(probs, perc_thresh)
45 self.weights[:, class_idx+1:][(probs >= lower_percentile)] = 0 # set the weights of the data points assigned to the class to 0
46 self.mu[class_idx] = (w/w.sum()) @ X # mean of the data points assigned to the class
47 if not fix_variance:
48 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
49 vv = self.sigma[class_idx].copy()
50 vv_det = abs(np.linalg.det(vv))
51 vv_side = (vv_det/3.14)**(1/vv.shape[0])
52 self.sigma[class_idx] = np.diag([vv_side]*vv.shape[0])
53
54 def fit(self, X, fix_variance = False, perc_thresh = 80):
55 self.initialize(X) # sets the means and covariance matrices
56 for _ in range(self.max_iter):
57 self.e_step(X) # E-step
58 self.m_step(X, fix_variance, perc_thresh) # M-step
59
60 def predict_proba(self, X):
61 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
62 return Px_Pk / Px_Pk.sum(axis=1)[:, np.newaxis] # return the probability of each data point belonging to each cluster
63
64 def predict(self, X):
65 return self.predict_proba(X).argmax(axis=1) # return the cluster to which each data point belongs
66
67
68def create_plot_with_coverage(im, prompt_text,
69 height_length = 500, width_length = 1500, #in meters,
70 update_radii = False,
71 perc_thresh = 80
72 ):
73 """
74 Creates a plot with circles based on the given radii and calculates the percentage of points covered by the circles.
75
76 Args:
77 mask_img: The mask image as a NumPy array.
78 radii: A list of radii for the circles.
79
80 Returns:
81 plot_array: A NumPy array representing the plot.
82 coverage_percentage: The percentage of points covered by the circles.
83 """
84 height_length, width_length = float(height_length), float(width_length)
85 # Parse the prompt text to get the radii
86 radii = list(map(float, prompt_text.split(',')))
87 img = im["background"][:,:,:3]
88 # print("img shape: ", img.shape)
89 mask_img = im["layers"][0][:,:,:3]
90
91 # Determine the current dimensions of the mask image
92 current_height, current_width = mask_img.shape[:2]
93
94 # print("current_height, current_width: ", current_height, current_width)
95
96 max_dim = max(current_height, current_width)
97 if max_dim == current_height:
98 px_to_len_ratio = current_height / height_length
99 # print("height ratio: ", px_to_len_ratio)
100 else:
101 px_to_len_ratio = current_width / width_length
102 # print("width ratio: ", px_to_len_ratio)
103
104 # mask_img = cv2.resize(mask_img, (new_width, new_height), interpolation=cv2.INTER_NEAREST)
105 data = np.argwhere(mask_img.sum(axis=2) != 0)
106 if len(data) > 500:
107 data = data[np.random.choice(data.shape[0], 500, replace=False)]
108 # data = data[~np.all(data == [0, 0], axis=1)]
109 # sample_size = int(0.03 * len(data))
110 # data = data[np.random.choice(data.shape[0], sample_size, replace=False)]
111 # print("*"*100)
112 # print(data)
113
114 variances = [ [(r*px_to_len_ratio/3)**2,
115 (r*px_to_len_ratio/3)**2] for r in radii ] # Convert radii to variances
116
117 def expectation_maximization(data=data, num_components=len(variances),
118 variances=variances,
119 max_iter = 100,
120 update_covariances= update_radii,
121 perc_thresh = 80
122 ):
123 """
124 """
125 # Initialize the GMM with the given number of components and variances
126 gmm = GMM(k=num_components,
127 max_iter=max_iter,
128 variances=variances)
129
130 # Fit the GMM to the data
131 gmm.fit(data, fix_variance=not(update_covariances),
132 perc_thresh=perc_thresh
133 )
134
135 # Extract the parameters
136 means = gmm.mu
137 covariances = gmm.sigma
138 weights = gmm.weights
139 responsibilities = gmm.predict_proba(data)
140
141 return means, covariances, weights, responsibilities
142 means, covariances, weights, responsibilities = expectation_maximization(data,
143 len(radii),
144 variances = variances,
145 max_iter = 1000,
146 update_covariances = update_radii,
147 perc_thresh = perc_thresh
148 )
149 # if update_radii:
150 radii = [(covariances[i][0][0]**(.5))*(3/px_to_len_ratio) for i in range(len(radii))]
151 # print("covariances: ", covariances)
152
153 # Plot circles based on the radii
154 # print("means: ", means)
155 covered_points = set()
156 # Calculate covered points
157 for i, radius in enumerate(radii):
158 for point in data:
159 if np.linalg.norm(point - means[i]) <= (radius*px_to_len_ratio):
160 covered_points.add(tuple(point))
161
162 # Calculate coverage percentage
163 coverage_percentage = (len(covered_points) / len(data)) * 100
164
165 # Overlay points, circles, and text directly on the image
166 # print("mask_img max min:",mask_img.max(), mask_img.min())
167 plot_array = (img * .6) + (mask_img * .4)
168 plot_array = plot_array.astype(np.uint8)
169
170 # # Plot the data points
171 # data_sample = data[np.random.choice(data.shape[0], int(0.1 * len(data)), replace=False)]
172 # for point in data_sample:
173 # cv2.circle(plot_array, (point[1], point[0]), 1, (152, 251, 152), -1) # palegreen color
174
175 # Plot circles and text based on the radii
176 for i, radius in enumerate(radii):
177 center = (int(means[i][1]), int(means[i][0]))
178 thickness = int(np.ceil(max_dim * .005))
179 cv2.circle(plot_array, center, int(radius * px_to_len_ratio), (0, 255, 255), thickness) # white color
180 cv2.circle(plot_array, center, thickness*2, (0, 255, 255), -1) # white color
181 coord_text = f"({means[i][1]:.1f},{means[i][0]:.1f})"
182 radius_text = f"r={np.ceil(radius)}"
183 text_color = (0, 255, 255) # green color
184 horizontal_margin = -50
185 vertical_margin = 15
186 cv2.putText(plot_array, coord_text, (center[0] + horizontal_margin,
187 center[1] + vertical_margin),
188 cv2.FONT_HERSHEY_SIMPLEX, 0.4,
189 text_color, 1, cv2.LINE_AA)
190 cv2.putText(plot_array, radius_text, (center[0] + horizontal_margin,
191 center[1] + (vertical_margin * 2)),
192 cv2.FONT_HERSHEY_SIMPLEX, 0.4,
193 text_color, 1, cv2.LINE_AA)
194
195 return plot_array, coverage_percentage
196
197
198with gr.Blocks() as demo:
199 selected_pixels = []
200 with gr.Row():
201 with gr.Column():
202 im = gr.ImageEditor(
203 type="numpy",
204 # canvas_size = (500,800),
205 # fixed_canvas = True,
206 # height = 800,
207 # crop_size="1:1",
208 )
209 examples_CD = gr.Examples(
210 examples=["./PEZ_map.png"],
211 inputs=[im],
212 )
213 with gr.Column():
214 output_img = gr.Image(label="Output")
215 percentage_covered_txt = gr.Textbox(lines=1, label="Percentage Covered", value="0.0")
216
217 with gr.Row():
218 prompt_text = gr.Textbox(lines=1, value= "50,75,75,100,100,150",
219 label="radii of the clusters (m)")
220 update_radii = gr.Checkbox(label="Update Radii", value=False)
221 perc_thresh = gr.Slider(minimum=0, maximum=100, value=80, label="Percentage Threshold")
222 with gr.Row():
223 height_length = gr.Textbox(lines=1, value= "500",
224 label="Height Length (m)")
225 width_length = gr.Textbox(lines=1, value= "500",
226 label="Width Length (m)")
227
228
229 with gr.Row():
230 submit = gr.Button("Submit")
231
232
233 submit.click(create_plot_with_coverage, [im, prompt_text,
234 height_length, width_length, update_radii,perc_thresh] ,
235 outputs = [output_img,percentage_covered_txt])
236
237demo.launch(
238 # debug=True
239)
240
241