diff --git a/application.py b/application.py
index 2b01876..ae9ee24 100644
--- a/application.py
+++ b/application.py
@@ -11,11 +11,11 @@ def __init__(self):
 
 obj = ClientApp()
 
-st.title('Kindney Disease Classification Using Deep Learning')
+st.title('Kindney Disease :health_worker: Classification Using Deep Learning')
 st.divider()
 st.subheader("Upload your image of kidney ct-scan to cheack for any disease")
 
-uploaded_file = st.file_uploader("Choose a file")
+uploaded_file = st.file_uploader("Choose a file", type=['jpg','png','jpeg'])
 if uploaded_file is not None:
     with open(obj.filename, "wb") as f:
         f.write(uploaded_file.getbuffer())
@@ -23,8 +23,12 @@ def __init__(self):
     prediction, result_raw = obj.classifier.predict()
     if(prediction=='Failed'):
         st.write("Something went wrong")
+    elif(prediction=='Normal'):
+        st.write("You have healthy kidney")
+        st.write("You don't have any disease")
     else:
-        st.write("You have ",prediction," kidney")
+        st.write("You have unhealthy kideny")
+        st.write("You have ",prediction," in your kidney")
         df = pd.DataFrame(data=result_raw, columns=['Cyst', 'Normal', 'Stone', 'Tumor'], index=None)
         st.write("Probablity for each case : ")
         st.dataframe(df)
diff --git a/inputImage.jpg b/inputImage.jpg
index b7fa0b6..89e3694 100644
Binary files a/inputImage.jpg and b/inputImage.jpg differ
diff --git a/src/CNN_Classifier/pipeline/prediction.py b/src/CNN_Classifier/pipeline/prediction.py
index b189a19..9fc8d03 100644
--- a/src/CNN_Classifier/pipeline/prediction.py
+++ b/src/CNN_Classifier/pipeline/prediction.py
@@ -16,19 +16,20 @@ def predict(self, ):
         test_image = np.expand_dims(test_image, axis=0)
 
         result_raw = model.predict(test_image)
-        result = np.argmax(result_raw, axis=0)
+        result = np.argmax(result_raw[0], axis=0)
         print(result_raw)
+        print(result)
 
-        if result[0] == 0:
+        if result == 0:
             prediction = "Cyst"
             return (prediction, result_raw)
-        elif result[1] ==  1:
+        elif result ==  1:
             prediction = "Normal"
             return (prediction, result_raw)
-        elif result[2] == 2:
+        elif result == 2:
             prediction = "Stone"
             return (prediction, result_raw)
-        elif result[3] == 3:
+        elif result == 3:
             prediction = "Tumor"
             return (prediction, result_raw)