Update README.md

README.md

@@ -1,3 +1,44 @@
----
-license: apache-2.0
----
+---
+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
+
+
+```