Package custom OWLv2 checkpoints for upload

README.md

@@ -223,6 +223,8 @@ Within the project object, you can retrieve metadata about the project, list ver
 
 Within the version object, you can download the dataset version in any model format, train the version on Roboflow, and deploy your own external model to Roboflow.
 
+For custom OWLv2 checkpoints and Auto Label, see [Upload a custom OWLv2 model](docs/owlv2-upload.md).
+
 ## 🏆 Contributing
 
 We would love your input on how we can improve the Roboflow Python package! Please see our [contributing guide](https://github.com/roboflow/roboflow-python/blob/main/CONTRIBUTING.md) to get started. Thank you 🙏 to all our contributors!

docs/owlv2-upload.md

@@ -0,0 +1,69 @@
+# Upload a custom OWLv2 model
+
+Use `model_type="owlv2"` to upload fine-tuned OWLv2 weights for Auto Label and Workflows.
+For a standard OWLv2 model, the upload folder needs only two files:
+
+```text
+weights-upload/
+  owlv2.pth
+  class_names.txt
+```
+
+`owlv2.pth` must contain the complete `Owlv2ForObjectDetection.state_dict()`, including the text encoder and detection heads.
+A training checkpoint that wraps the weights under another key is not accepted.
+`class_names.txt` contains unique, nonblank output labels, one per line. Class IDs follow this order.
+The labels also become the text prompts. The weights do not impose a fixed class count.
+
+The exporter assumes the standard `google/owlv2-base-patch16-ensemble` model configuration, image preprocessing, and tokenizer.
+It bundles these defaults at revision `cfd3195ba4ea9592eec887ded089f4c08eff231d` and uses them offline.
+Fine-tuning the weights does not require additional files if those components stay unchanged.
+The exporter saves the resolved configuration and processor in the inference package.
+
+## Upload
+
+Use an object-detection project that the workspace can access.
+Replace the workspace, project, folder, and model name with your values.
+Set `ROBOFLOW_API_KEY` in your environment before you run this example.
+
+```python
+from roboflow import Roboflow
+
+workspace = Roboflow().workspace("your-workspace")
+workspace.deploy_model(
+    model_type="owlv2",
+    model_path="/absolute/path/to/weights-upload",
+    filename="owlv2.pth",
+    project_ids=["your-project"],
+    model_name="custom-owlv2",
+)
+```
+
+After upload processing finishes, select the uploaded model under custom models in Auto Label.
+In Workflows, select its model ID in an Object Detection Model block.
+The saved prompts stay fixed for each model package.
+Inference reuses the existing OWLv2 implementation. Roboflow Instant embedding uploads keep their separate format and route.
+
+## Custom model or preprocessing
+
+If training changed a component, place its override files beside the weights.
+Each supplied component takes precedence over its default.
+
+| Optional files | When to supply them |
+| --- | --- |
+| `config.json` | A different model architecture, saved with `model.config.save_pretrained(folder)` |
+| `processor_config.json` or `preprocessor_config.json` | Different image preprocessing, such as resizing or normalization |
+| Tokenizer files | Different tokenization, saved with the original tokenizer or processor |
+| `owlv2_prompts.json` | Text prompts that differ from the output labels |
+
+A tokenizer override needs `tokenizer_config.json` and either `tokenizer.json` or both `vocab.json` and `merges.txt`.
+Include additional tokenizer files that the original processor saved, such as `added_tokens.json` and `special_tokens_map.json`.
+Partial tokenizer overrides produce an error. They are never mixed with the default tokenizer.
+You can save a complete processor override with `processor.save_pretrained(folder)`.
+
+`owlv2_prompts.json` is a JSON array with one nonblank prompt per label, in label order.
+For example, labels `dent` and `scratch` can use `["a dent in a metal panel", "a scratch in a metal panel"]`.
+
+The exporter strictly matches the weights to the selected model configuration and rejects partial or adapter-only checkpoints.
+If the weights do not match the default architecture, supply the original `config.json`.
+Weight shapes cannot identify custom resizing or normalization. Omitted image preprocessing files mean standard preprocessing.
+The SDK packages files without loading the checkpoint.

roboflow/util/model_processor.py

@@ -68,6 +68,7 @@
     "paligemma2",
     "florence-2",
     "rfdetr",
+    "owlv2",
 )
 
 SUPPORTED_HUGGINGFACE_TYPES = (
@@ -462,7 +463,11 @@ def _process_model(
             f"It must start with a supported family: {', '.join(SUPPORTED_MODELS)}."
         )
 
-    if model_type.startswith(("paligemma", "paligemma2", "florence-2")):
+    if model_type == "owlv2":
+        return _process_owlv2(model_path, filename, build_dir)
+    elif model_type.startswith("owlv2"):
+        raise UnsupportedModelError("Use model_type='owlv2' for OWLv2 checkpoints.")
+    elif model_type.startswith(("paligemma", "paligemma2", "florence-2")):
         return _process_huggingface(model_type, model_path, build_dir)
     if model_type.startswith("yolonas"):
         return _process_yolonas(model_type, model_path, filename, build_dir)
@@ -1125,6 +1130,79 @@ def _process_rfdetr(
     return archive_path, model_type, warnings
 
 
+def _process_owlv2(model_path: Path, filename: str, build_dir: Path) -> tuple[Path, str, list[str]]:
+    required = {"class_names.txt"}
+    tokenizer_files = {
+        "tokenizer_config.json",
+        "tokenizer.json",
+        "vocab.json",
+        "merges.txt",
+        "special_tokens_map.json",
+        "added_tokens.json",
+    }
+    optional = tokenizer_files | {
+        "config.json",
+        "preprocessor_config.json",
+        "processor_config.json",
+        "owlv2_prompts.json",
+    }
+    files = {name: _resolve_within_source(model_path, name) for name in required | optional}
+    missing = sorted(name for name in required if not files[name].is_file())
+    if missing:
+        raise MissingFileError(
+            f"Missing OWLv2 upload files: {', '.join(missing)}. Provide class_names.txt with one output label per line."
+        )
+    if any(files[name].is_file() for name in tokenizer_files) and (
+        not files["tokenizer_config.json"].is_file()
+        or not (
+            files["tokenizer.json"].is_file() or all(files[name].is_file() for name in ("vocab.json", "merges.txt"))
+        )
+    ):
+        raise MissingFileError(
+            "Provide a complete tokenizer override: tokenizer_config.json and either tokenizer.json "
+            "or both vocab.json and merges.txt. Omit all tokenizer files to use the standard OWLv2 tokenizer."
+        )
+    try:
+        config = json.loads(files["config.json"].read_text()) if files["config.json"].is_file() else None
+        class_names = files["class_names.txt"].read_text().splitlines()
+        prompts = (
+            json.loads(files["owlv2_prompts.json"].read_text())
+            if files["owlv2_prompts.json"].is_file()
+            else class_names
+        )
+    except (ValueError, UnicodeError) as error:
+        raise ModelPackagingError(f"Invalid OWLv2 metadata: {error}") from error
+    if files["config.json"].is_file() and (not isinstance(config, dict) or config.get("model_type") != "owlv2"):
+        raise ModelPackagingError("config.json must describe an OWLv2 model (model_type='owlv2').")
+    if (
+        not class_names
+        or any(not name.strip() or name != name.strip() for name in class_names)
+        or len(set(class_names)) != len(class_names)
+    ):
+        raise ModelPackagingError(
+            "class_names.txt must contain unique, nonblank labels, one per line, without surrounding spaces."
+        )
+    if (
+        not isinstance(prompts, list)
+        or len(prompts) != len(class_names)
+        or any(not isinstance(prompt, str) or not prompt.strip() for prompt in prompts)
+    ):
+        raise ModelPackagingError(
+            "owlv2_prompts.json must contain one nonblank text prompt per class, in class_names.txt order."
+        )
+    weights = _resolve_within_source(model_path, filename)
+    if not weights.is_file():
+        raise MissingFileError(
+            f"OWLv2 weights file '{weights}' was not found. Set filename to the checkpoint's relative path."
+        )
+    archive_path = build_dir / "roboflow_deploy.zip"
+    _write_zip(
+        archive_path,
+        [(weights, "weights.pt", True)] + [(path, name, name in required) for name, path in sorted(files.items())],
+    )
+    return archive_path, "owlv2", []
+
+
 def _process_huggingface(
     model_type: str,
     model_path: Path,

tests/util/test_model_processor.py

@@ -728,6 +728,92 @@ def test_rfdetr_multiple_discovered_checkpoints_warns_which_used(self):
                 bundle.cleanup()
 
 
+class PackageOwlV2Test(unittest.TestCase):
+    def setUp(self):
+        self.directory = tempfile.TemporaryDirectory()
+        self.addCleanup(self.directory.cleanup)
+        self.source = Path(self.directory.name)
+        for name, contents in {
+            "config.json": '{"model_type": "owlv2"}',
+            "processor_config.json": "{}",
+            "tokenizer_config.json": "{}",
+            "tokenizer.json": "{}",
+            "class_names.txt": "dent\nscratch\n",
+            "owlv2.pth": "checkpoint",
+        }.items():
+            (self.source / name).write_text(contents)
+
+    def test_packages_raw_weights_without_loading_them(self):
+        with mock.patch.object(model_processor, "import_module", side_effect=AssertionError("No model load")):
+            bundle = package_custom_weights("owlv2", str(self.source), "owlv2.pth")
+        self.addCleanup(bundle.cleanup)
+        with zipfile.ZipFile(bundle.archive_path) as archive:
+            self.assertEqual(archive.read("weights.pt"), b"checkpoint")
+            self.assertEqual(archive.read("class_names.txt").decode().splitlines(), ["dent", "scratch"])
+            self.assertIn("processor_config.json", archive.namelist())
+        self.assertFalse((self.source / "weights.pt").exists())
+
+    def test_rejects_prompt_count_mismatch(self):
+        (self.source / "owlv2_prompts.json").write_text('["a dent"]')
+        with self.assertRaisesRegex(ModelPackagingError, "one nonblank text prompt"):
+            package_custom_weights("owlv2", str(self.source), "owlv2.pth")
+
+    def test_omits_tokenizer_for_server_default(self):
+        (self.source / "tokenizer_config.json").unlink()
+        (self.source / "tokenizer.json").unlink()
+        bundle = package_custom_weights("owlv2", str(self.source), "owlv2.pth")
+        self.addCleanup(bundle.cleanup)
+        with zipfile.ZipFile(bundle.archive_path) as archive:
+            self.assertNotIn("tokenizer_config.json", archive.namelist())
+            self.assertNotIn("tokenizer.json", archive.namelist())
+            self.assertIn("processor_config.json", archive.namelist())
+
+    def test_rejects_partial_tokenizer_overrides(self):
+        (self.source / "tokenizer_config.json").unlink()
+        (self.source / "tokenizer.json").unlink()
+        for name in ("tokenizer_config.json", "tokenizer.json", "special_tokens_map.json", "vocab.json"):
+            with self.subTest(name=name):
+                (self.source / name).write_text("{}")
+                with self.assertRaisesRegex(MissingFileError, "complete tokenizer override"):
+                    package_custom_weights("owlv2", str(self.source), "owlv2.pth")
+                (self.source / name).unlink()
+
+    def test_packages_legacy_tokenizer_override(self):
+        (self.source / "tokenizer.json").unlink()
+        (self.source / "vocab.json").write_text("{}")
+        (self.source / "merges.txt").write_text("#version: 0.2\n")
+        bundle = package_custom_weights("owlv2", str(self.source), "owlv2.pth")
+        self.addCleanup(bundle.cleanup)
+        with zipfile.ZipFile(bundle.archive_path) as archive:
+            self.assertIn("tokenizer_config.json", archive.namelist())
+            self.assertIn("vocab.json", archive.namelist())
+            self.assertIn("merges.txt", archive.namelist())
+
+    def test_rejects_duplicate_labels(self):
+        (self.source / "class_names.txt").write_text("dent\ndent\n")
+        with self.assertRaisesRegex(ModelPackagingError, "unique"):
+            package_custom_weights("owlv2", str(self.source), "owlv2.pth")
+
+    def test_packages_only_weights_and_class_names(self):
+        for path in self.source.iterdir():
+            if path.name not in {"owlv2.pth", "class_names.txt"}:
+                path.unlink()
+        bundle = package_custom_weights("owlv2", str(self.source), "owlv2.pth")
+        self.addCleanup(bundle.cleanup)
+        with zipfile.ZipFile(bundle.archive_path) as archive:
+            self.assertEqual(set(archive.namelist()), {"weights.pt", "class_names.txt"})
+
+    def test_rejects_invalid_model_configuration_override(self):
+        (self.source / "config.json").write_text('{"model_type": "owlvit"}')
+        with self.assertRaisesRegex(ModelPackagingError, "config.json must describe an OWLv2 model"):
+            package_custom_weights("owlv2", str(self.source), "owlv2.pth")
+
+    def test_rejects_missing_class_names(self):
+        (self.source / "class_names.txt").unlink()
+        with self.assertRaisesRegex(MissingFileError, "class_names.txt"):
+            package_custom_weights("owlv2", str(self.source), "owlv2.pth")
+
+
 class ProcessHuggingfaceTest(unittest.TestCase):
     """Packaging for HuggingFace-backed models (paligemma / florence-2)."""