Update app.py

app.py

@@ -213,6 +213,27 @@ def res2net152_v1b_26w_4s(pretrained=False, **kwargs):
         model.load_state_dict(model_zoo.load_url(model_urls['res2net152_v1b_26w_4s']))
     return model
     
+    
+class mutil_model(nn.Module):
+
+    def __init__(self,category_num = 7):
+        super(mutil_model, self).__init__()
+        self.model1 = res2net50_v1b_26w_4s(pretrained=False)
+        self.model1.fc = nn.Sequential(
+    nn.Linear(in_features=2048, out_features=category_num, bias=True),
+) 
+        self.model2 = torch.load('./affectnet_emotions/'+'enet_b2_8'+'.pt',map_location=torch.device('cpu'))
+        self.model2.classifier = nn.Sequential(
+        nn.Linear(in_features=1408, out_features=category_num, bias=True),
+) 
+        self.fc = nn.Linear(in_features=category_num*2, out_features=category_num, bias=True)
+    def forward(self, x):
+        x1 = self.model1(x)
+        x2 = self.model2(x)
+        x = torch.cat((x1,x2),1)
+        x = self.fc(x)
+        return x
+            
 pth_path = './add_face_emotion_model_7.pt'
 category_num = 7
 
@@ -241,7 +262,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 #增加人脸识别模型
-model=torch.load(pth_path,map_location=torch.device('cpu'))
+model = mutil_model(category_num = 7)
+model_state = torch.load('./add_face_emotion_model_7.pt',map_location=torch.device('cpu')).state_dict()
+model.load_state_dict(model_state)  # 加载模型参数
 model.eval()
 
 labels = ['怀旧','伤感','快乐','激励','清新','浪漫','思念']