Update app.py

app.py

@@ -1,133 +1,64 @@
 """
-Simplified Enhanced Dataset Creator with Fixed Image Selection
+Simplified Dog Detection Demo with MegaDescriptor
 """
 import gradio as gr
 import cv2
 import numpy as np
-import pandas as pd
-import json
-import shutil
 import torch
 from pathlib import Path
-from typing import List, Dict, Optional, Tuple
-from datetime import datetime
-import zipfile
+from typing import Dict
 import gc
+import base64
+from io import BytesIO
+from PIL import Image
 
-# Import required modules
+# Import modules
 from detection import DogDetector
 from tracking import SimpleTracker
-from reid import MultiComponentReID
-from ultralytics import YOLO
+from reid import MegaDescriptorReID
 
 
-class ImageQualityAnalyzer:
-    """Simple image quality scoring"""
-    
-    def calculate_quality(self, image: np.ndarray, bbox: List[float]) -> float:
-        """Calculate overall quality score (0-100)"""
-        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-        
-        # Sharpness (Laplacian variance)
-        sharpness = min(100, cv2.Laplacian(gray, cv2.CV_64F).var())
-        
-        # Brightness (optimal around 127)
-        brightness = 100 - abs(np.mean(gray) - 127) * 0.78
-        
-        # Size score
-        h, w = image.shape[:2]
-        size_score = min(100, (h * w) / (224 * 224) * 100)
-        
-        # Combine scores
-        return (sharpness * 0.4 + brightness * 0.3 + size_score * 0.3)
-
-
-class SimpleDatasetCreator:
-    """Simplified dataset creator with intuitive interface"""
+class DogDetectionDemo:
+    """Simplified demo for dog detection and ReID"""
     
     def __init__(self):
-        # Core directories
-        self.temp_dir = Path("temp_dataset")
-        self.database_dir = Path("permanent_database")
-        self.export_dir = Path("export_dataset")
-        
-        # Create directories
-        for dir_path in [self.temp_dir, self.database_dir, self.export_dir]:
-            dir_path.mkdir(exist_ok=True)
-        
-        # Components
+        # Initialize components
         device = 'cuda' if torch.cuda.is_available() else 'cpu'
         self.detector = DogDetector(device=device)
         self.tracker = SimpleTracker()
-        self.reid = MultiComponentReID(device=device)
-        self.quality_analyzer = ImageQualityAnalyzer()
-        
-        # Session data
-        self.current_dogs = {}  # dog_id -> metadata
-        self.dog_images = {}    # dog_id -> list of image paths
-        self.selected_images = []  # Currently selected image paths
-        self.next_dog_id = 1
-        
-        # Load existing database
-        self.load_database()
+        self.reid = MegaDescriptorReID(device=device)
+        
+        # Temporary storage for current session
+        self.current_dogs = {}  # dog_id -> list of images
+        
+    def reset_session(self):
+        """Reset everything for new video or parameter change"""
+        self.current_dogs.clear()
+        self.tracker.reset()
+        self.reid.reset_all()
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        print("🔄 Session reset")
     
