implement operator

Author: harpreetsahota204

__init__.py

@@ -0,0 +1,235 @@
+import os
+
+os.environ['FIFTYONE_ALLOW_LEGACY_ORCHESTRATORS'] = 'true'
+
+import requests 
+
+from fiftyone.core.utils import add_sys_path
+import fiftyone.operators as foo
+from fiftyone.operators import types
+
+with add_sys_path(os.path.dirname(os.path.abspath(__file__))):
+    from moondream import (
+        run_moondream_model
+    )
+
+def _handle_calling(
+        uri, 
+        sample_collection, 
+        model_name,
+        bbox_field,
+        output_field,
+        confidence_threshold,
+        delegate=False
+        ):
+    ctx = dict(dataset=sample_collection)
+
+    params = dict(
+        model_name=model_name,
+        bbox_field=bbox_field,
+        output_field=output_field,
+        confidence_threshold=confidence_threshold,
+        delegate=delegate
+        )
+    return foo.execute_operator(uri, ctx, params=params)
+
+MOONDREAM_MODES = {
+    "caption": "Caption images", 
+    "query": "Visual question answering",
+    "detect": "Object detection",
+    "point": "Apply point on object",
+}
+
+
+# Load Moondream2 revisions from HuggingFace
+MOONDREAM_VERSIONS_URL = "https://huggingface.co/vikhyatk/moondream2/raw/main/versions.txt"
+
+def load_moondream_versions():
+    try:
+        response = requests.get(MOONDREAM_VERSIONS_URL)
+        response.raise_for_status()
+        # Split by newlines and remove any empty strings
+        versions = [v.strip() for v in response.text.splitlines() if v.strip()]
+        return versions
+    except Exception as e:
+        print(f"Failed to load Moondream versions: {e}")
+        return ["2025-01-09"]  # Return default version as fallback
+
+# Load available versions
+MOONDREAM_REVISIONS = load_moondream_versions()
+
+class MoondreamOperator(foo.Operator):
+    @property
+    def config(self):
+        return foo.OperatorConfig(
+            name="moondream",
+            label="Run Moondream2",
+            description="Run the Moondream model on your Dataset!",
+            dynamic=True,
+            icon="/assets/moon-phase-svgrepo-com.svg",
+            )
+
+    def resolve_input(self, ctx):
+        """Implement this method to collect user inputs as parameters
+        that are stored in `ctx.params`.
+
+        Returns:
+            a `types.Property` defining the form's components
+        """
+        inputs = types.Object()
+
+        mode_dropdown = types.Dropdown(label="What would you like to use Moondream2 for?")
+        
+        for arch_key, arch_value in MOONDREAM_MODES.items():
+            mode_dropdown.add_choice(arch_value, label=arch_key)
+
+        revision_dropdown = types.Dropdown(label="Which revision would you like to use?")
+
+        inputs.enum(
+            "revision",
+            values=revision_dropdown.values(),
+            label="Revision",
+            default="2025-01-09",
+            description="Select from one of the available revisions. Note: The model weights will be downloaded from Hugging Face.",
+            view=revision_dropdown,
+            required=False
+        )
+
+        inputs.enum(
+            "operation",
+            values=mode_dropdown.values(),
+            label="Moondream2 Tasks",
+            description="Select from one of the supported tasks.",
+            view=mode_dropdown,
+            required=True
+        )
+
+        length_radio = types.RadioGroup()
+        length_radio.add_choice("short", label="A short caption")
+        length_radio.add_choice("normal", label="A more descriptive caption")        
+        
+        chosen_task = ctx.params.get("operation")
+
+        if chosen_task == "caption":
+            inputs.enum(
+                "length",
+                label="Caption Length",
+                description="Which caption type would you like?",
+                required=True,
+                view=length_radio
+            )
+
+        if chosen_task == "query":
+            inputs.str(
+                "query_text",
+                label="Query",
+                description="What's your query?",
+                required=True,
+            )
+
+        if chosen_task == "detect":
+            inputs.str(
+                "object_type",
+                label="Detect",
+                description="What do you want to detect? Currently this model only supports passing one object.",
+                required=True,
+            )
+
+        if chosen_task == "point":
+            inputs.str(
+                "object_type",
+                label="Point",
+                description="What do you want to place a point on? Currently this model only supports passing one object",
+                required=True,
+            )
+       
+
+        inputs.str(
+            "output_field",            
+            required=True,
+            label="Output Field",
+            description="Name of the field to store the results in."
+            )
+        
+        inputs.bool(
+            "delegate",
+            default=False,
+            required=True,
+            label="Delegate execution?",
+            description=("If you choose to delegate this operation you must first have a delegated service running. "
+            "You can launch a delegated service by running `fiftyone delegated launch` in your terminal"),
+            view=types.CheckboxView(),
+        )
+
+        inputs.view_target(ctx)
+
+        return types.Property(inputs)
+
+    def resolve_delegation(self, ctx):
+        """Implement this method if you want to programmatically *force*
+        this operation to be delegated or executed immediately.
+
+        Returns:
+            whether the operation should be delegated (True), run
+            immediately (False), or None to defer to
+            `resolve_execution_options()` to specify the available options
+        """
+        return ctx.params.get("delegate", False)
+
+    def execute(self, ctx):
+        """Executes the actual operation based on the hydrated `ctx`.
+        All operators must implement this method.
+
+        This method can optionally be implemented as `async`.
+
+        Returns:
+            an optional dict of results values
+        """
+        view = ctx.target_view()
+        revision = ctx.params.get("revision")
+        operation = ctx.params.get("operation")
+        output_field = ctx.params.get("output_field")
+        
+        # Create kwargs dictionary with additional parameters based on operation
+        kwargs = {}
+        if operation == "caption":
+            kwargs["length"] = ctx.params.get("length")
+        elif operation == "query":
+            kwargs["query_text"] = ctx.params.get("query_text")
+        elif operation in ["detect", "point"]:
+            kwargs["object_type"] = ctx.params.get("object_type")
+      
+        run_moondream_model(
+            dataset=view,
+            revision=revision,
+            operation=operation,
+            output_field=output_field,
+            **kwargs
+            )
+        
+        ctx.ops.reload_dataset()
+
+    def __call__(
+            self, 
+            sample_collection, 
+            revision,
+            operation,
+            output_field,
+            delegate,
+            **kwargs
+            ):
+        return _handle_calling(
+            self.uri,
+            sample_collection,
+            revision,
+            operation,
+            output_field,
+            delegate,
+            **kwargs
+            )
+
+def register(p):
+    """Always implement this method and register() each operator that your
+    plugin defines.
+    """
+    p.register(MoondreamOperator)

