Update app.py

app.py

@@ -2,8 +2,6 @@ import torch
 import gradio as gr
 from torch import nn
 import torch.nn.functional as F
-#from torchsummary import summary
-from tqdm import tqdm
 import os
 from torchvision import transforms
 from torch.utils.data import DataLoader,random_split,Dataset
@@ -132,7 +130,7 @@ class CRNN(nn.Module):
         outputs = torch.stack(outputs, dim=1)  # Shape: (batch_size, 5, vocab_size)
 
         return outputs
-#output=CRNN(64,CRNN_KERNEL,CRNN_POOL_KERNEL,CRNN_DROPOUT,CRNN_LATENT,LSTM_HIDDEN_DIM,VOCAB_SIZE,OUTPUT_LENGTH).to(device)(torch.zeros((2,64,256,256)).to(device))
+output=CRNN(64,CRNN_KERNEL,CRNN_POOL_KERNEL,CRNN_DROPOUT,CRNN_LATENT,LSTM_HIDDEN_DIM,VOCAB_SIZE,OUTPUT_LENGTH).to(device)(torch.zeros((2,64,256,256)).to(device))
 
 
 class CaptchaCrackNet(nn.Module):
@@ -178,13 +176,6 @@ 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'])
-
-# 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)}
 
@@ -199,9 +190,19 @@ def predict_captcha(image):
         if image is None:
             return "No image provided"
 
-        # Directly use the image as a PIL.Image
-        if not isinstance(image, Image.Image):
-            return "Invalid image format"
+        # 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))
 
         # Process image
         transform = transforms.Compose([
@@ -223,12 +224,15 @@ def predict_captcha(image):
     except Exception as e:
         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)
+# Restore states
+print("Checkpoint keys:", checkpoint.keys())
 model.load_state_dict(checkpoint['model_state_dict'])
-model.eval()  # Set the model to evaluation mode after loading the state
 
+optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+model.eval()
 
 # Update Gradio interface
 iface = gr.Interface(