File size: 1,224 Bytes
44100ae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
---
license: apache-2.0
tags:
- vision
- image-classification
datasets:
- imagenet-1k
---
## How to Use
``` python
# Preprocess Image
def process_image(image, model):
preprocess = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(image).unsqueeze(0)
input_tensor = input_tensor.to(device)
with torch.no_grad():
output = model(input_tensor)
predicted_count = output.item()
print(f"Predicted Headcount: {predicted_count}")
return math.ceil(predicted_count)
# Load Model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def load_model(selected_model):
model = None
model_path = None
if selected_model == 'VGG16':
model = models.VGG16()
model_path = "vgg16_headcount.pth"
else:
model = models.ResNet50()
model_path = "resnet50_headcount.pth"
model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
model.to(device)
model.eval()
print(f"{selected_model}.Heavy Model loaded successfully")
return model
``` |