Upload registry.py

registry.py

@@ -0,0 +1,17 @@
+from typing import Callable, Dict, Any, List, Optional
+
+class ModelEntry:
+    def __init__(self, model: Any, preprocess: Callable, postprocess: Callable, class_names: List[str],
+                 display_name: Optional[str] = None, contributor: Optional[str] = None, model_path: Optional[str] = None):
+        self.model = model
+        self.preprocess = preprocess
+        self.postprocess = postprocess
+        self.class_names = class_names
+        self.display_name = display_name
+        self.contributor = contributor
+        self.model_path = model_path
+
+MODEL_REGISTRY: Dict[str, ModelEntry] = {}
+
+def register_model(model_id: str, model: Any, preprocess: Callable, postprocess: Callable, class_names: List[str]):
+    MODEL_REGISTRY[model_id] = ModelEntry(model, preprocess, postprocess, class_names)