Update app.py

app.py

@@ -118,12 +118,6 @@ class SmartImageSelector:
         if len(dog_data) <= max_images:
             return dog_data
         
-        # Calculate quality scores
-        for item in dog_data:
-            item['quality_score'] = self.quality_analyzer.calculate_overall_quality(
-                item['crop'], item['bbox'], item['frame'].shape
-            )
-        
         # Sort by quality
         dog_data.sort(key=lambda x: x['quality_score'], reverse=True)
         
@@ -235,7 +229,7 @@ class AdvancedHeadExtractor:
         edges = cv2.Canny(gray, 50, 150)
         
         # Find feature concentration (likely head area)
-        kernel_size = h // 10
+        kernel_size = max(1, h // 10)
         kernel = np.ones((kernel_size, kernel_size), np.float32)
         edge_density = cv2.filter2D(edges, -1, kernel)
         
@@ -276,6 +270,7 @@ class ResNetDatasetCreator:
     def __init__(self):
         self.temp_dir = Path("temp_dataset")
         self.final_dir = Path("resnet_finetune_dataset")
+        self.database_dir = Path("permanent_database")
         
         # Components
         self.detector = DogDetector(device='cuda' if torch.cuda.is_available() else 'cpu')
@@ -291,6 +286,47 @@ class ResNetDatasetCreator:
         # Create directories
         self.temp_dir.mkdir(exist_ok=True)
         self.final_dir.mkdir(exist_ok=True)
+        self.database_dir.mkdir(exist_ok=True)
+        
+        # Load existing database if exists
+        self.load_database()
+    
+    def load_database(self):
+        """Load existing permanent database"""
+        db_file = self.database_dir / "database.json"
+        if db_file.exists():
+            with open(db_file, 'r') as f:
+                data = json.load(f)
+                self.processed_dogs = {int(k): v for k, v in data.get('dogs', {}).items()}
+                print(f"Loaded {len(self.processed_dogs)} dogs from database")
+    
+    def save_to_database(self):
+        """Save current processed dogs to permanent database"""
+        db_file = self.database_dir / "database.json"
+        data = {
+            'dogs': {str(k): v for k, v in self.processed_dogs.items()},
+            'last_updated': datetime.now().isoformat()
+        }
+        with open(db_file, 'w') as f:
+            json.dump(data, f, indent=2)
+        
+        # Also save images to permanent location
+        for dog_id in self.processed_dogs:
+            src_dir = self.temp_dir / f"dog_{dog_id:03d}"
+            dst_dir = self.database_dir / f"dog_{dog_id:03d}"
+            if src_dir.exists():
+                if dst_dir.exists():
+                    shutil.rmtree(dst_dir)
+                shutil.copytree(src_dir, dst_dir)
+    
+    def clear_database(self):
+        """Clear all permanent database"""
+        if self.database_dir.exists():
+            shutil.rmtree(self.database_dir)
+        self.database_dir.mkdir(exist_ok=True)
+        self.processed_dogs = {}
+        self.current_session = None
+        print("Database cleared")
     
     def process_video(self, video_path: str, reid_threshold: float,
                      max_images_per_dog: int, sample_rate: int) -> Dict:
@@ -302,7 +338,7 @@ class ResNetDatasetCreator:
             max_images_per_dog: Maximum images to extract per dog
             sample_rate: Process every Nth frame
         """
-        # Clear temp directory
+        # Clear temp directory for new processing
         if self.temp_dir.exists():
             shutil.rmtree(self.temp_dir)
         self.temp_dir.mkdir()
@@ -344,7 +380,7 @@ class ResNetDatasetCreator:
                     dog_id = results['ResNet50']['dog_id']
                     confidence = results['ResNet50']['confidence']
                     
-                    if dog_id > 0 and confidence > 0.5:
+                    if dog_id > 0 and confidence > 0.3:  # Lower threshold for detection
                         # Get best detection
                         detection = None
                         for det in reversed(track.detections):
@@ -381,6 +417,8 @@ class ResNetDatasetCreator:
         
         # Select best images for each dog
         total_images = 0
+        new_dogs = {}
+        
         for dog_id, images in dog_data.items():
             # Use smart selector
             selected = self.image_selector.select_best_images(
@@ -410,20 +448,23 @@ class ResNetDatasetCreator:
             total_images += saved_count
             
             # Store metadata
-            self.processed_dogs[dog_id] = {
+            new_dogs[dog_id] = {
                 'num_images': saved_count,
                 'avg_confidence': np.mean([d['reid_confidence'] for d in selected]),
                 'quality_scores': [d['quality_score'] for d in selected]
             }
         
+        # Update processed dogs (append, don't replace)
+        self.processed_dogs.update(new_dogs)
+        
         # Save session info
         self.current_session = {
             'video': video_path,
             'timestamp': datetime.now().isoformat(),
-            'num_dogs': len(dog_data),
+            'num_dogs': len(new_dogs),
             'total_images': total_images,
             'reid_threshold': reid_threshold,
-            'dogs': {str(k): v for k, v in self.processed_dogs.items()}
+            'dogs': {str(k): v for k, v in new_dogs.items()}
         }
         
         # Save metadata
@@ -434,42 +475,65 @@ class ResNetDatasetCreator:
     
     def get_dog_images(self, dog_id: int) -> List:
         """Get images for verification interface"""
+        # Try temp dir first, then database dir
         dog_dir = self.temp_dir / f"dog_{dog_id:03d}"
-        full_dir = dog_dir / 'full'
+        if not dog_dir.exists():
+            dog_dir = self.database_dir / f"dog_{dog_id:03d}"
         
+        full_dir = dog_dir / 'full'
         if not full_dir.exists():
             return []
         
         images = []
         for img_path in sorted(full_dir.glob("*.jpg"))[:12]:
             img = cv2.imread(str(img_path))
-            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-            images.append(img_rgb)
+            if img is not None:
+                img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+                images.append(img_rgb)
         
         return images
     
     def remove_images(self, dog_id: int, image_indices: List[int]):
         """Remove specific images from a dog folder"""
-        dog_dir = self.temp_dir / f"dog_{dog_id:03d}"
-        full_dir = dog_dir / 'full'
-        head_dir = dog_dir / 'head'
-        
-        image_files = sorted(list(full_dir.glob("*.jpg")))
-        
-        for idx in image_indices:
-            if idx < len(image_files):
-                # Remove full image
-                image_files[idx].unlink()
-                # Remove corresponding head
-                head_file = head_dir / image_files[idx].name
-                if head_file.exists():
-                    head_file.unlink()
+        # Handle both temp and database directories
+        for base_dir in [self.temp_dir, self.database_dir]:
+            dog_dir = base_dir / f"dog_{dog_id:03d}"
+            if not dog_dir.exists():
+                continue
+                
+            full_dir = dog_dir / 'full'
+            head_dir = dog_dir / 'head'
+            
+            image_files = sorted(list(full_dir.glob("*.jpg")))
+            
+            # Extract actual indices from gallery selection
+            indices_to_remove = []
+            if isinstance(image_indices, list):
+                for item in image_indices:
+                    if isinstance(item, (list, tuple)) and len(item) > 0:
+                        indices_to_remove.append(item[0])
+                    elif isinstance(item, int):
+                        indices_to_remove.append(item)
+            
+            for idx in indices_to_remove:
+                if 0 <= idx < len(image_files):
+                    # Remove full image
+                    image_files[idx].unlink(missing_ok=True)
+                    # Remove corresponding head
+                    head_file = head_dir / image_files[idx].name
+                    if head_file.exists():
+                        head_file.unlink()
     
     def delete_dog(self, dog_id: int):
-        """Delete entire dog folder"""
-        dog_dir = self.temp_dir / f"dog_{dog_id:03d}"
-        if dog_dir.exists():
-            shutil.rmtree(dog_dir)
+        """Delete entire dog folder from both temp and database"""
+        for base_dir in [self.temp_dir, self.database_dir]:
+            dog_dir = base_dir / f"dog_{dog_id:03d}"
+            if dog_dir.exists():
+                shutil.rmtree(dog_dir)
+        
+        # Remove from processed dogs
+        if dog_id in self.processed_dogs:
+            del self.processed_dogs[dog_id]
     
     def save_final_dataset(self, format_type: str = 'folder') -> str:
         """
@@ -483,12 +547,24 @@ class ResNetDatasetCreator:
             shutil.rmtree(self.final_dir)
         self.final_dir.mkdir()
         
-        # Copy verified dogs
-        dog_dirs = sorted([d for d in self.temp_dir.iterdir() if d.is_dir()])
+        # Copy all dogs from both temp and database
+        all_dog_dirs = []
+        
+        # Get dogs from temp
+        for d in self.temp_dir.iterdir():
+            if d.is_dir() and d.name.startswith('dog_'):
+                all_dog_dirs.append(d)
+        
+        # Get dogs from database (if not already in temp)
+        temp_dogs = {d.name for d in all_dog_dirs}
+        for d in self.database_dir.iterdir():
+            if d.is_dir() and d.name.startswith('dog_') and d.name not in temp_dogs:
+                all_dog_dirs.append(d)
+        
         data_entries = []
         final_id = 1
         
-        for dog_dir in dog_dirs:
+        for dog_dir in sorted(all_dog_dirs):
             if not (dog_dir / 'full').exists():
                 continue
             
@@ -512,22 +588,30 @@ class ResNetDatasetCreator:
             # Create train/val split
             df = pd.DataFrame(data_entries)
             
-            # Stratified split by dog_id
-            from sklearn.model_selection import train_test_split
-            train_df, val_df = train_test_split(
-                df, test_size=0.2, stratify=df['dog_id'], random_state=42
-            )
-            
-            # Save CSV files
-            train_df.to_csv(self.final_dir / 'train.csv', index=False)
-            val_df.to_csv(self.final_dir / 'val.csv', index=False)
+            if len(df) > 0:
+                # Stratified split by dog_id
+                from sklearn.model_selection import train_test_split
+                
+                # Only split if we have enough samples
+                if len(df) > 5:
+                    train_df, val_df = train_test_split(
+                        df, test_size=0.2, stratify=df['dog_id'], random_state=42
+                    )
+                else:
+                    train_df = df
+                    val_df = pd.DataFrame()
+                
+                # Save CSV files
+                train_df.to_csv(self.final_dir / 'train.csv', index=False)
+                if len(val_df) > 0:
+                    val_df.to_csv(self.final_dir / 'val.csv', index=False)
         
         # Create metadata
         metadata = {
             'total_dogs': final_id - 1,
             'total_images': len(data_entries),
-            'train_images': len(train_df) if format_type in ['csv', 'both'] else 0,
-            'val_images': len(val_df) if format_type in ['csv', 'both'] else 0,
+            'train_images': len(train_df) if format_type in ['csv', 'both'] and 'train_df' in locals() else len(data_entries),
+            'val_images': len(val_df) if format_type in ['csv', 'both'] and 'val_df' in locals() else 0,
             'format': format_type,
             'created': datetime.now().isoformat()
         }
@@ -556,8 +640,6 @@ class ResNetDatasetCreator:
             
             # State to store processing results
             processing_state = gr.State(None)
-            # State for tab navigation
-            selected_tab = gr.State(0)
             
             # Step 1: Process Video
             with gr.Tabs() as tabs:
@@ -585,7 +667,7 @@ class ResNetDatasetCreator:
                     # Results display in formatted table
                     with gr.Column():
                         progress_bar = gr.Textbox(label="Progress", interactive=False)
-                        results_display = gr.HTML(label="Processing Results")
+                        results_display = gr.HTML(label="Processing Results", value="")
                         save_status = gr.Textbox(label="Save Status", interactive=False, visible=False)
                         
                         with gr.Row():
@@ -600,154 +682,152 @@ class ResNetDatasetCreator:
                                 variant="secondary",
                                 visible=False
                             )
-                
-                def format_results_table(session_data):
-                    """Format session data as HTML table"""
-                    if not session_data:
-                        return ""
-                    
-                    html = """
-                    <div style="padding: 20px; background-color: #f8f9fa; border-radius: 10px;">
-                        <h3 style="color: #2c3e50;">📊 Processing Results</h3>
-                        <table style="width: 100%; border-collapse: collapse; margin: 20px 0;">
-                            <tr style="background-color: #3498db; color: white;">
-                                <td style="padding: 10px; border: 1px solid #ddd;"><b>Metric</b></td>
-                                <td style="padding: 10px; border: 1px solid #ddd;"><b>Value</b></td>
+                    
+                    def format_results_table(session_data):
+                        """Format session data as HTML table"""
+                        if not session_data:
+                            return ""
+                        
+                        html = """
+                        <div style="padding: 20px; background-color: #f8f9fa; border-radius: 10px;">
+                            <h3 style="color: #2c3e50;">📊 Processing Results</h3>
+                            <table style="width: 100%; border-collapse: collapse; margin: 20px 0;">
+                                <tr style="background-color: #3498db; color: white;">
+                                    <td style="padding: 10px; border: 1px solid #ddd;"><b>Metric</b></td>
+                                    <td style="padding: 10px; border: 1px solid #ddd;"><b>Value</b></td>
+                                </tr>
+                        """
+                        
+                        # Basic info
+                        html += f"""
+                            <tr style="background-color: #ecf0f1;">
+                                <td style="padding: 10px; border: 1px solid #ddd;">Video File</td>
+                                <td style="padding: 10px; border: 1px solid #ddd;">{session_data['video'].split('/')[-1]}</td>
                             </tr>
