git commit -m "app.py"

app.py

@@ -1,20 +1,26 @@
 import torch
 from torchvision import transforms
+from torchvision.models import resnet18
 from PIL import Image
 import gradio as gr
 
-# Allow safe loading of torchvision ResNet class
-from torchvision.models import resnet18
-torch.serialization.add_safe_globals([resnet18])
-
 # Class labels
 class_names = ['Nu.1', 'Nu.10', 'Nu.100', 'Nu.1000', 'Nu.20', 'Nu.5', 'Nu.50', 'Nu.500']
 
-# Force CPU (Hugging Face Spaces do not support GPU)
+# Force CPU
 device = torch.device('cpu')
 
-# Load full model (code-safe if trusted model)
-model = torch.load("currency_model.pth", map_location=device, weights_only=False)
+# Step 1: Define model architecture
+model = resnet18(pretrained=False)
+
+# Step 2: Modify final layer (assuming 8 classes)
+model.fc = torch.nn.Linear(model.fc.in_features, len(class_names))
+
+# Step 3: Load weights
+model.load_state_dict(torch.load("currency_model.pth", map_location=device))
+
+# Step 4: Set to eval mode
+model.to(device)
 model.eval()
 
 # Image transform
@@ -23,7 +29,7 @@ transform = transforms.Compose([
     transforms.ToTensor(),
 ])
 
-# Prediction logic
+# Prediction function
 def predict(image):
     image = image.convert("RGB")
     image = transform(image).unsqueeze(0).to(device)
@@ -33,7 +39,7 @@ def predict(image):
         _, predicted = torch.max(outputs, 1)
         return class_names[predicted.item()]
 
-# Gradio Interface
+# Gradio interface
 interface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
@@ -42,4 +48,4 @@ interface = gr.Interface(
     description="Upload a currency note image to identify its value."
 )
 
-interface.launch() 
+interface.launch()