Update database.py

database.py

@@ -13,6 +13,7 @@ from typing import List, Dict, Optional, Tuple, Any
 from pathlib import Path
 import pandas as pd
 
+
 class DogDatabase:
     """SQLite database manager for dog monitoring system"""
     
@@ -38,54 +39,54 @@ class DogDatabase:
                 last_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                 total_sightings INTEGER DEFAULT 1,
                 notes TEXT,
-                merged_from TEXT,  -- JSON list of merged dog IDs
-                status TEXT DEFAULT 'active'  -- active, merged, deleted
+                merged_from TEXT,
+                status TEXT DEFAULT 'active'
             )
         """)
         
-        # Dog features table - stores extracted features
+        # Dog features table
         self.cursor.execute("""
             CREATE TABLE IF NOT EXISTS dog_features (
                 feature_id INTEGER PRIMARY KEY AUTOINCREMENT,
                 dog_id INTEGER,
-                resnet_features BLOB,  -- Pickled numpy array
-                color_histogram BLOB,  -- Pickled numpy array
+                resnet_features BLOB,
+                color_histogram BLOB,
                 timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                 confidence REAL,
                 FOREIGN KEY (dog_id) REFERENCES dogs(dog_id)
             )
         """)
         
-        # Dog images table - stores actual images
+        # Dog images table
         self.cursor.execute("""
             CREATE TABLE IF NOT EXISTS dog_images (
                 image_id INTEGER PRIMARY KEY AUTOINCREMENT,
                 dog_id INTEGER,
-                image_data BLOB,  -- Base64 encoded image
-                thumbnail BLOB,    -- Small preview
+                image_data BLOB,
+                thumbnail BLOB,
                 width INTEGER,
                 height INTEGER,
                 timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                 frame_number INTEGER,
                 video_source TEXT,
-                bbox TEXT,  -- JSON [x1, y1, x2, y2]
+                bbox TEXT,
                 confidence REAL,
-                pose_keypoints TEXT, -- JSON of keypoints
+                pose_keypoints TEXT,
                 is_validated BOOLEAN DEFAULT 0,
                 is_discarded BOOLEAN DEFAULT 0,
                 FOREIGN KEY (dog_id) REFERENCES dogs(dog_id)
             )
         """)
         
-        # Body parts table - stores cropped body parts
+        # Body parts table
         self.cursor.execute("""
             CREATE TABLE IF NOT EXISTS body_parts (
                 part_id INTEGER PRIMARY KEY AUTOINCREMENT,
                 dog_id INTEGER,
                 image_id INTEGER,
-                part_type TEXT,  -- 'head', 'torso', 'rear'
-                part_image BLOB,  -- Base64 encoded crop
-                crop_bbox TEXT,  -- JSON [x1, y1, x2, y2] relative to full image
+                part_type TEXT,
+                part_image BLOB,
+                crop_bbox TEXT,
                 confidence REAL,
                 is_validated BOOLEAN DEFAULT 0,
                 is_discarded BOOLEAN DEFAULT 0,
@@ -94,7 +95,7 @@ class DogDatabase:
             )
         """)
         
-        # Sightings table - tracks when/where dogs were seen
+        # Sightings table
         self.cursor.execute("""
             CREATE TABLE IF NOT EXISTS sightings (
                 sighting_id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -118,11 +119,11 @@ class DogDatabase:
                 video_path TEXT,
                 total_frames INTEGER,
                 dogs_detected INTEGER,
-                settings TEXT  -- JSON of processing settings
+                settings TEXT
             )
         """)
         
-        # Create indexes for performance
+        # Create indexes
         self.cursor.execute("CREATE INDEX IF NOT EXISTS idx_dog_features ON dog_features(dog_id)")
         self.cursor.execute("CREATE INDEX IF NOT EXISTS idx_dog_images ON dog_images(dog_id)")
         self.cursor.execute("CREATE INDEX IF NOT EXISTS idx_sightings ON sightings(dog_id)")
@@ -130,7 +131,6 @@ class DogDatabase:
         self.conn.commit()
     
     # ========== Dog Management ==========
-    
     def add_dog(self, dog_id: Optional[int] = None, name: Optional[str] = None) -> int:
         """Add a new dog to the database"""
         if dog_id:
@@ -144,7 +144,6 @@ class DogDatabase:
                 (name,)
             )
             dog_id = self.cursor.lastrowid
-        
         self.conn.commit()
         return dog_id
     
@@ -161,21 +160,15 @@ class DogDatabase:
     def merge_dogs(self, keep_id: int, merge_id: int) -> bool:
         """Merge two dogs, keeping keep_id"""
         try:
-            # Update all references
-            self.cursor.execute("UPDATE dog_features SET dog_id = ? WHERE dog_id = ?", 
-                              (keep_id, merge_id))
-            self.cursor.execute("UPDATE dog_images SET dog_id = ? WHERE dog_id = ?", 
-                              (keep_id, merge_id))
-            self.cursor.execute("UPDATE sightings SET dog_id = ? WHERE dog_id = ?", 
-                              (keep_id, merge_id))
+            self.cursor.execute("UPDATE dog_features SET dog_id = ? WHERE dog_id = ?", (keep_id, merge_id))
+            self.cursor.execute("UPDATE dog_images SET dog_id = ? WHERE dog_id = ?", (keep_id, merge_id))
+            self.cursor.execute("UPDATE sightings SET dog_id = ? WHERE dog_id = ?", (keep_id, merge_id))
             
-            # Get merged_from history
             self.cursor.execute("SELECT merged_from FROM dogs WHERE dog_id = ?", (merge_id,))
             row = self.cursor.fetchone()
             merged_history = json.loads(row['merged_from'] if row and row['merged_from'] else '[]')
             merged_history.append(merge_id)
             
-            # Update keep_id dog with merge history
             self.cursor.execute("""
                 UPDATE dogs 
                 SET merged_from = ?,
