Update README.md

chore(): update the model usage example

README.md

@@ -35,19 +35,51 @@ You can load and run this model **directly in PyTorch** **without** installing `
 
 ```python
 import torch
+import torchvision.transforms.v2 as v2
+from huggingface_hub import hf_hub_download
+from PIL import Image
 
-# 1. Load the model
-model = torch.jit.load("cosmoformer_traced.pt")
+label_mapping = {
+    0: 'barred_spiral',
+    1: 'edge_on_disk',
+    2: 'featured_without_bar_or_spiral',
+    3: 'irregular',
+    4: 'smooth_cigar',
+    5: 'smooth_inbetween',
+    6: 'smooth_round',
+    7: 'unbarred_spiral'
+}
+
+# 1. Define the path to the hugging face repo
+ts_path = hf_hub_download(
+    repo_id="artursultanov/cosmoformer-model",
+    filename="cosmoformer_traced_cpu.pt"
+)
+
+# 2. Load the model from the hugging face repo
+model = torch.jit.load(ts_path, map_location="cpu")
 model.eval()
 
-# 2. Inference
-# Suppose you have a 3-channel image tensor (1, 3, 224, 224)
-dummy_input = torch.randn(1, 3, 224, 224)
+# 3. Define image transform to match model's internal representation 
+transform = v2.Compose([
+    v2.Resize((224, 224)),
+    v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])
+])
+
+# 4. Load the image
+image_path = "test_image.jpg"
+image = Image.open(image_path).convert("RGB")
+
+tensor = transform(image)  # shape [3, 224, 224]
+tensor = tensor.unsqueeze(0)  shape [1, 3, 224, 224]
 
+# 5. Inference
 with torch.no_grad():
-    outputs = model(dummy_input)
+    output = model(tensor)
+    predicted_idx = torch.argmax(output, dim=1).item()
 
-print(outputs.shape)  # e.g., [1, num_classes]
+predicted_label = label_mapping[predicted_idx]
+print("Predicted class:", predicted_label)
 ```
 
 ```