-                    """
-                    
-                    # Basic info
-                    html += f"""
-                        <tr style="background-color: #ecf0f1;">
-                            <td style="padding: 10px; border: 1px solid #ddd;">Video File</td>
-                            <td style="padding: 10px; border: 1px solid #ddd;">{session_data['video'].split('/')[-1]}</td>
-                        </tr>
-                        <tr>
-                            <td style="padding: 10px; border: 1px solid #ddd;">Processing Time</td>
-                            <td style="padding: 10px; border: 1px solid #ddd;">{session_data['timestamp'].split('T')[1].split('.')[0]}</td>
-                        </tr>
-                        <tr style="background-color: #ecf0f1;">
-                            <td style="padding: 10px; border: 1px solid #ddd;">Number of Dogs Detected</td>
-                            <td style="padding: 10px; border: 1px solid #ddd;"><b>{session_data['num_dogs']}</b></td>
-                        </tr>
-                        <tr>
-                            <td style="padding: 10px; border: 1px solid #ddd;">Total Images Extracted</td>
-                            <td style="padding: 10px; border: 1px solid #ddd;"><b>{session_data['total_images']}</b></td>
-                        </tr>
-                        <tr style="background-color: #ecf0f1;">
-                            <td style="padding: 10px; border: 1px solid #ddd;">ReID Threshold Used</td>
-                            <td style="padding: 10px; border: 1px solid #ddd;">{session_data['reid_threshold']:.2f}</td>
-                        </tr>
-                    </table>
-                    """
-                    
-                    # Dog-specific details
-                    if session_data['dogs']:
-                        html += """
-                        <h4 style="color: #2c3e50; margin-top: 20px;">🐕 Dog Details</h4>
-                        <table style="width: 100%; border-collapse: collapse; margin: 10px 0;">
-                            <tr style="background-color: #27ae60; color: white;">
-                                <td style="padding: 10px; border: 1px solid #ddd;"><b>Dog ID</b></td>
-                                <td style="padding: 10px; border: 1px solid #ddd;"><b>Images</b></td>
-                                <td style="padding: 10px; border: 1px solid #ddd;"><b>Avg Confidence</b></td>
-                                <td style="padding: 10px; border: 1px solid #ddd;"><b>Avg Quality</b></td>
-                                <td style="padding: 10px; border: 1px solid #ddd;"><b>Quality Range</b></td>
+                            <tr>
+                                <td style="padding: 10px; border: 1px solid #ddd;">Processing Time</td>
+                                <td style="padding: 10px; border: 1px solid #ddd;">{session_data['timestamp'].split('T')[1].split('.')[0]}</td>
+                            </tr>
+                            <tr style="background-color: #ecf0f1;">
+                                <td style="padding: 10px; border: 1px solid #ddd;">Number of Dogs Detected</td>
+                                <td style="padding: 10px; border: 1px solid #ddd;"><b>{session_data['num_dogs']}</b></td>
+                            </tr>
+                            <tr>
+                                <td style="padding: 10px; border: 1px solid #ddd;">Total Images Extracted</td>
+                                <td style="padding: 10px; border: 1px solid #ddd;"><b>{session_data['total_images']}</b></td>
                             </tr>
