Update app.py

changed caching method, more suitable for the large amount of subgroups we have..

app.py

@@ -7,6 +7,8 @@ import requests
 import pickle
 from tqdm import tqdm
 from datetime import datetime
+from collections import OrderedDict
+
 
 # ==== CONFIG ==== #
 PERSISTENT_DIR = "/data" if os.path.exists("/data") else "."
@@ -16,6 +18,9 @@ MAPPING_FILE = "file_mapping.csv"
 DRIVE_LINK_TEMPLATE = "https://drive.google.com/uc?id={}"
 CACHE_FILE = os.path.join(PERSISTENT_DIR, "groups_cache.pkl")
 RESET = False  # Set to True to clear previous results and cache
+CACHE_LIMIT = 30  # Feel free to tweak this
+
+image_cache = OrderedDict()
 
 # ==== Optional Reset ====
 if RESET:
@@ -25,33 +30,32 @@ if RESET:
             os.remove(path)
             print(f"🗑️ Deleted {path}")
 
+def load_group_with_cache(group, resize=(256, 256)):
+    key = tuple(group)
+    if key in image_cache:
+        return image_cache[key]
+    
+    imgs = []
+    for file_name in group:
+        try:
+            url = get_drive_image_url(file_name)
+            response = requests.get(url, stream=True, timeout=10)
+            img = Image.open(response.raw).convert("RGB").resize(resize)
+            imgs.append(img)
+        except Exception as e:
+            print(f"❌ Error loading {file_name}: {e}")
+            imgs.append(None)
+    
+    image_cache[key] = imgs
+    if len(image_cache) > CACHE_LIMIT:
+        image_cache.popitem(last=False)  # Remove oldest group
+
+    return imgs
 # ==== Helpers ====
 def get_drive_image_url(file_name):
     file_id = file_dict.get(file_name)
     return DRIVE_LINK_TEMPLATE.format(file_id) if file_id else None
 
-def preload_groups(groups, resize=(256, 256)):
-    if os.path.exists(CACHE_FILE):
-        with open(CACHE_FILE, "rb") as f:
-            return pickle.load(f)
-
-    print("🖼️ Preloading groups from Drive...")
-    preloaded = []
-    for group in tqdm(groups):
-        try:
-            imgs = []
-            for file_name in group:
-                url = get_drive_image_url(file_name)
-                response = requests.get(url, stream=True, timeout=10)
-                img = Image.open(response.raw).convert("RGB").resize(resize)
-                imgs.append(img)
-            preloaded.append(imgs)
-        except Exception as e:
-            print(f"Error loading group {group}: {e}")
-    with open(CACHE_FILE, "wb") as f:
-        pickle.dump(preloaded, f)
-    return preloaded
-
 def load_reviewed_ids():
     try:
         reviewed = pd.read_csv(RESULTS_FILE).to_dict(orient="records")
@@ -62,13 +66,11 @@ def load_reviewed_ids():
 def get_remaining_groups():
     reviewed, reviewed_ids = load_reviewed_ids()
     remaining = [g for g in sample_names if tuple(g) not in reviewed_ids]
-    filtered_preloaded = [pg for g, pg in zip(sample_names, preloaded_groups) if tuple(g) not in reviewed_ids]
-    return reviewed, reviewed_ids, remaining, filtered_preloaded
+    return reviewed, reviewed_ids, remaining
 
 def review_group(decision, group):
     reviewed, reviewed_ids = load_reviewed_ids()
 
-    # Save updated results
     reviewed.append({
         "group": json.dumps(group),
         "decision": decision
@@ -81,14 +83,15 @@ def review_group(decision, group):
     except Exception as e:
         print(f"❌ Error saving results: {e}")
 
-    # Get updated remaining groups
-    _, _, remaining, filtered_preloaded = get_remaining_groups()
+    _, _, remaining = get_remaining_groups()
 
     if remaining:
-        return filtered_preloaded[0], remaining[0], f"Group {len(reviewed)+1} / {len(preloaded_groups)}"
+        group = remaining[0]
+        return load_group_with_cache(group), group, f"Group {len(reviewed)+1} / {len(sample_names)}"
     else:
         return [], None, "✅ All groups reviewed!"
 
+
 def prepare_download():
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     filename = f"review_results_{timestamp}.csv"
@@ -104,9 +107,10 @@ def prepare_download():
         return None
 
 def get_first_group():
-    reviewed, _, remaining, filtered_preloaded = get_remaining_groups()
+    reviewed, _, remaining = get_remaining_groups()
     if remaining:
-        return filtered_preloaded[0], remaining[0], f"Group {len(reviewed)+1} / {len(preloaded_groups)}"
+        group = remaining[0]
+        return load_group_with_cache(group), group, f"Group {len(reviewed)+1} / {len(sample_names)}"
     else:
         return [], None, "✅ All groups reviewed!"
 
@@ -114,7 +118,6 @@ def get_first_group():
 file_dict = pd.read_csv(MAPPING_FILE).set_index("name")["id"].to_dict()
 with open(GROUPS_FILE) as f:
     sample_names = json.load(f)
-preloaded_groups = preload_groups(sample_names)
 
 # ==== Gradio UI ====
 with gr.Blocks() as demo: