update

app.py

@@ -1,18 +1,23 @@
 import numpy as np
-
+import torch
 import gradio as gr
 from infer import detections
-
-'''
-'''
 import os
 os.system("mkdir data")
 os.system("mkdir data/models")
 os.system("wget https://www.cs.cmu.edu/~walt/models/walt_people.pth -O data/models/walt_people.pth")
 os.system("wget https://www.cs.cmu.edu/~walt/models/walt_vehicle.pth -O data/models/walt_vehicle.pth")
+'''
+'''
 def walt_demo(input_img):
     #detect_people = detections('configs/walt/walt_people.py', 'cuda:0', model_path='data/models/walt_people.pth')
-    detect = detections('configs/walt/walt_vehicle.py', 'cuda:0', model_path='data/models/walt_vehicle.pth')
+    if torch.cuda.is_available() == False:
+        device='cpu'
+    else:
+        device='cuda:0'
+    #detect_people = detections('configs/walt/walt_people.py', device, model_path='data/models/walt_people.pth')
+    detect = detections('configs/walt/walt_vehicle.py', device, model_path='data/models/walt_vehicle.pth', threshold=0.75)
+ 
     count = 0
     #img = detect_people.run_on_image(input_img)
     output_img = detect.run_on_image(input_img)
@@ -42,7 +47,6 @@ article="""
 examples = [
     'demo/images/img_1.jpg',
     'demo/images/img_2.jpg',
-    'demo/images/img_3.png',
     'demo/images/img_4.png',
 ]
 

infer.py

@@ -11,12 +11,12 @@ import os
 import cv2, glob
 
 class detections():
-    def __init__(self, cfg_path, device, model_path = 'data/models/walt_vehicle.pth'):
+    def __init__(self, cfg_path, device, model_path = 'data/models/walt_vehicle.pth', threshold=0.85):
         self.model = init_detector(cfg_path, model_path, device=device)
         self.all_preds = []
         self.all_scores = []
         self.index = []
-        self.score_thr = 0.9
+        self.score_thr = threshold
         self.result = []
         self.record_dict = {'model': cfg_path,'results': []}
         self.detect_count = []
@@ -75,8 +75,12 @@ class detections():
 
  
 def main():
-    detect_people = detections('configs/walt/walt_people.py', 'cuda:0', model_path='data/models/walt_people.pth')
-    detect = detections('configs/walt/walt_vehicle.py', 'cuda:0', model_path='data/models/walt_vehicle.pth')
+    if torch.cuda.is_available() == False:
+        device='cpu'
+    else:
+        device='cuda:0'
+    detect_people = detections('configs/walt/walt_people.py', device, model_path='data/models/walt_people.pth')
+    detect = detections('configs/walt/walt_vehicle.py', device, model_path='data/models/walt_vehicle.pth')
     filenames = sorted(glob.glob('demo/images/*'))
     count = 0
     for filename in filenames:

requirements.txt

@@ -1,10 +1,12 @@
-torch==1.9.0
-mmcv-full==1.4.0 -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.9.0/index.html
+--find-links https://download.pytorch.org/whl/torch_stable.html
+torch==1.9.0+cpu
+--find-links https://download.openmmlab.com/mmcv/dist/cpu/torch1.9.0/index.html
+mmcv-full==1.4.0
 gradio==3.0.20
 timm 
 scikit-image 
 imagesize
-torchvision==0.10.0
+torchvision==0.10.0+cpu 
 imantics
 terminaltables
 pycocotools