-    def load_database(self):
-        """Load existing dogs from database"""
-        db_file = self.database_dir / "database.json"
-        if db_file.exists():
-            with open(db_file, 'r') as f:
-                data = json.load(f)
-                self.current_dogs = {int(k): v for k, v in data.get('dogs', {}).items()}
-                self.next_dog_id = data.get('next_id', 1)
-                
-                # Load image paths
-                for dog_id in self.current_dogs:
-                    dog_dir = self.database_dir / f"dog_{dog_id:03d}"
-                    if dog_dir.exists():
-                        self.dog_images[dog_id] = sorted([str(p) for p in dog_dir.glob("*.jpg")])
-    
-    def save_database(self):
-        """Save current dogs to database"""
-        # Save metadata
-        db_file = self.database_dir / "database.json"
-        data = {
-            'dogs': {str(k): v for k, v in self.current_dogs.items()},
-            'next_id': self.next_dog_id,
-            'timestamp': datetime.now().isoformat()
-        }
-        with open(db_file, 'w') as f:
-            json.dump(data, f, indent=2)
-        
-        # Copy images from temp to permanent
-        for dog_id in self.current_dogs:
-            src_dir = self.temp_dir / f"dog_{dog_id:03d}"
-            if src_dir.exists():
-                dst_dir = self.database_dir / f"dog_{dog_id:03d}"
-                if dst_dir.exists():
-                    shutil.rmtree(dst_dir)
-                shutil.copytree(src_dir, dst_dir)
-    
-    def process_video(self, video_path: str, reid_threshold: float, 
-                     max_images: int, sample_rate: int) -> Dict:
-        """Simplified video processing"""
+    def process_video(self, video_path: str, reid_threshold: float, sample_rate: int):
+        """Process video and extract dog images"""
         if not video_path:
-            return {'status': 'error', 'message': 'No video provided'}
+            return None, "Please upload a video"
         
-        # Clear temp directory
-        if self.temp_dir.exists():
-            shutil.rmtree(self.temp_dir)
-        self.temp_dir.mkdir()
+        # Reset for new processing
+        self.reset_session()
         
         # Set ReID threshold
         self.reid.set_all_thresholds(reid_threshold)
         
         # Process video
         cap = cv2.VideoCapture(video_path)
-        fps = cap.get(cv2.CAP_PROP_FPS)
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         
-        dog_data = {}  # dog_id -> list of crops
+        dog_crops = {}  # dog_id -> list of crops
         frame_num = 0
+        processed_frames = 0
         
         while cap.isOpened():
             ret, frame = cap.read()
@@ -136,36 +67,33 @@ class SimpleDatasetCreator:
             
             # Process every N frames
             if frame_num % sample_rate == 0:
+                # Detect dogs
                 detections = self.detector.detect(frame)
+                
+                # Update tracks
                 tracks = self.tracker.update(detections)
                 
+                # Process each track
                 for track in tracks:
                     # Get ReID result
-                    results = self.reid.match_or_register_all(track)
-                    dog_id = results['ResNet50']['dog_id']
-                    confidence = results['ResNet50']['confidence']
+                    result = self.reid.match_or_register_all(track)
+                    dog_id = result['MegaDescriptor']['dog_id']
                     
-                    if dog_id > 0 and confidence > 0.3:
+                    if dog_id > 0:
                         # Get latest detection with crop
                         for det in reversed(track.detections[-3:]):
                             if det.image_crop is not None:
-                                if dog_id not in dog_data:
-                                    dog_data[dog_id] = []
-                                
-                                # Calculate quality score
-                                quality = self.quality_analyzer.calculate_quality(
-                                    det.image_crop, det.bbox
-                                )
+                                if dog_id not in dog_crops:
+                                    dog_crops[dog_id] = []
                                 
-                                dog_data[dog_id].append({
-                                    'crop': det.image_crop.copy(),
-                                    'frame_num': frame_num,
-                                    'quality': quality,
-                                    'confidence': confidence
-                                })
+                                # Store crop (max 10 per dog)
+                                if len(dog_crops[dog_id]) < 10:
+                                    dog_crops[dog_id].append(det.image_crop.copy())
                                 break
                 
-                # Memory cleanup
+                processed_frames += 1
+                
+                # Memory cleanup every 100 frames
                 if frame_num % 100 == 0:
                     gc.collect()
                     if torch.cuda.is_available():
@@ -173,541 +101,115 @@ class SimpleDatasetCreator:
             
             frame_num += 1
             
-            # Yield progress
+            # Show progress
             if frame_num % 30 == 0:
                 progress = int((frame_num / total_frames) * 100)
-                yield {'status': 'processing', 'progress': progress}
+                print(f"Processing: {progress}%")
         
         cap.release()
         
