Update app.py

app.py

@@ -12,10 +12,10 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # List of available model files
 MODEL_FILES = {
-    "lenet": "grayscale_lenet_state_dict.pt",
-    "cnn": "grayscale_custom_CNN_state_dict.pt",
-    "resnet": "grayscale_resnet_state_dict.pt",
-    "vgg": "grayscale_vgg_state_dict.pt"
+    "LeNet": "grayscale_lenet_state_dict.pt",
+    "CNN": "grayscale_custom_CNN_state_dict.pt",
+    "ResNet": "grayscale_resnet_state_dict.pt",
+    "VGG": "grayscale_vgg_state_dict.pt"
 }
 
 # Replace with your actual class names
@@ -157,14 +157,14 @@ def load_model(model_choice):
         raise FileNotFoundError(f"Model file {model_path} not found.")
     
 
-    if "cnn" in model_choice:
+    if "CNN" in model_choice:
         # Load custom model
         model = HandwrittenTextCNN()
         
-    elif "lenet" in model_choice:
+    elif "LeNet" in model_choice:
         model = LeNet5()
     
-    elif "vgg" in model_choice:
+    elif "VGG" in model_choice:
         model = torch.hub.load('pytorch/vision:v0.10.0','vgg11', pretrained=False)
         model.features[0] = nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
         model.classifier[-1] = nn.Linear(in_features=4096, out_features=68, bias=True)    
@@ -200,7 +200,7 @@ def predict(model_choice, image):
             predicted = torch.argmax(outputs,dim=1)
             predicted_class = class_names[f"{predicted.item():02d}"]
         
-        return f"Predicted class: {predicted_class}"
+        return f"{predicted_class}"
     
     except Exception as e:
         return f"Error: {str(e)}"
@@ -213,7 +213,7 @@ iface = gr.Interface(
         gr.Image(type="pil", label="Upload Image")
     ],
     outputs="text",
-    title="Image Classification with PyTorch Models",
+    title="Burapha-TH Character dataset classification",
     description="Select a custom or pre-trained model and upload an image to get a classification prediction."
 )