Spaces:
Build error
Build error
| """ | |
| Test script for the QuickDraw Classifier | |
| Run this to test the classification functionality locally | |
| """ | |
| import os | |
| import sys | |
| import torch | |
| from PIL import Image, ImageDraw | |
| import numpy as np | |
| # Add current directory to path to import app modules | |
| sys.path.append(os.path.dirname(os.path.abspath(__file__))) | |
| try: | |
| from app import classify_drawing, load_categories, CATEGORIES | |
| print("✅ Successfully imported app modules") | |
| except ImportError as e: | |
| print(f"❌ Error importing app modules: {e}") | |
| print("Make sure you've installed the requirements: pip install -r requirements.txt") | |
| sys.exit(1) | |
| def create_test_drawing(category="cat", size=(400, 400)): | |
| """ | |
| Create a simple test drawing for testing | |
| """ | |
| # Create a white canvas | |
| img = Image.new('RGB', size, 'white') | |
| draw = ImageDraw.Draw(img) | |
| # Draw simple shapes based on category | |
| if category == "cat": | |
| # Draw a simple cat face | |
| # Head circle | |
| draw.ellipse([150, 150, 250, 250], outline='black', width=3) | |
| # Ears | |
| draw.polygon([(140, 150), (160, 120), (180, 150)], outline='black', width=2) | |
| draw.polygon([(220, 150), (240, 120), (260, 150)], outline='black', width=2) | |
| # Eyes | |
| draw.ellipse([170, 170, 180, 180], fill='black') | |
| draw.ellipse([220, 170, 230, 180], fill='black') | |
| # Nose | |
| draw.polygon([(195, 190), (205, 190), (200, 200)], fill='black') | |
| # Mouth | |
| draw.arc([185, 200, 215, 220], 0, 180, fill='black', width=2) | |
| elif category == "house": | |
| # Draw a simple house | |
| # Base rectangle | |
| draw.rectangle([150, 200, 250, 300], outline='black', width=3) | |
| # Roof triangle | |
| draw.polygon([(140, 200), (200, 150), (260, 200)], outline='black', width=3) | |
| # Door | |
| draw.rectangle([180, 240, 220, 300], outline='black', width=2) | |
| # Windows | |
| draw.rectangle([160, 220, 180, 240], outline='black', width=2) | |
| draw.rectangle([220, 220, 240, 240], outline='black', width=2) | |
| elif category == "car": | |
| # Draw a simple car | |
| # Body | |
| draw.rectangle([120, 200, 280, 240], outline='black', width=3) | |
| # Roof | |
| draw.rectangle([150, 180, 250, 200], outline='black', width=3) | |
| # Wheels | |
| draw.ellipse([140, 230, 170, 260], outline='black', width=3) | |
| draw.ellipse([230, 230, 260, 260], outline='black', width=3) | |
| # Windows | |
| draw.rectangle([160, 185, 190, 195], outline='black', width=2) | |
| draw.rectangle([210, 185, 240, 195], outline='black', width=2) | |
| return img | |
| def test_classification(): | |
| """ | |
| Test the classification functionality | |
| """ | |
| print("\n🧪 Testing QuickDraw Classification") | |
| print("="*50) | |
| # Test 1: Check if categories are loaded | |
| print(f"📋 Loaded {len(CATEGORIES)} categories") | |
| print(f"First 10 categories: {CATEGORIES[:10]}") | |
| # Test 2: Test with simple drawings | |
| test_categories = ["cat", "house", "car"] | |
| for category in test_categories: | |
| print(f"\n🎨 Testing with {category} drawing...") | |
| # Create test drawing | |
| test_img = create_test_drawing(category) | |
| # Classify the drawing | |
| results = classify_drawing(test_img, top_k=5) | |
| if results: | |
| print(f"✅ Classification successful!") | |
| print(f"Top predictions:") | |
| for i, (pred_category, confidence) in enumerate(results, 1): | |
| percentage = confidence * 100 | |
| marker = "🎯" if pred_category == category else " " | |
| print(f"{marker} {i}. {pred_category.title()} - {percentage:.1f}%") | |
| else: | |
| print(f"❌ Classification failed for {category}") | |
| return True | |
| def test_model_loading(): | |
| """ | |
| Test if CLIP model loads correctly | |
| """ | |
| print("\n🤖 Testing Model Loading") | |
| print("="*30) | |
| try: | |
| import clip | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"📱 Using device: {device}") | |
| model, preprocess = clip.load("ViT-B/32", device=device) | |
| print("✅ CLIP model loaded successfully") | |
| # Test text encoding | |
| text_inputs = clip.tokenize(["a drawing of a cat", "a drawing of a dog"]).to(device) | |
| with torch.no_grad(): | |
| text_features = model.encode_text(text_inputs) | |
| print(f"✅ Text encoding working - shape: {text_features.shape}") | |
| return True | |
| except Exception as e: | |
| print(f"❌ Model loading failed: {e}") | |
| return False | |
| def main(): | |
| """ | |
| Run all tests | |
| """ | |
| print("🚀 QuickDraw Classifier Test Suite") | |
| print("="*40) | |
| # Test 1: Model loading | |
| if not test_model_loading(): | |
| print("\n❌ Model loading test failed. Please check your installation.") | |
| return False | |
| # Test 2: Classification | |
| if not test_classification(): | |
| print("\n❌ Classification test failed.") | |
| return False | |
| print("\n" + "="*50) | |
| print("🎉 All tests passed! Your QuickDraw Classifier is ready!") | |
| print("\n📖 Next steps:") | |
| print("1. Run 'python app.py' to start the local app") | |
| print("2. Deploy to Hugging Face Spaces following the README") | |
| print("3. Test the deployed app with real drawings") | |
| return True | |
| if __name__ == "__main__": | |
| success = main() | |
| sys.exit(0 if success else 1) |