Upload README.md with huggingface_hub

README.md

@@ -34,27 +34,27 @@ The model outputs a probability distribution over 360 angle bins (1 degree resol
 The inference code (`model_cgd.py`, `architectures.py`, `rotation_utils.py`) is included in this repo.
 
 ```python
-from huggingface_hub import hf_hub_download, snapshot_download
-from PIL import Image
+from huggingface_hub import snapshot_download
 
-# Download inference code and checkpoint
+# Download inference code
 snapshot_download(
     repo_id="maxwoe/image-rotation-angle-estimation",
     allow_patterns=["*.py", "*.json"],
     local_dir=".",
 )
-ckpt_path = hf_hub_download(
-    repo_id="maxwoe/image-rotation-angle-estimation",
-    filename="cgd_mambaout_base_coco2017.ckpt",
-)
 
-# Load model
+# Load model (defaults to COCO 2017 checkpoint)
 from model_cgd import CGDAngleEstimation
+model = CGDAngleEstimation.from_pretrained("maxwoe/image-rotation-angle-estimation")
 
-model = CGDAngleEstimation.try_load(checkpoint_path=ckpt_path)
-model.eval()
+# Or load a specific checkpoint by filename
+# model = CGDAngleEstimation.from_pretrained(
+#     "maxwoe/image-rotation-angle-estimation",
+#     model_name="cgd_mambaout_base_coco2014.ckpt",
+# )
 
 # Predict rotation angle
+from PIL import Image
 image = Image.open("your_image.jpg")
 angle = model.predict_angle(image)
 print(f"Predicted rotation: {angle:.2f} degrees")