@@ -185,12 +178,7 @@ class DogDatabase:
                 WHERE dog_id = ?
             """, (json.dumps(merged_history), merge_id, keep_id))
             
-            # Mark merge_id as merged
-            self.cursor.execute(
-                "UPDATE dogs SET status = 'merged' WHERE dog_id = ?", 
-                (merge_id,)
-            )
-            
+            self.cursor.execute("UPDATE dogs SET status = 'merged' WHERE dog_id = ?", (merge_id,))
             self.conn.commit()
             return True
         except Exception as e:
@@ -201,33 +189,25 @@ class DogDatabase:
     def delete_dog(self, dog_id: int, hard_delete: bool = False):
         """Delete or mark dog as deleted"""
         if hard_delete:
-            # Hard delete - remove all data
             self.cursor.execute("DELETE FROM dog_features WHERE dog_id = ?", (dog_id,))
             self.cursor.execute("DELETE FROM dog_images WHERE dog_id = ?", (dog_id,))
             self.cursor.execute("DELETE FROM sightings WHERE dog_id = ?", (dog_id,))
             self.cursor.execute("DELETE FROM dogs WHERE dog_id = ?", (dog_id,))
         else:
-            # Soft delete - mark as deleted
-            self.cursor.execute(
-                "UPDATE dogs SET status = 'deleted' WHERE dog_id = ?", 
-                (dog_id,)
-            )
+            self.cursor.execute("UPDATE dogs SET status = 'deleted' WHERE dog_id = ?", (dog_id,))
         self.conn.commit()
     
     # ========== Features Management ==========
-    
     def save_features(self, dog_id: int, resnet_features: np.ndarray, 
                      color_histogram: np.ndarray, confidence: float):
         """Save dog features to database"""
         resnet_blob = pickle.dumps(resnet_features)
         color_blob = pickle.dumps(color_histogram)
-        
         self.cursor.execute("""
             INSERT INTO dog_features 
             (dog_id, resnet_features, color_histogram, confidence)
             VALUES (?, ?, ?, ?)
         """, (dog_id, resnet_blob, color_blob, confidence))
-        
         self.conn.commit()
     
     def get_features(self, dog_id: int, limit: int = 20) -> List[Dict]:
@@ -238,7 +218,6 @@ class DogDatabase:
             ORDER BY timestamp DESC
             LIMIT ?
         """, (dog_id, limit))
