Update app.py

app.py

@@ -221,7 +221,7 @@ def res2net152_v1b_26w_4s(pretrained=False, **kwargs):
 
 class mutil_model(nn.Module):
 
-    def __init__(self, category_num=7):
+    def __init__(self, category_num=10):
         super(mutil_model, self).__init__()
         self.model1 = res2net50_v1b_26w_4s(pretrained=False)
         self.model1.fc = nn.Sequential(
@@ -241,39 +241,40 @@ class mutil_model(nn.Module):
         return x
 
 
-pth_path = './res2net_model_6_new.pt'
-category_num = 6
+pth_path = './res2net_pretrain_model_999.pt'
+category_num = 9
 
 # "cuda" only when GPUs are available.
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Initialize a model, and put it on the device specified.
+#device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cpu"
+#Initialize a model, and put it on the device specified.
 # 导入res2net预训练模型
-# pthfile = '/cbd_lixiaogang_lixianneng/morror_art/pre_train_model/res2net50_v1b.pth'
-# model = res2net50_v1b_26w_4s(pretrained=False)
+# pthfile = './res2net50_v1b.pth'
+model = res2net50_v1b_26w_4s(pretrained=False)
 # 修改全连接层，输出维度为预测 分类
-# num_ftrs = model.fc.in_features
-#     model.fc = nn.Sequential(
-#         nn.Linear(in_features=2048, out_features=1000, bias=True),
-#         nn.Dropout(0.5),
-#         nn.Linear(1000, out_features=category_num)
-#     )
-# model.fc = nn.Sequential(
-#    nn.Linear(in_features=2048, out_features=category_num, bias=True),
-# )
-# model = model.to(device)
-# model.device = device
-# model.load_state_dict(torch.load(pth_path,torch.device('cpu')))
-# model.eval()
+num_ftrs = model.fc.in_features
+model.fc = nn.Sequential(
+        nn.Linear(in_features=2048, out_features=1000, bias=True),
+        nn.Dropout(0.5),
+        nn.Linear(1000, out_features=category_num)
+    )
+model.fc = nn.Sequential(
+   nn.Linear(in_features=2048, out_features=category_num, bias=True),
+)
+
+model = model.to(device)
+model.device = device
+model.load_state_dict(torch.load(pth_path,torch.device('cpu')))
+model.eval()
 
 
 # 增加人脸识别模型
-model = mutil_model(category_num=7)
-model_state = torch.load('./res2net_model_6_new.pt', map_location=torch.device('cpu')).state_dict()
-model.load_state_dict(model_state)  # 加载模型参数
-model.eval()
+#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 = ['伤感', '开心', '励志', '宣泄', '平静', '感人']
+labels = ['中国风', '古典', '电子', '摇滚', '乡村', '说唱', '民谣', '动漫', '现代']
 
 import requests
 import torch
@@ -327,4 +328,5 @@ gr.Interface(
     # inputs='image',
     # outputs='label',
     # examples=[["images/cheetah1.jpg"], ["images/lion.jpg"]],
-).launch(debug=True, share=True)
+).launch(share=True)
+# share=True