remove pth from tracking

.gitignore

@@ -30,6 +30,7 @@ env.bak/
 venv.bak/
 
 # PyTorch specific
+*.pth
 *.pt
 *.pkl
 *.onnx

app.py

@@ -1,41 +1,64 @@
-import gradio as gr
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 import torch
-import torchvision.transforms as transforms
 from PIL import Image
-from torchvision.models import resnet50
-import torch.nn as nn
 
-# Load model
-model = resnet50(pretrained=False)
-model.fc = nn.Linear(model.fc.in_features, 10)
-model.load_state_dict(torch.load('best_model.pth'))
-model.eval()
-
-# Define classes (for CIFAR-10)
-classes = ['airplane', 'automobile', 'bird', 'cat', 'deer',
-           'dog', 'frog', 'horse', 'ship', 'truck']
+def load_model_from_hub(repo_id: str):
+    """
+    Load model from Hugging Face Hub
+    
+    Args:
+        repo_id: The repository ID (e.g., 'username/model-name')
+    Returns:
+        model: The loaded model
+        processor: The feature extractor/processor
+    """
+    # Load model and processor from Hub
+    model = AutoModelForImageClassification.from_pretrained(repo_id)
+    processor = AutoFeatureExtractor.from_pretrained(repo_id)
+    return model, processor
 
-def predict(image):
-    transform = transforms.Compose([
-        transforms.Resize(224),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ])
+def predict(image_path: str, model, processor):
+    """
+    Make prediction using the loaded model
     
-    img_tensor = transform(image).unsqueeze(0)
+    Args:
+        image_path: Path to input image
+        model: Loaded model
+        processor: Feature extractor/processor
+    Returns:
+        prediction: Model prediction
+    """
+    # Load and preprocess image
+    image = Image.open(image_path)
+    inputs = processor(images=image, return_tensors="pt")
     
+    # Make prediction
     with torch.no_grad():
-        outputs = model(img_tensor)
-        _, predicted = outputs.max(1)
-        
-    return classes[predicted.item()]
+        outputs = model(**inputs)
+        predictions = outputs.logits.softmax(-1)
+    
+    return predictions
+
+# Example usage in your Flask/FastAPI app
+from flask import Flask, request
+app = Flask(__name__)
 
-# Create Gradio interface
-iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=1),
-    examples=[["example1.jpg"], ["example2.jpg"]]
-)
+# Load model at startup
+model, processor = load_model_from_hub("srtangirala/resnet50-exp")
+
+@app.route('/predict', methods=['POST'])
+def predict_endpoint():
+    if 'file' not in request.files:
+        return {'error': 'No file provided'}, 400
+    
+    file = request.files['file']
+    image_path = "temp_image.jpg"  # You might want to generate a unique filename
+    file.save(image_path)
+    
+    predictions = predict(image_path, model, processor)
+    
+    # Convert predictions to list and return
+    return {'predictions': predictions.tolist()[0]}
 
-iface.launch() 
+if __name__ == '__main__':
+    app.run(debug=True) 

upload_to_hub.py

@@ -0,0 +1,122 @@
+from huggingface_hub import HfApi, login
+import torch
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
+import json
+import os
+
+def upload_model_to_hub(
+    model_path: str,
+    repo_name: str,
+    token: str,
+    num_labels: int,
+    label2id: dict,
+    id2label: dict,
+    model_architecture: str = "resnet50",
+    task: str = "image-classification",
+):
+    """
+    Upload a PyTorch model to Hugging Face Hub with proper configuration
+    """
+    # Login to Hugging Face
+    login(token=token)
+    api = HfApi()
+    
+    # Create the repository
+    repo_url = api.create_repo(
+        repo_id=repo_name,
+        exist_ok=True,
+        private=False
+    )
+
+    # Create config.json
+    config = {
+        "architectures": ["ResNetForImageClassification"],
+        "model_type": "resnet",
+        "num_labels": num_labels,
+        "id2label": id2label,
+        "label2id": label2id,
+        "num_channels": 3,
+        "hidden_sizes": [2048],
+        "image_size": [224, 224]
+    }
+
+    # Create feature extractor config
+    feature_extractor = {
+        "image_mean": [0.485, 0.456, 0.406],
+        "image_std": [0.229, 0.224, 0.225],
+        "do_normalize": True,
+        "do_resize": True,
+        "size": 224,
+        "resample": 2
+    }
+    
+    # Upload config files
+    api.upload_file(
+        path_or_fileobj=json.dumps(config).encode(),
+        path_in_repo="config.json",
+        repo_id=repo_name,
+        commit_message="Upload model config"
+    )
+
+    api.upload_file(
+        path_or_fileobj=json.dumps(feature_extractor).encode(),
+        path_in_repo="preprocessor_config.json",
+        repo_id=repo_name,
+        commit_message="Upload preprocessor config"
+    )
+
+    # Upload the model file
+    api.upload_file(
+        path_or_fileobj=model_path,
+        path_in_repo="pytorch_model.bin",
+        repo_id=repo_name,
+        commit_message="Upload model weights"
+    )
+    
+    # Create and upload model card
+    model_card = f"""
+    ---
+    language: en
+    tags:
+    - pytorch
+    - {model_architecture}
+    - {task}
+    ---
+    
+    # Model Card for {repo_name}
+    
+    This model is a fine-tuned version of {model_architecture} for {task}.
+    """
+    
+    api.upload_file(
+        path_or_fileobj=model_card.encode(),
+        path_in_repo="README.md",
+        repo_id=repo_name,
+        commit_message="Upload model card"
+    )
+    
+    print(f"Model uploaded successfully to: https://huggingface.co/{repo_name}")
+
+if __name__ == "__main__":
+    # Get Hugging Face token from environment variable
+    token = os.getenv("HF_TOKEN")
+    if not token:
+        raise ValueError("Please set the HF_TOKEN environment variable")
+    
+    # Example label mappings - replace with your actual labels
+    label2id = {
+        "class1": 0,
+        "class2": 1,
+        # ... add all your classes
+    }
+    id2label = {str(v): k for k, v in label2id.items()}
+    
+    # Upload the model
+    upload_model_to_hub(
+        model_path="best_model.pth",
+        repo_name="srtangirala/resnet50-exp",
+        token=token,
+        num_labels=len(label2id),
+        label2id=label2id,
+        id2label=id2label
+    )