diegokauer/segmentation-backend

1import os2import logging3import torch4import datetime5import requests6 7from google.cloud import storage8from transformers import AutoImageProcessor, AutoModelForObjectDetection, ViTImageProcessor, Swinv2ForImageClassification9from label_studio_ml.model import LabelStudioMLBase10from lxml import etree11from uuid import uuid412from PIL import Image13 14from creds import get_credentials15from io import BytesIO16 17 18def generate_download_signed_url_v4(blob_name):19    """Generates a v4 signed URL for downloading a blob.20 21    Note that this method requires a service account key file. You can not use22    this if you are using Application Default Credentials from Google Compute23    Engine or from the Google Cloud SDK.24    """25    bucket_name = os.getenv("bucket")26 27    storage_client = storage.Client()28    bucket = storage_client.bucket(bucket_name)29    blob = bucket.blob(blob_name.replace(f"gs://{bucket_name}/", ""))30 31    url = blob.generate_signed_url(32        version="v4",33        # This URL is valid for 15 minutes34        expiration=datetime.timedelta(minutes=15),35        # Allow GET requests using this URL.36        method="GET",37    )38    return url39 40 41class Model(LabelStudioMLBase):42    43    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = get_credentials()44    image_processor = AutoImageProcessor.from_pretrained("diegokauer/conditional-detr-coe-int-v2")45    model = AutoModelForObjectDetection.from_pretrained("diegokauer/conditional-detr-coe-int-v2")46    seg_image_processor = ViTImageProcessor.from_pretrained("diegokauer/int-pet-classifier-v2")47    seg_model = Swinv2ForImageClassification.from_pretrained("diegokauer/int-pet-classifier-v2")48    id2label = model.config.id2label49    seg_id2label = seg_model.config.id2label50 51    def predict(self, tasks, **kwargs):52        """ This is where inference happens: model returns 53            the list of predictions based on input list of tasks 54        """55        predictions = []56        for task in tasks:57 58 59            url = task["data"]["image"]60            response = requests.get(generate_download_signed_url_v4(url))61            print(response)62            image_data = BytesIO(response.content)63            image = Image.open(image_data)64            65            original_width, original_height = image.size66            with torch.no_grad():67                68                inputs = self.image_processor(images=image, return_tensors="pt")69                outputs = self.model(**inputs)70                target_sizes = torch.tensor([image.size[::-1]])71                results = self.image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[0]72 73            result_list = []74            for score, label, box in zip(results['scores'], results['labels'], results['boxes']):75                label_id = str(uuid4())76                x, y, x2, y2 = tuple(box)77 78                if self.id2label[label.item()] == 'Propuesta':79                    with torch.no_grad():80                        pred_label_id = str(uuid4())81                        image = image.crop((x.item(), y.item(), x2.item(), y2.item()))82                        inputs = self.seg_image_processor(images=image, return_tensors="pt")83                        logits = self.seg_model(**inputs).logits84                        logits = 1 / (1 + torch.exp(-logits))85                        print(logits)86                        preds = logits > 0.587                        preds = [self.seg_id2label[i] for i, pred in enumerate(preds.squeeze().tolist()) if pred]88                        preds = ["No Reportado"] if "No Reportado" in preds else preds89                        result_list.append({90                          "value": {91                            "choices": preds92                          },93                          "id": pred_label_id,94                          "from_name": "propuesta",95                          "to_name": "image",96                          "type": "choices"97                        })98                99                result_list.append({100                    'id': label_id,101                    'original_width': original_width,102                    'original_height': original_height,103                    'from_name': "bbox",104                    'to_name': "image",105                    'type': 'rectangle',106                    'score': score.item(),  # per-region score, visible in the editor 107                    'value': {108                        'x': x.item() * 100.0 / original_width,109                        'y': y.item() * 100.0 / original_height,110                        'width': (x2-x).item() * 100.0 / original_width,111                        'height': (y2-y).item() * 100.0 / original_height,112                        'rotation': 0,113                    }114                })115                result_list.append({116                    'id': label_id,117                    'original_width': original_width,118                    'original_height': original_height,119                    'from_name': "label",120                    'to_name': "image",121                    'type': 'labels',122                    'score': score.item(),  # per-region score, visible in the editor 123                    'value': {124                        'x': x.item() * 100.0 / original_width,125                        'y': y.item() * 100.0 / original_height,126                        'width': (x2-x).item() * 100.0 / original_width,127                        'height': (y2-y).item() * 100.0 / original_height,128                        'rotation': 0,129                        'labels': [self.id2label[label.item()]]130                    }131                })132            133            predictions.append({134                'score': results['scores'].mean().item(),  # prediction overall score, visible in the data manager columns135                'model_version': 'cdetr_v2.5',  # all predictions will be differentiated by model version136                'result': result_list137            })138        print(predictions)139        return predictions140 141    def fit(self, event, annotations, **kwargs):142        """ This is where training happens: train your model given list of annotations, 143            then returns dict with created links and resources144        """145        return {'path/to/created/model': 'my/model.bin'}