moondream.py

@@ -44,18 +44,14 @@ class Moondream2(Model):
     
     Args:
         operation (str): Type of operation to perform
-        model_name (str): Name of the model to load from HuggingFace
         revision (str, optional): Model revision/tag to use
-        device (str, optional): Device to run the model on ('cuda', 'mps', or 'cpu')
         **kwargs: Operation-specific parameters
     """
 
     def __init__(
         self, 
         operation: str,
-        model_name: str = "vikhyatk/moondream2",
         revision: Optional[str] = None,
-        device: Optional[str] = None,
         **kwargs
     ):
         self.operation = operation
@@ -73,11 +69,11 @@ def __init__(
         self.params = kwargs
 
         # Set device
-        self.device = device or get_device()
+        self.device = get_device()
 
         # Initialize model
         self.model = AutoModelForCausalLM.from_pretrained(
-            model_name,
+            "vikhyatk/moondream2",
             revision=revision,
             trust_remote_code=True,
             device_map={"": self.device},
@@ -233,9 +229,7 @@ def run_moondream_model(
     dataset: fo.Dataset,
     operation: str,
     output_field: str,
-    model_name: str = "vikhyatk/moondream2",
     revision: Optional[str] = None,
-    device: Optional[str] = None,
     **kwargs
 ) -> None:
     """Apply Moondream2 operations to a FiftyOne dataset.
@@ -244,16 +238,12 @@ def run_moondream_model(
         dataset: FiftyOne dataset to process
         operation: Type of operation to perform
         output_field: Field to store results in
-        model_name: Name of the model to load from HuggingFace
         revision: Model revision/tag to use
-        device: Device to run the model on
         **kwargs: Operation-specific parameters
     """
     model = Moondream2(
         operation=operation,
-        model_name=model_name,
         revision=revision,
-        device=device,
         **kwargs
     )
     dataset.apply_model(model, label_field=output_field)