PKUWilliamYang commited on
Commit
13706ea
1 Parent(s): 639ff14

Update vtoonify_model.py

Browse files
Files changed (1) hide show
  1. vtoonify_model.py +4 -1
vtoonify_model.py CHANGED
@@ -206,7 +206,10 @@ class Model():
206
  int(video_cap.get(4)*4)))
207
 
208
  batch_frames = []
209
- batch_size = max(1, int(round(4 * 256 * 256 / video_cap.get(3) / video_cap.get(4))))
 
 
 
210
  print('Using batch size of %d on %d frames'%(batch_size, num))
211
  with torch.no_grad():
212
  if self.color_transfer:
 
206
  int(video_cap.get(4)*4)))
207
 
208
  batch_frames = []
209
+ if video_cap.get(3) != 0:
210
+ batch_size = min(max(1, int(round(4 * 400 * 360 / video_cap.get(3) / video_cap.get(4)))), 8)
211
+ else:
212
+ batch_size = 4
213
  print('Using batch size of %d on %d frames'%(batch_size, num))
214
  with torch.no_grad():
215
  if self.color_transfer: