irfanhossainsust/simple-image-classifier

Simple Image Classifier

This model was pushed to the Hub using PyTorchModelHubMixin.

Usage

import requests
import torch
import torch.nn as nn
from PIL import Image
from torchvision import models, transforms
from huggingface_hub import hf_hub_download

class SimpleImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = models.resnet18(weights=None)
        self.transforms = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def forward(self, x):
        return self.model(x)

    @classmethod
    def from_pretrained(cls, repo_id):
        instance = cls()
        weights_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
        state_dict = torch.load(weights_path, map_location="cpu")
        instance.model.load_state_dict(state_dict, strict=False)
        return instance

# Load model
model = SimpleImageClassifier.from_pretrained("irfanhossainsust/simple-image-classifier")
model.eval()

# Run inference
url = "[https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg](https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg)"
image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
input_tensor = model.transforms(image).unsqueeze(0)

with torch.no_grad():
    output = model(input_tensor)

predicted_class = torch.argmax(output, dim=1).item()
print("Predicted Class ID:", predicted_class)