Update app.py

app.py

@@ -14,6 +14,27 @@ import re
 import matplotlib.pyplot as plt
 from transformers import AutoConfig
 from PIL import Image
+from facenet_pytorch import InceptionResnetV1
+vggface = InceptionResnetV1(pretrained='vggface2')
+
+
+class classifier_vggface(nn.Module):
+  def __init__(self):
+    super(classifier_vggface, self).__init__()
+    self.encoder= vggface
+    self.classifier= nn.Sequential(
+        nn.Linear(512, 512),
+        nn.BatchNorm1d(512),
+        nn.ReLU(inplace=True),
+        nn.Linear(512, 128),
+        nn.BatchNorm1d(128),
+        nn.ReLU(inplace=True),
+        nn.Linear(128, 2),
+    )
+  def forward(self, x):
+    x= self.encoder(x)
+    x= self.classifier(x)
+    return x
 
 
 # Set title
@@ -31,15 +52,17 @@ btn = st.button("Predict")
 # Create a function to load model
 @st.cache(allow_output_mutation=True)
 def load_model():
-    model_path = './model_1_updated.pth'  # Update with the correct file name
+    model1= classifier_vggface()
     
+    model_path = './model1.pt'  # Update with the correct file name
+    model1.load_state_dict(torch.load(model_path))
     # Load the model configuration
-    config = AutoConfig.from_pretrained(model_path)
+    # config = AutoConfig.from_pretrained(model_path)
     
     # Load the model and tokenizer from the specified file path
-    model = AutoModelForSequenceClassification.from_pretrained(model_path, config = config)
+    # model = AutoModelForSequenceClassification.from_pretrained(model_path, config = config)
     tokenizer = AutoTokenizer.from_pretrained(model_path)
-    return model, tokenizer
+    return model1, tokenizer
 
 # Create a function to predict the input image file using the loaded model
 def predict(image_file):