Update app.py

app.py

@@ -178,13 +178,13 @@ optimizer=torch.optim.Adam(model.parameters())
 
 
 
-checkpoint = torch.load(r'model/final.pth', map_location=device)
-# Restore states
-#print("Checkpoint keys:", checkpoint.keys())
-model.load_state_dict(checkpoint['model_state_dict'])
+# checkpoint = torch.load(r'model/final.pth', map_location=device)
+# # Restore states
+# #print("Checkpoint keys:", checkpoint.keys())
+# model.load_state_dict(checkpoint['model_state_dict'])
 
-optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
-model.eval()
+# optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+# model.eval()
 characters = string.ascii_letters + string.digits
 idx_to_char = {idx: char for idx, char in enumerate(characters)}
 
@@ -194,26 +194,14 @@ def to_text(arr):
     for c in arr:
         ans=ans+idx_to_char[c.item()]
     return ans
-
-    
 def predict_captcha(image):
     try:
         if image is None:
             return "No image provided"
 
-        # Handle Gradio's FileData input: dict with 'data' and 'meta'
-        if isinstance(image, dict) and 'data' in image:
-            image = image['data']
-
-        # Convert to PIL.Image
-        if isinstance(image, str) and image.startswith('data:image'):
-            import base64
-            from io import BytesIO
-            image_data = base64.b64decode(image.split(',')[1])
-            image = Image.open(BytesIO(image_data))
-        elif not isinstance(image, Image.Image):
-            from io import BytesIO
-            image = Image.open(BytesIO(image))
+        # Directly use the image as a PIL.Image
+        if not isinstance(image, Image.Image):
+            return "Invalid image format"
 
         # Process image
         transform = transforms.Compose([
@@ -236,6 +224,11 @@ def predict_captcha(image):
         print(f"Error details: {str(e)}")
         return f"Error processing image: {str(e)}"
 
+# Ensure model is loaded and set to eval mode
+checkpoint = torch.load(r'model/final.pth', map_location=device)
+model.load_state_dict(checkpoint['model_state_dict'])
+model.eval()  # Set the model to evaluation mode after loading the state
+
 
 # Update Gradio interface
 iface = gr.Interface(