|
6 | 6 | import os |
7 | 7 | import sys |
8 | 8 | import time |
9 | | -from typing import Any, Dict, List, Optional |
| 9 | +from typing import Any, Dict, Generator, List, Optional |
10 | 10 |
|
11 | 11 | import requests |
12 | 12 | from PIL import Image |
@@ -666,6 +666,112 @@ def _upload_zip( |
666 | 666 | except Exception as e: |
667 | 667 | print(f"An error occured when uploading the model: {e}") |
668 | 668 |
|
| 669 | + def search( |
| 670 | + self, |
| 671 | + query: str, |
| 672 | + page_size: int = 50, |
| 673 | + fields: Optional[List[str]] = None, |
| 674 | + continuation_token: Optional[str] = None, |
| 675 | + ) -> dict: |
| 676 | + """Search across all images in the workspace using RoboQL syntax. |
| 677 | +
|
| 678 | + Args: |
| 679 | + query: RoboQL search query (e.g. ``"tag:review"``, ``"project:false"`` |
| 680 | + for orphan images, or free-text for semantic CLIP search). |
| 681 | + page_size: Number of results per page (default 50). |
| 682 | + fields: Fields to include in each result. |
| 683 | + Defaults to ``["tags", "projects", "filename"]``. |
| 684 | + continuation_token: Token returned by a previous call for fetching |
| 685 | + the next page. |
| 686 | +
|
| 687 | + Returns: |
| 688 | + Dict with ``results`` (list), ``total`` (int), and |
| 689 | + ``continuationToken`` (str or None). |
| 690 | +
|
| 691 | + Example: |
| 692 | + >>> ws = rf.workspace() |
| 693 | + >>> page = ws.search("tag:review", page_size=10) |
| 694 | + >>> print(page["total"]) |
| 695 | + >>> for img in page["results"]: |
| 696 | + ... print(img["filename"]) |
| 697 | + """ |
| 698 | + if fields is None: |
| 699 | + fields = ["tags", "projects", "filename"] |
| 700 | + |
| 701 | + return rfapi.workspace_search( |
| 702 | + api_key=self.__api_key, |
| 703 | + workspace_url=self.url, |
| 704 | + query=query, |
| 705 | + page_size=page_size, |
| 706 | + fields=fields, |
| 707 | + continuation_token=continuation_token, |
| 708 | + ) |
| 709 | + |
| 710 | + def delete_images(self, image_ids: List[str]) -> dict: |
| 711 | + """Delete orphan images from the workspace. |
| 712 | +
|
| 713 | + Only deletes images not associated with any project. |
| 714 | + Images still in projects are skipped. |
| 715 | +
|
| 716 | + Args: |
| 717 | + image_ids: List of image IDs to delete. |
| 718 | +
|
| 719 | + Returns: |
| 720 | + Dict with ``deletedSources`` and ``skippedSources`` counts. |
| 721 | +
|
| 722 | + Example: |
| 723 | + >>> ws = rf.workspace() |
| 724 | + >>> result = ws.delete_images(["img_id_1", "img_id_2"]) |
| 725 | + >>> print(result["deletedSources"]) |
| 726 | + """ |
| 727 | + return rfapi.workspace_delete_images( |
| 728 | + api_key=self.__api_key, |
| 729 | + workspace_url=self.url, |
| 730 | + image_ids=image_ids, |
| 731 | + ) |
| 732 | + |
| 733 | + def search_all( |
| 734 | + self, |
| 735 | + query: str, |
| 736 | + page_size: int = 50, |
| 737 | + fields: Optional[List[str]] = None, |
| 738 | + ) -> Generator[List[dict], None, None]: |
| 739 | + """Paginated search across all images in the workspace. |
| 740 | +
|
| 741 | + Yields one page of results at a time, automatically following |
| 742 | + ``continuationToken`` until all results have been returned. |
| 743 | +
|
| 744 | + Args: |
| 745 | + query: RoboQL search query. |
| 746 | + page_size: Number of results per page (default 50). |
| 747 | + fields: Fields to include in each result. |
| 748 | + Defaults to ``["tags", "projects", "filename"]``. |
| 749 | +
|
| 750 | + Yields: |
| 751 | + A list of result dicts for each page. |
| 752 | +
|
| 753 | + Example: |
| 754 | + >>> ws = rf.workspace() |
| 755 | + >>> for page in ws.search_all("tag:review"): |
| 756 | + ... for img in page: |
| 757 | + ... print(img["filename"]) |
| 758 | + """ |
| 759 | + token = None |
| 760 | + while True: |
| 761 | + response = self.search( |
| 762 | + query=query, |
| 763 | + page_size=page_size, |
| 764 | + fields=fields, |
| 765 | + continuation_token=token, |
| 766 | + ) |
| 767 | + results = response.get("results", []) |
| 768 | + if not results: |
| 769 | + break |
| 770 | + yield results |
| 771 | + token = response.get("continuationToken") |
| 772 | + if not token: |
| 773 | + break |
| 774 | + |
669 | 775 | def search_export( |
670 | 776 | self, |
671 | 777 | query: str, |
|
0 commit comments