Upload app.py

app.py

@@ -5,30 +5,25 @@ import tensorflow as tf
 from PIL import Image
 import zipfile
 
-# Path to your model zip file in the Hugging Face Space
-MODEL_PATH = '/app/your_trained_model.keras.zip'  # Adjust this path for Hugging Face
+# Path to your zipped model file
+ZIP_MODEL_PATH = '/app/your_trained_model.keras.zip'  # Correct this path if needed
+UNZIPPED_MODEL_PATH = '/app/your_trained_model.keras'  # Path where the model will be extracted
 
-# Check if the model file exists and load the model
-model = None  # Initialize model variable outside the try block
-if os.path.exists(MODEL_PATH):
-    try:
-        # Unzip the model file in the Hugging Face Space directory if needed
-        with zipfile.ZipFile(MODEL_PATH, 'r') as zip_ref:
-            zip_ref.extractall("/app")  # Adjust the extraction directory in Hugging Face Space
-        
-        # Now load the model
-        model = tf.keras.models.load_model("/app/your_trained_model.keras")  # Adjust to the unzipped path
-        print("Model loaded successfully!")
-    except Exception as e:
-        print(f"Error loading model: {e}")
-else:
-    print(f"Model file not found at {MODEL_PATH}")
+# Unzip the model if it hasn't been unzipped already
+if not os.path.exists(UNZIPPED_MODEL_PATH):
+    with zipfile.ZipFile(ZIP_MODEL_PATH, 'r') as zip_ref:
+        zip_ref.extractall('/app')
+        print(f"Model unzipped to {UNZIPPED_MODEL_PATH}")
 
-# If model is not loaded, return an error in the prediction function
-def predict_decoration(image: Image.Image):
-    if model is None:
-        raise ValueError("Model is not loaded, cannot make predictions.")
+# Load the model
+try:
+    model = tf.keras.models.load_model(UNZIPPED_MODEL_PATH)
+    print("Model loaded successfully!")
+except Exception as e:
+    print(f"Error loading model: {e}")
 
+# Define the function to predict decoration
+def predict_decoration(image: Image.Image):
     # Preprocess the image to match the model input format
     image = image.resize((224, 224))  # Resize to match model's expected input size
     image_array = np.array(image) / 255.0  # Normalize the image to [0, 1]
@@ -96,11 +91,8 @@ with tab1:
         st.image(image, caption="Uploaded Tree Image", use_container_width=True)
 
         if st.button("Predict Decoration"):
-            try:
-                prediction = predict_decoration(image)
-                st.write(f"Prediction: {prediction}")
-            except ValueError as e:
-                st.error(str(e))
+            prediction = predict_decoration(image)
+            st.write(f"Prediction: {prediction}")
 
 # Tree Image URLs Tab
 with tab2: