Update utils.py

utils.py

@@ -1,41 +1,28 @@
-import os, zipfile
-import numpy as np
-from PIL import Image
-import random
-
-# Use local path instead of /content
-EXTRACTION_PATH = "./dataset"
-
-def extract_dataset(zip_path="test.zip"):
-    """Extracts test.zip into ./dataset/test"""
-    if not os.path.exists(EXTRACTION_PATH):
-        os.makedirs(EXTRACTION_PATH)
-
-    with zipfile.ZipFile(zip_path, "r") as zip_ref:
-        zip_ref.extractall(EXTRACTION_PATH)
-
-    return os.path.join(EXTRACTION_PATH, "test")
-
-def filter_images_with_buildings(base_path):
-    test_a = os.path.join(base_path, "A")
-    test_b = os.path.join(base_path, "B")
-    label = os.path.join(base_path, "label")
-
-    image_files = sorted(os.listdir(test_a))
-    valid_images = []
-
-    for img_name in image_files:
-        mask_path = os.path.join(label, img_name)
-        mask = Image.open(mask_path).convert("L")
-        mask_np = np.array(mask)
-
-        if np.any(mask_np > 0):  # keep only if mask has buildings
-            valid_images.append({
-                "image_a": os.path.join(test_a, img_name),
-                "image_b": os.path.join(test_b, img_name),
-                "mask": mask_path
-            })
-    return valid_images
-
-def get_random_valid_image(valid_images):
-    return random.choice(valid_images)
+import os
+import zipfile
+import shutil
+
+def setup_dataset(zip_path="test.zip", extract_path="./dataset"):
+    """Extract dataset if not already extracted"""
+    dataset_path = os.path.join(extract_path, "test")
+    
+    if not os.path.exists(dataset_path):
+        print("Extracting dataset...")
+        os.makedirs(extract_path, exist_ok=True)
+        
+        try:
+            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+                zip_ref.extractall(extract_path)
+            print("Dataset extracted successfully!")
+        except Exception as e:
+            print(f"Error extracting dataset: {e}")
+            return None
+    
+    return dataset_path
+
+def cleanup_cache():
+    """Clean up any temporary files"""
+    cache_dirs = ['__pycache__', '.ipynb_checkpoints']
+    for cache_dir in cache_dirs:
+        if os.path.exists(cache_dir):
+            shutil.rmtree(cache_dir)