Spaces:
Sleeping
Sleeping
File size: 1,785 Bytes
1970a5b | 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 | import os
import datetime
from pymongo import MongoClient
from pymongo.server_api import ServerApi
from dotenv import load_dotenv
load_dotenv()
MONGO_URI = os.getenv("MONGO_URI", "mongodb://localhost:27017/")
DB_NAME = "deepshield"
COLLECTION_NAME = "detections"
_db_client = None
def get_db():
global _db_client
if _db_client is None:
try:
# Connect to MongoDB Atlas or local MongoDB
_db_client = MongoClient(MONGO_URI, server_api=ServerApi('1') if "mongodb.net" in MONGO_URI else None)
# Send a ping to confirm a successful connection
_db_client.admin.command('ping')
print("Successfully connected to MongoDB!")
except Exception as e:
print(f"MongoDB connection error: {e}")
return None
return _db_client[DB_NAME]
def save_detection(session_id, filename, modality, verdict, confidence, details=None):
db = get_db()
if db is None:
return False
try:
record = {
"session_id": session_id,
"filename": filename,
"modality": modality,
"verdict": verdict,
"confidence": confidence,
"details": details or {},
"timestamp": datetime.datetime.utcnow()
}
db[COLLECTION_NAME].insert_one(record)
return True
except Exception as e:
print(f"Error saving to MongoDB: {e}")
return False
def get_recent_detections(limit=50):
db = get_db()
if db is None:
return []
try:
cursor = db[COLLECTION_NAME].find({}, {"_id": 0}).sort("timestamp", -1).limit(limit)
return list(cursor)
except Exception as e:
print(f"Error reading from MongoDB: {e}")
return []
|