-        
         features = []
         for row in self.cursor.fetchall():
             features.append({
@@ -247,27 +226,22 @@ class DogDatabase:
                 'confidence': row['confidence'],
                 'timestamp': row['timestamp']
             })
-        
         return features
     
     # ========== Images Management ==========
-    
     def save_image(self, dog_id: int, image: np.ndarray, 
                   frame_number: int, video_source: str,
                   bbox: List[float], confidence: float,
                   pose_keypoints: Optional[List] = None):
         """Save dog image to database"""
-        # Encode image as JPEG
         _, buffer = cv2.imencode('.jpg', image)
         image_data = base64.b64encode(buffer).decode('utf-8')
         
-        # Create thumbnail
         thumbnail = cv2.resize(image, (128, 128))
         _, thumb_buffer = cv2.imencode('.jpg', thumbnail, [cv2.IMWRITE_JPEG_QUALITY, 70])
         thumb_data = base64.b64encode(thumb_buffer).decode('utf-8')
         
         h, w = image.shape[:2]
-        
         self.cursor.execute("""
             INSERT INTO dog_images 
             (dog_id, image_data, thumbnail, width, height, 
@@ -276,21 +250,18 @@ class DogDatabase:
         """, (dog_id, image_data, thumb_data, w, h, 
               frame_number, video_source, json.dumps(bbox), 
               confidence, json.dumps(pose_keypoints) if pose_keypoints else None))
-        
         self.conn.commit()
-        return self.cursor.lastrowid  # THIS LINE - make sure it returns the image_id
+        return self.cursor.lastrowid
     
     def get_dog_images(self, dog_id: int, validated_only: bool = False,
                   include_discarded: bool = False) -> List[Dict]:
         """Get all images for a dog"""
         query = "SELECT * FROM dog_images WHERE dog_id = ?"
         params = [dog_id]
-        
         if validated_only:
             query += " AND is_validated = 1"
         if not include_discarded:
             query += " AND is_discarded = 0"
-        
         query += " ORDER BY timestamp DESC"
         
         self.cursor.execute(query, params)
@@ -298,11 +269,9 @@ class DogDatabase:
         
         images = []
         for row in rows:
-            # Decode image
             image_bytes = base64.b64decode(row['image_data'])
             nparr = np.frombuffer(image_bytes, np.uint8)
             image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
-            
             images.append({
                 'image_id': row['image_id'],
                 'image': image,
@@ -315,26 +284,17 @@ class DogDatabase:
                 'is_discarded': row['is_discarded'],
                 'pose_keypoints': json.loads(row['pose_keypoints']) if row['pose_keypoints'] else None
             })
-        
         return images
-
     
     def validate_image(self, image_id: int, is_valid: bool = True):
         """Mark image as validated or discarded"""
         if is_valid:
-            self.cursor.execute(
-                "UPDATE dog_images SET is_validated = 1 WHERE image_id = ?",
-                (image_id,)
-            )
+            self.cursor.execute("UPDATE dog_images SET is_validated = 1 WHERE image_id = ?", (image_id,))
         else:
-            self.cursor.execute(
-                "UPDATE dog_images SET is_discarded = 1 WHERE image_id = ?",
-                (image_id,)
-            )
+            self.cursor.execute("UPDATE dog_images SET is_discarded = 1 WHERE image_id = ?", (image_id,))
         self.conn.commit()
     
     # ========== Body Parts Management ==========
-    
     def save_body_parts(self, dog_id: int, image_id: int, 
                         head_crop: Optional[np.ndarray],
                         torso_crop: Optional[np.ndarray],
@@ -346,21 +306,16 @@ class DogDatabase:
             'torso': torso_crop,
             'rear': rear_crop
         }
-        
         for part_type, crop in parts.items():
             if crop is not None:
-                # Encode crop as JPEG
                 _, buffer = cv2.imencode('.jpg', crop)
                 crop_data = base64.b64encode(buffer).decode('utf-8')