-        # Select best images for each dog
-        total_images = 0
-        new_dogs = {}
+        # Convert crops to RGB for display
+        self.current_dogs = {}
+        for dog_id, crops in dog_crops.items():
+            self.current_dogs[dog_id] = [cv2.cvtColor(crop, cv2.COLOR_BGR2RGB) for crop in crops]
         
-        for temp_id, images in dog_data.items():
-            # Get new dog ID
-            dog_id = self.next_dog_id
-            self.next_dog_id += 1
-            
-            # Sort by quality and select top N
-            images.sort(key=lambda x: x['quality'], reverse=True)
-            selected = images[:max_images]
-            
-            # Save images
-            dog_dir = self.temp_dir / f"dog_{dog_id:03d}"
-            dog_dir.mkdir(exist_ok=True)
-            
-            saved_paths = []
-            for idx, img_data in enumerate(selected):
-                img_path = dog_dir / f"img_{idx:03d}.jpg"
-                cv2.imwrite(str(img_path), img_data['crop'])
-                saved_paths.append(str(img_path))
-            
-            # Update tracking
-            self.dog_images[dog_id] = saved_paths
-            new_dogs[dog_id] = {
-                'num_images': len(saved_paths),
-                'avg_confidence': np.mean([d['confidence'] for d in selected]),
-                'source': video_path
-            }
-            total_images += len(saved_paths)
+        # Create gallery HTML
+        gallery_html = self._create_gallery_html()
         
-        # Update current dogs
-        self.current_dogs.update(new_dogs)
+        stats_msg = f"✅ Found {len(self.current_dogs)} dogs | Processed {processed_frames} frames"
         
-        yield {
-            'status': 'complete',
-            'num_dogs': len(new_dogs),
-            'total_images': total_images,
-            'dogs': new_dogs
-        }
+        return gallery_html, stats_msg
     
-    def get_all_dog_galleries(self):
-        """Get all dog galleries for display"""
-        galleries = []
-        
-        for dog_id in sorted(self.current_dogs.keys()):
-            images = []
-            paths = self.dog_images.get(dog_id, [])
+    def _create_gallery_html(self):
+        """Create HTML gallery of detected dogs"""
+        if not self.current_dogs:
+            return "<p style='text-align:center; padding:20px;'>No dogs detected</p>"
+        
+        html = """
+        <div style='padding: 20px;'>
+            <h2 style='text-align:center; color:#2196F3;'>🐕 Detected Dogs</h2>
+            <div style='display: grid; grid-template-columns: repeat(auto-fit, minmax(400px, 1fr)); gap: 20px; margin-top: 20px;'>
+        """
+        
+        for dog_id, images in self.current_dogs.items():
+            html += f"""
+            <div style='border: 2px solid #2196F3; border-radius: 10px; padding: 15px; background: #f5f5f5;'>
+                <h3 style='margin: 0 0 10px 0; color: #1976D2;'>Dog ID: {dog_id}</h3>
+                <p style='margin: 5px 0; color: #666;'>Images captured: {len(images)}</p>
+                <div style='display: grid; grid-template-columns: repeat(5, 1fr); gap: 5px; margin-top: 10px;'>
+            """
             
-            for img_path in paths:
-                try:
-                    img = cv2.imread(img_path)
-                    if img is not None:
-                        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-                        images.append(img_rgb)
-                except:
-                    continue
+            for img in images:
+                img_base64 = self._img_to_base64(img)
+                html += f"""
+                <img src='data:image/jpeg;base64,{img_base64}' 
+                     style='width: 100%; aspect-ratio: 1; object-fit: cover; border-radius: 5px;'>
+                """
             
