Update app.py

app.py

@@ -25,7 +25,6 @@ TARGET_PER_PERSON = 20
 CONTACT_EMAIL = "ffallah@asu.edu"
 
 # --- Paths ---
-CAPTIONS_JSON_PATH = os.environ.get("CAPTIONS_JSON_PATH", "data/captions.json")
 
 GT_MASKED_DIR   = "data/gt_b"   # Image 1
 GT_UNMASKED_DIR = "data/adc_b"  # Image 2
@@ -147,42 +146,70 @@ def load_image(path: str) -> Image.Image:
     except Exception:
         return Image.new("RGB", (256, 256), color="gray")
 
-def load_dataset(captions_path: str, gt_masked_dir: str, gt_unmasked_dir: str, sr_dir: str, original_dir: str, image_5_dir: str) -> List[Sample]:
-    if not os.path.exists(captions_path):
-        print(f"Captions file not found at {captions_path}")
+def load_dataset(
+    gt_masked_dir: str,
+    gt_unmasked_dir: str,
+    sr_dir: str,
+    original_dir: str,
+    image_5_dir: str,
+) -> List[Sample]:
+    """
+    Build samples only from the 5 folders.
+    Each folder should have the same filenames.
+    Example layout:
+        data/gt_b/xxx.png
+        data/adc_b/xxx.png
+        data/sr_b/xxx.png
+        data/lr_b/xxx.png
+        data/see_b/xxx.png
+    """
+
+    def list_images(dir_path: str) -> set:
+        if not os.path.isdir(dir_path):
+            print(f"Warning: directory not found: {dir_path}")
+            return set()
+        files = []
+        for f in os.listdir(dir_path):
+            f_lower = f.lower()
+            if f_lower.endswith((".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp")):
+                files.append(f)
+        return set(files)
+
+    masked_files   = list_images(gt_masked_dir)
+    unmasked_files = list_images(gt_unmasked_dir)
+    sr_files       = list_images(sr_dir)
+    orig_files     = list_images(original_dir)
+    img5_files     = list_images(image_5_dir)
+
+    # Common filenames present in ALL 5 folders
+    common_files = masked_files & unmasked_files & sr_files & orig_files & img5_files
+
+    if not common_files:
+        print("No common image files found in all 5 folders.")
         return []
 
-    with open(captions_path, "r", encoding="utf-8") as f:
-        try:
-            captions_data = json.load(f)
-        except Exception:
-            print("Failed to parse captions JSON.")
-            return []
+    # Optional: simple debug info
+    print(f"Found {len(common_files)} common images.")
 
     samples: List[Sample] = []
-    for item in captions_data:
-        base_filename = item.get("image")
-        if not base_filename:
-            continue
-
+    for base_filename in sorted(common_files):
         sample_id = os.path.splitext(base_filename)[0]
+
         paths = {
-            "masked": os.path.join(gt_masked_dir, base_filename),
+            "masked":   os.path.join(gt_masked_dir, base_filename),
             "unmasked": os.path.join(gt_unmasked_dir, base_filename),
-            "sr": os.path.join(sr_dir, base_filename),
+            "sr":       os.path.join(sr_dir, base_filename),
             "original": os.path.join(original_dir, base_filename),
-            "img5": os.path.join(image_5_dir, base_filename)
+            "img5":     os.path.join(image_5_dir, base_filename),
         }
 
-        # If strict enforcement required, require all five files to exist.
+        # If STRICT_ENFORCEMENT is True, skip if any file missing
         if STRICT_ENFORCEMENT:
             if not all(os.path.exists(p) for p in paths.values()):
                 missing = [k for k, v in paths.items() if not os.path.exists(v)]
-                print(f"Skipping {base_filename}: Missing in folders {missing}")
+                print(f"Skipping {base_filename}: missing in folders {missing}")
                 continue
 
-        # In non-strict mode, it's okay to include samples even if some files missing;
-        # we will supply placeholders at load time.
         samples.append(
             Sample(
                 sample_id=sample_id,
@@ -190,12 +217,13 @@ def load_dataset(captions_path: str, gt_masked_dir: str, gt_unmasked_dir: str, s
                 unmasked_gt_path=paths["unmasked"],
                 sr_path=paths["sr"],
                 original_path=paths["original"],
-                image_5_path=paths["img5"]
+                image_5_path=paths["img5"],
             )
         )
 
     return samples
 
+
 # ----------------------
 # Progress & results I/O
 # ----------------------
@@ -258,7 +286,7 @@ def start_or_resume(name: str, email: str):
         raise gr.Error("Please enter your name and email to begin.")
 
     ensure_paths()
-    samples = load_dataset(CAPTIONS_JSON_PATH, GT_MASKED_DIR, GT_UNMASKED_DIR, SR_DIR, ORIGINAL_DIR, IMAGE_5_DIR)
+    samples = load_dataset(GT_MASKED_DIR, GT_UNMASKED_DIR, SR_DIR, ORIGINAL_DIR, IMAGE_5_DIR)
 
     if not samples:
         raise gr.Error("No images found. Please check dataset configuration.")
@@ -716,7 +744,7 @@ if __name__ == "__main__":
             print(f"Error reading from HF: {e}")
 
     ensure_paths()
-    _ = load_dataset(CAPTIONS_JSON_PATH, GT_MASKED_DIR, GT_UNMASKED_DIR, SR_DIR, ORIGINAL_DIR, IMAGE_5_DIR)
+    _ = load_dataset(GT_MASKED_DIR, GT_UNMASKED_DIR, SR_DIR, ORIGINAL_DIR, IMAGE_5_DIR)
 
     print("✅ Launching app.")
     demo.queue()