Update README.md

README.md

@@ -64,7 +64,42 @@ Users should carefully evaluate the model on their specific data to ensure compa
 Use the following code snippet to get started:
 
 ```python
-Coming soon
+from huggingface_hub import snapshot_download
+from transformers import AutoImageProcessor, AutoConfig, SwinForImageClassification
+from torchvision import transforms
+import torch
+
+# Download the model
+snapshot_download(repo_id='stefanodangelo/chartdet', local_dir='./', allow_patterns='*.pt')
+
+# Load the model and processor
+processor = AutoImageProcessor.from_pretrained("microsoft/swin-large-patch4-window7-224")
+config = AutoConfig.from_pretrained("microsoft/swin-large-patch4-window7-224", num_labels=2)
+model = SwinForImageClassification(config)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.load_state_dict(torch.load("models/chart_detection_model.pt", map_location=torch.device(device)))
+
+# Define transformations
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=processor.image_mean, std=processor.image_std)
+])
+
+id2class = {0: "Picture", 1: "Chart"}
+
+# Prepare an image
+from PIL import Image
+image = Image.open("YOUR_IMG_PATH")
+image_tensor = transform(image).unsqueeze(0).to(device)  # Add batch dimension and move to device
+
+# Make prediction
+with torch.no_grad():
+    outputs = model(image_tensor).logits  # Get logits from the model
+    predicted_class = torch.argmax(outputs, dim=1).item()  # Get the class index
+
+print(f"Predicted class: {id2class[predicted_class]}")
 ```
 
 ## Training Details
@@ -152,8 +187,9 @@ The model uses a Swin Transformer-based architecture adapted for binary image cl
 
 - Windows 11
 - Python 3.11
+- CUDA 11.7
 - HuggingFace Transformers Library
-- PyTorch
+- PyTorch & TorchVision
 
 ## Citation