Update app.py

app.py

@@ -23,7 +23,7 @@ def preprocess_image(image):
 def predict_image(image):
     # Save the image to a file-like object
     image_bytes = io.BytesIO()
-    image.save(image_bytes, format="jpg")
+    image.save(image_bytes, format="jpeg")
     image_bytes.seek(0)  # Reset file pointer to start
 
     # Load the image from the file-like object
@@ -32,9 +32,16 @@ def predict_image(image):
     img_array = preprocess_image(image)
     outputs = model.predict(img_array)
     predictions = tf.nn.softmax(outputs.logits, axis=-1)
-    predicted_class = np.argmax(predictions)
-    #confidence = float(np.max(predictions))
-    return{"prediction": predicted_class}
+    predicted_class_index = np.argmax(predictions)
+    
+    # Map class index to label
+    if predicted_class_index == 0:
+        predicted_label = "Clean"
+    else:
+        predicted_label = "Carries"
+    
+    return {"prediction": predicted_label}
+
     #return json.dumps(prediction_dict, indent=1)
 
 # Create the interface