Update app.py

app.py

@@ -1,36 +1,42 @@
 import gradio as gr
-import sys
-from huggingface_hub import snapshot_download
 import torch
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification, AutoConfig
 from PIL import Image
+from torchvision import transforms
+from transformers import AutoConfig, AutoModelForImageClassification
+import sys
+from huggingface_hub import snapshot_download
 
-# Step 1: Download the model files locally
+# ✅ Download and link custom model repo
 model_path = snapshot_download("shahad-alh/arabichar-finetuned-v2")
-
-# Step 2: Add to Python path
 sys.path.append(model_path)
 
-# Step 3: Load model and extractor
+# ✅ Load model with config
 config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
 model = AutoModelForImageClassification.from_pretrained(model_path, config=config, trust_remote_code=True)
-extractor = AutoFeatureExtractor.from_pretrained(model_path)
+model.eval()
+
+# ✅ Manual preprocessing (Grayscale → Resize → Tensor)
+transform = transforms.Compose([
+    transforms.Grayscale(num_output_channels=1),
+    transforms.Resize((32, 32)),
+    transforms.ToTensor()
+])
 
-# Step 4: Prediction logic
-def predict(img: Image.Image):
-    inputs = extractor(images=img, return_tensors="pt")
+# ✅ Prediction logic
+def predict(image: Image.Image):
+    tensor = transform(image).unsqueeze(0)  # Add batch dimension
     with torch.no_grad():
-        outputs = model(**inputs)
-        predicted = outputs.logits.argmax(-1).item()
-        label = model.config.id2label[str(predicted)]
+        logits = model(tensor).logits
+        predicted = torch.argmax(logits, dim=1).item()
+        label = config.id2label[str(predicted)]
     return label
 
-# Step 5: Gradio app
+# ✅ Gradio app
 gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", label="Upload Arabic Letter"),
-    outputs=gr.Textbox(label="Prediction"),
+    outputs=gr.Textbox(label="Predicted Letter"),
     title="Arabic Character Classifier",
-    allow_flagging="never",
-    allow_api=True
+    allow_api=True,
+    allow_flagging="never"
 ).launch()