Optimize ram usage

2026-07-06 09:23:36 +02:00
parent b728668440
commit 0c771f89c3
2 changed files with 118 additions and 58 deletions
@@ -81,32 +81,36 @@ def run_inference(model, data_loader, conf):

 def accumulate_hits(predictions: np.ndarray, reference_points: np.ndarray, 
                    binning_factor: int, number_of_subframes: int = 20):
-    ml_super_frames = np.zeros((number_of_subframes, NY * binning_factor, NX * binning_factor))
-    count_frame = np.zeros((NY, NX))
-    subpixel_dist = np.zeros((binning_factor, binning_factor))
+    ml_super_frames = np.zeros((number_of_subframes, NY * binning_factor, NX * binning_factor), dtype=np.uint32)
+    count_frame = np.zeros((NY, NX), dtype=np.uint32)
+    subpixel_dist = np.zeros((binning_factor, binning_factor), dtype=np.uint32)
    
-    # absolute coordinate = predicted subpixel + reference point
-    absolute_positions = predictions + reference_points[:, :2]
-    
-    # super resolution frames (binning)
-    hit_x = np.floor(absolute_positions[:, 0] * binning_factor).astype(int)
-    hit_y = np.floor(absolute_positions[:, 1] * binning_factor).astype(int)
+    chunk_size = len(predictions) // number_of_subframes
    
    for i in range(number_of_subframes):
-        start_idx = i * len(predictions) // number_of_subframes
-        end_idx = min((i + 1) * len(predictions) // number_of_subframes, len(predictions))
-        np.add.at(ml_super_frames[i], (hit_y[start_idx:end_idx], hit_x[start_idx:end_idx]), 1)
-    
-    # count frame (by reference point pixel index)
-    ref_x = (reference_points[:, 0] + 1).astype(int)  # reference point is lower-left corner, +1 to get pixel index
-    ref_y = (reference_points[:, 1] + 1).astype(int)
-    np.add.at(count_frame, (ref_y, ref_x), 1)
-    
-    # subpixel distribution
-    sub_x = np.floor((absolute_positions[:, 0] % 1) * binning_factor).astype(int)
-    sub_y = np.floor((absolute_positions[:, 1] % 1) * binning_factor).astype(int)
-    np.add.at(subpixel_dist, (sub_y, sub_x), 1)
-    
+        start_idx = i * chunk_size
+        end_idx = (i + 1) * chunk_size if i < number_of_subframes - 1 else len(predictions)
+        
+        pred_chunk = predictions[start_idx:end_idx]
+        ref_chunk = reference_points[start_idx:end_idx]
+        
+        abs_pos_chunk = pred_chunk + ref_chunk[:, :2]
+        
+        # --- Super resolution frames ---
+        hit_x = np.floor(abs_pos_chunk[:, 0] * binning_factor).astype(np.int32)
+        hit_y = np.floor(abs_pos_chunk[:, 1] * binning_factor).astype(np.int32)
+        np.add.at(ml_super_frames[i], (hit_y, hit_x), 1)
+        
+        # --- Count frame ---
+        ref_x = (ref_chunk[:, 0] + 1).astype(np.int32)
+        ref_y = (ref_chunk[:, 1] + 1).astype(np.int32)
+        np.add.at(count_frame, (ref_y, ref_x), 1)
+        
+        # --- Subpixel distribution ---
+        sub_x = np.floor((abs_pos_chunk[:, 0] % 1) * binning_factor).astype(np.int32)
+        sub_y = np.floor((abs_pos_chunk[:, 1] % 1) * binning_factor).astype(np.int32)
+        np.add.at(subpixel_dist, (sub_y, sub_x), 1)
+        
    return ml_super_frames, count_frame, subpixel_dist

 def save_results(ml_super_frames, ml_super_frames_eta,