oskarastrom commited on
Commit
f0af381
1 Parent(s): 72d14ec

Update inference.py

Browse files
Files changed (1) hide show
  1. inference.py +4 -6
inference.py CHANGED
@@ -268,8 +268,8 @@ def do_confidence_boost(inference, safe_preds, gp=None, batch_size=BATCH_SIZE, c
268
  if len(safe_frame) == 0:
269
  continue
270
 
271
- has_next_batch = batch_i+1 < len(inference)
272
- has_prev_batch = batch_i-1 >= 0
273
 
274
 
275
  for dt in range(-boost_range, boost_range+1):
@@ -278,11 +278,9 @@ def do_confidence_boost(inference, safe_preds, gp=None, batch_size=BATCH_SIZE, c
278
  temp_frame = None
279
  if idx >= 0 and idx < len(infer):
280
  temp_frame = infer[idx]
281
- elif idx < 0 and has_prev_batch:
282
- prev_batch = inference[batch_i - 1]
283
  temp_frame = prev_batch[idx]
284
- elif idx >= len(infer) and has_next_batch:
285
- next_batch = inference[batch_i + 1]
286
  temp_frame = next_batch[idx - len(infer)]
287
 
288
  if temp_frame is not None:
 
268
  if len(safe_frame) == 0:
269
  continue
270
 
271
+ next_batch = inference[batch_i + 1] if batch_i+1 < len(inference) else None
272
+ prev_batch = inference[batch_i - 1] if batch_i-1 >= 0 else None
273
 
274
 
275
  for dt in range(-boost_range, boost_range+1):
 
278
  temp_frame = None
279
  if idx >= 0 and idx < len(infer):
280
  temp_frame = infer[idx]
281
+ elif idx < 0 and prev_batch is not None and -idx >= len(prev_batch):
 
282
  temp_frame = prev_batch[idx]
283
+ elif idx >= len(infer) and next_batch is not None and idx - len(infer) < len(next_batch):
 
284
  temp_frame = next_batch[idx - len(infer)]
285
 
286
  if temp_frame is not None: