Update app.py

app.py

@@ -22,27 +22,48 @@ import seaborn as sns
 from collections import Counter
 import pickle
 from transformers import ResNetForImageClassification
+from datasets import load_dataset, Dataset  # Hugging Face datasets library
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Set the working directory to a writable location
 WORKING_DIR = "/app"  # Use /app for temporary files
-PERSISTENT_DIR = "/data"  # Use /data for persistent storage
 os.makedirs(WORKING_DIR, exist_ok=True)
-os.makedirs(PERSISTENT_DIR, exist_ok=True)
 os.chdir(WORKING_DIR)
 
 # Verify the current directory
 logging.info(f"Current working directory: {os.getcwd()}")
 
-# Cache file to store extracted text
-CACHE_FILE = os.path.join(PERSISTENT_DIR, "image_text_cache.pkl")  # Save cache in persistent storage
+# Hugging Face dataset name for persistent storage
+DATASET_NAME = "Fred808/data"  # Replace with your Hugging Face dataset name
 
-# Load cache if it exists
-if os.path.exists(CACHE_FILE):
-    with open(CACHE_FILE, "rb") as f:
-        cache = pickle.load(f)
+# Function to save data to Hugging Face Hub
+def save_to_hub(data, dataset_name):
+    """Save data to Hugging Face Hub as a dataset."""
+    try:
+        dataset = Dataset.from_dict(data)
+        dataset.push_to_hub(dataset_name)
+        logging.info(f"Data saved to Hugging Face Hub: {dataset_name}")
+    except Exception as e:
+        logging.error(f"Error saving data to Hugging Face Hub: {e}")
+
+# Function to load data from Hugging Face Hub
+def load_from_hub(dataset_name):
+    """Load data from Hugging Face Hub."""
+    try:
+        dataset = load_dataset(dataset_name)
+        logging.info(f"Data loaded from Hugging Face Hub: {dataset_name}")
+        return dataset["train"].to_dict()  # Convert to dictionary
+    except Exception as e:
+        logging.error(f"Error loading data from Hugging Face Hub: {e}")
+        return None
+
+# Load cache from Hugging Face Hub
+CACHE_KEY = "image_text_cache"
+cache = load_from_hub(DATASET_NAME)
+if cache and CACHE_KEY in cache:
+    cache = cache[CACHE_KEY]
 else:
     cache = {}
 
@@ -180,20 +201,13 @@ new_images = process_new_images(new_images)
 # Update the combined DataFrame with extracted text
 combined_df.update(new_images)
 
-# Save the updated dataset to persistent storage
-output_file = os.path.join(PERSISTENT_DIR, "data_with_extracted_text.csv")
-combined_df.to_csv(output_file, index=False)
-logging.info(f"Updated dataset saved to {output_file}")
-
-# Save the cache to persistent storage
-with open(CACHE_FILE, "wb") as f:
-    pickle.dump(cache, f)
-logging.info(f"Cache saved to {CACHE_FILE}")
+# Save the updated dataset to Hugging Face Hub
+save_to_hub({"data_with_extracted_text": combined_df.to_dict(), CACHE_KEY: cache}, DATASET_NAME)
 
 logging.info("Incremental processing complete!")
 
 # Load or download the ResNet model from Hugging Face
-MODEL_DIR = os.path.join(PERSISTENT_DIR, "models")
+MODEL_DIR = os.path.join(WORKING_DIR, "models")
 os.makedirs(MODEL_DIR, exist_ok=True)
 MODEL_PATH = os.path.join(MODEL_DIR, "resnet50-0676ba61.pth")