+                            <tr style="background-color: #ecf0f1;">
+                                <td style="padding: 10px; border: 1px solid #ddd;">ReID Threshold Used</td>
+                                <td style="padding: 10px; border: 1px solid #ddd;">{session_data['reid_threshold']:.2f}</td>
+                            </tr>
+                        </table>
                         """
                         
-                        for dog_id, dog_info in session_data['dogs'].items():
-                            avg_quality = np.mean(dog_info['quality_scores'])
-                            min_quality = min(dog_info['quality_scores'])
-                            max_quality = max(dog_info['quality_scores'])
-                            
-                            row_style = "background-color: #ecf0f1;" if int(dog_id) % 2 == 0 else ""
-                            html += f"""
-                            <tr style="{row_style}">
-                                <td style="padding: 10px; border: 1px solid #ddd;">Dog {dog_id}</td>
-                                <td style="padding: 10px; border: 1px solid #ddd;">{dog_info['num_images']}</td>
-                                <td style="padding: 10px; border: 1px solid #ddd;">{dog_info['avg_confidence']:.2%}</td>
-                                <td style="padding: 10px; border: 1px solid #ddd;">{avg_quality:.1f}</td>
-                                <td style="padding: 10px; border: 1px solid #ddd;">{min_quality:.1f} - {max_quality:.1f}</td>
-                            </tr>
+                        # Dog-specific details
+                        if session_data['dogs']:
+                            html += """
+                            <h4 style="color: #2c3e50; margin-top: 20px;">🐕 Dog Details</h4>
+                            <table style="width: 100%; border-collapse: collapse; margin: 10px 0;">
+                                <tr style="background-color: #27ae60; color: white;">
+                                    <td style="padding: 10px; border: 1px solid #ddd;"><b>Dog ID</b></td>
+                                    <td style="padding: 10px; border: 1px solid #ddd;"><b>Images</b></td>
+                                    <td style="padding: 10px; border: 1px solid #ddd;"><b>Avg Confidence</b></td>
+                                    <td style="padding: 10px; border: 1px solid #ddd;"><b>Avg Quality</b></td>
+                                    <td style="padding: 10px; border: 1px solid #ddd;"><b>Quality Range</b></td>
+                                </tr>
                             """
