Update app.py

app.py

@@ -13,39 +13,32 @@ def load_model():
 model = load_model()
 
 # Function to apply transforms to the image (update as per your model's requirement)
-def get_top_predictions(prediction, class_indices, top_k=5):
-    probabilities = torch.nn.functional.softmax(prediction, dim=1)
-    top_probs, top_catids = torch.topk(probabilities, top_k)
-    top_classes = [class_indices[str(catid)] for catid in top_catids[0].tolist()]
-    top_probs = top_probs[0].tolist()
-    return top_classes, top_probs
+def transform_image(image):
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),  # Resize to the input size that your model expects
+        transforms.ToTensor(),
+        # Add other transformations as needed
+    ])
+    return transform(image).unsqueeze(0)  # Add batch dimension
 
-# Function to display the prediction results
-def display_prediction(top_classes, top_probs):
-    for i, (cls, prob) in enumerate(zip(top_classes, top_probs)):
-        st.write(f"{i+1}: {cls} - {prob:.2%}")
+st.title("Animal Facial Expression Recognition")
 
-# Load the class indices (you will need to have a class_indices.json file)
-class_indices = load_class_indices()
 
-# Streamlit interface
-st.title("Petn Emotion Recognition")
-st.write("Upload a pet facial image.")
 
 # File uploader
-uploaded_file = st.file_uploader("", type=["jpg", "png", "jpeg"])
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
 if uploaded_file is not None:
-    # Display the uploaded image
     image = Image.open(uploaded_file).convert('RGB')
-    st.image(image, use_column_width=True)
-    
-    # Classify button
-    if st.button('Classify Image'):
-        st.write("Classifying...")
-        # Transform the image and make prediction
-        input_tensor = transform_image(image)
-        with torch.no_grad():
-            prediction = model(input_tensor)
-        # Get and display top predictions
-        top_classes, top_probs = get_top_predictions(prediction, class_indices)
-        display_prediction(top_classes, top_probs)
+    st.image(image, caption='Uploaded Image.', use_column_width=True)
+    st.write("")
+    st.write("Classifying...")
+
+    # Transform the image
+    input_tensor = transform_image(image)
+
+    # Make prediction
+    with torch.no_grad():
+        prediction = model(input_tensor)
+
+    # Display the prediction (modify as per your output)
+    st.write('Predicted class:', prediction.argmax().item())