Update database.py

database.py

@@ -326,25 +326,64 @@ class DogDatabase:
                   frame_number: int, video_source: str,
                   bbox: List[float], confidence: float,
                   pose_keypoints: Optional[List] = None):
-        """Save dog image to database"""
-        _, buffer = cv2.imencode('.jpg', image)
-        image_data = base64.b64encode(buffer).decode('utf-8')
-        
-        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, 
-             frame_number, video_source, bbox, confidence, pose_keypoints)
-            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-        """, (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
+            """Save dog image to database"""
+            _, buffer = cv2.imencode('.jpg', image)
+            image_data = base64.b64encode(buffer).decode('utf-8')
+            
+            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, 
+                 frame_number, video_source, bbox, confidence, pose_keypoints)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """, (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
+            def add_dog_image(self, dog_id: int, image: np.ndarray, 
+                      timestamp: float, confidence: float, bbox: tuple):
+        """Add a dog image to database (simplified for dataset collection)"""
+        try:
+            # Convert image to base64 for storage
+            _, 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]
+            
+            # Insert into database
+            self.cursor.execute("""
+                INSERT INTO dog_images 
+                (dog_id, image_data, thumbnail, width, height, 
+                 timestamp, bbox, confidence)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+            """, (dog_id, image_data, thumb_data, w, h, 
+                  timestamp, json.dumps(bbox), confidence))
+            
+            # Update dog's last_seen and sighting count
+            self.cursor.execute("""
+                UPDATE dogs 
+                SET last_seen = datetime('now'),
+                    total_sightings = total_sightings + 1
+                WHERE dog_id = ?
+            """, (dog_id,))
+            
+            self.conn.commit()
+            return self.cursor.lastrowid
+            
+    except Exception as e:
+        self.conn.rollback()
+        print(f"Error adding dog image: {e}")
+        raise
     
     def get_dog_images(self, dog_id: int, validated_only: bool = False,
                   include_discarded: bool = False) -> List[Dict]: