Update main.py

main.py

@@ -1,27 +1,27 @@
-import gradio as gr
+import streamlit as st
 from tensorflow import keras
 import numpy as np
 from huggingface_hub import HfApi
-import h5py
-from io import BytesIO
 
-# Authenticate and read the custom model from Hugging Face Spaces
+# Authenticate and download the custom model from Hugging Face Spaces
 hf_api = HfApi()
 model_url = hf_api.presigned_url('dhhd255', 'idk_test', filename='best_model.h5', token='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')
 r = requests.get(model_url)
-model_file = h5py.File(BytesIO(r.content), 'r')
+with open('best_model.h5', 'wb') as f:
+    f.write(r.content)
 
 # Load your custom model
-model = keras.models.load_model(model_file)
+model = keras.models.load_model('best_model.h5')
 
-def image_classifier(inp):
+# Define a function that takes an image as input and uses the model for inference
+def image_classifier(image):
     # Preprocess the input image
-    inp = np.array(inp)
-    inp = inp / 255.0
-    inp = np.expand_dims(inp, axis=0)
+    image = np.array(image)
+    image = image / 255.0
+    image = np.expand_dims(image, axis=0)
     
     # Use your custom model for inference
-    predictions = model.predict(inp)
+    predictions = model.predict(image)
     
     # Process the predictions and return the result
     result = {}
@@ -31,5 +31,8 @@ def image_classifier(inp):
     
     return result
 
-demo = gr.Interface(fn=image_classifier, inputs='image', outputs='label')
-demo.launch()
+# Create a Streamlit app with an image upload input
+image = st.file_uploader('Upload an image')
+if image is not None:
+    result = image_classifier(image)
+    st.write(result)