Upload README.md with huggingface_hub

README.md

@@ -6,12 +6,12 @@ tags:
   - angle-detection
   - circular-gaussian-distribution
   - mambaout
+  - pytorch
 datasets:
   - coco
 metrics:
   - mae
 pipeline_tag: image-classification
-library_name: pytorch
 ---
 
 # Image Rotation Angle Estimation
@@ -31,12 +31,18 @@ The model outputs a probability distribution over 360 angle bins (1 degree resol
 
 ## Usage
 
+The inference code (`model_cgd.py`, `architectures.py`, `rotation_utils.py`) is included in this repo.
+
 ```python
-import torch
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, snapshot_download
 from PIL import Image
 
-# Download checkpoint
+# Download inference code and checkpoint
+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",
@@ -45,7 +51,7 @@ ckpt_path = hf_hub_download(
 # Load model
 from model_cgd import CGDAngleEstimation
 
-model = CGDAngleEstimation.load_from_checkpoint(ckpt_path, image_size=224)
+model = CGDAngleEstimation.try_load(checkpoint_path=ckpt_path)
 model.eval()
 
 # Predict rotation angle