+                            
+                            for dog_id, dog_info in session_data['dogs'].items():
+                                avg_quality = np.mean(dog_info['quality_scores'])
+                                min_quality = min(dog_info['quality_scores'])
+                                max_quality = max(dog_info['quality_scores'])
+                                
+                                row_style = "background-color: #ecf0f1;" if int(dog_id) % 2 == 0 else ""
+                                html += f"""
+                                <tr style="{row_style}">
+                                    <td style="padding: 10px; border: 1px solid #ddd;">Dog {dog_id}</td>
+                                    <td style="padding: 10px; border: 1px solid #ddd;">{dog_info['num_images']}</td>
+                                    <td style="padding: 10px; border: 1px solid #ddd;">{dog_info['avg_confidence']:.2%}</td>
+                                    <td style="padding: 10px; border: 1px solid #ddd;">{avg_quality:.1f}</td>
+                                    <td style="padding: 10px; border: 1px solid #ddd;">{min_quality:.1f} - {max_quality:.1f}</td>
+                                </tr>
+                                """
+                            
+                            html += "</table>"
                         
-                        html += "</table>"
-                    
-                    html += """
-                        <div style="margin-top: 20px; padding: 10px; background-color: #d4edda; border-radius: 5px;">
-                            <p style="margin: 0; color: #155724;">
-                                ✅ <b>Processing Complete!</b> Click "Save Results & Proceed" to continue to verification step.
-                            </p>
+                        html += """
+                            <div style="margin-top: 20px; padding: 10px; background-color: #d4edda; border-radius: 5px;">
+                                <p style="margin: 0; color: #155724;">
+                                    ✅ <b>Processing Complete!</b> Click "Save Results & Proceed" to continue to verification step.
+                                </p>
+                            </div>
                         </div>
-                    </div>
-                    """
-                    
-                    return html
-                
-                def process_wrapper(video, threshold, max_img, sample):
-                    """Process video and format results"""
-                    if not video:
-                        return None, "", "Please upload a video", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-                    
-                    # Clear previous session
-                    self.current_session = None
-                    self.processed_dogs = {}
-                    
-                    # Process video
-                    for update in self.process_video(video, threshold, int(max_img), int(sample)):
-                        if 'progress' in update:
-                            yield None, "", update['status'], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-                        else:
-                            # Store session data
-                            self.current_session = update['session']
-                            # Format results as table
-                            formatted_table = format_results_table(update['session'])
-                            yield update['session'], formatted_table, "Complete! ✅", gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
-                
-                def save_and_proceed():
-                    """Save current results and notify user"""
-                    if self.current_session and self.processed_dogs:
-                        # Debug info
-                        dog_count = len(self.processed_dogs)
-                        img_count = sum(d['num_images'] for d in self.processed_dogs.values())
-                        
-                        message = f"""✅ Results saved successfully! 
-                        
-                        📊 Summary:
-                        - Dogs saved: {dog_count}
-                        - Total images: {img_count}
-                        - Data location: {self.temp_dir}
+                        """
                         
-                        You can now proceed to Step 2: Verify & Clean
-                        Click the 'Refresh List' button in Step 2 to load the dogs."""
+                        return html
+                    
+                    def process_wrapper(video, threshold, max_img, sample):
+                        """Process video and format results"""
+                        if not video:
+                            return None, "", "Please upload a video", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
                         
