Update app.py

app.py

@@ -4,42 +4,51 @@ from PIL import Image
 from tensorflow.keras.models import load_model
 from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
 import joblib
-import gdown
+from huggingface_hub import hf_hub_url, cached_download
 
-# Google Drive model URLs
-KNN_MODEL_URL = 'https://drive.google.com/uc?id=1TJ0KbzFw-2NfuJf67xvp-32uaYLIqpj3'
-EXTRACTOR_URL = 'https://drive.google.com/uc?id=1HR2Qc8Fji6RzbtG_K_sqSoiG0AQnvyZa'
+# Replace with your Space name (from the link)
+SPACE_NAME = "engrharis/Throat_Image_Classifier"
 
-# Download the model files
-st.write("Downloading models...")
-gdown.download(KNN_MODEL_URL, 'knn_pharyngitis_model.pkl', quiet=False)
-gdown.download(EXTRACTOR_URL, 'mobilenetv2_feature_extractor.h5', quiet=False)
-st.write("Models downloaded successfully!")
+# Assuming the filenames are the same as before 
+KNN_MODEL_FILE = "knn_pharyngitis_model.pkl"
+EXTRACTOR_FILE = "mobilenetv2_feature_extractor.h5"
+
+
+def download_models(url, filename):
+    """Downloads model files from Hugging Face space if not cached locally."""
+    model_path = hf_hub_url(SPACE_NAME, filename=filename)
+    if not cached_download(model_path):
+        st.write(f"Downloading {filename}...")
+        cached_download(model_path)
+        st.write(f"{filename} downloaded successfully!")
+
+
+# Load the saved models (download if not cached)
+download_models(SPACE_NAME, KNN_MODEL_FILE)
+download_models(SPACE_NAME, EXTRACTOR_FILE)
+
+knn = joblib.load(KNN_MODEL_FILE)
+feature_extractor = load_model(EXTRACTOR_FILE)
 
-# Load the saved models
-knn = joblib.load('knn_pharyngitis_model.pkl')
-feature_extractor = load_model('mobilenetv2_feature_extractor.h5')
 
-# Function to preprocess the uploaded image
 def preprocess_image(image):
     img = image.resize((224, 224))  # Resize to match MobileNetV2 input size
     img_array = np.array(img)
     img_array = preprocess_input(img_array)  # Apply MobileNetV2 preprocessing
     return np.expand_dims(img_array, axis=0)
 
-# Function to classify the image
+
 def classify_image(image):
     processed_image = preprocess_image(image)
     features = feature_extractor.predict(processed_image)
     prediction = knn.predict(features)
     return "Pharyngitis" if prediction[0] == 1 else "No Pharyngitis"
 
+
 # Streamlit app UI
 st.title("Pharyngitis Classification App")
 st.write("Upload an image to classify it as 'Pharyngitis' or 'No Pharyngitis'.")
-
 uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
-
 if uploaded_file is not None:
     # Load the uploaded image
     image = Image.open(uploaded_file)
@@ -48,4 +57,4 @@ if uploaded_file is not None:
     # Classify the image
     st.write("Classifying...")
     prediction = classify_image(image)
-    st.write(f"Prediction: **{prediction}**")
+    st.write(f"Prediction: **{prediction}**")