Spaces:
Sleeping
Sleeping
| import os | |
| import logging | |
| import torch | |
| import datetime | |
| import requests | |
| from google.cloud import storage | |
| from transformers import AutoImageProcessor, AutoModelForObjectDetection, ViTImageProcessor, Swinv2ForImageClassification | |
| from label_studio_ml.model import LabelStudioMLBase | |
| from lxml import etree | |
| from uuid import uuid4 | |
| from PIL import Image | |
| from creds import get_credentials | |
| from io import BytesIO | |
| def generate_download_signed_url_v4(blob_name): | |
| """Generates a v4 signed URL for downloading a blob. | |
| Note that this method requires a service account key file. You can not use | |
| this if you are using Application Default Credentials from Google Compute | |
| Engine or from the Google Cloud SDK. | |
| """ | |
| bucket_name = os.getenv("bucket") | |
| storage_client = storage.Client() | |
| bucket = storage_client.bucket(bucket_name) | |
| blob = bucket.blob(blob_name.replace(f"gs://{bucket_name}/", "")) | |
| url = blob.generate_signed_url( | |
| version="v4", | |
| # This URL is valid for 15 minutes | |
| expiration=datetime.timedelta(minutes=15), | |
| # Allow GET requests using this URL. | |
| method="GET", | |
| ) | |
| return url | |
| class Model(LabelStudioMLBase): | |
| os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = get_credentials() | |
| image_processor = AutoImageProcessor.from_pretrained("diegokauer/conditional-detr-coe-int-v2") | |
| model = AutoModelForObjectDetection.from_pretrained("diegokauer/conditional-detr-coe-int-v2") | |
| seg_image_processor = ViTImageProcessor.from_pretrained("diegokauer/int-pet-classifier-v2") | |
| seg_model = Swinv2ForImageClassification.from_pretrained("diegokauer/int-pet-classifier-v2") | |
| id2label = model.config.id2label | |
| seg_id2label = seg_model.config.id2label | |
| def predict(self, tasks, **kwargs): | |
| """ This is where inference happens: model returns | |
| the list of predictions based on input list of tasks | |
| """ | |
| predictions = [] | |
| for task in tasks: | |
| url = task["data"]["image"] | |
| response = requests.get(generate_download_signed_url_v4(url)) | |
| print(response) | |
| image_data = BytesIO(response.content) | |
| image = Image.open(image_data) | |
| original_width, original_height = image.size | |
| with torch.no_grad(): | |
| inputs = self.image_processor(images=image, return_tensors="pt") | |
| outputs = self.model(**inputs) | |
| target_sizes = torch.tensor([image.size[::-1]]) | |
| results = self.image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[0] | |
| result_list = [] | |
| for score, label, box in zip(results['scores'], results['labels'], results['boxes']): | |
| label_id = str(uuid4()) | |
| x, y, x2, y2 = tuple(box) | |
| if self.id2label[label.item()] == 'Propuesta': | |
| with torch.no_grad(): | |
| pred_label_id = str(uuid4()) | |
| image = image.crop((x.item(), y.item(), x2.item(), y2.item())) | |
| inputs = self.seg_image_processor(images=image, return_tensors="pt") | |
| logits = self.seg_model(**inputs).logits | |
| logits = 1 / (1 + torch.exp(-logits)) | |
| print(logits) | |
| preds = logits > 0.5 | |
| preds = [self.seg_id2label[i] for i, pred in enumerate(preds.squeeze().tolist()) if pred] | |
| preds = ["No Reportado"] if "No Reportado" in preds else preds | |
| result_list.append({ | |
| "value": { | |
| "choices": preds | |
| }, | |
| "id": pred_label_id, | |
| "from_name": "propuesta", | |
| "to_name": "image", | |
| "type": "choices" | |
| }) | |
| result_list.append({ | |
| 'id': label_id, | |
| 'original_width': original_width, | |
| 'original_height': original_height, | |
| 'from_name': "bbox", | |
| 'to_name': "image", | |
| 'type': 'rectangle', | |
| 'score': score.item(), # per-region score, visible in the editor | |
| 'value': { | |
| 'x': x.item() * 100.0 / original_width, | |
| 'y': y.item() * 100.0 / original_height, | |
| 'width': (x2-x).item() * 100.0 / original_width, | |
| 'height': (y2-y).item() * 100.0 / original_height, | |
| 'rotation': 0, | |
| } | |
| }) | |
| result_list.append({ | |
| 'id': label_id, | |
| 'original_width': original_width, | |
| 'original_height': original_height, | |
| 'from_name': "label", | |
| 'to_name': "image", | |
| 'type': 'labels', | |
| 'score': score.item(), # per-region score, visible in the editor | |
| 'value': { | |
| 'x': x.item() * 100.0 / original_width, | |
| 'y': y.item() * 100.0 / original_height, | |
| 'width': (x2-x).item() * 100.0 / original_width, | |
| 'height': (y2-y).item() * 100.0 / original_height, | |
| 'rotation': 0, | |
| 'labels': [self.id2label[label.item()]] | |
| } | |
| }) | |
| predictions.append({ | |
| 'score': results['scores'].mean().item(), # prediction overall score, visible in the data manager columns | |
| 'model_version': 'cdetr_v2.5', # all predictions will be differentiated by model version | |
| 'result': result_list | |
| }) | |
| print(predictions) | |
| return predictions | |
| def fit(self, event, annotations, **kwargs): | |
| """ This is where training happens: train your model given list of annotations, | |
| then returns dict with created links and resources | |
| """ | |
| return {'path/to/created/model': 'my/model.bin'} |