Spaces:
Runtime error
Runtime error
| from transformers import DPTImageProcessor, DPTForDepthEstimation | |
| from segment_anything import SamAutomaticMaskGenerator, sam_model_registry, SamPredictor | |
| import gradio as gr | |
| import supervision as sv | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| import requests | |
| import open3d as o3d | |
| import pandas as pd | |
| import plotly.express as px | |
| class DepthPredictor: | |
| def __init__(self): | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-large") | |
| self.model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large") | |
| self.model.eval() | |
| def predict(self, image): | |
| # prepare image for the model | |
| encoding = self.feature_extractor(image, return_tensors="pt") | |
| self.img = image | |
| # forward pass | |
| with torch.no_grad(): | |
| outputs = self.model(**encoding) | |
| predicted_depth = outputs.predicted_depth | |
| # interpolate to original size | |
| prediction = torch.nn.functional.interpolate( | |
| predicted_depth.unsqueeze(1), | |
| size=image.size[::-1], | |
| mode="bicubic", | |
| align_corners=False, | |
| ).squeeze() | |
| output = prediction.cpu().numpy() | |
| formatted = (output * 255 / np.max(output)).astype('uint8') | |
| #img = Image.fromarray(formatted) | |
| return formatted | |
| def generate_pcl(self, image): | |
| depth = self.predict(image) | |
| # Step 2: Create an RGBD image from the RGB and depth image | |
| depth_o3d = o3d.geometry.Image(depth) | |
| image_o3d = o3d.geometry.Image(image) | |
| rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(image_o3d, depth_o3d, convert_rgb_to_intensity=False) | |
| # Step 3: Create a PointCloud from the RGBD image | |
| pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, o3d.camera.PinholeCameraIntrinsic(o3d.camera.PinholeCameraIntrinsicParameters.PrimeSenseDefault)) | |
| # Step 4: Convert PointCloud data to a NumPy array | |
| points = np.asarray(pcd.points) | |
| colors = np.asarray(pcd.colors) | |
| return points, colors | |
| def generate_fig(self, image): | |
| points, colors = self.generate_pcl(image) | |
| data = {'x': points[:, 0], 'y': points[:, 1], 'z': points[:, 2], | |
| 'red': colors[:, 0], 'green': colors[:, 1], 'blue': colors[:, 2]} | |
| df = pd.DataFrame(data) | |
| size = np.zeros(len(df)) | |
| size[:] = 0.01 | |
| # Step 6: Create a 3D scatter plot using Plotly Express | |
| fig = px.scatter_3d(df, x='x', y='y', z='z', color='red', size=size) | |
| return fig | |
| class SegmentPredictor: | |
| def __init__(self): | |
| MODEL_TYPE = "vit_b" | |
| checkpoint = "sam_vit_b_01ec64.pth" | |
| sam = sam_model_registry[MODEL_TYPE](checkpoint=checkpoint) | |
| # Select device | |
| self.device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| sam.to(device=self.device) | |
| self.mask_generator = SamAutomaticMaskGenerator(sam) | |
| self.conditioned_pred = SamPredictor(sam) | |
| def encode(self, image): | |
| image = np.array(image) | |
| self.conditioned_pred.set_image(image) | |
| def cond_pred(self, pts, lbls): | |
| masks, _, _ = self.conditioned_pred.predict( | |
| point_coords=pts, | |
| point_labels=lbls, | |
| multimask_output=True | |
| ) | |
| return masks | |
| def segment_everything(self, image): | |
| image = np.array(image) | |
| sam_result = self.mask_generator.generate(image) | |
| mask_annotator = sv.MaskAnnotator() | |
| detections = sv.Detections.from_sam(sam_result=sam_result) | |
| annotated_image = mask_annotator.annotate(scene=image.copy(), detections=detections) | |
| return annotated_image |