-                        return message, gr.update(visible=True)
-                    return "❌ No results to save. Please process a video first.", gr.update(visible=False)
-                
-                def clear_results():
-                    """Clear current processing results"""
-                    self.current_session = None
-                    self.processed_dogs = {}
-                    if self.temp_dir.exists():
-                        shutil.rmtree(self.temp_dir)
-                        self.temp_dir.mkdir()
-                    return None, "", "", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-                
-                process_btn.click(
-                    process_wrapper,
-                    inputs=[video_input, reid_threshold, max_images, sample_rate],
-                    outputs=[processing_state, results_display, progress_bar, save_status, save_proceed_btn, clear_btn]
-                )
-                
-                save_proceed_btn.click(
-                    save_and_proceed,
-                    outputs=[save_status, save_status]  # Use save_status for both message and visibility
-                )
-                
-                clear_btn.click(
-                    clear_results,
-                    outputs=[processing_state, results_display, progress_bar, save_status, save_proceed_btn, clear_btn]
-                )
+                        # Process video
+                        for update in self.process_video(video, threshold, int(max_img), int(sample)):
+                            if 'progress' in update:
+                                yield None, "", update['status'], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+                            else:
+                                # Store session data
+                                self.current_session = update['session']
+                                # Format results as table
+                                formatted_table = format_results_table(update['session'])
+                                yield update['session'], formatted_table, "Complete! ✅", gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
+                    
+                    def save_and_proceed():
+                        """Save current results and notify user"""
+                        if self.current_session and self.processed_dogs:
+                            # Save to permanent database
+                            self.save_to_database()
+                            
+                            # Debug info
+                            dog_count = len(self.processed_dogs)
+                            img_count = sum(d.get('num_images', 0) for d in self.processed_dogs.values())
+                            
+                            message = f"""✅ Results saved successfully to database! 
+                            
+                            📊 Summary:
+                            - Total dogs in database: {dog_count}
+                            - Total images: {img_count}
+                            - Data location: {self.database_dir}
+                            
+                            You can now proceed to Step 2: Verify & Clean
+                            Click the 'Refresh List' button in Step 2 to load all dogs."""
+                            
+                            return message, gr.update(visible=True)
+                        return "❌ No results to save. Please process a video first.", gr.update(visible=False)
+                    
+                    def clear_results():
+                        """Clear current processing results (not database)"""
+                        self.current_session = None
+                        if self.temp_dir.exists():
+                            shutil.rmtree(self.temp_dir)
+                            self.temp_dir.mkdir()
+                        return None, "", "", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+                    
+                    process_btn.click(
+                        process_wrapper,
+                        inputs=[video_input, reid_threshold, max_images, sample_rate],
+                        outputs=[processing_state, results_display, progress_bar, save_status, save_proceed_btn, clear_btn]
+                    )
+                    
+                    save_proceed_btn.click(
+                        save_and_proceed,
+                        outputs=[save_status, save_status]
+                    )
+                    
+                    clear_btn.click(
+                        clear_results,
+                        outputs=[processing_state, results_display, progress_bar, save_status, save_proceed_btn, clear_btn]
+                    )
                 
                 # Step 2: Verify & Clean
                 with gr.Tab("✅ Step 2: Verify & Clean", id=1):
@@ -759,7 +839,14 @@ class ResNetDatasetCreator:
                             choices=[],
                             interactive=True
                         )
+                    
+                    # Add diagnostic and management buttons
+                    with gr.Row():
                         refresh_btn = gr.Button("🔄 Refresh List")
+                        diagnose_btn = gr.Button("🔍 Diagnose Data", variant="secondary")
+                        clear_db_btn = gr.Button("⚠️ Clear All Database", variant="stop")
+                    
+                    diagnostic_output = gr.Textbox(label="Diagnostic Info", visible=False)
                     
                     image_gallery = gr.Gallery(
                         label="Dog Images (Click to select for removal)",
@@ -769,18 +856,27 @@ class ResNetDatasetCreator:
                         rows=3,
                         object_fit="contain",
                         height="auto",
-                        interactive=True
+                        type="numpy",
+                        interactive=False
                     )
                     
                     with gr.Row():
