Spaces:
Sleeping
Sleeping
| from samgeo import tms_to_geotiff | |
| from samgeo.text_sam import LangSAM | |
| sam = LangSAM() | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| import torch | |
| from torchvision import transforms | |
| from matplotlib import pyplot as plt | |
| from samgeo.text_sam import LangSAM | |
| import cv2 | |
| import matplotlib.patches as patches | |
| from transformers import SamModel, SamConfig, SamProcessor | |
| from math import floor, ceil | |
| from matplotlib.colors import LinearSegmentedColormap | |
| from samgeo import tms_to_geotiff | |
| from samgeo.text_sam import LangSAM | |
| # Load the SAM model | |
| sam = LangSAM() | |
| # methods for sidewalk inferences | |
| def get_input_image(image_file, processor, bbox=None): | |
| # img = torch.tensor(np.array(Image.open(image_file))).permute(2, 0, 1) | |
| img = torch.tensor(np.array(image_file)).permute(2, 0, 1) | |
| ''' | |
| image = Image.open(image_file).convert('RGB') | |
| img = np.array(image) | |
| ''' | |
| if bbox is None: | |
| bbox = [0, 0, img.shape[1], img.shape[0]] # Use image dimensions as bounding box | |
| # prepare image and prompt for the model | |
| inputs = processor(img, input_boxes=[[bbox]], return_tensors="pt") | |
| # remove batch dimension which the processor adds by default | |
| inputs = {k: v.squeeze(0) for k, v in inputs.items()} | |
| inputs["org_img"] = img | |
| return inputs | |
| def process_image(inputs): | |
| model.eval() | |
| with torch.no_grad(): | |
| outputs = model(pixel_values=inputs["pixel_values"].unsqueeze(0).to(device), | |
| input_boxes=inputs["input_boxes"].unsqueeze(0).to(device), | |
| multimask_output=False) | |
| medsam_seg_prob = torch.sigmoid(outputs.pred_masks.squeeze(1)) | |
| medsam_seg_prob = medsam_seg_prob.cpu().numpy().squeeze() | |
| orig = inputs["org_img"].permute(1, 2, 0).cpu().numpy() | |
| return orig, medsam_seg_prob | |
| def display_image(medsam_seg_prob, threshold=0.5): | |
| medsam_seg = (medsam_seg_prob > threshold).astype(np.uint8) | |
| return medsam_seg | |
| # output sidewalk with original photo | |
| def output_sidewalk(image, medsam_seg, alpha=0.7): | |
| # Color for 0: transparent, for 1: blue | |
| colors = [(0, 0, 0, 0), (0, 0, 1, 1)] # RGBA tuples | |
| cmap = LinearSegmentedColormap.from_list("custom_cmap", colors) | |
| fig, axes = plt.subplots(1, 1, figsize=(8, 8)) | |
| axes.imshow(np.array(image)) | |
| axes.imshow(np.array(medsam_seg), cmap=cmap, alpha=alpha) | |
| axes.axis('off') | |
| # Ensure the figure canvas is drawn | |
| fig.canvas.draw() | |
| # Now convert it to a NumPy array | |
| data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) | |
| data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) | |
| return data | |
| # methods for smoother sidewalk mask | |
| def filter_weak(medsam_seg, size_threshold=10): | |
| num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(medsam_seg, connectivity=8, | |
| ltype=cv2.CV_32S) | |
| result = np.zeros_like(medsam_seg) | |
| for i in range(1, num_labels): | |
| if stats[i, cv2.CC_STAT_AREA] >= size_threshold: | |
| result[labels == i] = 1 | |
| return result | |
| def smoothing(mask, kernel_size=(6, 6)): | |
| kernel = np.ones(kernel_size, np.uint8) | |
| opening = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) | |
| closing = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) | |
| return closing | |
| def pipeline(data, size_threshold=25, kernel_size=(9, 9)): | |
| result = filter_weak(data, size_threshold) | |
| result = smoothing(result, kernel_size) | |
| return result | |
| # methods for occlusion handling | |
| def create_boundary_mask_from_bbox(bbox, array_size, thickness=1): | |
| # Create an empty mask with the same dimensions as the array_size | |
| mask = np.zeros(array_size, dtype=np.uint8) | |
| # Calculate xmin, ymin, xmax, ymax from the bbox | |
| xmin, ymin, xmax, ymax = bbox | |
| # Ensure the bbox coordinates are within the array bounds to avoid IndexErrors | |
| xmin = floor(max(xmin, 0)) | |
| xmax = ceil(min(xmax, array_size[1] - 1)) | |
| ymin = floor(max(ymin, 0)) | |
| ymax = ceil(min(ymax, array_size[0] - 1)) | |
| # Draw top and bottom horizontal lines | |
| mask[ymin:ymin + thickness, xmin:xmax] = 2 | |
| mask[ymax - thickness + 1:ymax + 1, xmin:xmax] = 2 | |
| # Draw left and right vertical lines | |
| mask[ymin:ymax, xmin:xmin + thickness] = 2 | |
| mask[ymin:ymax, xmax - thickness + 1:xmax + 1] = 2 | |
| return mask | |
| def check_boundary(m1, m2, radius=1): | |
| # Initialize an output mask of the same shape as m2, filled with zeros | |
| boundary_mask = np.zeros_like(m2) | |
| # Get the dimensions of the masks | |
| rows, cols = m2.shape | |
| # Iterate through each pixel in the m2 mask | |
| for r in range(rows): | |
| for c in range(cols): | |
| # Check if the current pixel is a 'tree' pixel | |
| if m2[r, c] == 2: | |
| # Initialize a flag to check for at least one adjacent 'sidewalk' | |
| found_sidewalk = 0 | |
| # Check the square around the current pixel with given radius | |
| for dr in range(-radius, radius + 1): | |
| for dc in range(-radius, radius + 1): | |
| # Calculate the neighbor's position | |
| nr, nc = r + dr, c + dc | |
| # Ensure we're not out of bounds and we're not checking the center pixel itself | |
| if 0 <= nr < rows and 0 <= nc < cols and (dr != 0 or dc != 0): | |
| if m1[nr, nc] == 1: | |
| found_sidewalk += 1 | |
| boundary_mask[r, c] = found_sidewalk | |
| return boundary_mask | |
| def linear_regression_two_points(point1, point2): | |
| # Create arrays of x and y values | |
| x = np.array([point1[0], point2[0]]) | |
| y = np.array([point1[1], point2[1]]) | |
| # Perform linear regression: np.polyfit returns the slope and intercept | |
| m, b = np.polyfit(x, y, 1) | |
| return m, b, x, y | |
| def generate_road_mask(x1, x2, slope, intercept, road_width=5, image_size=(256, 256)): | |
| # Create a blank black image (all zeros) | |
| image = np.zeros(image_size, dtype=np.uint8) | |
| # Define x values within the specified range x1 to x2 | |
| x_values = np.array(range(x1, x2 + 1)) | |
| # Calculate corresponding y values using the slope and intercept | |
| y_values = (slope * x_values + intercept).astype(int) | |
| # Draw the road line with the specified width | |
| for i in range(len(x_values)): | |
| if 0 <= y_values[i] < image_size[0]: # Check if the y-value is within the image boundaries | |
| cv2.circle(image, (x_values[i], y_values[i]), road_width // 2, 1, -1) # Draw circles to create a thick line | |
| return image | |
| def get_road_mask_per_bbox(filtered_med_seg, bbox, radius=1): | |
| array_size = (256, 256) # Define the size of the 2D mask | |
| mask = create_boundary_mask_from_bbox(bbox, array_size, thickness=1) | |
| # get intersection | |
| output = check_boundary(filtered_med_seg, mask, radius) | |
| # get connected component and centriods | |
| num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(output, 8, cv2.CV_32S) | |
| centroids = centroids[1:] | |
| centroids = sorted(centroids, key=lambda x: x[0]) | |
| # check if we have two 2 centriods | |
| if len(centroids) == 2: | |
| # linear regression | |
| slope, intercept, x, y = linear_regression_two_points(centroids[0], centroids[1]) | |
| # get road mask inferred from tree bbox intersection points | |
| road_mask = generate_road_mask(int(x[0]), int(x[1]), slope, intercept, 3) | |
| else: | |
| return None | |
| return road_mask | |
| def analyze_sidewalk(sam, filtered_med_seg, image, alpha=0.7): | |
| # Using SAM model to predict on the image with a specific prompt | |
| text_prompt = "tree" | |
| masks, boxes, labels, logits = sam.predict(image, text_prompt, box_threshold=0.24, text_threshold=0.24, | |
| return_results=True) | |
| # Setting up custom color maps for overlays | |
| colors = [(0, 0, 0, 0), (0, 0, 1, 1)] # Blue color | |
| cmap = LinearSegmentedColormap.from_list("custom_cmap", colors) | |
| colors_alt = [(0, 0, 0, 0), (0, 1, 0, 1)] # Green color | |
| cmap_alt = LinearSegmentedColormap.from_list("custom_cmap", colors_alt) | |
| # Plotting the results | |
| # fig, axes = plt.subplots(1, 3, figsize=(18, 6)) | |
| fig, axes = plt.subplots(1, 1, figsize=(8, 8)) | |
| # fig.suptitle(f"Sidewalk Detection with SAM Model \n{image}", fontsize=16) | |
| ''' | |
| axes[0].imshow(image) | |
| axes[0].set_title("Original Image") | |
| axes[0].axis('off') | |
| axes[1].imshow(image) | |
| axes[1].imshow(filtered_med_seg, cmap=cmap, alpha=0.7) | |
| axes[1].axis('off') | |
| axes[1].set_title("Sidewalk Mask - Initial") | |
| ''' | |
| axes.imshow(image) | |
| axes.imshow(filtered_med_seg, cmap=cmap, alpha=alpha) | |
| axes.axis('off') | |
| # axes.set_title("Sidewalk Mask - Refined with Occlusion Handling") | |
| for bbox in boxes: | |
| road_mask = get_road_mask_per_bbox(filtered_med_seg, bbox.tolist(), 1) | |
| if road_mask is not None: | |
| axes.imshow(road_mask, cmap=cmap_alt, alpha=alpha) | |
| rect = patches.Rectangle((bbox[0], bbox[1]), bbox[2] - bbox[0], bbox[3] - bbox[1], linewidth=1, | |
| edgecolor='r', facecolor='none') | |
| axes.add_patch(rect) | |
| # Ensure the figure canvas is drawn | |
| fig.canvas.draw() | |
| # Now convert it to a NumPy array | |
| data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) | |
| data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) | |
| return data | |
| # Load pretrained model | |
| model_config = SamConfig.from_pretrained("facebook/sam-vit-base") | |
| model = SamModel(config=model_config) | |
| model.load_state_dict(torch.load("model_checkpoint_final1.pth", map_location=torch.device('cpu'))) | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model.to(device) # Move model to device once here instead of in the function | |
| # special methods for gradio | |
| partial_results = {} | |
| def process_pipeline(image, threshold, alpha): | |
| processor = SamProcessor.from_pretrained("facebook/sam-vit-base") | |
| processed_inputs = get_input_image(image, processor, bbox=[0, 0, 256, 256]) | |
| orig, medsam_seg_prob = process_image(processed_inputs) | |
| medsam_seg = display_image(medsam_seg_prob, threshold) | |
| filtered_med_seg = pipeline(medsam_seg) | |
| output_image = output_sidewalk(orig, filtered_med_seg, alpha) | |
| filled_image = analyze_sidewalk(sam, filtered_med_seg, image, alpha=alpha) | |
| partial_results["prob"] = medsam_seg_prob | |
| partial_results["orig"] = orig | |
| partial_results["filtered_med_seg"] = filtered_med_seg | |
| return output_image, filled_image | |
| def update_output(image, threshold, alpha): | |
| if "prob" in partial_results and "orig" in partial_results: | |
| medsam_seg_prob = partial_results['prob'] | |
| orig = partial_results['orig'] | |
| medsam_seg = display_image(medsam_seg_prob, threshold) | |
| filtered_med_seg = pipeline(medsam_seg) | |
| output_image = output_sidewalk(orig, filtered_med_seg, alpha) | |
| filled_image = analyze_sidewalk(sam, filtered_med_seg, image, alpha=alpha) | |
| partial_results["filtered_med_seg"] = filtered_med_seg | |
| return output_image, filled_image | |
| def update_output_alpha(image, threshold, alpha): | |
| if "prob" in partial_results and "filtered_med_seg" in partial_results: | |
| medsam_seg_prob = partial_results['prob'] | |
| orig = partial_results['orig'] | |
| filtered_med_seg = partial_results["filtered_med_seg"] | |
| output_image = output_sidewalk(orig, filtered_med_seg, alpha=alpha) | |
| filled_image = analyze_sidewalk(sam, filtered_med_seg, image, alpha=alpha) | |
| return output_image, filled_image | |
| with gr.Blocks() as app: | |
| gr.Markdown("# Sidewalk Detection with SAM Model") | |
| gr.Markdown("#### by Dan Mao, Kevin Tan") | |
| with gr.Row(): | |
| with gr.Column(): | |
| img_in = gr.Image(type="pil", label="Upload Image") | |
| threshold = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.5, label="Threshold") | |
| alpha = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.7, label="Alpha for Mask Overlay") | |
| submit_button = gr.Button("Process Image") | |
| with gr.Column(): | |
| img_out1 = gr.Image(label="Sidewalk Mask - Initial") | |
| img_out2 = gr.Image(label="Sidewalk Mask - Refine with Occlusion Handling") | |
| gr.ClearButton(components=[img_in, img_out1, img_out2]) | |
| # Setting up triggers for changes and button clicks | |
| threshold.change(fn=update_output, inputs=[img_in, threshold, alpha], outputs=[img_out1, img_out2]) | |
| alpha.change(fn=update_output_alpha, inputs=[img_in, threshold, alpha], outputs=[img_out1, img_out2]) | |
| submit_button.click( | |
| fn=process_pipeline, | |
| inputs=[img_in, threshold, alpha], | |
| outputs=[img_out1, img_out2] | |
| ) | |
| app.launch(debug=True) | |