Update database.py

database.py

@@ -321,31 +321,33 @@ class DogDatabase:
             })
         return features
     
+    # ========== Images Management ==========
     # ========== 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"""
-            _, 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, 
+        """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:
@@ -380,13 +382,13 @@ class DogDatabase:
             self.conn.commit()
             return self.cursor.lastrowid
             
-    except Exception as e:
-        self.conn.rollback()
-        print(f"Error adding dog image: {e}")
-        raise
+        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]:
+                       include_discarded: bool = False) -> List[Dict]:
         """Get all images for a dog"""
         query = "SELECT * FROM dog_images WHERE dog_id = ?"
         params = [dog_id]
@@ -408,13 +410,13 @@ class DogDatabase:
                 'image_id': row['image_id'],
                 'image': image,
                 'thumbnail': row['thumbnail'],
-                'bbox': json.loads(row['bbox']),
+                'bbox': json.loads(row['bbox']) if row['bbox'] else [],
                 'confidence': row['confidence'],
-                'frame_number': row['frame_number'],
-                'video_source': row['video_source'],
+                'frame_number': row.get('frame_number'),
+                'video_source': row.get('video_source'),
                 'is_validated': row['is_validated'],
                 'is_discarded': row['is_discarded'],
-                'pose_keypoints': json.loads(row['pose_keypoints']) if row['pose_keypoints'] else None
+                'pose_keypoints': json.loads(row['pose_keypoints']) if row.get('pose_keypoints') else None
             })
         return images