+                        selected_images = gr.Textbox(
+                            label="Selected Image Indices (comma-separated)",
+                            placeholder="e.g., 0,2,5",
+                            interactive=True
+                        )
                         remove_selected_btn = gr.Button("🗑 Remove Selected Images", variant="secondary")
                         delete_dog_btn = gr.Button("❌ Delete Entire Dog", variant="stop")
                     
                     status_text = gr.Textbox(label="Status", interactive=False)
                     
                     def refresh_dogs():
-                        """Refresh the dog list from processed data"""
-                        if not self.current_session or not self.processed_dogs:
+                        """Refresh the dog list from all available data"""
+                        # Load from database
+                        self.load_database()
+                        
+                        if not self.processed_dogs:
                             return gr.update(choices=[], value=None)
                         
                         choices = [f"Dog {dog_id}" for dog_id in sorted(self.processed_dogs.keys())]
@@ -788,6 +884,51 @@ class ResNetDatasetCreator:
                             return gr.update(choices=choices, value=choices[0])
                         return gr.update(choices=[], value=None)
                     
+                    def diagnose_data():
+                        """Show diagnostic information about saved data"""
+                        info = []
+                        info.append("=== DIAGNOSTIC INFORMATION ===\n")
+                        
+                        # Check session
+                        if self.current_session:
+                            info.append(f"✅ Session exists: {self.current_session['num_dogs']} dogs, {self.current_session['total_images']} images")
+                        else:
+                            info.append("❌ No current session data")
+                        
+                        # Check processed dogs
+                        if self.processed_dogs:
+                            info.append(f"✅ Processed dogs dict: {len(self.processed_dogs)} dogs")
+                            for dog_id, data in self.processed_dogs.items():
+                                info.append(f"  - Dog {dog_id}: {data.get('num_images', 0)} images, conf={data.get('avg_confidence', 0):.2f}")
+                        else:
+                            info.append("❌ No processed dogs data")
+                        
+                        # Check temp directory
+                        if self.temp_dir.exists():
+                            info.append(f"✅ Temp directory exists: {self.temp_dir}")
+                            dog_dirs = list(self.temp_dir.glob("dog_*"))
+                            info.append(f"  - Found {len(dog_dirs)} dog directories")
+                            for dog_dir in sorted(dog_dirs):
+                                if (dog_dir / 'full').exists():
+                                    img_count = len(list((dog_dir / 'full').glob("*.jpg")))
+                                    info.append(f"    • {dog_dir.name}: {img_count} full images")
+                        else:
+                            info.append("❌ Temp directory not found")
+                        
+                        # Check database directory
+                        if self.database_dir.exists():
+                            info.append(f"✅ Database directory exists: {self.database_dir}")
+                            dog_dirs = list(self.database_dir.glob("dog_*"))
+                            info.append(f"  - Found {len(dog_dirs)} dog directories")
+                            for dog_dir in sorted(dog_dirs):
+                                if (dog_dir / 'full').exists():
+                                    img_count = len(list((dog_dir / 'full').glob("*.jpg")))
+                                    info.append(f"    • {dog_dir.name}: {img_count} full images")
+                        else:
+                            info.append("❌ Database directory not found")
+                        
+                        return "\n".join(info), gr.update(visible=True)
+                    
                     def show_dog_images(dog_selection):
                         """Display images for selected dog"""
                         if not dog_selection:
@@ -803,31 +944,59 @@ class ResNetDatasetCreator:
                             print(f"Error loading images: {e}")
                             return []
                     
-                    def remove_selected(dog_selection, selected_indices):
-                        if not dog_selection or not selected_indices:
-                            return "No images selected", gr.update()
+                    def remove_selected(dog_selection, indices_str):
+                        """Remove selected images based on text input"""
+                        if not dog_selection or not indices_str:
+                            return "No images selected", []
                         
-                        dog_id = int(dog_selection.split()[1])
-                        self.remove_images(dog_id, selected_indices)
-                        return f"Removed {len(selected_indices)} images", self.get_dog_images(dog_id)
+                        try:
+                            # Parse comma-separated indices
+                            indices = [int(i.strip()) for i in indices_str.split(',')]
+                            dog_id = int(dog_selection.split()[1])
+                            
+                            self.remove_images(dog_id, indices)
+                            
+                            # Update database
+                            self.save_to_database()
+                            
+                            return f"Removed {len(indices)} images", self.get_dog_images(dog_id)
+                        except Exception as e:
+                            return f"Error: {str(e)}", []
                     
                     def delete_dog(dog_selection):
                         if not dog_selection:
-                            return "No dog selected"
+                            return "No dog selected", []
                         
                         dog_id = int(dog_selection.split()[1])
                         self.delete_dog(dog_id)
