-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_pretrained_detector.py
More file actions
265 lines (206 loc) Β· 8.84 KB
/
test_pretrained_detector.py
File metadata and controls
265 lines (206 loc) Β· 8.84 KB
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
#!/usr/bin/env python3
"""
Test script for the Pretrained Damage Detector
This script demonstrates how to use the pretrained damage detection system
with sample images or test cases.
Author: AI Assistant
Date: June 2025
"""
import os
import sys
import json
from pathlib import Path
import logging
# Add the project root to the path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from pretrained_damage_detector import PretrainedDamageDetector
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
def test_single_image():
"""Test detection on a single image."""
print("\n" + "="*60)
print("TESTING SINGLE IMAGE DETECTION")
print("="*60)
# Initialize detector
detector = PretrainedDamageDetector(device='auto')
# Create a sample test image path (you'll need to provide an actual image)
test_image = "test_images/sample_poster.jpg"
if not os.path.exists(test_image):
print(f"β οΈ Test image not found: {test_image}")
print("Please create a 'test_images' directory and add some poster images for testing.")
return
# Analyze the image
result = detector.detect_damage(test_image)
# Display results
print(f"\nπ ANALYSIS RESULTS for {test_image}")
print("-" * 40)
print(f"Damaged: {'YES' if result['is_damaged'] else 'NO'}")
print(f"Confidence: {result['confidence']:.3f}")
print(f"Primary Damage Type: {result.get('primary_damage_type', 'none')}")
print(f"\nπ DAMAGE SCORES:")
for damage_type, score in result['damage_scores'].items():
if score > 0.1: # Only show significant scores
print(f" {damage_type}: {score:.3f}")
print(f"\nπ BLIP ANALYSIS:")
print(f"Caption: {result['detailed_analysis']['blip_caption']}")
print(f"\nβ Q&A RESPONSES:")
for question, answer in result['detailed_analysis']['qa_responses'].items():
print(f" Q: {question}")
print(f" A: {answer}")
def test_batch_processing():
"""Test batch processing of multiple images."""
print("\n" + "="*60)
print("TESTING BATCH PROCESSING")
print("="*60)
# Initialize detector
detector = PretrainedDamageDetector(device='auto')
# Set up test directory
test_dir = "test_images"
if not os.path.exists(test_dir):
print(f"β οΈ Test directory not found: {test_dir}")
print("Please create a 'test_images' directory and add some poster images for testing.")
return
# Get all images in test directory
image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
image_paths = [
os.path.join(test_dir, f)
for f in os.listdir(test_dir)
if Path(f).suffix.lower() in image_extensions
]
if not image_paths:
print(f"β οΈ No images found in {test_dir}")
return
print(f"Found {len(image_paths)} images to process")
# Process batch
results = detector.batch_analyze(image_paths, "batch_results.json")
# Generate report
report = detector.generate_report(results)
# Display summary
print(f"\nπ BATCH ANALYSIS SUMMARY")
print("-" * 40)
print(f"Total Images Processed: {report['summary']['total_images']}")
print(f"Damaged Images Found: {report['summary']['damaged_images']}")
print(f"Damage Rate: {report['summary']['damage_rate']:.2%}")
print(f"Average Confidence: {report['summary']['average_confidence']:.3f}")
print(f"\nπ·οΈ DAMAGE TYPE BREAKDOWN:")
for damage_type, count in report['damage_breakdown'].items():
if count > 0:
print(f" {damage_type}: {count} images")
print(f"\nπΎ Detailed results saved to: batch_results.json")
def create_sample_test_structure():
"""Create sample directory structure for testing."""
print("\n" + "="*60)
print("CREATING SAMPLE TEST STRUCTURE")
print("="*60)
# Create test directories
test_dir = Path("test_images")
test_dir.mkdir(exist_ok=True)
# Create subdirectories for different damage types
damage_dirs = [
"undamaged", "torn", "stained", "faded",
"creased", "water_damaged", "burned"
]
for damage_dir in damage_dirs:
(test_dir / damage_dir).mkdir(exist_ok=True)
# Create a README file with instructions
readme_content = """# Test Images Directory
This directory structure is for testing the Pretrained Damage Detector.
## Directory Structure:
- `undamaged/` - Place clean, undamaged poster images here
- `torn/` - Place images of posters with tears here
- `stained/` - Place images of stained posters here
- `faded/` - Place images of faded posters here
- `creased/` - Place images of creased/wrinkled posters here
- `water_damaged/` - Place images of water-damaged posters here
- `burned/` - Place images of burned posters here
## Usage:
1. Add poster images to the appropriate subdirectories
2. Run the test script: `python test_pretrained_detector.py`
3. Or test individual features using the main detector
## Supported Formats:
- JPEG (.jpg, .jpeg)
- PNG (.png)
- BMP (.bmp)
- TIFF (.tiff)
- WebP (.webp)
## Note:
The detector can analyze any poster image and will automatically
determine the damage type and severity.
"""
with open(test_dir / "README.md", "w") as f:
f.write(readme_content)
print(f"β
Created test directory structure in: {test_dir}")
print("π Please add poster images to the subdirectories for testing")
print(f"π See {test_dir}/README.md for detailed instructions")
def demonstrate_features():
"""Demonstrate key features of the detector without requiring images."""
print("\n" + "="*60)
print("PRETRAINED DAMAGE DETECTOR FEATURES")
print("="*60)
print("π€ LOADED MODELS:")
print(" β’ BLIP2 - Visual understanding and captioning")
print(" β’ ViT (Vision Transformer) - Defect classification")
print(" β’ DinoV2 - Advanced anomaly detection")
print(" β’ DETR - Object detection for structural analysis")
print("\nπ DETECTION CAPABILITIES:")
damage_types = [
'tear', 'crease', 'stain', 'fade', 'burn',
'water_damage', 'scratches', 'holes', 'wrinkles'
]
for i, damage_type in enumerate(damage_types, 1):
print(f" {i:2d}. {damage_type.replace('_', ' ').title()}")
print("\nπ ANALYSIS METHODS:")
print(" β’ Visual Question Answering (VQA) using BLIP2")
print(" β’ Anomaly detection with pretrained classifiers")
print(" β’ Structural integrity analysis")
print(" β’ Visual feature extraction (edges, texture, color)")
print(" β’ Rules-based damage classification")
print("\nπ‘ KEY ADVANTAGES:")
print(" β
No training required - uses pretrained models")
print(" β
Multi-modal analysis (visual + text understanding)")
print(" β
Detailed damage type classification")
print(" β
Confidence scoring for each prediction")
print(" β
Batch processing capabilities")
print(" β
Comprehensive reporting features")
print("\nπ USAGE EXAMPLES:")
print(" # Single image analysis:")
print(" python pretrained_damage_detector.py image.jpg")
print()
print(" # Batch processing with report:")
print(" python pretrained_damage_detector.py test_images/ --batch --report")
print()
print(" # Save results to JSON:")
print(" python pretrained_damage_detector.py image.jpg --output results.json")
def main():
"""Main test function."""
print("π― PRETRAINED DAMAGE DETECTOR - TEST SUITE")
# Check if test images exist
test_dir = Path("test_images")
has_test_images = test_dir.exists() and any(test_dir.iterdir())
if not has_test_images:
print("\nβ οΈ No test images found. Setting up test structure...")
create_sample_test_structure()
print("\nπ Please add some poster images to the test directories and run again.")
demonstrate_features()
return
try:
# Demonstrate features first
demonstrate_features()
# Test single image detection
test_single_image()
# Test batch processing
test_batch_processing()
print("\n" + "="*60)
print("β
ALL TESTS COMPLETED SUCCESSFULLY!")
print("="*60)
except Exception as e:
logger.error(f"Test failed with error: {e}")
print(f"\nβ Test failed: {e}")
print("\nPossible solutions:")
print("1. Install missing dependencies: pip install -r requirements.txt")
print("2. Ensure you have internet connection for model downloads")
print("3. Check that you have sufficient disk space (~2GB for models)")
print("4. Try running on CPU if GPU issues: --device cpu")
if __name__ == "__main__":
main()