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'))