Update README.md

README.md

@@ -73,25 +73,12 @@ You can download the model and load it using `torch`.
 import torch
 from huggingface_hub import hf_hub_download
 
-# Download the model
-model_path = hf_hub_download("Ayamohamed/DiaNoneclassi", "diagram_classifier_full.pth")
+# Download model from Hugging Face Hub
+model_path = hf_hub_download(repo_id="Ayamohamed/DiaClassification", filename="dia_none_classifier_full.pth")
 
-# Load the model
-from torchvision import models
-import torch.nn as nn
-
-class DiagramClassifier(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.model = models.resnet18(pretrained=False)
-        self.model.fc = nn.Linear(self.model.fc.in_features, 2)
-
-    def forward(self, x):
-        return self.model(x)
-
-model = DiagramClassifier()
-model.load_state_dict(torch.load(model_path, map_location="cpu"))
-model.eval()  # Set to evaluation mode
+# Load model
+model_hg = torch.load(model_path)
+model_hg.eval()  # Set to evaluation mode
 
 ```
 ### **2️⃣ Preprocess and Classify an Image**
@@ -105,20 +92,18 @@ transform = transforms.Compose([
     transforms.ToTensor(),
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
-
-def classify_image(image_path):
-    image = Image.open(image_path).convert("RGB")
-    image = transform(image).unsqueeze(0)  # Add batch dimension
-
+def predict(image_path):
+    image = Image.open(image_path).convert("RGB")  
+    image = transform(image).unsqueeze(0) 
     with torch.no_grad():
-        output = model(image)
-        predicted_class = output.argmax(1).item()
+        output = model_hg(image)
+        class_idx = torch.argmax(output, dim=1).item()
+
+    return "Diagram" if class_idx == 0 else "Not Diagram"
 
-    labels = ["diagram", "none"]
-    return labels[predicted_class]
+# Example usage
+print(predict("my-diagram-classifier/31188_1536932698.jpg"))
 
-# Test
-print(classify_image("test_image.jpg"))
 
 ```