nnilayy/segment-anything

1import os2 3from label_studio_converter import brush4from typing import List, Dict, Optional5from uuid import uuid46from sam_predictor import SAMPredictor7from label_studio_ml.model import LabelStudioMLBase8 9SAM_CHOICE = os.environ.get("SAM_CHOICE", "MobileSAM")  # other option is just SAM10PREDICTOR = SAMPredictor(SAM_CHOICE)11 12 13class SamMLBackend(LabelStudioMLBase):14 15    def predict(self, tasks: List[Dict], context: Optional[Dict] = None, **kwargs) -> List[Dict]:16        """ Returns the predicted mask for a smart keypoint that has been placed."""17 18        from_name, to_name, value = self.get_first_tag_occurence('BrushLabels', 'Image')19 20        if not context or not context.get('result'):21            # if there is no context, no interaction has happened yet22            return []23 24        image_width = context['result'][0]['original_width']25        image_height = context['result'][0]['original_height']26 27        # collect context information28        point_coords = []29        point_labels = []30        input_box = None31        selected_label = None32        for ctx in context['result']:33            x = ctx['value']['x'] * image_width / 10034            y = ctx['value']['y'] * image_height / 10035            ctx_type = ctx['type']36            selected_label = ctx['value'][ctx_type][0]37            if ctx_type == 'keypointlabels':38                point_labels.append(int(ctx['is_positive']))39                point_coords.append([int(x), int(y)])40            elif ctx_type == 'rectanglelabels':41                box_width = ctx['value']['width'] * image_width / 10042                box_height = ctx['value']['height'] * image_height / 10043                input_box = [int(x), int(y), int(box_width + x), int(box_height + y)]44 45        print(f'Point coords are {point_coords}, point labels are {point_labels}, input box is {input_box}')46 47        img_path = tasks[0]['data'][value]48        predictor_results = PREDICTOR.predict(49            img_path=img_path,50            point_coords=point_coords or None,51            point_labels=point_labels or None,52            input_box=input_box53        )54 55        predictions = self.get_results(56            masks=predictor_results['masks'],57            probs=predictor_results['probs'],58            width=image_width,59            height=image_height,60            from_name=from_name,61            to_name=to_name,62            label=selected_label)63 64        return predictions65 66    def get_results(self, masks, probs, width, height, from_name, to_name, label):67        results = []68        for mask, prob in zip(masks, probs):69            # creates a random ID for your label everytime so no chance for errors70            label_id = str(uuid4())[:4]71            # converting the mask from the model to RLE format which is usable in Label Studio72            mask = mask * 25573            rle = brush.mask2rle(mask)74 75            results.append({76                'id': label_id,77                'from_name': from_name,78                'to_name': to_name,79                'original_width': width,80                'original_height': height,81                'image_rotation': 0,82                'value': {83                    'format': 'rle',84                    'rle': rle,85                    'brushlabels': [label],86                },87                'score': prob,88                'type': 'brushlabels',89                'readonly': False90            })91 92        return [{93            'result': results,94            'model_version': PREDICTOR.model_name95        }]96 97 98if __name__ == '__main__':99    # test the model100    model = SamMLBackend()101    model.use_label_config('''102    <View>103        <Image name="image" value="$image" zoom="true"/>104        <BrushLabels name="tag" toName="image">105            <Label value="Banana" background="#FF0000"/>106            <Label value="Orange" background="#0d14d3"/>107        </BrushLabels>108        <KeyPointLabels name="tag2" toName="image" smart="true" >109            <Label value="Banana" background="#000000" showInline="true"/>110            <Label value="Orange" background="#000000" showInline="true"/>111        </KeyPointLabels>112        <RectangleLabels name="tag3" toName="image"  >113            <Label value="Banana" background="#000000" showInline="true"/>114            <Label value="Orange" background="#000000" showInline="true"/>115        </RectangleLabels>116    </View>117    ''')118    results = model.predict(119        tasks=[{120            'data': {121                'image': 'https://s3.amazonaws.com/htx-pub/datasets/images/125245483_152578129892066_7843809718842085333_n.jpg'122            }}],123        context={124            'result': [{125                'original_width': 1080,126                'original_height': 1080,127                'image_rotation': 0,128                'value': {129                    'x': 49.441786283891545,130                    'y': 59.96810207336522,131                    'width': 0.3189792663476874,132                    'labels': ['Banana'],133                    'keypointlabels': ['Banana']134                },135                'is_positive': True,136                'id': 'fBWv1t0S2L',137                'from_name': 'tag2',138                'to_name': 'image',139                'type': 'keypointlabels',140                'origin': 'manual'141            }]}142    )143    import json144    results[0]['result'][0]['value']['rle'] = f'...{len(results[0]["result"][0]["value"]["rle"])} integers...'145    print(json.dumps(results, indent=2))