File size: 8,922 Bytes
105944b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 | """
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!")
|