Update app.py

app.py

@@ -74,7 +74,16 @@ if uploaded_file is not None:
     # Preprocess the image and perform inference
     image_tensor, final_image = preprocess_image(image, device)
     predicted_class, probabilities = predict(image_tensor, model, device)
-    
+    # {'Dry': 0, 'Normal': 1, 'Oily': 2}
     st.write(f"Predicted class: {predicted_class.item()}")
+
+    if predicted_class.item()==2:
+        st.write(f"Oily Skin : {probabilities[2]}")
+    elif predicted_class.item()==1:
+        st.write("Normal Skin : {probabilities[1]}")
+    elif predicted_class.item()==0:
+        st.write("Dry Skin : {probabilities[0]}") 
+    st.write(f"probabilities: {probabilities}")
+    
     # Display the final processed image
     st.image(final_image, caption='Processed Image', use_column_width=True)