Update inference.py

inference.py

@@ -268,23 +268,20 @@ def do_confidence_boost(inference, safe_preds, gp=None, batch_size=BATCH_SIZE, c
 
                 frames = [None, None]
                 
-                next_frame = None
-                if i+1 < len(infer):
-                    next_frame = infer[i+1]
-                elif has_next_batch:
-                    next_frame = inference[batch_i + 1][0]
-
-                if next_frame != None:
-                    boost_frame(safe_frame, next_frame, 1, decay=conf_decay)
-
-                prev_frame = None
-                if i-1 >= 0:
-                    prev_frame = infer[i-1]
-                elif has_prev_batch:
-                    prev_frame = inference[batch_i - 1][len(inference[batch_i - 1]) - 1]
-
-                if prev_frame != None:
-                    boost_frame(safe_frame, prev_frame, -1, power=conf_power, decay=conf_decay)
+                for dt in [-1, 1]:
+                    idx = i+dt 
+                    temp_frame = None
+                    if idx >= 0 and idx < len(infer):
+                        temp_frame = infer[idx]
+                    elif idx < 0 and has_prev_batch:
+                        prev_batch = inference[batch_i - 1]
+                        temp_frame = prev_batch[idx]
+                    elif idx >= len(infer) and has_next_batch:
+                        next_batch = inference[batch_i + 1]
+                        temp_frame = next_batch[idx - len(infer)]
+                    
+                    if temp_frame is not None:
+                        boost_frame(safe_frame, temp_frame, dt, power=conf_power, decay=conf_decay)
 
             pbar.update(1*batch_size)