test_speed.py

"""Test optimized model speed"""
import sys, os, time
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from model import GarbageClassifier
from PIL import Image

# Create test image
img = Image.new("RGB", (512, 512), color=(200, 150, 80))
test_path = "test_speed.jpg"
img.save(test_path)

# Test inference speed
clf = GarbageClassifier()
t0 = time.time()
result = clf.classify(test_path)
t = time.time() - t0

print(f"\n=== Inference time: {t:.1f}s ===")
print(f"Item: {result['item_zh']}")
print(f"Category: {result['category']}")
print(f"Confidence: {result['confidence']:.2%}")
print(f"Has DNA: {result['dna'] is not None}")

# Cleanup
os.remove(test_path)