evelynmitchell
1 year ago
committed by
GitHub
commit
fb0cfaefb4
@ -1,11 +1,3 @@
|
|||||||
from swarms.memory.pinecone import PineconeVector
|
# from swarms.memory.pinecone import PineconeVector
|
||||||
from swarms.memory.base import BaseVectorStore
|
# from swarms.memory.base import BaseVectorStore
|
||||||
from swarms.memory.pg import PgVectorVectorStore
|
# from swarms.memory.pg import PgVectorVectorStore
|
||||||
from swarms.memory.ocean import OceanDB
|
|
||||||
|
|
||||||
__all__ = [
|
|
||||||
"BaseVectorStore",
|
|
||||||
"PineconeVector",
|
|
||||||
"PgVectorVectorStore",
|
|
||||||
"OceanDB",
|
|
||||||
]
|
|
||||||
|
|||||||
@ -1,177 +0,0 @@
|
|||||||
import uuid
|
|
||||||
from abc import ABC
|
|
||||||
from typing import Any, Dict, List, Optional
|
|
||||||
|
|
||||||
from swarms.memory.schemas import Artifact, Status
|
|
||||||
from swarms.memory.schemas import Step as APIStep
|
|
||||||
from swarms.memory.schemas import Task as APITask
|
|
||||||
|
|
||||||
|
|
||||||
class Step(APIStep):
|
|
||||||
additional_properties: Optional[Dict[str, str]] = None
|
|
||||||
|
|
||||||
|
|
||||||
class Task(APITask):
|
|
||||||
steps: List[Step] = []
|
|
||||||
|
|
||||||
|
|
||||||
class NotFoundException(Exception):
|
|
||||||
"""
|
|
||||||
Exception raised when a resource is not found.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, item_name: str, item_id: str):
|
|
||||||
self.item_name = item_name
|
|
||||||
self.item_id = item_id
|
|
||||||
super().__init__(f"{item_name} with {item_id} not found.")
|
|
||||||
|
|
||||||
|
|
||||||
class TaskDB(ABC):
|
|
||||||
async def create_task(
|
|
||||||
self,
|
|
||||||
input: Optional[str],
|
|
||||||
additional_input: Any = None,
|
|
||||||
artifacts: Optional[List[Artifact]] = None,
|
|
||||||
steps: Optional[List[Step]] = None,
|
|
||||||
) -> Task:
|
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
async def create_step(
|
|
||||||
self,
|
|
||||||
task_id: str,
|
|
||||||
name: Optional[str] = None,
|
|
||||||
input: Optional[str] = None,
|
|
||||||
is_last: bool = False,
|
|
||||||
additional_properties: Optional[Dict[str, str]] = None,
|
|
||||||
) -> Step:
|
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
async def create_artifact(
|
|
||||||
self,
|
|
||||||
task_id: str,
|
|
||||||
file_name: str,
|
|
||||||
relative_path: Optional[str] = None,
|
|
||||||
step_id: Optional[str] = None,
|
|
||||||
) -> Artifact:
|
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
async def get_task(self, task_id: str) -> Task:
|
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
async def get_step(self, task_id: str, step_id: str) -> Step:
|
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
async def get_artifact(self, task_id: str, artifact_id: str) -> Artifact:
|
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
async def list_tasks(self) -> List[Task]:
|
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
async def list_steps(
|
|
||||||
self, task_id: str, status: Optional[Status] = None
|
|
||||||
) -> List[Step]:
|
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
|
||||||
class InMemoryTaskDB(TaskDB):
|
|
||||||
_tasks: Dict[str, Task] = {}
|
|
||||||
|
|
||||||
async def create_task(
|
|
||||||
self,
|
|
||||||
input: Optional[str],
|
|
||||||
additional_input: Any = None,
|
|
||||||
artifacts: Optional[List[Artifact]] = None,
|
|
||||||
steps: Optional[List[Step]] = None,
|
|
||||||
) -> Task:
|
|
||||||
if not steps:
|
|
||||||
steps = []
|
|
||||||
if not artifacts:
|
|
||||||
artifacts = []
|
|
||||||
task_id = str(uuid.uuid4())
|
|
||||||
task = Task(
|
|
||||||
task_id=task_id,
|
|
||||||
input=input,
|
|
||||||
steps=steps,
|
|
||||||
artifacts=artifacts,
|
|
||||||
additional_input=additional_input,
|
|
||||||
)
|
|
||||||
self._tasks[task_id] = task
|
|
||||||
return task
|
|
||||||
|
|
||||||
async def create_step(
|
|
||||||
self,
|
|
||||||
task_id: str,
|
|
||||||
name: Optional[str] = None,
|
|
||||||
input: Optional[str] = None,
|
|
||||||
is_last=False,
|
|
||||||
additional_properties: Optional[Dict[str, Any]] = None,
|
|
||||||
) -> Step:
|
|
||||||
step_id = str(uuid.uuid4())
|
|
||||||
step = Step(
|
|
||||||
task_id=task_id,
|
|
||||||
step_id=step_id,
|
|
||||||
name=name,
|
|
||||||
input=input,
|
|
||||||
status=Status.created,
|
|
||||||
is_last=is_last,
|
|
||||||
additional_properties=additional_properties,
|
|
||||||
)
|
|
||||||
task = await self.get_task(task_id)
|
|
||||||
task.steps.append(step)
|
|
||||||
return step
|
|
||||||
|
|
||||||
async def get_task(self, task_id: str) -> Task:
|
|
||||||
task = self._tasks.get(task_id, None)
|
|
||||||
if not task:
|
|
||||||
raise NotFoundException("Task", task_id)
|
|
||||||
return task
|
|
||||||
|
|
||||||
async def get_step(self, task_id: str, step_id: str) -> Step:
|
|
||||||
task = await self.get_task(task_id)
|
|
||||||
step = next(filter(lambda s: s.task_id == task_id, task.steps), None)
|
|
||||||
if not step:
|
|
||||||
raise NotFoundException("Step", step_id)
|
|
||||||
return step
|
|
||||||
|
|
||||||
async def get_artifact(self, task_id: str, artifact_id: str) -> Artifact:
|
|
||||||
task = await self.get_task(task_id)
|
|
||||||
artifact = next(
|
|
||||||
filter(lambda a: a.artifact_id == artifact_id, task.artifacts), None
|
|
||||||
)
|
|
||||||
if not artifact:
|
|
||||||
raise NotFoundException("Artifact", artifact_id)
|
|
||||||
return artifact
|
|
||||||
|
|
||||||
async def create_artifact(
|
|
||||||
self,
|
|
||||||
task_id: str,
|
|
||||||
file_name: str,
|
|
||||||
relative_path: Optional[str] = None,
|
|
||||||
step_id: Optional[str] = None,
|
|
||||||
) -> Artifact:
|
|
||||||
artifact_id = str(uuid.uuid4())
|
|
||||||
artifact = Artifact(
|
|
||||||
artifact_id=artifact_id,
|
|
||||||
file_name=file_name,
|
|
||||||
relative_path=relative_path,
|
|
||||||
)
|
|
||||||
task = await self.get_task(task_id)
|
|
||||||
task.artifacts.append(artifact)
|
|
||||||
|
|
||||||
if step_id:
|
|
||||||
step = await self.get_step(task_id, step_id)
|
|
||||||
step.artifacts.append(artifact)
|
|
||||||
|
|
||||||
return artifact
|
|
||||||
|
|
||||||
async def list_tasks(self) -> List[Task]:
|
|
||||||
return [task for task in self._tasks.values()]
|
|
||||||
|
|
||||||
async def list_steps(
|
|
||||||
self, task_id: str, status: Optional[Status] = None
|
|
||||||
) -> List[Step]:
|
|
||||||
task = await self.get_task(task_id)
|
|
||||||
steps = task.steps
|
|
||||||
if status:
|
|
||||||
steps = list(filter(lambda s: s.status == status, steps))
|
|
||||||
return steps
|
|
||||||
@ -1,157 +0,0 @@
|
|||||||
import logging
|
|
||||||
from typing import List
|
|
||||||
|
|
||||||
import oceandb
|
|
||||||
from oceandb.utils.embedding_function import MultiModalEmbeddingFunction
|
|
||||||
|
|
||||||
|
|
||||||
class OceanDB:
|
|
||||||
"""
|
|
||||||
A class to interact with OceanDB.
|
|
||||||
|
|
||||||
...
|
|
||||||
|
|
||||||
Attributes
|
|
||||||
----------
|
|
||||||
client : oceandb.Client
|
|
||||||
a client to interact with OceanDB
|
|
||||||
|
|
||||||
Methods
|
|
||||||
-------
|
|
||||||
create_collection(collection_name: str, modality: str):
|
|
||||||
Creates a new collection in OceanDB.
|
|
||||||
append_document(collection, document: str, id: str):
|
|
||||||
Appends a document to a collection in OceanDB.
|
|
||||||
add_documents(collection, documents: List[str], ids: List[str]):
|
|
||||||
Adds multiple documents to a collection in OceanDB.
|
|
||||||
query(collection, query_texts: list[str], n_results: int):
|
|
||||||
Queries a collection in OceanDB.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, client: oceandb.Client = None):
|
|
||||||
"""
|
|
||||||
Constructs all the necessary attributes for the OceanDB object.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
client : oceandb.Client, optional
|
|
||||||
a client to interact with OceanDB (default is None, which creates a new client)
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
self.client = client if client else oceandb.Client()
|
|
||||||
print(self.client.heartbeat())
|
|
||||||
except Exception as e:
|
|
||||||
logging.error(f"Failed to initialize OceanDB client. Error: {e}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
def create_collection(self, collection_name: str, modality: str):
|
|
||||||
"""
|
|
||||||
Creates a new collection in OceanDB.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
collection_name : str
|
|
||||||
the name of the new collection
|
|
||||||
modality : str
|
|
||||||
the modality of the new collection
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
collection
|
|
||||||
the created collection
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
embedding_function = MultiModalEmbeddingFunction(modality=modality)
|
|
||||||
collection = self.client.create_collection(
|
|
||||||
collection_name, embedding_function=embedding_function
|
|
||||||
)
|
|
||||||
return collection
|
|
||||||
except Exception as e:
|
|
||||||
logging.error(f"Failed to create collection. Error {e}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
def append_document(self, collection, document: str, id: str):
|
|
||||||
"""
|
|
||||||
Appends a document to a collection in OceanDB.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
collection
|
|
||||||
the collection to append the document to
|
|
||||||
document : str
|
|
||||||
the document to append
|
|
||||||
id : str
|
|
||||||
the id of the document
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
result
|
|
||||||
the result of the append operation
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
return collection.add(documents=[document], ids=[id])
|
|
||||||
except Exception as e:
|
|
||||||
logging.error(
|
|
||||||
f"Failed to append document to the collection. Error {e}"
|
|
||||||
)
|
|
||||||
raise
|
|
||||||
|
|
||||||
def add_documents(self, collection, documents: List[str], ids: List[str]):
|
|
||||||
"""
|
|
||||||
Adds multiple documents to a collection in OceanDB.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
collection
|
|
||||||
the collection to add the documents to
|
|
||||||
documents : List[str]
|
|
||||||
the documents to add
|
|
||||||
ids : List[str]
|
|
||||||
the ids of the documents
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
result
|
|
||||||
the result of the add operation
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
return collection.add(documents=documents, ids=ids)
|
|
||||||
except Exception as e:
|
|
||||||
logging.error(f"Failed to add documents to collection. Error: {e}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
def query(self, collection, query_texts: list[str], n_results: int):
|
|
||||||
"""
|
|
||||||
Queries a collection in OceanDB.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
collection
|
|
||||||
the collection to query
|
|
||||||
query_texts : list[str]
|
|
||||||
the texts to query
|
|
||||||
n_results : int
|
|
||||||
the number of results to return
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
results
|
|
||||||
the results of the query
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
results = collection.query(
|
|
||||||
query_texts=query_texts, n_results=n_results
|
|
||||||
)
|
|
||||||
return results
|
|
||||||
except Exception as e:
|
|
||||||
logging.error(f"Failed to query the collection. Error {e}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
|
|
||||||
# Example
|
|
||||||
# ocean = OceanDB()
|
|
||||||
# collection = ocean.create_collection("test", "text")
|
|
||||||
# ocean.append_document(collection, "hello world", "1")
|
|
||||||
# ocean.add_documents(collection, ["hello world", "hello world"], ["2", "3"])
|
|
||||||
# results = ocean.query(collection, ["hello world"], 3)
|
|
||||||
# print(results)
|
|
||||||
@ -0,0 +1,291 @@
|
|||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
from PIL import Image
|
||||||
|
from transformers import SamImageProcessor, SamModel, SamProcessor, pipeline
|
||||||
|
|
||||||
|
try:
|
||||||
|
import cv2
|
||||||
|
import supervision as sv
|
||||||
|
except ImportError:
|
||||||
|
print("Please install supervision and cv")
|
||||||
|
|
||||||
|
|
||||||
|
from enum import Enum
|
||||||
|
|
||||||
|
|
||||||
|
class FeatureType(Enum):
|
||||||
|
"""
|
||||||
|
An enumeration to represent the types of features for mask adjustment in image
|
||||||
|
segmentation.
|
||||||
|
"""
|
||||||
|
|
||||||
|
ISLAND = "ISLAND"
|
||||||
|
HOLE = "HOLE"
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def list(cls):
|
||||||
|
return list(map(lambda c: c.value, cls))
|
||||||
|
|
||||||
|
|
||||||
|
def compute_mask_iou_vectorized(masks: np.ndarray) -> np.ndarray:
|
||||||
|
"""
|
||||||
|
Vectorized computation of the Intersection over Union (IoU) for all pairs of masks.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
masks (np.ndarray): A 3D numpy array with shape `(N, H, W)`, where `N` is the
|
||||||
|
number of masks, `H` is the height, and `W` is the width.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
np.ndarray: A 2D numpy array of shape `(N, N)` where each element `[i, j]` is
|
||||||
|
the IoU between masks `i` and `j`.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If any of the masks is found to be empty.
|
||||||
|
"""
|
||||||
|
if np.any(masks.sum(axis=(1, 2)) == 0):
|
||||||
|
raise ValueError(
|
||||||
|
"One or more masks are empty. Please filter out empty masks before"
|
||||||
|
" using `compute_iou_vectorized` function."
|
||||||
|
)
|
||||||
|
|
||||||
|
masks_bool = masks.astype(bool)
|
||||||
|
masks_flat = masks_bool.reshape(masks.shape[0], -1)
|
||||||
|
intersection = np.logical_and(masks_flat[:, None], masks_flat[None, :]).sum(
|
||||||
|
axis=2
|
||||||
|
)
|
||||||
|
union = np.logical_or(masks_flat[:, None], masks_flat[None, :]).sum(axis=2)
|
||||||
|
iou_matrix = intersection / union
|
||||||
|
return iou_matrix
|
||||||
|
|
||||||
|
|
||||||
|
def mask_non_max_suppression(
|
||||||
|
masks: np.ndarray, iou_threshold: float = 0.6
|
||||||
|
) -> np.ndarray:
|
||||||
|
"""
|
||||||
|
Performs Non-Max Suppression on a set of masks by prioritizing larger masks and
|
||||||
|
removing smaller masks that overlap significantly.
|
||||||
|
|
||||||
|
When the IoU between two masks exceeds the specified threshold, the smaller mask
|
||||||
|
(in terms of area) is discarded. This process is repeated for each pair of masks,
|
||||||
|
effectively filtering out masks that are significantly overlapped by larger ones.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
masks (np.ndarray): A 3D numpy array with shape `(N, H, W)`, where `N` is the
|
||||||
|
number of masks, `H` is the height, and `W` is the width.
|
||||||
|
iou_threshold (float): The IoU threshold for determining significant overlap.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
np.ndarray: A 3D numpy array of filtered masks.
|
||||||
|
"""
|
||||||
|
num_masks = masks.shape[0]
|
||||||
|
areas = masks.sum(axis=(1, 2))
|
||||||
|
sorted_idx = np.argsort(-areas)
|
||||||
|
keep_mask = np.ones(num_masks, dtype=bool)
|
||||||
|
iou_matrix = compute_mask_iou_vectorized(masks)
|
||||||
|
for i in range(num_masks):
|
||||||
|
if not keep_mask[sorted_idx[i]]:
|
||||||
|
continue
|
||||||
|
|
||||||
|
overlapping_masks = iou_matrix[sorted_idx[i]] > iou_threshold
|
||||||
|
overlapping_masks[sorted_idx[i]] = False
|
||||||
|
keep_mask[sorted_idx] = np.logical_and(
|
||||||
|
keep_mask[sorted_idx], ~overlapping_masks
|
||||||
|
)
|
||||||
|
|
||||||
|
return masks[keep_mask]
|
||||||
|
|
||||||
|
|
||||||
|
def filter_masks_by_relative_area(
|
||||||
|
masks: np.ndarray, minimum_area: float = 0.01, maximum_area: float = 1.0
|
||||||
|
) -> np.ndarray:
|
||||||
|
"""
|
||||||
|
Filters masks based on their relative area within the total area of each mask.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
masks (np.ndarray): A 3D numpy array with shape `(N, H, W)`, where `N` is the
|
||||||
|
number of masks, `H` is the height, and `W` is the width.
|
||||||
|
minimum_area (float): The minimum relative area threshold. Must be between `0`
|
||||||
|
and `1`.
|
||||||
|
maximum_area (float): The maximum relative area threshold. Must be between `0`
|
||||||
|
and `1`.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
np.ndarray: A 3D numpy array containing masks that fall within the specified
|
||||||
|
relative area range.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If `minimum_area` or `maximum_area` are outside the `0` to `1`
|
||||||
|
range, or if `minimum_area` is greater than `maximum_area`.
|
||||||
|
"""
|
||||||
|
|
||||||
|
if not (isinstance(masks, np.ndarray) and masks.ndim == 3):
|
||||||
|
raise ValueError("Input must be a 3D numpy array.")
|
||||||
|
|
||||||
|
if not (0 <= minimum_area <= 1) or not (0 <= maximum_area <= 1):
|
||||||
|
raise ValueError(
|
||||||
|
"`minimum_area` and `maximum_area` must be between 0 and 1."
|
||||||
|
)
|
||||||
|
|
||||||
|
if minimum_area > maximum_area:
|
||||||
|
raise ValueError(
|
||||||
|
"`minimum_area` must be less than or equal to `maximum_area`."
|
||||||
|
)
|
||||||
|
|
||||||
|
total_area = masks.shape[1] * masks.shape[2]
|
||||||
|
relative_areas = masks.sum(axis=(1, 2)) / total_area
|
||||||
|
return masks[
|
||||||
|
(relative_areas >= minimum_area) & (relative_areas <= maximum_area)
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def adjust_mask_features_by_relative_area(
|
||||||
|
mask: np.ndarray,
|
||||||
|
area_threshold: float,
|
||||||
|
feature_type: FeatureType = FeatureType.ISLAND,
|
||||||
|
) -> np.ndarray:
|
||||||
|
"""
|
||||||
|
Adjusts a mask by removing small islands or filling small holes based on a relative
|
||||||
|
area threshold.
|
||||||
|
|
||||||
|
!!! warning
|
||||||
|
|
||||||
|
Running this function on a mask with small islands may result in empty masks.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
mask (np.ndarray): A 2D numpy array with shape `(H, W)`, where `H` is the
|
||||||
|
height, and `W` is the width.
|
||||||
|
area_threshold (float): Threshold for relative area to remove or fill features.
|
||||||
|
feature_type (FeatureType): Type of feature to adjust (`ISLAND` for removing
|
||||||
|
islands, `HOLE` for filling holes).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
np.ndarray: A 2D numpy array containing mask.
|
||||||
|
"""
|
||||||
|
height, width = mask.shape
|
||||||
|
total_area = width * height
|
||||||
|
|
||||||
|
mask = np.uint8(mask * 255)
|
||||||
|
operation = (
|
||||||
|
cv2.RETR_EXTERNAL
|
||||||
|
if feature_type == FeatureType.ISLAND
|
||||||
|
else cv2.RETR_CCOMP
|
||||||
|
)
|
||||||
|
contours, _ = cv2.findContours(mask, operation, cv2.CHAIN_APPROX_SIMPLE)
|
||||||
|
|
||||||
|
for contour in contours:
|
||||||
|
area = cv2.contourArea(contour)
|
||||||
|
relative_area = area / total_area
|
||||||
|
if relative_area < area_threshold:
|
||||||
|
cv2.drawContours(
|
||||||
|
image=mask,
|
||||||
|
contours=[contour],
|
||||||
|
contourIdx=-1,
|
||||||
|
color=(0 if feature_type == FeatureType.ISLAND else 255),
|
||||||
|
thickness=-1,
|
||||||
|
)
|
||||||
|
return np.where(mask > 0, 1, 0).astype(bool)
|
||||||
|
|
||||||
|
|
||||||
|
def masks_to_marks(masks: np.ndarray) -> sv.Detections:
|
||||||
|
"""
|
||||||
|
Converts a set of masks to a marks (sv.Detections) object.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
masks (np.ndarray): A 3D numpy array with shape `(N, H, W)`, where `N` is the
|
||||||
|
number of masks, `H` is the height, and `W` is the width.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
sv.Detections: An object containing the masks and their bounding box
|
||||||
|
coordinates.
|
||||||
|
"""
|
||||||
|
return sv.Detections(mask=masks, xyxy=sv.mask_to_xyxy(masks=masks))
|
||||||
|
|
||||||
|
|
||||||
|
def refine_marks(
|
||||||
|
marks: sv.Detections,
|
||||||
|
maximum_hole_area: float = 0.01,
|
||||||
|
maximum_island_area: float = 0.01,
|
||||||
|
minimum_mask_area: float = 0.02,
|
||||||
|
maximum_mask_area: float = 1.0,
|
||||||
|
) -> sv.Detections:
|
||||||
|
"""
|
||||||
|
Refines a set of masks by removing small islands and holes, and filtering by mask
|
||||||
|
area.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
marks (sv.Detections): An object containing the masks and their bounding box
|
||||||
|
coordinates.
|
||||||
|
maximum_hole_area (float): The maximum relative area of holes to be filled in
|
||||||
|
each mask.
|
||||||
|
maximum_island_area (float): The maximum relative area of islands to be removed
|
||||||
|
from each mask.
|
||||||
|
minimum_mask_area (float): The minimum relative area for a mask to be retained.
|
||||||
|
maximum_mask_area (float): The maximum relative area for a mask to be retained.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
sv.Detections: An object containing the masks and their bounding box
|
||||||
|
coordinates.
|
||||||
|
"""
|
||||||
|
result_masks = []
|
||||||
|
for mask in marks.mask:
|
||||||
|
mask = adjust_mask_features_by_relative_area(
|
||||||
|
mask=mask,
|
||||||
|
area_threshold=maximum_island_area,
|
||||||
|
feature_type=FeatureType.ISLAND,
|
||||||
|
)
|
||||||
|
mask = adjust_mask_features_by_relative_area(
|
||||||
|
mask=mask,
|
||||||
|
area_threshold=maximum_hole_area,
|
||||||
|
feature_type=FeatureType.HOLE,
|
||||||
|
)
|
||||||
|
if np.any(mask):
|
||||||
|
result_masks.append(mask)
|
||||||
|
result_masks = np.array(result_masks)
|
||||||
|
result_masks = filter_masks_by_relative_area(
|
||||||
|
masks=result_masks,
|
||||||
|
minimum_area=minimum_mask_area,
|
||||||
|
maximum_area=maximum_mask_area,
|
||||||
|
)
|
||||||
|
return sv.Detections(
|
||||||
|
mask=result_masks, xyxy=sv.mask_to_xyxy(masks=result_masks)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class SegmentAnythingMarkGenerator:
|
||||||
|
"""
|
||||||
|
A class for performing image segmentation using a specified model.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
device (str): The device to run the model on (e.g., 'cpu', 'cuda').
|
||||||
|
model_name (str): The name of the model to be loaded. Defaults to
|
||||||
|
'facebook/sam-vit-huge'.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self, device: str = "cpu", model_name: str = "facebook/sam-vit-huge"
|
||||||
|
):
|
||||||
|
self.model = SamModel.from_pretrained(model_name).to(device)
|
||||||
|
self.processor = SamProcessor.from_pretrained(model_name)
|
||||||
|
self.image_processor = SamImageProcessor.from_pretrained(model_name)
|
||||||
|
self.pipeline = pipeline(
|
||||||
|
task="mask-generation",
|
||||||
|
model=self.model,
|
||||||
|
image_processor=self.image_processor,
|
||||||
|
device=device,
|
||||||
|
)
|
||||||
|
|
||||||
|
def run(self, image: np.ndarray) -> sv.Detections:
|
||||||
|
"""
|
||||||
|
Generate image segmentation marks.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
image (np.ndarray): The image to be marked in BGR format.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
sv.Detections: An object containing the segmentation masks and their
|
||||||
|
corresponding bounding box coordinates.
|
||||||
|
"""
|
||||||
|
image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||||||
|
outputs = self.pipeline(image, points_per_batch=64)
|
||||||
|
masks = np.array(outputs["masks"])
|
||||||
|
return masks_to_marks(masks=masks)
|
||||||
@ -0,0 +1,11 @@
|
|||||||
|
# System prompt
|
||||||
|
FLOW_SYSTEM_PROMPT = """
|
||||||
|
You are an autonomous agent granted autonomy in a autonomous loop structure.
|
||||||
|
Your role is to engage in multi-step conversations with your self or the user,
|
||||||
|
generate long-form content like blogs, screenplays, or SOPs,
|
||||||
|
and accomplish tasks bestowed by the user.
|
||||||
|
|
||||||
|
You can have internal dialogues with yourself or can interact with the user
|
||||||
|
to aid in these complex tasks. Your responses should be coherent, contextually relevant, and tailored to the task at hand.
|
||||||
|
|
||||||
|
"""
|
||||||
@ -0,0 +1,99 @@
|
|||||||
|
VISION_PROMPT = """
|
||||||
|
You are a Self-Operating Computer. You use the same operating system as a human.
|
||||||
|
|
||||||
|
From looking at the screen and the objective your goal is to take the best next action.
|
||||||
|
|
||||||
|
To operate the computer you have the four options below.
|
||||||
|
|
||||||
|
1. CLICK - Move mouse and click
|
||||||
|
2. TYPE - Type on the keyboard
|
||||||
|
3. SEARCH - Search for a program on Mac and open it
|
||||||
|
4. DONE - When you completed the task respond with the exact following phrase content
|
||||||
|
|
||||||
|
Here are the response formats below.
|
||||||
|
|
||||||
|
1. CLICK
|
||||||
|
Response: CLICK {{ "x": "percent", "y": "percent", "description": "~description here~", "reason": "~reason here~" }}
|
||||||
|
|
||||||
|
2. TYPE
|
||||||
|
Response: TYPE "value you want to type"
|
||||||
|
|
||||||
|
2. SEARCH
|
||||||
|
Response: SEARCH "app you want to search for on Mac"
|
||||||
|
|
||||||
|
3. DONE
|
||||||
|
Response: DONE
|
||||||
|
|
||||||
|
Here are examples of how to respond.
|
||||||
|
__
|
||||||
|
Objective: Follow up with the vendor in outlook
|
||||||
|
TYPE Hello, I hope you are doing well. I wanted to follow up
|
||||||
|
__
|
||||||
|
Objective: Open Spotify and play the beatles
|
||||||
|
SEARCH Spotify
|
||||||
|
__
|
||||||
|
Objective: Find a image of a banana
|
||||||
|
CLICK {{ "x": "50%", "y": "60%", "description": "Click: Google Search field", "reason": "This will allow me to search for a banana" }}
|
||||||
|
__
|
||||||
|
Objective: Go buy a book about the history of the internet
|
||||||
|
TYPE https://www.amazon.com/
|
||||||
|
__
|
||||||
|
|
||||||
|
A few important notes:
|
||||||
|
|
||||||
|
- Default to opening Google Chrome with SEARCH to find things that are on the internet.
|
||||||
|
- Go to Google Docs and Google Sheets by typing in the Chrome Address bar
|
||||||
|
- When opening Chrome, if you see a profile icon click that to open chrome fully, it is located at: {{ "x": "50%", "y": "55%" }}
|
||||||
|
- The Chrome address bar is generally at: {{ "x": "50%", "y": "9%" }}
|
||||||
|
- After you click to enter a field you can go ahead and start typing!
|
||||||
|
|
||||||
|
{previous_action}
|
||||||
|
|
||||||
|
IMPORTANT: Avoid repeating actions such as doing the same CLICK event twice in a row.
|
||||||
|
|
||||||
|
Objective: {objective}
|
||||||
|
"""
|
||||||
|
|
||||||
|
USER_QUESTION = "Hello, I can help you with anything. What would you like done?"
|
||||||
|
|
||||||
|
SUMMARY_PROMPT = """
|
||||||
|
You are a Self-Operating Computer. You just completed a request from a user by operating the computer. Now you need to share the results.
|
||||||
|
|
||||||
|
You have three pieces of key context about the completed request.
|
||||||
|
|
||||||
|
1. The original objective
|
||||||
|
2. The steps you took to reach the objective that are available in the previous messages
|
||||||
|
3. The screenshot you are looking at.
|
||||||
|
|
||||||
|
Now you need to summarize what you did to reach the objective. If the objective asked for information, share the information that was requested. IMPORTANT: Don't forget to answer a user's question if they asked one.
|
||||||
|
|
||||||
|
Thing to note: The user can not respond to your summary. You are just sharing the results of your work.
|
||||||
|
|
||||||
|
The original objective was: {objective}
|
||||||
|
|
||||||
|
Now share the results!
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
def format_summary_prompt(objective):
|
||||||
|
"""
|
||||||
|
Format the summary prompt
|
||||||
|
"""
|
||||||
|
prompt = SUMMARY_PROMPT.format(objective=objective)
|
||||||
|
return prompt
|
||||||
|
|
||||||
|
|
||||||
|
def format_vision_prompt(objective, previous_action):
|
||||||
|
"""
|
||||||
|
Format the vision prompt
|
||||||
|
"""
|
||||||
|
if previous_action:
|
||||||
|
previous_action = (
|
||||||
|
f"Here was the previous action you took: {previous_action}"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
previous_action = ""
|
||||||
|
prompt = VISION_PROMPT.format(
|
||||||
|
objective=objective, previous_action=previous_action
|
||||||
|
)
|
||||||
|
return prompt
|
||||||
@ -0,0 +1,60 @@
|
|||||||
|
# Prompts
|
||||||
|
DYNAMIC_STOP_PROMPT = """
|
||||||
|
|
||||||
|
Now, when you 99% sure you have completed the task, you may follow the instructions below to escape the autonomous loop.
|
||||||
|
|
||||||
|
When you have finished the task from the Human, output a special token: <DONE>
|
||||||
|
This will enable you to leave the autonomous loop.
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
# Make it able to handle multi input tools
|
||||||
|
DYNAMICAL_TOOL_USAGE = """
|
||||||
|
You have access to the following tools:
|
||||||
|
Output a JSON object with the following structure to use the tools
|
||||||
|
commands: {
|
||||||
|
"tools": {
|
||||||
|
tool1: "tool_name",
|
||||||
|
"params": {
|
||||||
|
"tool1": "inputs",
|
||||||
|
"tool1": "inputs"
|
||||||
|
}
|
||||||
|
"tool2: "tool_name",
|
||||||
|
"params": {
|
||||||
|
"tool1": "inputs",
|
||||||
|
"tool1": "inputs"
|
||||||
|
}
|
||||||
|
"tool3: "tool_name",
|
||||||
|
"params": {
|
||||||
|
"tool1": "inputs",
|
||||||
|
"tool1": "inputs"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
-------------TOOLS---------------------------
|
||||||
|
{tools}
|
||||||
|
"""
|
||||||
|
|
||||||
|
SCENARIOS = """
|
||||||
|
commands: {
|
||||||
|
"tools": {
|
||||||
|
tool1: "tool_name",
|
||||||
|
"params": {
|
||||||
|
"tool1": "inputs",
|
||||||
|
"tool1": "inputs"
|
||||||
|
}
|
||||||
|
"tool2: "tool_name",
|
||||||
|
"params": {
|
||||||
|
"tool1": "inputs",
|
||||||
|
"tool1": "inputs"
|
||||||
|
}
|
||||||
|
"tool3: "tool_name",
|
||||||
|
"params": {
|
||||||
|
"tool1": "inputs",
|
||||||
|
"tool1": "inputs"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
"""
|
||||||
@ -1,97 +0,0 @@
|
|||||||
from swarms.models import OpenAIChat
|
|
||||||
from swarms.structs.agent import Agent
|
|
||||||
|
|
||||||
import concurrent.futures
|
|
||||||
from typing import Callable, List, Dict, Any, Sequence
|
|
||||||
|
|
||||||
|
|
||||||
class Task:
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
id: str,
|
|
||||||
task: str,
|
|
||||||
flows: Sequence[Agent],
|
|
||||||
dependencies: List[str] = [],
|
|
||||||
):
|
|
||||||
self.id = id
|
|
||||||
self.task = task
|
|
||||||
self.flows = flows
|
|
||||||
self.dependencies = dependencies
|
|
||||||
self.results = []
|
|
||||||
|
|
||||||
def execute(self, parent_results: Dict[str, Any]):
|
|
||||||
args = [parent_results[dep] for dep in self.dependencies]
|
|
||||||
for agent in self.flows:
|
|
||||||
result = agent.run(self.task, *args)
|
|
||||||
self.results.append(result)
|
|
||||||
args = [
|
|
||||||
result
|
|
||||||
] # The output of one agent becomes the input to the next
|
|
||||||
|
|
||||||
|
|
||||||
class Workflow:
|
|
||||||
def __init__(self):
|
|
||||||
self.tasks: Dict[str, Task] = {}
|
|
||||||
self.executor = concurrent.futures.ThreadPoolExecutor()
|
|
||||||
|
|
||||||
def add_task(self, task: Task):
|
|
||||||
self.tasks[task.id] = task
|
|
||||||
|
|
||||||
def run(self):
|
|
||||||
completed_tasks = set()
|
|
||||||
while len(completed_tasks) < len(self.tasks):
|
|
||||||
futures = []
|
|
||||||
for task in self.tasks.values():
|
|
||||||
if task.id not in completed_tasks and all(
|
|
||||||
dep in completed_tasks for dep in task.dependencies
|
|
||||||
):
|
|
||||||
future = self.executor.submit(
|
|
||||||
task.execute,
|
|
||||||
{
|
|
||||||
dep: self.tasks[dep].results
|
|
||||||
for dep in task.dependencies
|
|
||||||
},
|
|
||||||
)
|
|
||||||
futures.append((future, task.id))
|
|
||||||
|
|
||||||
for future, task_id in futures:
|
|
||||||
future.result() # Wait for task completion
|
|
||||||
completed_tasks.add(task_id)
|
|
||||||
|
|
||||||
def get_results(self):
|
|
||||||
return {task_id: task.results for task_id, task in self.tasks.items()}
|
|
||||||
|
|
||||||
|
|
||||||
# create flows
|
|
||||||
llm = OpenAIChat(openai_api_key="sk-")
|
|
||||||
|
|
||||||
flow1 = Agent(llm, max_loops=1)
|
|
||||||
flow2 = Agent(llm, max_loops=1)
|
|
||||||
flow3 = Agent(llm, max_loops=1)
|
|
||||||
flow4 = Agent(llm, max_loops=1)
|
|
||||||
|
|
||||||
|
|
||||||
# Create tasks with their respective Agents and task strings
|
|
||||||
task1 = Task("task1", "Generate a summary on Quantum field theory", [flow1])
|
|
||||||
task2 = Task(
|
|
||||||
"task2",
|
|
||||||
"Elaborate on the summary of topic X",
|
|
||||||
[flow2, flow3],
|
|
||||||
dependencies=["task1"],
|
|
||||||
)
|
|
||||||
task3 = Task(
|
|
||||||
"task3", "Generate conclusions for topic X", [flow4], dependencies=["task1"]
|
|
||||||
)
|
|
||||||
|
|
||||||
# Create a workflow and add tasks
|
|
||||||
workflow = Workflow()
|
|
||||||
workflow.add_task(task1)
|
|
||||||
workflow.add_task(task2)
|
|
||||||
workflow.add_task(task3)
|
|
||||||
|
|
||||||
# Run the workflow
|
|
||||||
workflow.run()
|
|
||||||
|
|
||||||
# Get results
|
|
||||||
results = workflow.get_results()
|
|
||||||
print(results)
|
|
||||||
@ -0,0 +1,49 @@
|
|||||||
|
from swarms.structs.agent import Agent
|
||||||
|
|
||||||
|
from typing import List, Dict, Any, Sequence
|
||||||
|
|
||||||
|
|
||||||
|
class Task:
|
||||||
|
"""
|
||||||
|
Task is a unit of work that can be executed by a set of agents.
|
||||||
|
|
||||||
|
A task is defined by a task name and a set of agents that can execute the task.
|
||||||
|
The task can also have a set of dependencies, which are the names of other tasks
|
||||||
|
that must be executed before this task can be executed.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
id (str): A unique identifier for the task
|
||||||
|
task (str): The name of the task
|
||||||
|
agents (Sequence[Agent]): A list of agents that can execute the task
|
||||||
|
dependencies (List[str], optional): A list of task names that must be executed before this task can be executed. Defaults to [].
|
||||||
|
|
||||||
|
Methods:
|
||||||
|
execute(parent_results: Dict[str, Any]): Executes the task by passing the results of the parent tasks to the agents.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
id: str,
|
||||||
|
task: str,
|
||||||
|
agents: Sequence[Agent],
|
||||||
|
dependencies: List[str] = [],
|
||||||
|
):
|
||||||
|
self.id = id
|
||||||
|
self.task = task
|
||||||
|
self.agents = agents
|
||||||
|
self.dependencies = dependencies
|
||||||
|
self.results = []
|
||||||
|
|
||||||
|
def execute(self, parent_results: Dict[str, Any]):
|
||||||
|
"""Executes the task by passing the results of the parent tasks to the agents.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
parent_results (Dict[str, Any]): _description_
|
||||||
|
"""
|
||||||
|
args = [parent_results[dep] for dep in self.dependencies]
|
||||||
|
for agent in self.agents:
|
||||||
|
result = agent.run(self.task, *args)
|
||||||
|
self.results.append(result)
|
||||||
|
args = [
|
||||||
|
result
|
||||||
|
] # The output of one agent becomes the input to the next
|
||||||
@ -0,0 +1,107 @@
|
|||||||
|
import os
|
||||||
|
from unittest.mock import Mock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
|
||||||
|
from swarms.models.gpt4_vision_api import GPT4VisionAPI
|
||||||
|
from swarms.prompts.multi_modal_autonomous_instruction_prompt import (
|
||||||
|
MULTI_MODAL_AUTO_AGENT_SYSTEM_PROMPT_1,
|
||||||
|
)
|
||||||
|
from swarms.structs.agent import Agent
|
||||||
|
from swarms.structs.task import Task
|
||||||
|
|
||||||
|
load_dotenv()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def llm():
|
||||||
|
return GPT4VisionAPI()
|
||||||
|
|
||||||
|
|
||||||
|
def test_agent_run_task(llm):
|
||||||
|
task = (
|
||||||
|
"Analyze this image of an assembly line and identify any issues such as"
|
||||||
|
" misaligned parts, defects, or deviations from the standard assembly"
|
||||||
|
" process. IF there is anything unsafe in the image, explain why it is"
|
||||||
|
" unsafe and how it could be improved."
|
||||||
|
)
|
||||||
|
img = "assembly_line.jpg"
|
||||||
|
|
||||||
|
agent = Agent(
|
||||||
|
llm=llm,
|
||||||
|
max_loops="auto",
|
||||||
|
sop=MULTI_MODAL_AUTO_AGENT_SYSTEM_PROMPT_1,
|
||||||
|
dashboard=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
result = agent.run(task=task, img=img)
|
||||||
|
|
||||||
|
# Add assertions here to verify the expected behavior of the agent's run method
|
||||||
|
assert isinstance(result, dict)
|
||||||
|
assert "response" in result
|
||||||
|
assert "dashboard_data" in result
|
||||||
|
# Add more assertions as needed
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def task():
|
||||||
|
agents = [Agent(llm=llm, id=f"Agent_{i}") for i in range(5)]
|
||||||
|
return Task(id="Task_1", task="Task_Name", agents=agents, dependencies=[])
|
||||||
|
|
||||||
|
|
||||||
|
# Basic tests
|
||||||
|
|
||||||
|
|
||||||
|
def test_task_init(task):
|
||||||
|
assert task.id == "Task_1"
|
||||||
|
assert task.task == "Task_Name"
|
||||||
|
assert isinstance(task.agents, list)
|
||||||
|
assert len(task.agents) == 5
|
||||||
|
assert isinstance(task.dependencies, list)
|
||||||
|
|
||||||
|
|
||||||
|
def test_task_execute(task, mocker):
|
||||||
|
mocker.patch.object(Agent, "run", side_effect=[1, 2, 3, 4, 5])
|
||||||
|
parent_results = {}
|
||||||
|
task.execute(parent_results)
|
||||||
|
assert isinstance(task.results, list)
|
||||||
|
assert len(task.results) == 5
|
||||||
|
for result in task.results:
|
||||||
|
assert isinstance(result, int)
|
||||||
|
|
||||||
|
|
||||||
|
# Parameterized tests
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("num_agents", [1, 3, 5, 10])
|
||||||
|
def test_task_num_agents(task, num_agents, mocker):
|
||||||
|
task.agents = [Agent(id=f"Agent_{i}") for i in range(num_agents)]
|
||||||
|
mocker.patch.object(Agent, "run", return_value=1)
|
||||||
|
parent_results = {}
|
||||||
|
task.execute(parent_results)
|
||||||
|
assert len(task.results) == num_agents
|
||||||
|
|
||||||
|
|
||||||
|
# Exception testing
|
||||||
|
|
||||||
|
|
||||||
|
def test_task_execute_with_dependency_error(task, mocker):
|
||||||
|
task.dependencies = ["NonExistentTask"]
|
||||||
|
mocker.patch.object(Agent, "run", return_value=1)
|
||||||
|
parent_results = {}
|
||||||
|
with pytest.raises(KeyError):
|
||||||
|
task.execute(parent_results)
|
||||||
|
|
||||||
|
|
||||||
|
# Mocking and monkeypatching tests
|
||||||
|
|
||||||
|
|
||||||
|
def test_task_execute_with_mocked_agents(task, mocker):
|
||||||
|
mock_agents = [Mock(spec=Agent) for _ in range(5)]
|
||||||
|
mocker.patch.object(task, "agents", mock_agents)
|
||||||
|
for mock_agent in mock_agents:
|
||||||
|
mock_agent.run.return_value = 1
|
||||||
|
parent_results = {}
|
||||||
|
task.execute(parent_results)
|
||||||
|
assert len(task.results) == 5
|
||||||
Loading…
Reference in new issue