Update app.py

Update predict_text() with print debugging.

app.py

@@ -28,16 +28,17 @@ def decode_prediction(prediction):
     Expects prediction to be a numpy array of shape (1, 78).
     It returns the drug name corresponding to the highest probability.
     """
-    # Get the index of the highest probability class
-    predicted_index = np.argmax(prediction, axis=-1)[0]
-    # Return the corresponding drug name
-    return CLASS_NAMES[predicted_index]
+    if prediction.shape != (1, 78):
+        return "Error: Unexpected model output shape"
+
+    predicted_index = np.argmax(prediction, axis=-1)[0]  # Get the index of the highest probability
+    return CLASS_NAMES[predicted_index]  # Return the corresponding drug name
 
 # Function to preprocess the uploaded image
 def preprocess_image(image):
-    image = image.convert("RGB")  # Convert to grayscale
-    image = image.resize((64, 64))  # Resize to match model input
-    image = np.array(image) / 255.0  # Normalize
+    image = image.convert("RGB")  # Ensure 3 channels
+    image = image.resize((64, 64))  # Match model input size
+    image = np.array(image) / 255.0  # Normalize to [0,1]
     image = np.expand_dims(image, axis=0)  # Add batch dimension
     return image
 
@@ -51,48 +52,17 @@ def predict_text(image):
     segment_width = image.shape[1] // num_chars  # Split image into equal parts
 
 def predict_text(image):
-    processed_image = preprocess_image(image)  # Make sure the image is preprocessed to (64, 64, 3)
+    processed_image = preprocess_image(image)  # Ensure input is (64, 64, 3)
     prediction = model.predict(processed_image)
     
     print("Model output shape:", prediction.shape)  # Should be (1, 78)
-    print("Model output values:", prediction)        # Check the raw probabilities
+    print("Model output values:", prediction)        # Print raw probabilities
     
     # Decode the prediction to get the drug name
     predicted_text = decode_prediction(prediction)
     return predicted_text
 
 
-import numpy as np
-import string
-
-# Define the possible characters in prescription handwriting
-CHARACTER_SET = string.ascii_letters + string.digits + " .,-/()"
-
-def decode_prediction(prediction):
-    # Ensure prediction is iterable
-    if len(prediction.shape) == 2:  # (1, num_classes), meaning single character classification
-        indices = np.argmax(prediction, axis=-1)  # Pick the most likely character
-        text = CHARACTER_SET[indices[0]]  # Convert to actual character
-        return text
-
-    elif len(prediction.shape) == 3:  # (1, sequence_length, num_classes), meaning character sequence classification
-        prediction = prediction[0]  # Remove batch dimension
-        indices = np.argmax(prediction, axis=-1)  # Get character indices at each step
-        
-        # Convert indices to characters while removing duplicates
-        decoded_text = []
-        prev_char = None
-        for i in indices:
-            if i != prev_char and i < len(CHARACTER_SET):  # Avoid duplicate characters
-                decoded_text.append(CHARACTER_SET[i])
-            prev_char = i  # Update previous character
-        
-        return "".join(decoded_text)
-
-    else:
-        return "Error: Unexpected output shape!"
-
-
 # Gradio UI
 interface = gr.Interface(
     fn=predict_text,