Views
No views yet
1import torch
2import torch.nn as nn
3
4class VGG16(nn.Module):
5 def __init__(self):
6 super(VGG16, self).__init__()
7 self.features = nn.Sequential(
8 nn.Conv2d(3, 64, kernel_size=3, padding=1),
9 nn.ReLU(inplace=True),
10 nn.Conv2d(64, 64, kernel_size=3, padding=1),
11 nn.ReLU(inplace=True),
12 nn.MaxPool2d(kernel_size=2, stride=2),
13
14 nn.Conv2d(64, 128, kernel_size=3, padding=1),
15 nn.ReLU(inplace=True),
16 nn.Conv2d(128, 128, kernel_size=3, padding=1),
17 nn.ReLU(inplace=True),
18 nn.MaxPool2d(kernel_size=2, stride=2),
19
20 nn.Conv2d(128, 256, kernel_size=3, padding=1),
21 nn.ReLU(inplace=True),
22 nn.Conv2d(256, 256, kernel_size=3, padding=1),
23 nn.ReLU(inplace=True),
24 nn.Conv2d(256, 256, kernel_size=3, padding=1),
25 nn.ReLU(inplace=True),
26 nn.MaxPool2d(kernel_size=2, stride=2),
27
28 nn.Conv2d(256, 512, kernel_size=3, padding=1),
29 nn.ReLU(inplace=True),
30 nn.Conv2d(512, 512, kernel_size=3, padding=1),
31 nn.ReLU(inplace=True),
32 nn.Conv2d(512, 512, kernel_size=3, padding=1),
33 nn.ReLU(inplace=True),
34 nn.MaxPool2d(kernel_size=2, stride=2),
35
36 nn.Conv2d(512, 512, kernel_size=3, padding=1),
37 nn.ReLU(inplace=True),
38 nn.Conv2d(512, 512, kernel_size=3, padding=1),
39 nn.ReLU(inplace=True),
40 nn.Conv2d(512, 512, kernel_size=3, padding=1),
41 nn.ReLU(inplace=True),
42 nn.MaxPool2d(kernel_size=2, stride=2),
43 )
44 self.classifier = nn.Sequential(
45 nn.Linear(512 * 7 * 7, 4096),
46 nn.ReLU(inplace=True),
47 nn.Dropout(),
48 nn.Linear(4096, 4096),
49 nn.ReLU(inplace=True),
50 nn.Dropout(),
51 nn.Linear(4096, 1) # Outputting head count as a single value
52 )
53
54 def forward(self, x):
55 x = self.features(x)
56 x = torch.flatten(x, 1)
57 x = self.classifier(x)
58 return x
591# Preprocessing function
2def preprocess_image(image, channels=6):
3 transform = transforms.Compose([
4 transforms.Resize((224, 224)),
5 transforms.ToTensor()
6 ])
7 image_tensor = transform(image)
8 # Simulating 6-channel input if required
9 if channels == 6:
10 image_tensor = torch.cat([image_tensor, image_tensor], dim=0)
11 return image_tensor.unsqueeze(0).to(device)
12
13device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
14# Load model
15def load_model(selected_model):
16 model = None
17 model_path = None
18 if selected_model == 'VGG16':
19 model = models.VGG16()
20 model_path = "vgg16_headcount.pth"
21 else:
22 model = models.ResNet50()
23 model_path = "resnet50_headcount.pth"
24 model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
25 model.to(device)
26 model.eval()
27 print(f"{selected_model}.Heavy Model loaded successfully")
28 return model
29
30# Prediction Function
31def process_image(image, model):
32 preprocess = transforms.Compose([
33 transforms.Resize((224, 224)),
34 transforms.ToTensor(),
35 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
36 ])
37 input_tensor = preprocess(image).unsqueeze(0)
38 input_tensor = input_tensor.to(device)
39 with torch.no_grad():
40 output = model(input_tensor)
41 predicted_count = output.item()
42 print(f"Predicted Headcount: {predicted_count}")
43 return math.ceil(predicted_count)