Spaces:
Running
Running
File size: 5,230 Bytes
887f5f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import os
import json
import firebase_admin
from firebase_admin import credentials, firestore
class FirebaseManager:
"""
Manages connection to Firebase Firestore for caching deepfake detection results
and tracking async task states. Falls back to in-memory dictionaries if Firebase is offline.
Uses a single collection 'deepfake_analyses' for all documents.
"""
def __init__(self):
self.db = None
self.local_cache = {}
cred_json_str = os.environ.get("FIREBASE_CREDENTIALS_JSON")
cred_file_path = os.environ.get("FIREBASE_SERVICE_ACCOUNT_KEY")
cred = None
try:
if cred_json_str:
cred_info = json.loads(cred_json_str)
cred = credentials.Certificate(cred_info)
elif cred_file_path and os.path.exists(cred_file_path):
cred = credentials.Certificate(cred_file_path)
if not firebase_admin._apps:
if cred:
firebase_admin.initialize_app(cred)
else:
print("No Firebase credentials provided. Skipping Firebase initialization.")
raise ValueError("Missing Firebase credentials")
self.db = firestore.client()
print("Connected to Firebase Firestore successfully!")
except Exception as e:
print(f"Warning: Firebase is not reachable ({e}). Falling back to in-memory cache.")
self.db = None
def get_cached_result(self, file_hash: str) -> dict | None:
if self.db:
try:
doc_ref = self.db.collection("deepfake_analyses").document(file_hash)
doc = doc_ref.get()
if doc.exists:
return doc.to_dict()
except Exception as e:
print(f"Firebase get cache error: {e}")
return self.local_cache.get(file_hash)
def cache_result(self, file_hash: str, doc_data: dict) -> None:
self.local_cache[file_hash] = doc_data
if self.db:
try:
doc_ref = self.db.collection("deepfake_analyses").document(file_hash)
doc_ref.set(doc_data, merge=True)
except Exception as e:
print(f"Firebase set cache error: {e}")
def get_task_status(self, task_id: str) -> dict | None:
if self.db:
try:
docs = self.db.collection("deepfake_analyses").where("task_id", "==", task_id).limit(1).get()
if docs:
return docs[0].to_dict()
except Exception as e:
print(f"Firebase get task error: {e}")
# Local search
for k, v in self.local_cache.items():
if v.get("task_id") == task_id:
return v
return None
def update_task_status(self, file_hash: str, status: dict) -> None:
existing = self.local_cache.get(file_hash, {})
existing.update(status)
self.local_cache[file_hash] = existing
if self.db:
try:
doc_ref = self.db.collection("deepfake_analyses").document(file_hash)
doc_ref.set(status, merge=True)
except Exception as e:
print(f"Firebase set task error: {e}")
def get_history(self) -> list[dict]:
history = []
if self.db:
try:
docs = self.db.collection("deepfake_analyses").stream()
for doc in docs:
res = doc.to_dict()
if res and "result" in res:
history.append(self._strip_heavy_payload(res))
history.sort(key=lambda x: x.get("timestamp", 0), reverse=True)
return history
except Exception as e:
print(f"Firebase get history error: {e}")
for file_hash, res in self.local_cache.items():
if "result" in res:
history.append(self._strip_heavy_payload(res))
history.sort(key=lambda x: x.get("timestamp", 0), reverse=True)
return history
def _strip_heavy_payload(self, res: dict) -> dict:
light_res = res.copy()
if "image" in light_res:
light_res["image"] = None
if "result" in light_res and isinstance(light_res["result"], dict):
if "frames" in light_res["result"]:
for frame in light_res["result"]["frames"]:
if "image" in frame:
frame["image"] = None
if "faces" in frame:
for face in frame["faces"]:
if "crop_b64" in face:
face["crop_b64"] = None
return light_res
def clear_cache(self) -> None:
self.local_cache.clear()
if self.db:
try:
results_ref = self.db.collection("deepfake_analyses")
docs = results_ref.list_documents()
for doc in docs:
doc.delete()
print("Firebase cache cleared.")
except Exception as e:
print(f"Firebase clear cache error: {e}")
|