codeby-hp/human-pose-classification

1"""Hugging Face utilities for downloading ML models."""
2
3import os
4import logging
5from transformers import AutoImageProcessor, AutoModelForImageClassification
6from huggingface_hub.utils import RepositoryNotFoundError
7
8logger = logging.getLogger(__name__)
9
10HF_MODEL_ID = "codeby-hp/finetune-VIT-HumanPoseClassification"
11
12
13def download_model_from_huggingface(local_path: str) -> bool:
14    """Download model from Hugging Face Hub.
15    
16    Args:
17        local_path: Local directory path to save model
18        
19    Returns:
20        True if successful, False otherwise
21    """
22    try:
23        logger.info(f"Downloading model from Hugging Face: {HF_MODEL_ID}")
24        os.makedirs(local_path, exist_ok=True)
25        
26        # Download image processor
27        logger.info("Downloading image processor...")
28        image_processor = AutoImageProcessor.from_pretrained(
29            HF_MODEL_ID,
30            cache_dir=local_path,
31        )
32        image_processor.save_pretrained(local_path)
33        
34        # Download model
35        logger.info("Downloading model weights...")
36        model = AutoModelForImageClassification.from_pretrained(
37            HF_MODEL_ID,
38            cache_dir=local_path,
39        )
40        model.save_pretrained(local_path)
41        
42        logger.info(f"Successfully downloaded model to {local_path}")
43        return True
44        
45    except RepositoryNotFoundError as e:
46        logger.error(f"Model not found on Hugging Face Hub: {e}")
47        return False
48    except Exception as e:
49        logger.error(f"Error downloading model from Hugging Face: {e}")
50        return False
51