Update media info schema to include extension (#494)

* update media info schema to include extension Signed-off-by: Igor-Davidyuk <igor.davidyuk@intel.com> * make extension field optional * fix field type Signed-off-by: Igor-Davidyuk <igor.davidyuk@intel.com> * add training framework field Signed-off-by: Igor-Davidyuk <igor.davidyuk@intel.com> --------- Signed-off-by: Igor-Davidyuk <igor.davidyuk@intel.com>
openvinotoolkit · Sep 12, 2024 · 88be98a · 88be98a
1 parent b53c8c7
commit 88be98a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/geti_sdk/data_models/media.py b/geti_sdk/data_models/media.py
@@ -55,6 +55,7 @@ class MediaInformation:
     display_url: str
     height: int
     width: int
+    extension: Optional[str] = None  # Added in Geti v2.5
     size: Optional[int] = None  # Added in Geti v1.2
 
 

diff --git a/geti_sdk/data_models/model.py b/geti_sdk/data_models/model.py
@@ -58,6 +58,16 @@ class ModelPurgeInfo:
     user_uid: Optional[str] = None
 
 
+@attr.define
+class TrainingFramework:
+    """
+    Representation of the training framework used to train the model.
+    """
+
+    type: str
+    version: str
+
+
 @attr.define
 class OptimizationConfigurationParameter:
     """
@@ -97,6 +107,7 @@ class BaseModel:
         default=None
     )  # Added in Intel Geti 1.1
     total_disk_size: Optional[int] = None  # Added in Intel Geti 2.3
+    training_framework: Optional[TrainingFramework] = None  # Added in Intel Geti 2.5
 
     def __attrs_post_init__(self):
         """