def count(img):
model = CSRNet()
if C["cuda"]:
model = model.cuda()
pth = torch.load(C["pth"])
else:
model = model.cpu()
pth = torch.load(C["pth"], map_location="cpu")
model.load_state_dict(pth["state_dict"])
# img = 255.0 * to_tensor(Image.open(img_path).convert("RGB"))
# for i in range(3):
# img[i,:,:] = img[i,:,:] + C["img_corr"][i]
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=C["mean"], std=C["std"])
])
img = transform(img.convert("RGB"))
if C["cuda"]:
img = img.cuda()
else:
img = img.cpu()
output = model(img.unsqueeze(0)).detach().cpu()
dmap = np.asarray(output.reshape(output.shape[2], output.shape[3]))
count = int(np.sum(dmap))
return count
import re
#import json
from matplotlib import cm as c
from model import CSRNet
import torch
from torchvision import datasets, transforms
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224,
0.225]),
])
#creating model
model = CSRNet()
model.cpu()
#loading the trained weights
checkpoint = torch.load('0model_best.pth.tar',
map_location=lambda storage, location: storage)
model.load_state_dict(checkpoint['state_dict'])
#valid image types
valid_types = ['image/jpeg', '/image/png']
#decoding an image from base64 into raw representation
def convertImage(imgData1):
imgstr = re.search(r'base64,(.*)', imgData1).group(1)
#print(imgstr)
with open('output.png', 'wb') as output:
output.write(imgstr.decode('base64'))
