Method segment_image_with_coordinate

visual_chatgpt.py:947–988

Args: img (numpy.ndarray): the given image, shape: H x W x 3. is_positive: whether the click is positive, if want to add mask use True else False. coordinate: the position of the click If the position is (x,y), me

(self, img, is_positive: bool,
                            coordinate: tuple)

        return img

    def segment_image_with_coordinate(self, img, is_positive: bool,
                            coordinate: tuple):
        '''
            Args:
                img (numpy.ndarray): the given image, shape: H x W x 3.
                is_positive: whether the click is positive, if want to add mask use True else False.
                coordinate: the position of the click
                          If the position is (x,y), means click at the x-th column and y-th row of the pixel matrix.
                          So x correspond to W, and y correspond to H.
            Output:
                img (PLI.Image.Image): the result image
                result_mask (numpy.ndarray): the result mask, shape: H x W

            Other parameters:
                transparency (float): the transparenccy of the mask
                                      to control he degree of transparency after the mask is superimposed.
                                      if transparency=1, then the masked part will be completely replaced with other colors.
        '''
        self.sam_predictor.set_image(img)
        self.saved_points.append([coordinate[0], coordinate[1]])
        self.saved_labels.append(1 if is_positive else 0)
        input_point = np.array(self.saved_points)
        input_label = np.array(self.saved_labels)

        # Predict the mask
        with torch.cuda.amp.autocast():
            masks, scores, logits = self.sam_predictor.predict(
                point_coords=input_point,
                point_labels=input_label,
                multimask_output=False,
            )


        img = self.show_mask(masks[0], img, random_color=False, transparency=0.3)

        img = self.show_points(input_point, input_label, img)

        img = Image.fromarray(img)
        
        result_mask = masks[0]

        return img, result_mask

    @prompts(name="Segment the Image",
             description="useful when you want to segment all the part of the image, but not segment a certain object."