Update schema to allow metadata for versions and add micro sam EM models

FrancisCrickInstitute · Mar 1, 2024 · 535d35d · 535d35d
1 parent aecb2bd
commit 535d35d
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 32 deletions.
diff --git a/aiod_registry/manifests/sam.json b/aiod_registry/manifests/sam.json
@@ -99,6 +99,42 @@
           "location": "https://syncandshare.desy.de/index.php/s/yLfdFbpfEGSHJWY/download/medsam_20230423_vit_b_0.0.1.pth"
         }
       }
+    },
+    "MicroSAM-Boundaries": {
+      "metadata": {
+        "description": "Segment Anything for Microscopy finetuned models (for)",
+        "url": "https://computational-cell-analytics.github.io/micro-sam/micro_sam.html",
+        "repo": "https://github.com/computational-cell-analytics/micro-sam",
+        "pubs": [
+          {
+            "info": "Segment Anything for Microscopy paper",
+            "url": "https://doi.org/10.1101/2023.08.21.554208"
+          }
+        ]
+      },
+      "tasks": {
+        "everything": {
+          "location": "https://zenodo.org/records/10524894/files/vit_b_em_boundaries.pth?download=1"
+        }
+      }
+    },
+    "MicroSAM-Organelles": {
+      "metadata": {
+        "description": "Segment Anything for Microscopy finetuned models (for)",
+        "url": "https://computational-cell-analytics.github.io/micro-sam/micro_sam.html",
+        "repo": "https://github.com/computational-cell-analytics/micro-sam",
+        "pubs": [
+          {
+            "info": "Segment Anything for Microscopy paper",
+            "url": "https://doi.org/10.1101/2023.08.21.554208"
+          }
+        ]
+      },
+      "tasks": {
+        "everything": {
+          "location": "https://zenodo.org/records/10524828/files/vit_b_em_organelles.pth?download=1"
+        }
+      }
     }
   },
   "params": [

diff --git a/aiod_registry/schema.py b/aiod_registry/schema.py
@@ -66,38 +66,6 @@ def extract_arg_type(self):
         return self
 
 
-class ModelVersionTask(StrictModel):
-    location: str = Field(
-        ...,
-        description="Either a url or a filepath (will be skipped if the path does not exist/cannot be read!)",
-    )
-    config_path: Optional[Union[Path, str]] = None
-    params: Optional[list[ModelParam]] = None
-    location_type: Optional[str] = None
-
-    @model_validator(mode="after")
-    def get_location_type(self):
-        # Skip if provided
-        if self.location_type is not None:
-            return self
-        # Otherwise, determine the type
-        res = urlparse(self.location)
-        if res.scheme in ("http", "https"):
-            self.location_type = "url"
-        elif res.scheme in ("file", ""):
-            self.location_type = "file"
-        else:
-            # NOTE: Because of including "" above, it is unlikely this will be reached
-            raise TypeError(
-                f"Cannot determine type (file/url) of location: {self.location}!"
-            )
-        return self
-
-
-class ModelVersion(StrictModel):
-    tasks: dict[Task, ModelVersionTask]
-
-
 class Author(StrictModel):
     name: str
     affiliation: str
@@ -136,6 +104,40 @@ class Metadata(StrictModel):
     repo: Optional[AnyUrl] = None
 
 
+class ModelVersionTask(StrictModel):
+    location: str = Field(
+        ...,
+        description="Either a url or a filepath (will be skipped if the path does not exist/cannot be read!)",
+    )
+    config_path: Optional[Union[Path, str]] = None
+    params: Optional[list[ModelParam]] = None
+    location_type: Optional[str] = None
+    metadata: Optional[Metadata] = None
+
+    @model_validator(mode="after")
+    def get_location_type(self):
+        # Skip if provided
+        if self.location_type is not None:
+            return self
+        # Otherwise, determine the type
+        res = urlparse(self.location)
+        if res.scheme in ("http", "https"):
+            self.location_type = "url"
+        elif res.scheme in ("file", ""):
+            self.location_type = "file"
+        else:
+            # NOTE: Because of including "" above, it is unlikely this will be reached
+            raise TypeError(
+                f"Cannot determine type (file/url) of location: {self.location}!"
+            )
+        return self
+
+
+class ModelVersion(StrictModel):
+    tasks: dict[Task, ModelVersionTask]
+    metadata: Optional[Metadata] = None
+
+
 class ModelManifest(StrictModel):
     name: str = Field(..., min_length=1, max_length=50)
     short_name: Optional[str] = None