-                        del self.processed_dogs[dog_id]
-                        return f"Deleted Dog {dog_id}"
+                        
+                        # Update database
+                        self.save_to_database()
+                        
+                        return f"Deleted Dog {dog_id}", []
+                    
+                    def clear_all_database():
+                        """Clear entire database"""
+                        self.clear_database()
+                        return "Database cleared successfully", gr.update(choices=[], value=None), []
                     
                     refresh_btn.click(refresh_dogs, outputs=dog_selector)
+                    diagnose_btn.click(diagnose_data, outputs=[diagnostic_output, diagnostic_output])
                     dog_selector.change(show_dog_images, inputs=dog_selector, outputs=image_gallery)
                     remove_selected_btn.click(
                         remove_selected,
-                        inputs=[dog_selector, image_gallery],
+                        inputs=[dog_selector, selected_images],
                         outputs=[status_text, image_gallery]
                     )
-                    delete_dog_btn.click(delete_dog, inputs=dog_selector, outputs=status_text)
+                    delete_dog_btn.click(
+                        delete_dog, 
+                        inputs=dog_selector, 
+                        outputs=[status_text, image_gallery]
+                    )
+                    clear_db_btn.click(
+                        clear_all_database,
+                        outputs=[status_text, dog_selector, image_gallery]
+                    )
                 
                 # Step 3: Export Dataset
                 with gr.Tab("💾 Step 3: Export Dataset", id=2):
@@ -845,36 +1014,64 @@ class ResNetDatasetCreator:
                         label="Export Format"
                     )
                     
-                    export_btn = gr.Button("📦 Export Final Dataset", variant="primary", size="lg")
+                    with gr.Row():
+                        export_btn = gr.Button("📦 Export Final Dataset", variant="primary", size="lg")
+                        export_status = gr.Button("📊 Check Export Status", variant="secondary")
+                    
                     export_output = gr.Textbox(label="Export Path", interactive=False)
                     download_file = gr.File(label="Download Dataset", interactive=False)
                     stats_display = gr.Markdown()
                     
                     def export_dataset(format_type):
-                        zip_path = self.save_final_dataset(format_type)
-                        
-                        # Get statistics
-                        with open(self.final_dir / 'metadata.json', 'r') as f:
-                            metadata = json.load(f)
+                        try:
+                            zip_path = self.save_final_dataset(format_type)
+                            
+                            # Get statistics
+                            with open(self.final_dir / 'metadata.json', 'r') as f:
+                                metadata = json.load(f)
+                            
+                            stats = f"""
+                            ### ✅ Dataset Exported Successfully!
+                            
+                            - **Total Dogs**: {metadata['total_dogs']}
+                            - **Total Images**: {metadata['total_images']}
+                            - **Training Images**: {metadata.get('train_images', 'N/A')}
+                            - **Validation Images**: {metadata.get('val_images', 'N/A')}
+                            
+                            Dataset is ready for ResNet fine-tuning!
+                            Download the ZIP file below.
+                            """
+                            
+                            return zip_path, zip_path, stats
+                        except Exception as e:
+                            return "", None, f"### ❌ Export Error\n{str(e)}"
+                    
+                    def check_export_status():
+                        """Check what data is available for export"""
+                        total_dogs = len(self.processed_dogs)
+                        total_images = sum(d.get('num_images', 0) for d in self.processed_dogs.values())
                         
-                        stats = f"""
-                        ### ✅ Dataset Exported Successfully!
+                        status = f"""
+                        ### 📊 Export Status
                         
-                        - **Total Dogs**: {metadata['total_dogs']}
-                        - **Total Images**: {metadata['total_images']}
-                        - **Training Images**: {metadata.get('train_images', 'N/A')}
-                        - **Validation Images**: {metadata.get('val_images', 'N/A')}
+                        **Available Data:**
+                        - Dogs in database: {total_dogs}
+                        - Total images: {total_images}
                         
-                        Dataset is ready for ResNet fine-tuning!
+                        {'✅ Ready to export!' if total_dogs > 0 else '❌ No data available. Process videos first.'}
                         """
-                        
-                        return zip_path, zip_path, stats
+                        return status
                     
                     export_btn.click(
                         export_dataset,
                         inputs=format_selector,
                         outputs=[export_output, download_file, stats_display]
                     )
+                    
+                    export_status.click(
+                        check_export_status,
+                        outputs=stats_display
+                    )
             
             return app