R2bEEaton/SidewalkSegmentation

1import torch2from transformers import SamProcessor3import numpy as np4from PIL import Image5 6class ModelInference:7    def __init__(self):8        self.device = "cuda" if torch.cuda.is_available() else "cpu"9        torch.device(self.device)10        self.model = torch.load('sidewalk_model.pth', map_location=self.device)11        self.processor = SamProcessor.from_pretrained("facebook/sam-vit-base", cache_dir='.cache')12        self.model.eval()13 14    def infer(self, input_img):15        inputs = self.processor(input_img, input_points=get_grid_points(), return_tensors="pt")16 17        # Move the input tensor to the GPU if it's not already there18        inputs = {k: v.to(self.device) for k, v in inputs.items()}19 20        # forward pass21        with torch.no_grad():22            outputs = self.model(**inputs, multimask_output=False)23 24        # apply sigmoid25        seg_prob = torch.sigmoid(outputs.pred_masks.squeeze(1))26        # convert soft mask to hard mask27        seg_prob = seg_prob.cpu().numpy().squeeze()28        seg_pred = (seg_prob > 0.5).astype(np.uint8)29 30        return seg_pred31 32def get_grid_points():33    # FROM EXAMPLE VIDEO REPO34 35    """36    input_points (torch.FloatTensor of shape (batch_size, num_points, 2)) —37    Input 2D spatial points, this is used by the prompt encoder to encode the prompt.38    Generally yields to much better results. The points can be obtained by passing a39    list of list of list to the processor that will create corresponding torch tensors40    of dimension 4. The first dimension is the image batch size, the second dimension41    is the point batch size (i.e. how many segmentation masks do we want the model to42    predict per input point), the third dimension is the number of points per segmentation43    mask (it is possible to pass multiple points for a single mask), and the last dimension44    is the x (vertical) and y (horizontal) coordinates of the point. If a different number45    of points is passed either for each image, or for each mask, the processor will create46    “PAD” points that will correspond to the (0, 0) coordinate, and the computation of the47    embedding will be skipped for these points using the labels.48 49    """50    # Define the size of your array51    array_size = 25652 53    # Define the size of your grid54    grid_size = 1055 56    # Generate the grid points57    x = np.linspace(0, array_size-1, grid_size)58    y = np.linspace(0, array_size-1, grid_size)59 60    # Generate a grid of coordinates61    xv, yv = np.meshgrid(x, y)62 63    # Convert the numpy arrays to lists64    xv_list = xv.tolist()65    yv_list = yv.tolist()66 67    # Combine the x and y coordinates into a list of list of lists68    input_points = [[[int(x), int(y)] for x, y in zip(x_row, y_row)] for x_row, y_row in zip(xv_list, yv_list)]69 70    #We need to reshape our nxn grid to the expected shape of the input_points tensor71    # (batch_size, point_batch_size, num_points_per_image, 2),72    # where the last dimension of 2 represents the x and y coordinates of each point.73    #batch_size: The number of images you're processing at once.74    #point_batch_size: The number of point sets you have for each image.75    #num_points_per_image: The number of points in each set.76    input_points = torch.tensor(input_points).view(1, 1, grid_size*grid_size, 2)77    return input_points78 79if __name__ == "__main__":80    # Test this module81 82    model = ModelInference()83    o = model.infer(Image.open("examples/example_input_1.tif"))84 85    img = Image.fromarray(o * 255, 'L')86    img.show()87