Update app.py

app.py

@@ -1,14 +1,26 @@
 import streamlit as st
 import pickle
 from Prediction import  Prediction
-from transformers import TFAutoModel
-
+from huggingface_hub import hf_hub_download
 
 
 # resnet_model = pickle.load(open('models/ResNet01.pkl','rb'))
 # cnn_model = pickle.load(open('models/CNNModel2.pkl','rb'))
 # inc_model = pickle.load(open('models/Inception01.pkl','rb'))
-resnet_model = TFAutoModel.from_pretrained("yashpat85/ResNet01")
+# resnet_model = TFAutoModel.from_pretrained("yashpat85/ResNet01")
+
+REPO_ID1 = "yashpat85/ResNet01"
+MODEL_DIR1 = hf_hub_download(repo_id=REPO_ID1, filename="ResNet01.pkl", repo_type="model")
+
+# REPO_ID2= "yashpat85/Inception01"
+# MODEL_DIR2 = hf_hub_download(repo_id=REPO_ID2, filename="Inception01.pkl", repo_type="model")
+
+REPO_ID3 = "yashpat85/CNNModel2"
+MODEL_DIR3 = hf_hub_download(repo_id=REPO_ID3, filename="CNNModel2.pkl", repo_type="model")
+
+
+resnet_model = pickle.load(open(MODEL_DIR1,'rb'))
+cnn_model = pickle.load(open(MODEL_DIR3,'rb'))
 
 
 def show_error_popup(message):
@@ -42,8 +54,8 @@ if uploaded_file is not None:
             p = pm.predict_image(cnn_model, image_data)
         elif option=="ResNet":
             p = pm.predict_image(resnet_model, image_data)
-        elif option=="InceptionNet":
-            p = pm.predict_image(inc_model, image_data)
+        # elif option=="InceptionNet":
+        #     p = pm.predict_image(inc_model, image_data)
         else:
             p = "Other Models are still under training due to overfitting"