Harinivas-28/ResNet50_head_count

How to Use

# 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