-            if images:
-                galleries.append({
-                    'dog_id': dog_id,
-                    'images': images,
-                    'paths': paths,
-                    'num_images': len(images)
-                })
-        
-        return galleries
-    
-    def reassign_images(self, from_dog: int, to_dog: int, image_indices: List[int]):
-        """Move selected images from one dog to another"""
-        if from_dog not in self.dog_images:
-            return f"Dog {from_dog} not found"
-        
-        from_paths = self.dog_images[from_dog]
-        moved_paths = []
-        
-        # Get paths to move
-        for idx in sorted(image_indices, reverse=True):
-            if 0 <= idx < len(from_paths):
-                moved_paths.append(from_paths.pop(idx))
-        
-        if not moved_paths:
-            return "No valid images to move"
-        
-        # Create target dog if needed
-        if to_dog not in self.current_dogs:
-            self.current_dogs[to_dog] = {
-                'num_images': 0,
-                'avg_confidence': 0.5,
-                'source': 'reassigned'
-            }
-            self.dog_images[to_dog] = []
-        
-        # Move files
-        to_dir = self.temp_dir / f"dog_{to_dog:03d}"
-        to_dir.mkdir(exist_ok=True)
-        
-        for old_path in moved_paths:
-            old_path = Path(old_path)
-            if old_path.exists():
-                new_path = to_dir / f"img_{len(self.dog_images[to_dog]):03d}.jpg"
-                shutil.move(str(old_path), str(new_path))
-                self.dog_images[to_dog].append(str(new_path))
-        
-        # Update metadata
-        self.current_dogs[from_dog]['num_images'] = len(self.dog_images[from_dog])
-        self.current_dogs[to_dog]['num_images'] = len(self.dog_images[to_dog])
-        
-        # Remove empty dogs
-        if len(self.dog_images[from_dog]) == 0:
-            del self.current_dogs[from_dog]
-            del self.dog_images[from_dog]
+            html += "</div></div>"
         
-        return f"✅ Moved {len(moved_paths)} images from Dog {from_dog} to Dog {to_dog}"
-    
-    def delete_dog_images(self, dog_id: int, image_indices: List[int]):
-        """Delete selected images from a dog"""
-        if dog_id not in self.dog_images:
-            return f"Dog {dog_id} not found"
-        
-        paths = self.dog_images[dog_id]
-        deleted_count = 0
-        
-        # Delete in reverse order to maintain indices
-        for idx in sorted(image_indices, reverse=True):
-            if 0 <= idx < len(paths):
-                img_path = Path(paths.pop(idx))
-                if img_path.exists():
-                    img_path.unlink()
-                deleted_count += 1
-        
-        # Update metadata
-        self.current_dogs[dog_id]['num_images'] = len(paths)
-        
-        # Remove dog if no images left
-        if len(paths) == 0:
-            del self.current_dogs[dog_id]
-            del self.dog_images[dog_id]
-        
-        return f"🗑️ Deleted {deleted_count} images from Dog {dog_id}"
-    
-    def export_dataset(self, include_csv: bool = True):
-        """Export the dataset"""
-        # Clear export directory
-        if self.export_dir.exists():
-            shutil.rmtree(self.export_dir)
-        self.export_dir.mkdir()
-        
-        # Copy all dog directories
-        total_images = 0
-        for dog_id in self.current_dogs:
-            # Try permanent first, then temp
-            src_dir = self.database_dir / f"dog_{dog_id:03d}"
-            if not src_dir.exists():
-                src_dir = self.temp_dir / f"dog_{dog_id:03d}"
-            
-            if src_dir.exists():
-                dst_dir = self.export_dir / f"dog_{dog_id:03d}"
-                shutil.copytree(src_dir, dst_dir)
-                total_images += len(list(dst_dir.glob("*.jpg")))
-        
-        # Create CSV if requested
-        if include_csv:
-            csv_data = []
-            for dog_id in self.current_dogs:
-                dog_dir = self.export_dir / f"dog_{dog_id:03d}"
-                if dog_dir.exists():
-                    for img_path in dog_dir.glob("*.jpg"):
-                        csv_data.append({
-                            'dog_id': dog_id,
-                            'image_path': str(img_path.relative_to(self.export_dir)),
-                            'filename': img_path.name
-                        })
-            
-            df = pd.DataFrame(csv_data)
-            df.to_csv(self.export_dir / "dataset.csv", index=False)
-        
-        # Create zip
-        zip_path = Path("dog_dataset.zip")
-        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
-            for file_path in self.export_dir.rglob("*"):
-                if file_path.is_file():
-                    zipf.write(file_path, file_path.relative_to(self.export_dir))
-        
-        return str(zip_path), len(self.current_dogs), total_images
+        html += "</div></div>"
+        return html
     
     def _img_to_base64(self, img):
