asdasdasdasd commited on
Commit
26d4625
1 Parent(s): e71c28e

Update detect_from_videos.py

Browse files
Files changed (1) hide show
  1. detect_from_videos.py +3 -3
detect_from_videos.py CHANGED
@@ -12,7 +12,7 @@ from model_core import Two_Stream_Net
12
  from torchvision import transforms
13
 
14
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
15
-
16
 
17
  xception_default_data_transforms_256 = {
18
  'train': transforms.Compose([
@@ -115,7 +115,7 @@ def predict_with_model(image, model, post_function=nn.Softmax(dim=1),
115
 
116
 
117
  def test_full_image_network(video_path, model_path, output_path,
118
- start_frame=0, end_frame=None, cuda=True):
119
  """
120
  Reads a video and evaluates a subset of frames with the a detection network
121
  that takes in a full frame. Outputs are only given if a face is present
@@ -150,7 +150,7 @@ def test_full_image_network(video_path, model_path, output_path,
150
  # Load model
151
  # model, *_ = model_selection(modelname='xception', num_out_classes=2)
152
  model = Two_Stream_Net()
153
- model.load_state_dict(torch.load(model_path))
154
  model = model.to(device)
155
  model.eval()
156
 
 
12
  from torchvision import transforms
13
 
14
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
15
+ map_location=torch.device('cpu')
16
 
17
  xception_default_data_transforms_256 = {
18
  'train': transforms.Compose([
 
115
 
116
 
117
  def test_full_image_network(video_path, model_path, output_path,
118
+ start_frame=0, end_frame=None, cuda=False):
119
  """
120
  Reads a video and evaluates a subset of frames with the a detection network
121
  that takes in a full frame. Outputs are only given if a face is present
 
150
  # Load model
151
  # model, *_ = model_selection(modelname='xception', num_out_classes=2)
152
  model = Two_Stream_Net()
153
+ model.load_state_dict(torch.load(model_path,map_location))
154
  model = model.to(device)
155
  model.eval()
156