-                
                 confidence = confidences.get(part_type, 0.0)
-                
                 self.cursor.execute("""
                     INSERT INTO body_parts 
                     (dog_id, image_id, part_type, part_image, confidence)
                     VALUES (?, ?, ?, ?, ?)
                 """, (dog_id, image_id, part_type, crop_data, confidence))
-        
         self.conn.commit()
     
     def get_body_parts(self, dog_id: int, part_type: Optional[str] = None,
@@ -368,26 +323,20 @@ class DogDatabase:
         """Get body part crops for a dog"""
         query = "SELECT * FROM body_parts WHERE dog_id = ?"
         params = [dog_id]
-        
         if part_type:
             query += " AND part_type = ?"
             params.append(part_type)
-        
         if validated_only:
             query += " AND is_validated = 1"
-        
         if not include_discarded:
             query += " AND is_discarded = 0"
-        
         self.cursor.execute(query, params)
         
         parts = []
         for row in self.cursor.fetchall():
-            # Decode image
             image_bytes = base64.b64decode(row['part_image'])
             nparr = np.frombuffer(image_bytes, np.uint8)
             image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
-            
             parts.append({
                 'part_id': row['part_id'],
                 'part_type': row['part_type'],
@@ -396,21 +345,14 @@ class DogDatabase:
                 'is_validated': row['is_validated'],
                 'image_id': row['image_id']
             })
-        
         return parts
     
     def validate_body_part(self, part_id: int, is_valid: bool = True):
         """Mark body part as validated or discarded"""
         if is_valid:
-            self.cursor.execute(
-                "UPDATE body_parts SET is_validated = 1 WHERE part_id = ?",
-                (part_id,)
-            )
+            self.cursor.execute("UPDATE body_parts SET is_validated = 1 WHERE part_id = ?", (part_id,))
         else:
-            self.cursor.execute(
-                "UPDATE body_parts SET is_discarded = 1 WHERE part_id = ?",
-                (part_id,)
-            )
+            self.cursor.execute("UPDATE body_parts SET is_discarded = 1 WHERE part_id = ?", (part_id,))
         self.conn.commit()
     
     def add_sighting(self, dog_id: int, position: Tuple[float, float],
@@ -421,41 +363,32 @@ class DogDatabase:
             (dog_id, position_x, position_y, video_source, frame_number, confidence)
             VALUES (?, ?, ?, ?, ?, ?)
         """, (dog_id, position[0], position[1], video_source, frame_number, confidence))
-        
         self.conn.commit()
     
     # ========== Query Methods ==========
-    
     def get_all_dogs(self, active_only: bool = True) -> pd.DataFrame:
         """Get all dogs as DataFrame"""
         query = "SELECT * FROM dogs"
         if active_only:
             query += " WHERE status = 'active'"
         query += " ORDER BY dog_id"
-        
         return pd.read_sql_query(query, self.conn)
     
     def get_dog_statistics(self) -> Dict:
         """Get overall statistics"""
         stats = {}
-        
-        # Total dogs
         self.cursor.execute("SELECT COUNT(*) FROM dogs WHERE status = 'active'")
         stats['total_active_dogs'] = self.cursor.fetchone()[0]
         
-        # Total images
         self.cursor.execute("SELECT COUNT(*) FROM dog_images WHERE is_discarded = 0")
         stats['total_images'] = self.cursor.fetchone()[0]
         
-        # Validated images
         self.cursor.execute("SELECT COUNT(*) FROM dog_images WHERE is_validated = 1")
         stats['validated_images'] = self.cursor.fetchone()[0]
         
-        # Total sightings
         self.cursor.execute("SELECT COUNT(*) FROM sightings")
         stats['total_sightings'] = self.cursor.fetchone()[0]
         
-        # Most seen dog
         self.cursor.execute("""
             SELECT d.dog_id, d.name, d.total_sightings
             FROM dogs d
@@ -470,47 +403,33 @@ class DogDatabase:
                 'name': row[1] or f"Dog #{row[0]}",
                 'sightings': row[2]
             }
-        
         return stats
     
     # ========== Export Methods ==========
-    
     def export_training_dataset(self, output_dir: str, validated_only: bool = True) -> Dict:
         """Export dataset with body parts for fine-tuning"""
         output_path = Path(output_dir)
         output_path.mkdir(parents=True, exist_ok=True)
         
-        # Create directories
         images_dir = output_path / "images"
         images_dir.mkdir(exist_ok=True)
         
-        # Export data
         dataset = []
-        
         dogs = self.get_all_dogs()
         for _, dog in dogs.iterrows():
             dog_id = dog['dog_id']
-            
-            # Create directories for each dog
             dog_dir = images_dir / f"dog_{dog_id}"
             dog_dir.mkdir(exist_ok=True)
             
-            # Subdirectories for body parts
             for part in ['full', 'head', 'torso', 'rear']:
-                part_dir = dog_dir / part
-                part_dir.mkdir(exist_ok=True)
+                (dog_dir / part).mkdir(exist_ok=True)
             
-            # Get full images
             images = self.get_dog_images(dog_id, validated_only=validated_only)
-            
             for idx, img_data in enumerate(images):
-                # Save full image
                 full_path = dog_dir / 'full' / f"img_{idx:04d}.jpg"
                 cv2.imwrite(str(full_path), img_data['image'])
                 
-                # Get and save body parts for this image
                 parts = self.get_body_parts(dog_id, validated_only=validated_only)
-                
                 part_paths = {}
                 for part_data in parts:
                     if part_data['image_id'] == img_data['image_id']:
@@ -519,25 +438,19 @@ class DogDatabase:
                         cv2.imwrite(str(part_path), part_data['image'])
                         part_paths[part_type] = str(part_path.relative_to(output_path))
                 
-                # Add to dataset
                 dataset_entry = {
                     'dog_id': dog_id,
                     'full_image': str(full_path.relative_to(output_path)),
                     'bbox': img_data['bbox'],
                     'confidence': img_data['confidence']
                 }
-                
-                # Add body part paths if available
                 for part_type in ['head', 'torso', 'rear']:
                     dataset_entry[f'{part_type}_image'] = part_paths.get(part_type, None)
-                
                 dataset.append(dataset_entry)
         
-        # Save dataset info
         dataset_df = pd.DataFrame(dataset)
         dataset_df.to_csv(output_path / "dataset.csv", index=False)
         
-        # Save metadata
         metadata = {
             'total_dogs': len(dogs),
             'total_images': len(dataset),
@@ -545,37 +458,27 @@ class DogDatabase:
             'validated_only': validated_only,
             'includes_body_parts': True
         }
-        
         with open(output_path / "metadata.json", 'w') as f:
             json.dump(metadata, f, indent=2)
         
-        # Create training splits
         from sklearn.model_selection import train_test_split
-        
-        train_df, test_df = train_test_split(dataset_df, test_size=0.2, 
-                                            stratify=dataset_df['dog_id'])
+        train_df, test_df = train_test_split(dataset_df, test_size=0.2, stratify=dataset_df['dog_id'])
         train_df.to_csv(output_path / "train.csv", index=False)
         test_df.to_csv(output_path / "test.csv", index=False)
         
         metadata['train_samples'] = len(train_df)
         metadata['test_samples'] = len(test_df)
-        
         return metadata
     
     # ========== Cleanup Methods ==========
-    
     def reset_database(self, confirm: bool = False):
         """Reset entire database"""
         if not confirm:
             return False
-        
         tables = ['sightings', 'dog_images', 'dog_features', 'dogs', 'sessions']
         for table in tables:
             self.cursor.execute(f"DELETE FROM {table}")
-        
-        # Reset autoincrement
         self.cursor.execute("DELETE FROM sqlite_sequence")
-        
         self.conn.commit()
         return True
     
@@ -590,4 +493,102 @@ class DogDatabase:
     def __del__(self):
         """Ensure connection is closed"""
         if hasattr(self, 'conn'):
-            self.conn.close()
+            self.conn.close()
+    
+    # ========== Health Methods (NEW) ==========
+    def save_health_assessment(self, dog_id: int, health_score: float, status: str,
+                              posture_score: float = None, gait_score: float = None,
+                              body_condition_score: float = None, activity_score: float = None,
+                              alerts: List[str] = None, recommendations: List[str] = None,
+                              confidence: float = 0.5, video_source: str = None,
+                              frame_number: int = None):
+        """Save health assessment to database"""
+        self.cursor.execute("""
+            INSERT INTO health_assessments 
+            (dog_id, health_score, status, posture_score, gait_score,
+             body_condition_score, activity_score, alerts, recommendations,
+             confidence, video_source, frame_number)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+        """, (dog_id, health_score, status, posture_score, gait_score,
+              body_condition_score, activity_score,
+              json.dumps(alerts) if alerts else None,
+              json.dumps(recommendations) if recommendations else None,
+              confidence, video_source, frame_number))
+        
+        self.cursor.execute("""
+            UPDATE dogs 
+            SET last_health_score = ?,
+                health_status = ?
+            WHERE dog_id = ?
+        """, (health_score, status, dog_id))
+        self.conn.commit()
+    
+    def get_health_history(self, dog_id: int, limit: int = 50) -> List[Dict]:
+        """Get health assessment history for a dog"""
+        self.cursor.execute("""
+            SELECT * FROM health_assessments 
+            WHERE dog_id = ?
+            ORDER BY timestamp DESC
+            LIMIT ?
+        """, (dog_id, limit))
+        assessments = []
+        for row in self.cursor.fetchall():
+            assessments.append({
+                'timestamp': row['timestamp'],
+                'health_score': row['health_score'],
+                'status': row['status'],
+                'posture_score': row['posture_score'],
+                'gait_score': row['gait_score'],
+                'body_condition_score': row['body_condition_score'],
+                'activity_score': row['activity_score'],
+                'alerts': json.loads(row['alerts']) if row['alerts'] else [],
+                'recommendations': json.loads(row['recommendations']) if row['recommendations'] else [],
+                'confidence': row['confidence']
+            })
+        return assessments
+    
+    def get_health_statistics(self) -> Dict:
+        """Get overall health statistics"""
+        stats = {}
+        self.cursor.execute("""
+            SELECT AVG(last_health_score) as avg_health,
+                   COUNT(CASE WHEN last_health_score >= 8 THEN 1 END) as healthy_count,
+                   COUNT(CASE WHEN last_health_score < 6 THEN 1 END) as unhealthy_count,
+                   COUNT(*) as total_dogs
+            FROM dogs 
+            WHERE status = 'active'
+        """)
+        row = self.cursor.fetchone()
+        if row:
+            stats['average_health'] = round(row['avg_health'] or 5.0, 1)
+            stats['healthy_dogs'] = row['healthy_count'] or 0
+            stats['unhealthy_dogs'] = row['unhealthy_count'] or 0
+            stats['total_dogs'] = row['total_dogs'] or 0
+        
+        self.cursor.execute("""
+            SELECT dog_id, name, last_health_score, health_status
+            FROM dogs 
+            WHERE status = 'active' AND last_health_score < 6
+            ORDER BY last_health_score ASC
+            LIMIT 10
+        """)
+        stats['dogs_needing_attention'] = []
+        for row in self.cursor.fetchall():
+            stats['dogs_needing_attention'].append({
+                'dog_id': row['dog_id'],
+                'name': row['name'] or f"Dog #{row['dog_id']}",
+                'health_score': row['last_health_score'],
+                'status': row['health_status']
+            })
+        return stats
+    
+    def save_pose_keypoints(self, dog_id: int, keypoints: np.ndarray,
+                           frame_number: int, video_source: str):
+        """Save pose keypoints for a dog"""
+        keypoints_json = json.dumps(keypoints.tolist()) if keypoints is not None else None
+        self.cursor.execute("""
+            UPDATE sightings 
+            SET pose_keypoints = ?
+            WHERE dog_id = ? AND frame_number = ? AND video_source = ?
+        """, (keypoints_json, dog_id, frame_number, video_source))
+        self.conn.commit()