plastic_client.py

"""
Plastic Detection API Client Library

A simple Python client for interacting with the Plastic Detection API.

Usage:
    from plastic_client import PlasticDetectionClient
    
    client = PlasticDetectionClient("http://localhost:8000")
    results = client.predict("path/to/image.jpg")
    print(results)
"""

import requests
import json
from typing import Dict, List, Optional
from pathlib import Path


class PlasticDetectionClient:
    """Client for Plastic Detection API"""
    
    # Hardcoded API server URL
    DEFAULT_API_URL = "http://localhost:8000"
    
    def __init__(self, base_url: str = None):
        """
        Initialize the client
        
        Args:
            base_url: Base URL of the API server (defaults to localhost:8000)
        """
        if base_url is None:
            base_url = self.DEFAULT_API_URL
        self.base_url = base_url.rstrip('/')
        self.session = requests.Session()
    
    def health_check(self) -> bool:
        """
        Check if the API server is running
        
        Returns:
            True if server is running, False otherwise
        """
        try:
            response = self.session.get(f"{self.base_url}/")
            return response.status_code == 200
        except requests.exceptions.ConnectionError:
            return False
    
    def get_classes(self) -> Dict[int, str]:
        """
        Get all detectable plastic classes
        
        Returns:
            Dictionary mapping class IDs to class names
        """
        response = self.session.get(f"{self.base_url}/classes")
        response.raise_for_status()
        data = response.json()
        return data.get('classes', {})
    
    def predict(self, image_path: str, confidence_threshold: float = 0.25) -> Dict:
        """
        Predict plastic types in an image
        
        Args:
            image_path: Path to the image file
            confidence_threshold: Minimum confidence score (0.0-1.0)
        
        Returns:
            Dictionary containing detection results
        
        Raises:
            FileNotFoundError: If image file doesn't exist
            requests.exceptions.RequestException: If API call fails
        """
        image_path = Path(image_path)
        
        if not image_path.exists():
            raise FileNotFoundError(f"Image not found: {image_path}")
        
        with open(image_path, 'rb') as f:
            files = {'file': f}
            response = self.session.post(f"{self.base_url}/predict", files=files)
        
        response.raise_for_status()
        return response.json()
    
    def predict_bytes(self, image_bytes: bytes, filename: str = "image.jpg") -> Dict:
        """
        Predict plastic types from image bytes
        
        Args:
            image_bytes: Image file content as bytes
            filename: Original filename for reference
        
        Returns:
            Dictionary containing detection results
        """
        files = {'file': (filename, image_bytes)}
        response = self.session.post(f"{self.base_url}/predict", files=files)
        response.raise_for_status()
        return response.json()
    
    def predict_multiple(self, image_paths: List[str]) -> List[Dict]:
        """
        Predict on multiple images
        
        Args:
            image_paths: List of image file paths
        
        Returns:
            List of prediction results
        """
        results = []
        for image_path in image_paths:
            try:
                result = self.predict(image_path)
                results.append({
                    'image': str(image_path),
                    'success': True,
                    'data': result
                })
            except Exception as e:
                results.append({
                    'image': str(image_path),
                    'success': False,
                    'error': str(e)
                })
        
        return results
    
    def filter_detections(self, predictions: Dict, min_confidence: float = 0.5, 
                         classes: Optional[List[str]] = None) -> Dict:
        """
        Filter detections by confidence and class
        
        Args:
            predictions: Prediction result from predict()
            min_confidence: Minimum confidence threshold
            classes: List of class names to filter by (None = all)
        
        Returns:
            Filtered prediction results
        """
        filtered_detections = []
        
        for det in predictions.get('detections', []):
            # Filter by confidence
            if det['confidence'] < min_confidence:
                continue
            
            # Filter by class
            if classes and det['class'] not in classes:
                continue
            
            filtered_detections.append(det)
        
        predictions_copy = predictions.copy()
        predictions_copy['detections'] = filtered_detections
        predictions_copy['count'] = len(filtered_detections)
        
        return predictions_copy
    
    def get_statistics(self, predictions: Dict) -> Dict:
        """
        Get statistics from predictions
        
        Args:
            predictions: Prediction result from predict()
        
        Returns:
            Dictionary with statistics
        """
        detections = predictions.get('detections', [])
        
        if not detections:
            return {
                'total_detections': 0,
                'classes': {},
                'average_confidence': 0.0,
                'max_confidence': 0.0,
                'min_confidence': 0.0
            }
        
        classes_count = {}
        confidences = []
        
        for det in detections:
            classes_count[det['class']] = classes_count.get(det['class'], 0) + 1
            confidences.append(det['confidence'])
        
        return {
            'total_detections': len(detections),
            'classes': classes_count,
            'average_confidence': sum(confidences) / len(confidences),
            'max_confidence': max(confidences),
            'min_confidence': min(confidences)
        }
    
    def save_results_json(self, predictions: Dict, output_path: str):
        """
        Save prediction results to JSON file
        
        Args:
            predictions: Prediction result from predict()
            output_path: Path to save JSON file
        """
        with open(output_path, 'w') as f:
            json.dump(predictions, f, indent=2)
    
    def close(self):
        """Close the session"""
        self.session.close()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


# Example usage
if __name__ == "__main__":
    import sys
    
    print("Plastic Detection API Client Example")
    print("=" * 50)
    
    # Initialize client
    client = PlasticDetectionClient("http://localhost:8000")
    
    # Check health
    print("\n✓ Checking server health...")
    if client.health_check():
        print("  Server is running!")
    else:
        print("  ✗ Server is not responding")
        sys.exit(1)
    
    # Get classes
    print("\n✓ Available classes:")
    classes = client.get_classes()
    for cls_id, cls_name in classes.items():
        print(f"  {cls_id}: {cls_name}")
    
    # Test prediction
    test_image = "test/images/00000062_jpg.rf.e4f04e59310a48547b50edcdebc1c286.jpg"
    
    try:
        print(f"\n✓ Running prediction on {test_image}...")
        results = client.predict(test_image)
        
        print(f"  Detections: {results['count']}")
        
        if results['detections']:
            print("  Objects found:")
            for det in results['detections']:
                print(f"    - {det['class']}: {det['confidence']:.2%}")
        
        # Get statistics
        stats = client.get_statistics(results)
        print(f"\n  Statistics:")
        print(f"    Total: {stats['total_detections']}")
        print(f"    Avg Confidence: {stats['average_confidence']:.2%}")
        print(f"    Classes: {stats['classes']}")
        
        # Filter by confidence
        filtered = client.filter_detections(results, min_confidence=0.8)
        print(f"\n  After filtering (conf > 0.8): {filtered['count']} detections")
        
    except FileNotFoundError:
        print(f"  ✗ Test image not found: {test_image}")
    except Exception as e:
        print(f"  ✗ Error: {e}")
    
    finally:
        client.close()
    
    print("\n" + "=" * 50)
    print("Example completed!")