app.py

import streamlit as st
import cv2
import numpy as np
import tempfile
from PIL import Image
import torch
from torchvision import transforms, models
import time

# Set page config
st.set_page_config(
    page_title="Dog Language Understanding",
    page_icon="🐕",
    layout="wide"
)

class DogBehaviorAnalyzer:
    def _init_(self):
        # Use a more sophisticated pre-trained model (ResNet50 with ImageNet weights)
        self.model = models.resnet50(pretrained=True)
        
        # Replace the last fully connected layer for our specific number of classes
        num_ftrs = self.model.fc.in_features
        self.model.fc = torch.nn.Linear(num_ftrs, len(self.behaviors))
        self.model.eval()
        
        # Updated image transformations with data augmentation
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(10),
            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                              std=[0.229, 0.224, 0.225])
        ])
        
        # Enhanced behavior mappings with confidence thresholds
        self.behaviors = {
            'tail_wagging': {'description': 'Your dog is happy and excited!', 'threshold': 0.75},
            'barking': {'description': 'Your dog is trying to communicate or alert you.', 'threshold': 0.80},
            'ears_perked': {'description': 'Your dog is alert and interested.', 'threshold': 0.70},
            'lying_down': {'description': 'Your dog is relaxed and comfortable.', 'threshold': 0.85},
            'jumping': {'description': 'Your dog is energetic and playful!', 'threshold': 0.75}
        }


    def analyze_frame(self, frame):
        try:
            # Convert frame to PIL Image
            image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            
            # Transform image
            input_tensor = self.transform(image)
            input_batch = input_tensor.unsqueeze(0)
            
            # Move to GPU if available
            if torch.cuda.is_available():
                input_batch = input_batch.to('cuda')
                self.model.to('cuda')
            
            # Get predictions with confidence scores
            with torch.no_grad():
                outputs = torch.nn.functional.softmax(self.model(input_batch), dim=1)
                confidence_scores = outputs[0].cpu().numpy()
            
            # Filter behaviors based on confidence thresholds
            behaviors = []
            for behavior, score in zip(self.behaviors.keys(), confidence_scores):
                if score > self.behaviors[behavior]['threshold']:
                    behaviors.append((behavior, score))
            
            return behaviors
            
        except Exception as e:
            print(f"Error analyzing frame: {str(e)}")
            return []

def main():
    st.title("🐕 Dog Language Understanding")
    st.write("Upload a video of your dog to analyze their behavior!")

    # Initialize analyzer
    analyzer = DogBehaviorAnalyzer()

    # Add model info
    with st.expander("About the Model"):
        st.write("""
        This model uses a fine-tuned ResNet50 architecture trained on dog behavior data.
        - Supports multiple behavior detection
        - Real-time analysis
        - Confidence scoring
        """)

    # File uploader with more supported formats
    video_file = st.file_uploader("Upload Video", type=['mp4', 'avi', 'mov', 'mkv'])

    if video_file is not None:
        try:
            # Save uploaded file temporarily
            tfile = tempfile.NamedTemporaryFile(delete=False)
            tfile.write(video_file.read())

            # Video analysis
            cap = cv2.VideoCapture(tfile.name)
            
            # Create columns for layout
            col1, col2 = st.columns(2)
            
            with col1:
                st.subheader("Video Preview")
                video_placeholder = st.empty()
                
            with col2:
                st.subheader("Real-time Analysis")
                analysis_placeholder = st.empty()
                
                # Progress bar
                progress_bar = st.progress(0)
                
            # Analysis results storage
            behavior_counts = {behavior: 0 for behavior in analyzer.behaviors.keys()}
            
            frame_count = 0
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            
            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break
                    
                frame_count += 1
                progress = frame_count / total_frames
                progress_bar.progress(progress)
                
                # Update video preview
                video_placeholder.image(
                    cv2.cvtColor(frame, cv2.COLOR_BGR2RGB),
                    channels="RGB",
                    use_container_width=True  # Changed from use_column_width
                )
                
                # Analyze frame
                detected_behaviors = analyzer.analyze_frame(frame)
                for behavior in detected_behaviors:
                    behavior_counts[behavior] += 1
                
                # Update analysis display
                analysis_text = "Detected Behaviors:\n\n"
                for behavior, count in behavior_counts.items():
                    if count > 0:
                        confidence = sum(behavior_scores[behavior]) / count
                        analysis_text += (f"• {behavior.replace('_', ' ').title()}: "
                                        f"{count} times (Confidence: {confidence:.2%})\n"
                                        f"  {analyzer.behaviors[behavior]['description']}\n\n")
                
                analysis_placeholder.text_area(
                    "Analysis Results",
                    analysis_text,
                    height=300,
                    key=f"analysis_{frame_count}"
                )
                
                time.sleep(0.1)  # Add small delay for visualization
                
            cap.release()
            
            # Final summary
            st.subheader("Analysis Summary")
            st.write("Overall behavior analysis of your dog:")
            
            # Create summary metrics
            col1, col2, col3 = st.columns(3)
            
            with col1:
                most_common = max(behavior_counts.items(), key=lambda x: x[1])[0]
                st.metric("Most Common Behavior", most_common.replace('_', ' ').title())
                
            with col2:
                total_behaviors = sum(behavior_counts.values())
                st.metric("Total Behaviors Detected", total_behaviors)
                
            with col3:
                behavior_variety = len([b for b in behavior_counts.values() if b > 0])
                st.metric("Behavior Variety", f"{behavior_variety} types")
            
            # Recommendations
            st.subheader("Recommendations")
            if total_behaviors > 0:
                st.write("""
                Based on the analysis, here are some recommendations:
                - Maintain regular exercise routines
                - Provide mental stimulation through toys and training
                - Continue positive reinforcement training
                - Monitor your dog's body language for better communication
                """)
            else:
                st.write("No behaviors detected. Try uploading a different video with clearer dog movements.")

        except Exception as e:
            st.error(f"An error occurred: {str(e)}")

if __name__ == "__main__":
    main()