Expand schema to include metadata

FrancisCrickInstitute · Dec 18, 2023 · 68024cb · 68024cb
1 parent 785d670
commit 68024cb
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 0 deletions.
diff --git a/aiod_registry/manifests/mitonet.json b/aiod_registry/manifests/mitonet.json
@@ -1,5 +1,24 @@
 {
     "name": "Mitonet",
+    "metadata": {
+      "description": "MitoNet is a deep learning model for mitochondria segmentation in EM images.",
+      "authors": [
+        {
+          "name": "Ryan Conrad",
+          "affiliation": "Center for Molecular Microscopy, Center for Cancer Research, National Cancer Institute, National Institutes of Health, Bethesda, MD 20892, USA"
+        },
+        {
+          "name": "Kedar Narayan",
+          "affiliation": "Center for Molecular Microscopy, Center for Cancer Research, National Cancer Institute, National Institutes of Health, Bethesda, MD 20892, USA"
+        }
+      ],
+      "pubs": [
+        {
+          "info": "Main paper that describes model & data",
+          "url": "https://doi.org/10.1016/j.cels.2022.12.006"
+        }
+      ]
+    },
     "versions": {
       "MitoNet v1": {
         "tasks": {

diff --git a/aiod_registry/manifests/sam.json b/aiod_registry/manifests/sam.json
@@ -1,6 +1,17 @@
 {
   "name": "Segment Anything",
   "short_name": "sam",
+  "metadata": {
+    "description": "Segment Anything is a vision foundation model with flexible prompting.",
+    "url": "https://segment-anything.com/",
+    "repo": "https://github.com/facebookresearch/segment-anything",
+    "pubs": [
+      {
+        "info": "Main paper that describes model & data",
+        "url": "https://arxiv.org/abs/2304.02643"
+      }
+    ]
+  },
   "versions": {
     "default": {
       "tasks": {

diff --git a/aiod_registry/manifests/unet_seai.json b/aiod_registry/manifests/unet_seai.json
@@ -1,6 +1,9 @@
 {
   "name": "SEAI U-Net",
   "short_name": "seai_unet",
+  "metadata": {
+    "description": "SEAI U-Net developed on internal Crick EM data"
+  },
   "versions": {
     "U-Net": {
       "tasks": {

diff --git a/aiod_registry/schema.py b/aiod_registry/schema.py
@@ -98,12 +98,49 @@ class ModelVersion(StrictModel):
     tasks: dict[Task, ModelVersionTask]
 
 
+class Author(StrictModel):
+    name: str
+    affiliation: str
+    email: Optional[str] = None
+    url: Optional[AnyUrl] = None
+    github: Optional[str] = None
+    orcid: Optional[str] = None
+
+
+class Publication(StrictModel):
+    info: Annotated[
+        str,
+        Field(
+            ...,
+            description="Information on publication, whether it pertains to the model or the underlying data or something else.",
+        ),
+    ]
+    url: AnyUrl
+    doi: Optional[str] = None
+    authors: Optional[list[Author]] = None
+
+
+class Metadata(StrictModel):
+    description: Annotated[
+        str,
+        Field(
+            ...,
+            description="A short description of the model to provide context.",
+        ),
+    ]
+    authors: Optional[list[Author]] = None
+    pubs: Optional[list[Publication]] = None
+    url: Optional[AnyUrl] = None
+    repo: Optional[AnyUrl] = None
+
+
 class ModelManifest(StrictModel):
     name: str = Field(..., min_length=1, max_length=50)
     short_name: Optional[str] = None
     versions: dict[ModelName, ModelVersion]
     params: Optional[list[ModelParam]] = None
     config: Optional[Path] = None
+    metadata: Metadata
 
     @model_validator(mode="after")
     def create_short_name(self):