Spaces:
Sleeping
Sleeping
| """ | |
| MongoDB database setup and models for storing classification results. | |
| Uses Motor (async MongoDB driver) for FastAPI integration. | |
| """ | |
| import time | |
| import os | |
| import base64 | |
| from datetime import datetime, timedelta | |
| from typing import Optional, List, Dict, Any | |
| from motor.motor_asyncio import AsyncIOMotorClient | |
| from pymongo.errors import ConnectionFailure | |
| import logging | |
| from .config import settings | |
| logger = logging.getLogger(__name__) | |
| # MongoDB connection settings | |
| MONGODB_URL = settings.MONGODB_URL | |
| DATABASE_NAME = settings.MONGODB_DATABASE | |
| # Global MongoDB client | |
| mongodb_client: Optional[AsyncIOMotorClient] = None | |
| mongodb_db = None | |
| gridfs_bucket = None | |
| async def connect_to_mongodb(): | |
| """Connect to MongoDB database.""" | |
| global mongodb_client, mongodb_db, gridfs_bucket | |
| if not MONGODB_URL: | |
| raise ValueError( | |
| "MONGODB_URL environment variable is not set. " | |
| "Please set it to your MongoDB connection string." | |
| ) | |
| try: | |
| from motor.motor_asyncio import AsyncIOMotorGridFSBucket | |
| mongodb_client = AsyncIOMotorClient(MONGODB_URL) | |
| # Test connection | |
| await mongodb_client.admin.command('ping') | |
| mongodb_db = mongodb_client[DATABASE_NAME] | |
| # Initialize GridFS bucket for large file storage (videos) | |
| gridfs_bucket = AsyncIOMotorGridFSBucket(mongodb_db) | |
| logger.info(f"Connected to MongoDB database: {DATABASE_NAME}") | |
| return mongodb_db | |
| except ConnectionFailure as e: | |
| logger.error(f"Failed to connect to MongoDB: {e}") | |
| raise | |
| async def close_mongodb_connection(): | |
| """Close MongoDB connection.""" | |
| global mongodb_client | |
| if mongodb_client: | |
| mongodb_client.close() | |
| logger.info("MongoDB connection closed") | |
| def get_database(): | |
| """Get MongoDB database instance.""" | |
| if mongodb_db is None: | |
| raise RuntimeError("MongoDB database not initialized. Call connect_to_mongodb() first.") | |
| return mongodb_db | |
| def get_gridfs_bucket(): | |
| """Get GridFS bucket for large file storage.""" | |
| if gridfs_bucket is None: | |
| raise RuntimeError("GridFS bucket not initialized. Call connect_to_mongodb() first.") | |
| return gridfs_bucket | |
| async def create_indexes(): | |
| """Create database indexes for better query performance.""" | |
| db = get_database() | |
| # Image classifications collection | |
| collection = db["classifications"] | |
| await collection.create_index("id", unique=True) | |
| await collection.create_index("timestamp") | |
| await collection.create_index([("latitude", 1), ("longitude", 1)]) | |
| # Video classifications collection | |
| video_collection = db["video_classifications"] | |
| await video_collection.create_index("id", unique=True) | |
| await video_collection.create_index("timestamp") | |
| await video_collection.create_index([("latitude", 1), ("longitude", 1)]) | |
| # Job tracking collection | |
| jobs_collection = db["jobs"] | |
| await jobs_collection.create_index("job_id", unique=True) | |
| await jobs_collection.create_index("status") | |
| await jobs_collection.create_index("created_at") | |
| logger.info("MongoDB indexes created successfully") | |
| class ClassificationRecord: | |
| """ | |
| MongoDB document model for classification results. | |
| """ | |
| def encode_image(image_bytes: bytes) -> str: | |
| """Encode image bytes to base64 string for MongoDB storage.""" | |
| return base64.b64encode(image_bytes).decode('utf-8') | |
| def decode_image(image_base64: str) -> bytes: | |
| """Decode base64 string back to image bytes.""" | |
| return base64.b64decode(image_base64.encode('utf-8')) | |
| async def create( | |
| record_id: str, | |
| image_bytes: bytes, | |
| latitude: Optional[float], | |
| longitude: Optional[float], | |
| flowers: List[Dict[str, Any]], | |
| filename: str = "image.jpg" | |
| ) -> Dict[str, Any]: | |
| """Create a new classification record in MongoDB. | |
| Note: Only stores the annotated image (with bounding boxes), not the original. | |
| """ | |
| db = get_database() | |
| collection = db["classifications"] | |
| # Encode annotated image to base64 | |
| image_base64 = ClassificationRecord.encode_image(image_bytes) | |
| document = { | |
| "id": record_id, | |
| "image_base64": image_base64, | |
| "image_filename": filename, | |
| "image_content_type": "image/jpeg", | |
| "latitude": latitude, | |
| "longitude": longitude, | |
| "timestamp": datetime.utcnow(), | |
| "flowers": flowers, | |
| } | |
| result = await collection.insert_one(document) | |
| return document | |
| async def get_by_id(record_id: str) -> Optional[Dict[str, Any]]: | |
| """Get classification record by ID.""" | |
| db = get_database() | |
| collection = db["classifications"] | |
| return await collection.find_one({"id": record_id}) | |
| async def get_all() -> List[Dict[str, Any]]: | |
| db = get_database() | |
| collection = db["classifications"] | |
| cursor = ( | |
| collection | |
| .find( | |
| {}, | |
| { | |
| "id": 1, | |
| "flowers": 1, | |
| "timestamp": 1, | |
| "_id": 0, # exclude Mongo ObjectId | |
| }, | |
| ) | |
| .sort("timestamp", -1) | |
| ) | |
| return await cursor.to_list(length=None) | |
| async def delete_by_id(record_id: str) -> bool: | |
| """Delete classification record by ID.""" | |
| db = get_database() | |
| collection = db["classifications"] | |
| result = await collection.delete_one({"id": record_id}) | |
| return result.deleted_count > 0 | |
| async def delete_all() -> int: | |
| """Delete all classification records.""" | |
| db = get_database() | |
| collection = db["classifications"] | |
| result = await collection.delete_many({}) | |
| return result.deleted_count | |
| async def count() -> int: | |
| """Get total number of classification records.""" | |
| db = get_database() | |
| collection = db["classifications"] | |
| return await collection.count_documents({}) | |
| class VideoClassificationRecord: | |
| """ | |
| MongoDB document model for video classification results. | |
| Uses GridFS for storing large video files. | |
| """ | |
| async def create( | |
| record_id: str, | |
| video_bytes: bytes, | |
| latitude: Optional[float], | |
| longitude: Optional[float], | |
| frame_results: List[List[Dict[str, Any]]], | |
| video_metadata: Dict[str, Any], | |
| filename: str = "video.mp4" | |
| ) -> Dict[str, Any]: | |
| """ | |
| Create a new video classification record in MongoDB. | |
| Stores video in GridFS (for large files) and metadata in regular collection. | |
| Args: | |
| record_id: Unique identifier for the record | |
| video_bytes: Raw video bytes (annotated) | |
| latitude: GPS latitude | |
| longitude: GPS longitude | |
| frame_results: List of detection lists (one per frame) | |
| video_metadata: Dict with 'total_frames', 'fps', 'duration' | |
| filename: Original filename | |
| Returns: | |
| Created document | |
| """ | |
| db = get_database() | |
| bucket = get_gridfs_bucket() | |
| collection = db["video_classifications"] | |
| # Store video in GridFS | |
| video_id = await bucket.upload_from_stream( | |
| filename, | |
| video_bytes, | |
| metadata={ | |
| "record_id": record_id, | |
| "content_type": "video/mp4", | |
| "uploaded_at": datetime.utcnow() | |
| } | |
| ) | |
| # Aggregate statistics across all frames | |
| stage_summary = {0: 0, 1: 0, 2: 0} | |
| total_detections = 0 | |
| for frame_detections in frame_results: | |
| total_detections += len(frame_detections) | |
| for det in frame_detections: | |
| stage = det['stage'] | |
| stage_summary[stage] = stage_summary.get(stage, 0) + 1 | |
| # Convert integer keys to strings for MongoDB compatibility | |
| stage_summary = {str(k): v for k, v in stage_summary.items()} | |
| # Calculate average flowers per frame | |
| total_frames = video_metadata.get('total_frames', len(frame_results)) | |
| avg_flowers = total_detections / total_frames if total_frames > 0 else 0 | |
| # Store metadata in regular collection | |
| document = { | |
| "id": record_id, | |
| "video_gridfs_id": video_id, | |
| "video_filename": filename, | |
| "video_content_type": "video/mp4", | |
| "latitude": latitude, | |
| "longitude": longitude, | |
| "timestamp": datetime.utcnow(), | |
| "total_frames": total_frames, | |
| "fps": video_metadata.get('fps', 0), | |
| "duration_seconds": video_metadata.get('duration', 0), | |
| "frame_results": frame_results, # Store all frame-by-frame detections | |
| "stage_summary": stage_summary, | |
| "total_detections": total_detections, | |
| "average_flowers_per_frame": avg_flowers, | |
| } | |
| await collection.insert_one(document) | |
| logger.info(f"Video classification record created: {record_id}") | |
| return document | |
| async def get_by_id(record_id: str) -> Optional[Dict[str, Any]]: | |
| """Get video classification record by ID.""" | |
| db = get_database() | |
| collection = db["video_classifications"] | |
| return await collection.find_one({"id": record_id}) | |
| async def get_all() -> List[Dict[str, Any]]: | |
| """Get all video classification records, sorted by timestamp.""" | |
| db = get_database() | |
| collection = db["video_classifications"] | |
| cursor = collection.find({}).sort("timestamp", -1) | |
| return await cursor.to_list(length=None) | |
| async def get_video_bytes(record_id: str) -> Optional[bytes]: | |
| """ | |
| Retrieve video bytes from GridFS. | |
| Args: | |
| record_id: Record identifier | |
| Returns: | |
| Video bytes or None if not found | |
| """ | |
| record = await VideoClassificationRecord.get_by_id(record_id) | |
| if not record: | |
| return None | |
| bucket = get_gridfs_bucket() | |
| video_id = record.get('video_gridfs_id') | |
| if not video_id: | |
| logger.error(f"No video_gridfs_id found for record {record_id}") | |
| return None | |
| try: | |
| # Download from GridFS | |
| grid_out = await bucket.open_download_stream(video_id) | |
| video_bytes = await grid_out.read() | |
| return video_bytes | |
| except Exception as e: | |
| logger.error(f"Error downloading video from GridFS: {e}") | |
| return None | |
| async def delete_by_id(record_id: str) -> bool: | |
| """ | |
| Delete video classification record by ID. | |
| Also deletes the video file from GridFS. | |
| """ | |
| db = get_database() | |
| bucket = get_gridfs_bucket() | |
| collection = db["video_classifications"] | |
| # Get record first to get GridFS ID | |
| record = await VideoClassificationRecord.get_by_id(record_id) | |
| if not record: | |
| return False | |
| # Delete video from GridFS | |
| video_id = record.get('video_gridfs_id') | |
| if video_id: | |
| try: | |
| await bucket.delete(video_id) | |
| logger.info(f"Deleted video from GridFS: {video_id}") | |
| except Exception as e: | |
| logger.error(f"Error deleting video from GridFS: {e}") | |
| # Delete metadata document | |
| result = await collection.delete_one({"id": record_id}) | |
| return result.deleted_count > 0 | |
| async def delete_all() -> int: | |
| """Delete all video classification records and their GridFS video files.""" | |
| db = get_database() | |
| bucket = get_gridfs_bucket() | |
| collection = db["video_classifications"] | |
| records = await collection.find({}, {"video_gridfs_id": 1}).to_list(length=None) | |
| for record in records: | |
| video_id = record.get("video_gridfs_id") | |
| if video_id: | |
| try: | |
| await bucket.delete(video_id) | |
| except Exception as e: | |
| logger.error(f"Error deleting video from GridFS: {e}") | |
| result = await collection.delete_many({}) | |
| return result.deleted_count | |
| async def count() -> int: | |
| """Get total number of video classification records.""" | |
| db = get_database() | |
| collection = db["video_classifications"] | |
| return await collection.count_documents({}) | |
| class JobRecord: | |
| """ | |
| MongoDB document model for async job tracking. | |
| """ | |
| async def create( | |
| job_id: str, | |
| job_type: str, | |
| status: str = "queued", | |
| metadata: Optional[Dict[str, Any]] = None | |
| ) -> Dict[str, Any]: | |
| """Create a new job record.""" | |
| db = get_database() | |
| collection = db["jobs"] | |
| document = { | |
| "job_id": job_id, | |
| "job_type": job_type, | |
| "status": status, | |
| "progress": 0, | |
| "message": None, | |
| "created_at": datetime.utcnow(), | |
| "updated_at": datetime.utcnow(), | |
| "completed_at": None, | |
| "metadata": metadata or {}, | |
| "result": None, | |
| "error": None, | |
| } | |
| await collection.insert_one(document) | |
| logger.info(f"Job record created: {job_id}") | |
| return document | |
| async def update_status( | |
| job_id: str, | |
| status: str, | |
| progress: Optional[int] = None, | |
| message: Optional[str] = None, | |
| result: Optional[Dict[str, Any]] = None, | |
| error: Optional[str] = None | |
| ) -> bool: | |
| """Update job status.""" | |
| db = get_database() | |
| collection = db["jobs"] | |
| update_data = { | |
| "status": status, | |
| "updated_at": datetime.utcnow(), | |
| } | |
| if progress is not None: | |
| update_data["progress"] = progress | |
| if message is not None: | |
| update_data["message"] = message | |
| if result is not None: | |
| update_data["result"] = result | |
| if error is not None: | |
| update_data["error"] = error | |
| if status in ["completed", "failed"]: | |
| update_data["completed_at"] = datetime.utcnow() | |
| result = await collection.update_one( | |
| {"job_id": job_id}, | |
| {"$set": update_data} | |
| ) | |
| return result.modified_count > 0 | |
| async def get_by_id(job_id: str) -> Optional[Dict[str, Any]]: | |
| """Get job record by ID.""" | |
| db = get_database() | |
| collection = db["jobs"] | |
| return await collection.find_one({"job_id": job_id}) | |
| async def get_active_jobs() -> List[Dict[str, Any]]: | |
| """Get all active (not completed/failed) jobs.""" | |
| db = get_database() | |
| collection = db["jobs"] | |
| cursor = collection.find({ | |
| "status": {"$in": ["queued", "processing"]} | |
| }).sort("created_at", -1) | |
| return await cursor.to_list(length=None) | |
| async def delete_by_id(job_id: str) -> bool: | |
| """Delete job record by ID.""" | |
| db = get_database() | |
| collection = db["jobs"] | |
| result = await collection.delete_one({"job_id": job_id}) | |
| return result.deleted_count > 0 | |
| async def delete_all() -> int: | |
| """Delete all job records.""" | |
| db = get_database() | |
| collection = db["jobs"] | |
| result = await collection.delete_many({}) | |
| return result.deleted_count | |
| async def cleanup_old_jobs(days: int = 7) -> int: | |
| """Delete job records older than specified days.""" | |
| db = get_database() | |
| collection = db["jobs"] | |
| cutoff_date = datetime.utcnow() - timedelta(days=days) | |
| result = await collection.delete_many({ | |
| "created_at": {"$lt": cutoff_date}, | |
| "status": {"$in": ["completed", "failed"]} | |
| }) | |
| logger.info(f"Cleaned up {result.deleted_count} old job records") | |
| return result.deleted_count | |