Spaces:
Build error
Build error
File size: 5,482 Bytes
625c0a2 | 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 | """
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) |