promptsai commited on
Commit
480d0bf
·
verified ·
1 Parent(s): 230ee18

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +54 -0
app.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from models import vgg19
3
+ import gdown
4
+ from PIL import Image
5
+ from torchvision import transforms
6
+ import gradio as gr
7
+ import cv2
8
+ import numpy as np
9
+ import scipy
10
+
11
+ model_path = "pretrained_models/model_qnrf.pth"
12
+ url = "https://drive.google.com/uc?id=1nnIHPaV9RGqK8JHL645zmRvkNrahD9ru"
13
+ gdown.download(url, model_path, quiet=False)
14
+
15
+ device = torch.device('cpu') # device can be "cpu" or "gpu"
16
+
17
+ model = vgg19()
18
+ model.to(device)
19
+ model.load_state_dict(torch.load(model_path, device))
20
+ model.eval()
21
+
22
+
23
+ def predict(inp):
24
+ inp = Image.fromarray(inp.astype('uint8'), 'RGB')
25
+ inp = transforms.ToTensor()(inp).unsqueeze(0)
26
+ inp = inp.to(device)
27
+ with torch.set_grad_enabled(False):
28
+ outputs, _ = model(inp)
29
+ count = torch.sum(outputs).item()
30
+ vis_img = outputs[0, 0].cpu().numpy()
31
+ # normalize density map values from 0 to 1, then map it to 0-255.
32
+ vis_img = (vis_img - vis_img.min()) / (vis_img.max() - vis_img.min() + 1e-5)
33
+ vis_img = (vis_img * 255).astype(np.uint8)
34
+ vis_img = cv2.applyColorMap(vis_img, cv2.COLORMAP_JET)
35
+ vis_img = cv2.cvtColor(vis_img, cv2.COLOR_BGR2RGB)
36
+ return vis_img, int(count)
37
+
38
+
39
+ inputs = gr.Image(label="Image of Crowd")
40
+ outputs = [
41
+ gr.Image(label="Predicted Density Map"),
42
+ gr.Label(label="Predicted Count")
43
+ ]
44
+
45
+ # Assuming `title`, `desc`, and `examples` variables are defined elsewhere in your code.
46
+ title = "Your App Title"
47
+ desc = "Your App Description"
48
+
49
+ gr.Interface(fn=predict,
50
+ inputs=inputs,
51
+ outputs=outputs,
52
+ title=title,
53
+ description=desc,
54
+ allow_flagging="never").launch(share=True)