-        """Convert image to base64 for HTML display"""
-        import base64
-        from io import BytesIO
-        from PIL import Image
-        
+        """Convert image to base64"""
         pil_img = Image.fromarray(img)
         buffered = BytesIO()
-        pil_img.save(buffered, format="JPEG", quality=70)
-        img_str = base64.b64encode(buffered.getvalue()).decode()
-        return img_str
+        pil_img.save(buffered, format="JPEG", quality=85)
+        return base64.b64encode(buffered.getvalue()).decode()
     
     def create_interface(self):
         """Create simplified Gradio interface"""
-        with gr.Blocks(title="Dog Dataset Creator", theme=gr.themes.Soft()) as app:
-            gr.Markdown("# 🐕 Dog Dataset Creator")
+        with gr.Blocks(title="Dog Detection Demo", theme=gr.themes.Soft()) as app:
+            gr.Markdown(
+                """
+                # 🐕 Dog Detection & Tracking Demo
+                ### Using MegaDescriptor for Individual Dog Recognition
+                Upload a video to detect and track individual dogs. Each dog gets a unique ID.
+                """
+            )
             
-            with gr.Tabs():
-                # ========== STEP 1: Process Video ==========
-                with gr.Tab("📹 Process Video"):
-                    with gr.Row():
-                        with gr.Column():
-                            video_input = gr.Video(label="Upload Video")
-                            
-                            reid_threshold = gr.Slider(
-                                0.3, 0.8, 0.6, step=0.05,
-                                label="ReID Threshold",
-                                info="Lower = more lenient matching"
-                            )
-                            
-                            max_images = gr.Slider(
-                                10, 50, 30, step=5,
-                                label="Max Images per Dog"
-                            )
-                            
-                            sample_rate = gr.Slider(
-                                1, 5, 2, step=1,
-                                label="Sample Rate (process every N frames)"
-                            )
-                            
-                            process_btn = gr.Button("🚀 Process Video", variant="primary")
-                        
-                        with gr.Column():
-                            progress_text = gr.Textbox(label="Progress", interactive=False)
-                            results_html = gr.HTML()
-                    
-                    def process_video_wrapper(video, threshold, max_img, sample):
-                        if not video:
-                            return "No video uploaded", ""
-                        
-                        for update in self.process_video(video, threshold, int(max_img), int(sample)):
-                            if update['status'] == 'processing':
-                                yield f"Processing: {update['progress']}%", ""
-                            else:
-                                html = f"""
-                                <div style="padding: 15px; background: #e8f5e9; border-radius: 8px;">
-                                    <h3>✅ Processing Complete!</h3>
-                                    <p>Dogs detected: <b>{update['num_dogs']}</b></p>
-                                    <p>Total images: <b>{update['total_images']}</b></p>
-                                </div>
-                                """
-                                yield "Complete!", html
-                    
-                    process_btn.click(
-                        process_video_wrapper,
-                        inputs=[video_input, reid_threshold, max_images, sample_rate],
-                        outputs=[progress_text, results_html]
-                    )
-                
-                # ========== STEP 2: Verify & Edit ==========
-                with gr.Tab("✏️ Verify & Edit"):
-                    gr.Markdown("""
-                    ### 📋 How to Edit Dogs:
-                    1. Click **Refresh Galleries** to see all dogs
-                    2. Note the image numbers (0, 1, 2...) shown on each image
-                    3. Enter the Dog ID and image numbers to move or delete
-                    """)
-                    
-                    refresh_btn = gr.Button("🔄 Refresh Galleries", variant="primary", size="lg")
-                    
-                    # Gallery display
-                    gallery_html = gr.HTML(label="Dog Galleries")
-                    
-                    gr.Markdown("---")
-                    
-                    # Simplified controls in clear sections
-                    with gr.Row():
-                        with gr.Column(scale=1):
-                            gr.Markdown("### 🔄 Move Images Between Dogs")
-                            from_dog = gr.Number(label="From Dog ID", value=1, precision=0)
-                            image_indices = gr.Textbox(
-                                label="Image Numbers",
-                                placeholder="0,2,5",
-                                info="Enter image numbers shown on thumbnails"
-                            )
-                            to_dog = gr.Number(label="To Dog ID", value=2, precision=0)
-                            move_btn = gr.Button("Move Images →", variant="primary")
-                        
-                        with gr.Column(scale=1):
-                            gr.Markdown("### 🗑️ Delete Images")
-                            del_dog = gr.Number(label="Dog ID", value=1, precision=0)
-                            del_indices = gr.Textbox(
-                                label="Image Numbers to Delete",
-                                placeholder="0,1,2",
-                                info="Enter image numbers to remove"
-                            )
-                            delete_btn = gr.Button("Delete Images", variant="stop")
-                    
-                    gr.Markdown("---")
-                    
-                    # Status and save
-                    status_text = gr.Textbox(label="Status", interactive=False)
-                    save_btn = gr.Button("💾 Save All Dogs to Database", variant="primary", size="lg")
-                    
-                    def refresh_galleries():
-                        """Create HTML grid of dog galleries"""
-                        galleries = self.get_all_dog_galleries()
-                        
-                        if not galleries:
-                            return "<p style='text-align:center; color:#666;'>No dogs found. Process a video first.</p>"
-                        
-                        html = """
-                        <div style='max-width: 1200px; margin: 0 auto;'>
-                        <div style='display: grid; grid-template-columns: repeat(auto-fit, minmax(500px, 1fr)); gap: 20px;'>
-                        """
-                        
-                        for gal in galleries:
-                            dog_id = gal['dog_id']
-                            num_images = gal['num_images']
-                            
-                            html += f"""
-                            <div style='border: 2px solid #2196F3; border-radius: 10px; padding: 15px; background: #f5f5f5;'>
-                                <h3 style='margin: 0 0 10px 0; color: #1976D2;'>🐕 Dog {dog_id}</h3>
-                                <p style='margin: 5px 0; color: #666;'>Total: {num_images} images</p>
-                                <div style='display: grid; grid-template-columns: repeat(4, 1fr); gap: 8px; margin-top: 10px;'>
-                            """
-                            
-                            # Show first 12 images as thumbnails
-                            for i, img in enumerate(gal['images'][:12]):
-                                html += f"""
-                                <div style='position: relative; aspect-ratio: 1/1; overflow: hidden; 
-                                          border: 2px solid #ddd; border-radius: 5px;'>
-                                    <img src='data:image/jpeg;base64,{self._img_to_base64(img)}' 
-                                         style='width: 100%; height: 100%; object-fit: cover;'
-                                         title='Image {i}'>
-                                    <div style='position: absolute; top: 4px; left: 4px; 
-                                              background: #2196F3; color: white; 
-                                              padding: 2px 6px; font-size: 12px; font-weight: bold;
-                                              border-radius: 3px;'>{i}</div>
-                                </div>
-                                """
-                            
-                            if num_images > 12:
-                                html += f"""
-                                <div style='grid-column: span 4; text-align: center; 
-                                          padding: 10px; color: #666; font-style: italic;'>
-                                    ... and {num_images - 12} more images
-                                </div>
-                                """
-                            
-                            html += "</div></div>"
-                        
-                        html += "</div></div>"
-                        
-                        return html
-                    
-                    refresh_btn.click(
-                        refresh_galleries,
-                        outputs=gallery_html
+            with gr.Row():
+                with gr.Column(scale=1):
+                    video_input = gr.Video(label="Upload Video")
+                    
+                    reid_threshold = gr.Slider(
+                        0.3, 0.8, 0.6, step=0.05,
+                        label="ReID Matching Threshold",
+                        info="Lower = more lenient matching"
                     )
                     
-                    def move_images_wrapper(from_id, indices, to_id):
-                        try:
-                            if not indices:
-                                return "Please enter image numbers to move"
-                            indices_list = [int(x.strip()) for x in indices.split(',') if x.strip()]
-                            return self.reassign_images(int(from_id), int(to_id), indices_list)
-                        except ValueError:
-                            return "Invalid input. Use numbers like: 0,1,2"
-                        except Exception as e:
-                            return f"Error: {str(e)}"
-                    
-                    def delete_images_wrapper(dog_id, indices):
-                        try:
-                            if not indices:
-                                return "Please enter image numbers to delete"
-                            indices_list = [int(x.strip()) for x in indices.split(',') if x.strip()]
-                            return self.delete_dog_images(int(dog_id), indices_list)
-                        except ValueError:
-                            return "Invalid input. Use numbers like: 0,1,2"
-                        except Exception as e:
-                            return f"Error: {str(e)}"
-                    
-                    move_btn.click(
-                        move_images_wrapper,
-                        inputs=[from_dog, image_indices, to_dog],
-                        outputs=status_text
-                    )
-                    
-                    delete_btn.click(
-                        delete_images_wrapper,
-                        inputs=[del_dog, del_indices],
-                        outputs=status_text
+                    sample_rate = gr.Slider(
+                        1, 5, 2, step=1,
+                        label="Frame Sample Rate",
+                        info="Process every N frames (higher = faster)"
                     )
                     
-                    save_btn.click(
-                        lambda: (self.save_database(), "✅ All dogs saved to database!")[1],
-                        outputs=status_text
-                    )
+                    process_btn = gr.Button("🚀 Process Video", variant="primary", size="lg")
                 
-                # ========== STEP 3: Export ==========
-                with gr.Tab("📦 Export Dataset"):
-                    gr.Markdown("""
-                    ### 📥 Export Your Dataset
-                    Download all dogs as a ZIP file for training your model.
-                    """)
-                    
-                    with gr.Row():
-                        with gr.Column():
-                            include_csv = gr.Checkbox(
-                                label="Include CSV metadata file",
-                                value=True,
-                                info="Creates a CSV with image paths and dog IDs"
-                            )
-                            
-                            export_btn = gr.Button("📥 Create Export ZIP", variant="primary", size="lg")
-                        
-                        with gr.Column():
-                            download_file = gr.File(
-                                label="Download Dataset",
-                                interactive=False,
-                                visible=False
-                            )
-                            export_status = gr.HTML()
-                    
-                    def export_wrapper(csv):
-                        try:
-                            zip_path, num_dogs, num_images = self.export_dataset(csv)
-                            html = f"""
-                            <div style='padding: 20px; background: #e3f2fd; border-radius: 10px;
-                                      border: 2px solid #2196F3;'>
-                                <h3 style='color: #1976D2; margin-top: 0;'>✅ Dataset Ready!</h3>
-                                <div style='display: grid; grid-template-columns: repeat(2, 1fr); gap: 10px;'>
-                                    <div style='background: white; padding: 10px; border-radius: 5px;'>
-                                        <p style='margin: 0; color: #666;'>Dogs</p>
-                                        <p style='margin: 0; font-size: 24px; font-weight: bold;'>{num_dogs}</p>
-                                    </div>
-                                    <div style='background: white; padding: 10px; border-radius: 5px;'>
-                                        <p style='margin: 0; color: #666;'>Images</p>
-                                        <p style='margin: 0; font-size: 24px; font-weight: bold;'>{num_images}</p>
-                                    </div>
-                                </div>
-                                <p style='margin-top: 15px; color: #666;'>
-                                    Click the download button below to get your dataset ZIP file.
-                                </p>
-                            </div>
-                            """
-                            return gr.update(value=zip_path, visible=True), html
-                        except Exception as e:
-                            html = f"""
-                            <div style='padding: 15px; background: #ffebee; border-radius: 10px;
-                                      border: 2px solid #f44336;'>
-                                <h3 style='color: #c62828; margin-top: 0;'>❌ Export Failed</h3>
-                                <p style='color: #666;'>{str(e)}</p>
-                            </div>
-                            """
-                            return gr.update(visible=False), html
-                    
-                    export_btn.click(
-                        export_wrapper,
-                        inputs=include_csv,
-                        outputs=[download_file, export_status]
-                    )
-                    
-                    def move_images(from_dog, to_dog, indices_str):
-                        try:
-                            indices = [int(x.strip()) for x in indices_str.split(',')]
-                            return self.reassign_images(int(from_dog), int(to_dog), indices)
-                        except:
-                            return "Invalid input. Use format: 0,1,2"
-                    
-                    def delete_images(dog_id, indices_str):
-                        try:
-                            indices = [int(x.strip()) for x in indices_str.split(',')]
-                            return self.delete_dog_images(int(dog_id), indices)
-                        except:
-                            return "Invalid input. Use format: 0,1,2"
-                    
-                    move_btn.click(
-                        move_images,
-                        inputs=[selected_dog, target_dog, selected_indices],
-                        outputs=status_text
-                    )
-                    
-                    delete_btn.click(
-                        delete_images,
-                        inputs=[selected_dog, selected_indices],
-                        outputs=status_text
-                    )
-                    
-                    save_btn.click(
-                        lambda: (self.save_database(), "✅ Saved to database!")[1],
-                        outputs=status_text
-                    )
-                
-                # ========== STEP 3: Export ==========
-                with gr.Tab("📦 Export Dataset"):
-                    gr.Markdown("### Export your dataset for training")
-                    
-                    include_csv = gr.Checkbox(label="Include CSV file", value=True)
-                    
-                    export_btn = gr.Button("📥 Export Dataset", variant="primary", size="lg")
-                    
-                    download_file = gr.File(label="Download", interactive=False)
-                    export_status = gr.HTML()
-                    
-                    def export_wrapper(csv):
-                        zip_path, num_dogs, num_images = self.export_dataset(csv)
-                        html = f"""
-                        <div style='padding: 15px; background: #e3f2fd; border-radius: 8px;'>
-                            <h3>✅ Export Complete!</h3>
-                            <p>Dogs: <b>{num_dogs}</b></p>
-                            <p>Images: <b>{num_images}</b></p>
-                            <p>Ready to download!</p>
-                        </div>
-                        """
-                        return zip_path, html
-                    
-                    export_btn.click(
-                        export_wrapper,
-                        inputs=include_csv,
-                        outputs=[download_file, export_status]
-                    )
+                with gr.Column(scale=2):
+                    status_text = gr.Textbox(label="Status", interactive=False)
+                    gallery_output = gr.HTML(label="Detected Dogs")
+            
+            # Process video on button click
+            process_btn.click(
+                self.process_video,
+                inputs=[video_input, reid_threshold, sample_rate],
+                outputs=[gallery_output, status_text]
+            )
+            
+            # Auto-reset when parameters change
+            video_input.change(fn=lambda: (None, "Ready for new video"), outputs=[gallery_output, status_text])
+            reid_threshold.change(fn=lambda: (None, "Parameters changed - upload video to process"), outputs=[gallery_output, status_text])
+            sample_rate.change(fn=lambda: (None, "Parameters changed - upload video to process"), outputs=[gallery_output, status_text])
             
-            return app
+        return app
 
 
 # Main entry point
 if __name__ == "__main__":
-    creator = SimpleDatasetCreator()
-    app = creator.create_interface()
+    demo = DogDetectionDemo()
+    app = demo.create_interface()
     app.launch(
         server_name="0.0.0.0",
         server_port=7860,