test_zeroshot.py

"""
Test script to verify zero-shot models are working properly.
"""
import cv2
import numpy as np
import sys
import os

# Add parent directory to path
sys.path.insert(0, os.path.dirname(__file__))

from app import run_clip_anomaly_inference, run_owlvit_inference

def create_test_image_with_defect():
    """Create a simple test image with a visible defect."""
    # Create white background
    img = np.ones((640, 640, 3), dtype=np.uint8) * 255

    # Draw a normal grid pattern
    for i in range(0, 640, 80):
        cv2.line(img, (i, 0), (i, 640), (200, 200, 200), 2)
        cv2.line(img, (0, i), (640, i), (200, 200, 200), 2)

    # Draw a "defect" - irregular shapes
    cv2.circle(img, (320, 320), 50, (0, 0, 0), -1)  # Black circle (defect)
    cv2.rectangle(img, (100, 100), (150, 180), (50, 50, 50), -1)  # Dark rectangle (scratch)

    # Save the test image
    cv2.imwrite("test_defect_image.jpg", img)

    # Convert to bytes
    _, img_encoded = cv2.imencode('.jpg', img)
    return img_encoded.tobytes()

def create_normal_test_image():
    """Create a simple test image without defects."""
    # Create white background
    img = np.ones((640, 640, 3), dtype=np.uint8) * 255

    # Draw a normal grid pattern only
    for i in range(0, 640, 80):
        cv2.line(img, (i, 0), (i, 640), (200, 200, 200), 2)
        cv2.line(img, (0, i), (640, i), (200, 200, 200), 2)

    # Save the test image
    cv2.imwrite("test_normal_image.jpg", img)

    # Convert to bytes
    _, img_encoded = cv2.imencode('.jpg', img)
    return img_encoded.tobytes()

def test_clip():
    """Test CLIP anomaly detection."""
    print("\n" + "="*60)
    print("Testing CLIP Anomaly Detection")
    print("="*60)

    # Test with defect image
    print("\n1. Testing with DEFECT image (should detect anomaly)...")
    defect_image = create_test_image_with_defect()
    detections, score = run_clip_anomaly_inference(defect_image, confidence=0.3)
    print(f"   Anomaly Score: {score:.4f}")
    print(f"   Detections: {len(detections)}")
    if detections:
        for i, det in enumerate(detections):
            print(f"   Detection {i+1}: {det}")
    else:
        print("   ⚠️ NO DETECTIONS (this is the problem!)")

    # Test with normal image
    print("\n2. Testing with NORMAL image (should NOT detect anomaly)...")
    normal_image = create_normal_test_image()
    detections, score = run_clip_anomaly_inference(normal_image, confidence=0.3)
    print(f"   Anomaly Score: {score:.4f}")
    print(f"   Detections: {len(detections)}")
    if detections:
        print("   ⚠️ False positive detected!")
    else:
        print("   ✓ Correctly identified as normal")

def test_owlvit():
    """Test OWL-ViT object detection."""
    print("\n" + "="*60)
    print("Testing OWL-ViT Object Detection")
    print("="*60)

    # Test with defect image
    print("\n1. Testing with DEFECT image...")
    defect_image = create_test_image_with_defect()
    detections = run_owlvit_inference(defect_image, confidence=0.05)
    print(f"   Detections: {len(detections)}")
    if detections:
        for i, det in enumerate(detections):
            print(f"   Detection {i+1}: bbox={det['bbox']}, conf={det['confidence']:.4f}, class={det['class_name']}")
    else:
        print("   ⚠️ NO DETECTIONS (this is the problem!)")

if __name__ == "__main__":
    print("Testing Zero-Shot Models")
    print("This will create test images and run inference")

    try:
        test_clip()
        test_owlvit()

        print("\n" + "="*60)
        print("Test Complete!")
        print("="*60)
        print("\nTest images saved:")
        print("  - test_defect_image.jpg (has defects)")
        print("  - test_normal_image.jpg (normal)")

    except Exception as e:
        print(f"\n❌ ERROR: {e}")
        import traceback
        traceback.print_exc()