Upload 4 files

ModelClass.py

@@ -3,22 +3,25 @@ from torch import nn
 from torchvision import transforms, models
 
 class ActionClassifier(nn.Module):
-    def __init__(self, ntargets):
+    def __init__(self, train_last_nlayer, hidden_size, dropout, ntargets):
         super().__init__()
-        resnet = models.resnet50(pretrained=True, progress=True)
+        resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT, progress=True)
         modules = list(resnet.children())[:-1] # delete last layer
+
         self.resnet = nn.Sequential(*modules)
-        for param in self.resnet.parameters():
+        for param in self.resnet[:-train_last_nlayer].parameters():
             param.requires_grad = False
+            
         self.fc = nn.Sequential(
             nn.Flatten(),
             nn.BatchNorm1d(resnet.fc.in_features),
-            nn.Dropout(0.2),
-            nn.Linear(resnet.fc.in_features, 256),
+            nn.Dropout(dropout),
+            nn.Linear(resnet.fc.in_features, hidden_size),
             nn.ReLU(),
-            nn.BatchNorm1d(256),
-            nn.Dropout(0.2),
-            nn.Linear(256, ntargets)
+            nn.BatchNorm1d(hidden_size),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size, ntargets),
+            nn.Sigmoid()
         )
     
     def forward(self, x):
@@ -27,22 +30,28 @@ class ActionClassifier(nn.Module):
         return x
 
 
-
 def get_transform():
-    transform = transforms.Compose([ 
-                transforms.Resize([224, 244]), 
-                transforms.ToTensor(),
-                # std multiply by 255 to convert img of [0, 255]
-                # to img of [0, 1]
-                transforms.Normalize((0.485, 0.456, 0.406), 
-                                     (0.229*255, 0.224*255, 0.225*255))]
-            )
+    transform = transforms.Compose([
+        transforms.Resize([224, 244]),
+        models.ResNet50_Weights.DEFAULT.transforms()
+    ])
     return transform
 
+# def get_transform():
+#     transform = transforms.Compose([ 
+#                 transforms.Resize([224, 244]), 
+#                 transforms.ToTensor(),
+#                 # std multiply by 255 to convert img of [0, 255]
+#                 # to img of [0, 1]
+#                 transforms.Normalize((0.485, 0.456, 0.406), 
+#                                      (0.229*255, 0.224*255, 0.225*255))]
+#             )
+#     return transform
+
 
 def get_model():
-    model = ActionClassifier(15)
-    model.load_state_dict(torch.load('./classifier_weights.pth', map_location=torch.device('cpu')))
+    model = ActionClassifier(0, 512, 0.2, 15)
+    model.load_state_dict(torch.load('./model_weights.pth', map_location=torch.device('cpu')))
     return model
 
 

app.py

@@ -35,7 +35,7 @@ def infer(img):
         
 
 st.set_page_config(
-    page_title="Whale Identification",
+    page_title="ActionNet",
     page_icon="🧊",
     layout="centered",
     initial_sidebar_state="expanded",
@@ -86,7 +86,7 @@ hide_st_style = """
             header {visibility: hidden;}
             </style>
             """
-#st.markdown(hide_st_style, unsafe_allow_html=True)
+st.markdown(hide_st_style, unsafe_allow_html=True)
 
 
 
@@ -129,10 +129,11 @@ def app():
             
             res = infer(image)
             prob = res.numpy()
-            idx = np.argpartition(prob, -4)[-4:]
+            idx = np.argpartition(prob, -6)[-6:]
             right_column.markdown('#### Results')
             
             idx = list(idx)
+            idx.sort(key=lambda x: prob[x].astype(float), reverse=True)
             for i in idx:
                 
                 class_name = ModelClass.get_class(i).replace('_', ' ').capitalize()

model_weights.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d300834d5794b294533827f8f7200c7f5fa29fb984fa17075ae0b87b8e4c7e6
+size 98624253