Spaces:
Runtime error
Runtime error
| import tempfile | |
| from typing import Dict, List, Union | |
| import numpy as np | |
| from dds_cloudapi_sdk import ( | |
| DetectionTask, | |
| Client, | |
| Config, | |
| TextPrompt, | |
| DetectionModel, | |
| DetectionTarget, | |
| ) | |
| from PIL import Image | |
| import concurrent.futures | |
| class GroundingDINOAPIWrapper: | |
| """API wrapper for Grounding DINO 1.5 | |
| Args: | |
| token (str): The token for Grounding DINO 1.5 API. We are now opening free API access to Grounding DINO 1.5. For | |
| educators, students, and researchers, we offer an API with extensive usage times to | |
| support your educational and research endeavors. You can get free API token at here: | |
| https://deepdataspace.com/request_api | |
| """ | |
| def __init__(self, token: str): | |
| self.client = Client(Config(token=token)) | |
| def inference(self, prompt: Dict, return_mask:bool=False): | |
| """Main inference function of Grounding DINO 1.5. We take batch as input and | |
| each image is a dict. N. We do not support batch inference for now. | |
| Args: | |
| prompts (dict): Annotations with the following keys: | |
| - "image" (str): Path to image. E.g. "test1.jpg", | |
| - "prompt" (str): Text prompt sepearted by '.' E.g. 'cate1 . cate2 . cate3' | |
| return_mask (bool): Whether to return mask. Defaults to False. | |
| Returns: | |
| (Dict): Detection results in dict format with keys:: | |
| - "scores": (List[float]): A list of scores for each object in the batch | |
| - "labels": (List[int]): A list of labels for each object in the batch | |
| - "boxes": (List[List[int]]): A list of boxes for each object in the batch, | |
| in format [xmin, ymin, xmax, ymax] | |
| - "masks": (List[np.ndarray]): A list of segmentations for each object in the batch | |
| """ | |
| # construct input prompts | |
| image=self.get_image_url(prompt["image"]), | |
| task=DetectionTask( | |
| image_url=image[0], | |
| prompts=[TextPrompt(text=prompt['prompt'])], | |
| targets=[DetectionTarget.Mask, DetectionTarget.BBox] if return_mask else [DetectionTarget.BBox], | |
| model=DetectionModel.GDino1_5_Pro, | |
| ) | |
| self.client.run_task(task) | |
| result = task.result | |
| return self.postprocess(result, task, return_mask) | |
| def postprocess(self, result, task, return_mask): | |
| """Postprocess the result from the API call | |
| Args: | |
| result (TaskResult): Task result with the following keys: | |
| - objects (List[DetectionObject]): Each DetectionObject has the following keys: | |
| - bbox (List[float]): Box in xyxy format | |
| - category (str): Detection category | |
| - score (float): Detection score | |
| - mask (DetectionObjectMask): Use mask.counts to parse RLE mask | |
| task (DetectionTask): The task object | |
| return_mask (bool): Whether to return mask | |
| Returns: | |
| (Dict): Return dict in format: | |
| { | |
| "scores": (List[float]): A list of scores for each object | |
| "categorys": (List[str]): A list of categorys for each object | |
| "boxes": (List[List[int]]): A list of boxes for each object | |
| "masks": (List[PIL.Image]): A list of masks in the format of PIL.Image | |
| } | |
| """ | |
| def process_object_with_mask(object): | |
| box = object.bbox | |
| score = object.score | |
| category = object.category | |
| mask = task.rle2rgba(object.mask) | |
| return box, score, category, mask | |
| def process_object_without_mask(object): | |
| box = object.bbox | |
| score = object.score | |
| category = object.category | |
| mask = None | |
| return box, score, category, mask | |
| boxes, scores, categorys, masks = [], [], [], [] | |
| with concurrent.futures.ThreadPoolExecutor() as executor: | |
| if return_mask: | |
| process_object = process_object_with_mask | |
| else: | |
| process_object = process_object_without_mask | |
| futures = [executor.submit(process_object, obj) for obj in result.objects] | |
| for future in concurrent.futures.as_completed(futures): | |
| box, score, category, mask = future.result() | |
| boxes.append(box) | |
| scores.append(score) | |
| categorys.append(category) | |
| if mask is not None: | |
| masks.append(mask) | |
| return dict(boxes=boxes, categorys=categorys, scores=scores, masks=masks) | |
| def get_image_url(self, image: Union[str, np.ndarray]): | |
| """Upload Image to server and return the url | |
| Args: | |
| image (Union[str, np.ndarray]): The image to upload. Can be a file path or np.ndarray. | |
| If it is a np.ndarray, it will be saved to a temporary file. | |
| Returns: | |
| str: The url of the image | |
| """ | |
| if isinstance(image, str): | |
| url = self.client.upload_file(image) | |
| else: | |
| with tempfile.NamedTemporaryFile(delete=True, suffix=".png") as tmp_file: | |
| # image is in numpy format, convert to PIL Image | |
| image = Image.fromarray(image) | |
| image.save(tmp_file, format="PNG") | |
| tmp_file_path = tmp_file.name | |
| url = self.client.upload_file(tmp_file_path) | |
| return url |