Keep 915 inference viewer images after 100 results

NN_server/Models/ensemble_915_v44.py

@@ -6,6 +6,8 @@ import torch
 import cv2
 import gc
 import io
+import os
+import re
 
 
 def _as_display_image(image):
@@ -22,6 +24,42 @@ def _as_display_image(image):
     return arr
 
 
+def _prune_old_inference_images(src, model_type, model_id, keep_last=200):
+    try:
+        keep_last = int(os.getenv("INFERENCE_IMAGE_KEEP_LAST", str(keep_last)))
+    except ValueError:
+        keep_last = keep_last
+
+    if keep_last <= 0 or not src or not os.path.isdir(src):
+        return
+
+    pattern = re.compile(
+        r"_inference_(\d+)_.*_"
+        + re.escape(str(model_id))
+        + "_"
+        + re.escape(str(model_type))
+        + r"\.png$"
+    )
+    grouped = {}
+    for name in os.listdir(src):
+        match = pattern.match(name)
+        if match is None:
+            continue
+        grouped.setdefault(int(match.group(1)), []).append(name)
+
+    if len(grouped) <= keep_last:
+        return
+
+    for old_result_id in sorted(grouped)[: len(grouped) - keep_last]:
+        for name in grouped[old_result_id]:
+            try:
+                os.remove(os.path.join(src, name))
+            except FileNotFoundError:
+                pass
+            except OSError as exc:
+                print(f"failed to remove old inference image {name}: {exc}")
+
+
 def _render_plot(values, figsize=(16, 16), dpi=16):
     import matplotlib.pyplot as plt
 
@@ -178,7 +216,7 @@ def post_func_ensemble(src="", model_type="", prediction="", model_id=0, ind_inf
         matplotlib.use("Agg")
         plt.ioff()
 
-        if int(ind_inference) <= 100 and isinstance(data, (list, tuple)) and len(data) >= 2:
+        if isinstance(data, (list, tuple)) and len(data) >= 2:
             fig, ax = plt.subplots()
             ax.imshow(_as_display_image(data[0]))
             plt.savefig(src + "_inference_" + str(ind_inference) + "_" + prediction + "_real_" + str(model_id) + "_" + model_type + ".png")
@@ -197,6 +235,8 @@ def post_func_ensemble(src="", model_type="", prediction="", model_id=0, ind_inf
             cv2.destroyAllWindows()
             gc.collect()
 
+            _prune_old_inference_images(src, model_type, model_id)
+
         plt.clf()
         plt.cla()
         plt.close()