feat: Added new TV models to demo

frgfm · Sep 16, 2022 · f6ed0b5 · f6ed0b5
1 parent c9d244a
commit f6ed0b5
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/demo/app.py b/demo/app.py
@@ -18,7 +18,15 @@
 from torchcam.utils import overlay_mask
 
 CAM_METHODS = ["CAM", "GradCAM", "GradCAMpp", "SmoothGradCAMpp", "ScoreCAM", "SSCAM", "ISCAM", "XGradCAM", "LayerCAM"]
-TV_MODELS = ["resnet18", "resnet50", "mobilenet_v2", "mobilenet_v3_small", "mobilenet_v3_large"]
+TV_MODELS = [
+    "resnet18",
+    "resnet50",
+    "mobilenet_v3_small",
+    "mobilenet_v3_large",
+    "regnet_y_400mf",
+    "convnext_tiny",
+    "convnext_small",
+]
 LABEL_MAP = requests.get(
     "https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json"
 ).json()

diff --git a/pyproject.toml b/pyproject.toml
@@ -65,7 +65,7 @@ docs = [
 ]
 demo = [
     "streamlit>=0.65.0,<2.0.0",
-    "torchvision>=0.8.0,<1.0.0",
+    "torchvision>=0.12.0,<1.0.0",
 ]
 dev = [
     # test