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:

    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 __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:
        # 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.")

if __name__ == "__main__":
    main()