bensonsantos commited on
Commit
c494b78
1 Parent(s): e699e28

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -0
app.py CHANGED
@@ -10,11 +10,18 @@ checkpoint = torch.load('part_B_pre.pth.tar',map_location=torch.device('cpu'))
10
  model.load_state_dict(checkpoint['state_dict'])
11
  model.eval()
12
 
 
13
  transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])])
14
 
 
15
  def crowd(img):
 
16
  img = transform(img)
 
 
17
  img = rearrange(img, "c h w -> 1 c h w")
 
 
18
  h = img.shape[2]
19
  w = img.shape[3]
20
  h_d = int(h/2)
@@ -23,11 +30,15 @@ def crowd(img):
23
  img_2 = img[:,:,:h_d,w_d:]
24
  img_3 = img[:,:,h_d:,:w_d]
25
  img_4 = img[:,:,h_d:,w_d:]
 
 
26
  with torch.no_grad():
27
  density_1 = model(img_1).numpy().sum()
28
  density_2 = model(img_2).numpy().sum()
29
  density_3 = model(img_3).numpy().sum()
30
  density_4 = model(img_4).numpy().sum()
 
 
31
  pred = density_1 + density_2 + density_3 + density_4
32
  pred = int(pred.round())
33
  return pred
 
10
  model.load_state_dict(checkpoint['state_dict'])
11
  model.eval()
12
 
13
+ ## Defining the transform function
14
  transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])])
15
 
16
+
17
  def crowd(img):
18
+ ## Transforming the image
19
  img = transform(img)
20
+
21
+ ## Adding batch dimension
22
  img = rearrange(img, "c h w -> 1 c h w")
23
+
24
+ ## Slicing the image into four parts
25
  h = img.shape[2]
26
  w = img.shape[3]
27
  h_d = int(h/2)
 
30
  img_2 = img[:,:,:h_d,w_d:]
31
  img_3 = img[:,:,h_d:,:w_d]
32
  img_4 = img[:,:,h_d:,w_d:]
33
+
34
+ ## Inputting the 4 images into the model, converting it to numpy array, and summing to get the density
35
  with torch.no_grad():
36
  density_1 = model(img_1).numpy().sum()
37
  density_2 = model(img_2).numpy().sum()
38
  density_3 = model(img_3).numpy().sum()
39
  density_4 = model(img_4).numpy().sum()
40
+
41
+ ## Summing up the estimated density and rounding the result to get an integer
42
  pred = density_1 + density_2 + density_3 + density_4
43
  pred = int(pred.round())
44
  return pred