File size: 29,086 Bytes
84b458c
e6535be
84b458c
 
 
 
51aa863
e6535be
84b458c
 
 
 
51aa863
84b458c
 
 
23f6757
e6535be
 
 
84b458c
 
 
e6535be
 
84b458c
 
 
 
 
 
 
 
 
 
 
e6535be
51aa863
e6535be
 
84b458c
23f6757
51aa863
23f6757
 
68b9325
 
 
 
4181aa6
5cc72f5
68b9325
 
 
 
46a7dd6
4181aa6
51aa863
23f6757
 
26f20f9
23f6757
 
68b9325
 
 
 
 
 
 
fac11ae
68b9325
fac11ae
26f20f9
 
23f6757
 
68b9325
 
 
 
fac11ae
 
68b9325
23f6757
 
 
 
5cc72f5
23f6757
5cc72f5
68b9325
23f6757
 
5cc72f5
 
 
23f6757
 
5cc72f5
68b9325
23f6757
 
 
 
 
5cc72f5
23f6757
 
5cc72f5
23f6757
 
 
 
5cc72f5
23f6757
 
5cc72f5
23f6757
 
 
 
e6535be
 
23f6757
26f20f9
 
23f6757
 
26f20f9
23f6757
26f20f9
23f6757
26f20f9
51aa863
 
23f6757
 
51aa863
23f6757
 
26f20f9
51aa863
26f20f9
e6535be
 
 
 
 
 
 
 
 
 
 
 
 
 
26f20f9
 
 
 
5cc72f5
23f6757
 
e6535be
 
23f6757
 
 
 
 
68b9325
5cc72f5
 
 
e6535be
23f6757
51aa863
 
23f6757
5cc72f5
e6535be
23f6757
26f20f9
 
23f6757
84b458c
ceb9e66
 
 
 
84b458c
 
 
 
ceb9e66
84b458c
 
5cc72f5
84b458c
e6535be
 
5cc72f5
 
e6535be
 
5cc72f5
 
e6535be
 
 
 
5cc72f5
 
23f6757
5cc72f5
 
 
 
 
 
 
 
 
23f6757
5cc72f5
84b458c
 
23f6757
e6535be
84b458c
23f6757
5cc72f5
 
84b458c
 
 
 
 
 
5cc72f5
 
 
 
e6535be
 
 
5cc72f5
 
 
 
 
 
 
 
 
 
23f6757
 
e6535be
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
84b458c
 
51aa863
 
 
23f6757
 
51aa863
e6535be
84b458c
51aa863
5045d14
23f6757
e6535be
51aa863
e6535be
23f6757
51aa863
 
 
 
 
 
 
 
 
 
 
 
 
 
23f6757
51aa863
 
 
 
5cc72f5
 
 
51aa863
 
5cc72f5
 
51aa863
 
 
 
 
 
 
e6535be
51aa863
84b458c
e6535be
 
51aa863
e6535be
e83afa6
e6535be
 
51aa863
 
 
e6535be
51aa863
 
ceb9e66
e6535be
 
ceb9e66
51aa863
e6535be
 
 
 
 
e83afa6
 
e6535be
 
 
 
 
51aa863
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e6535be
51aa863
5cc72f5
 
51aa863
 
 
e6535be
5cc72f5
 
e6535be
51aa863
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e6535be
51aa863
 
 
 
84b458c
 
 
 
ceb9e66
e6535be
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
84b458c
 
ceb9e66
51aa863
 
 
 
 
 
ceb9e66
51aa863
 
 
 
 
 
 
 
23f6757
84b458c
 
e6535be
 
84b458c
 
51aa863
5045d14
23f6757
e6535be
51aa863
e6535be
23f6757
51aa863
84b458c
 
 
e6535be
 
 
51aa863
84b458c
 
 
 
 
 
 
 
e6535be
 
 
 
 
 
 
e83afa6
e6535be
 
 
 
 
 
 
 
 
 
 
 
 
 
 
84b458c
e6535be
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
51aa863
 
 
 
e6535be
51aa863
 
 
 
e6535be
 
 
51aa863
 
e6535be
51aa863
84b458c
 
 
51aa863
84b458c
51aa863
 
84b458c
23f6757
84b458c
51aa863
e6535be
84b458c
 
 
51aa863
 
ceb9e66
84b458c
 
 
 
e6535be
 
 
 
 
 
5cc72f5
 
 
84b458c
 
 
 
5cc72f5
 
51aa863
e6535be
 
84b458c
 
 
 
51aa863
 
 
 
 
e6535be
84b458c
 
51aa863
e6535be
51aa863
 
 
 
 
 
 
84b458c
ceb9e66
84b458c
51aa863
84b458c
51aa863
 
 
5cc72f5
e6535be
84b458c
 
 
e6535be
 
 
 
84b458c
 
 
 
 
 
 
 
11f621a
 
84b458c
e6535be
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
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
import io
from typing import List, Dict
import uvicorn
import numpy as np
import uuid
from datetime import datetime
from fastapi import FastAPI, UploadFile, File, HTTPException, Form
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from PIL import Image
import cv2
import yaml
from src.detection import YOLOv11Detector
from src.comparison import DamageComparator
from src.visualization import DamageVisualizer
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
import torch
import gc

app = FastAPI(
    title="Car Damage Detection API",
    description="YOLOv11-based car damage detection with DINOv2 ReID (Memory Optimized)",
    version="1.3.0"
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# GLOBAL COMPONENTS - Load once at startup
detector = None
comparator = None
visualizer = None

# Model paths mapping - PT and ONNX versions
MODEL_PATHS = {
    # PT models (original)
    0: "models_small/best.pt",  # Small v1 PT
    1: "models_small_version_2/best.pt",  # Small v2 PT,
    2: "models_small_version3/best.pt",  # Small v3 PT
    3: "models_medium/best.pt",  # Medium v1 PT
    4: "models_medium_version_2/best.pt",  # Medium v2 PT
    5: "model_medium_version3/best.pt",  # Medium v3 PT
    # ONNX models (optimized with v1.19 + opset 21)
    6: "models_small/best.onnx",  # Small v1 ONNX
    7: "models_small_version_2/best.onnx",  # Small v2 ONNX,
    8: "models_small_version3/best.onnx",  # Small v3 ONNX
    9: "models_medium/best.onnx",  # Medium v1 ONNX
    10: "models_medium_version_2/best.onnx",  # Medium v2 ONNX,
    11: "model_medium_version3/best.onnx"  # Medium v3 ONNX,
}

# Config paths - ONNX uses same config as PT version
CONFIG_PATHS = {
    0: "config.yaml",  # Small v1 PT
    1: "config_version2.yaml",  # Small v2 PT
    2: "config_version3.yaml",  # Small v3 PT
    3: "config.yaml",  # Medium v1 PT
    4: "config_version2.yaml",  # Medium v2 PT
    5: "config_version3.yaml",  # Medium v3 PT
    6: "config.yaml",  # Small v1 ONNX
    7: "config_version2.yaml",  # Small v2 ONNX
    8: "config_version3.yaml",  # Small v3 ONNX
    9: "config.yaml",  # Medium v1 ONNX
    10: "config_version2.yaml",  # Medium v2 ONNX
    11: "config_version3.yaml"  # Medium v3 ONNX
}

# Mapping from PT index to ONNX index
PT_TO_ONNX_MAPPING = {
    0: 5,  # Small v1 -> ONNX
    1: 6,  # Small v2 -> ONNX
    2: 7,  # Medium v1 -> ONNX
    3: 8,  # Medium v2 -> ONNX
    4: 9,   # Medium v3 -> ONNX
    5: 10,   # Medium v3 -> ONNX
    6: 11   # Medium v3 -> ONNX
}

def get_optimal_model_index(select_models: int, prefer_onnx: bool = True) -> int:
    """
    Enhanced model selection with performance optimization info
    """
    # If user explicitly selects ONNX index (5..8) => use that ONNX with optimizations
    if select_models in (6, 7, 8, 9, 10, 11):
        onnx_path = Path(MODEL_PATHS.get(select_models, ""))
        if not onnx_path.exists():
            raise FileNotFoundError(
                f"Requested ONNX model index {select_models} not found at {MODEL_PATHS.get(select_models)}")
        print(f"πŸš€ Selected ONNX model with MAXIMUM optimizations: {MODEL_PATHS[select_models]}")
        return select_models

    # Normalize to valid PT indices
    if select_models not in (0, 1, 2, 3, 4, 5):
        select_models = 2  # default to medium v1

    # PT preferred for 0..4
    pt_path = Path(MODEL_PATHS.get(select_models, ""))
    if pt_path.exists():
        print(f"πŸ“¦ Selected PyTorch model: {MODEL_PATHS[select_models]}")
        return select_models

    # If PT not found and prefer_onnx: fallback to ONNX with optimizations
    onnx_index = PT_TO_ONNX_MAPPING.get(select_models)
    if prefer_onnx and onnx_index is not None:
        onnx_path = Path(MODEL_PATHS.get(onnx_index, ""))
        if onnx_path.exists():
            print(f"PT not found at {pt_path}, falling back to optimized ONNX {MODEL_PATHS[onnx_index]}")
            return onnx_index

    # No suitable file found
    raise FileNotFoundError(f"Requested PT model index {select_models} not found at {MODEL_PATHS.get(select_models)}")

def load_detector(select_models: int = 2, prefer_onnx: bool = True):
    """
    Load detector with optimized ONNX Runtime v1.19 support
    IMPORTANT: This loads GLOBAL instances that are shared across threads
    """
    global detector, comparator, visualizer

    actual_model_index = get_optimal_model_index(select_models, prefer_onnx)

    # Get appropriate config file
    config_file = CONFIG_PATHS.get(actual_model_index, "config.yaml")

    # Load config
    with open(config_file, 'r') as f:
        config = yaml.safe_load(f)

    # Update config with selected model path
    config['model']['path'] = MODEL_PATHS[actual_model_index]

    # Save updated config to temp file
    temp_config = f'temp_config_{actual_model_index}.yaml'
    with open(temp_config, 'w') as f:
        yaml.dump(config, f, default_flow_style=False)

    # Clear previous models from memory before loading new ones
    if detector is not None:
        del detector
    if comparator is not None:
        del comparator
    if visualizer is not None:
        del visualizer

    # Force garbage collection
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # Load all components with new config
    detector = YOLOv11Detector(config_path=temp_config)
    comparator = DamageComparator(config_path=temp_config)
    visualizer = DamageVisualizer(config_path=temp_config)

    # Log model info with optimization status
    model_type = "ONNX" if MODEL_PATHS[actual_model_index].endswith('.onnx') else "PyTorch"
    model_labels = [
        "Small v1", "Small v2", "Small v3", "Medium v1", "Medium v2", "Medium v3",
        "Small v1 ONNX", "Small v2 ONNX", "Small v3 ONNX", "Medium v1 ONNX", "Medium v2 ONNX", "Medium v3 ONNX"
    ]

    if 0 <= select_models < len(model_labels):
        model_size = model_labels[select_models]
    else:
        raise ValueError(f"select_models={select_models} must be 0-11")

    optimization_status = "πŸš€ MAXIMUM OPTIMIZATIONS" if model_type == "ONNX" else "πŸ“¦ Standard PyTorch"
    print(f"Loaded {model_size} model in {model_type} format - {optimization_status}")
    print(f"βœ… DINOv2 ReID enabled for damage comparison")

    return detector


# Initialize default detector with medium model (preferring ONNX for performance)
print("πŸš€ Initializing API with optimized ONNX Runtime and DINOv2 ReID support...")
detector = load_detector(2, prefer_onnx=True)
comparator = DamageComparator(config_path=CONFIG_PATHS[2])
visualizer = DamageVisualizer(config_path=CONFIG_PATHS[2])

# Create necessary directories
UPLOADS_DIR = Path("uploads")
RESULTS_DIR = Path("results")
UPLOADS_DIR.mkdir(exist_ok=True)
RESULTS_DIR.mkdir(exist_ok=True)

# Mount static files directory
app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")


@app.get("/")
async def root():
    """Root endpoint with enhanced model info"""
    return {
        "message": "Car Damage Detection API with YOLOv11 + DINOv2 ReID (Memory Optimized)",
        "version": "1.3.0",
        "optimizations": {
            "onnx_runtime": "v1.19+ with opset 21 support",
            "reid_model": "DINOv2 (Meta) - Superior visual feature extraction",
            "memory_management": "Global model loading with ThreadPoolExecutor",
            "performance_features": [
                "Graph optimizations (ALL level)",
                "DINOv2 ReID for cross-view damage matching",
                "Memory-efficient threading",
                "torch.no_grad() context for inference",
                "Automatic CUDA cache clearing"
            ]
        },
        "model_options": {
            "0": "Small model v1 (PyTorch)",
            "1": "Small model v2 (PyTorch)",
            "2": "Medium model v1 (PyTorch)",
            "3": "Medium model v2 (PyTorch)",
            "4": "Large model (PyTorch only)",
            "5": "Small model v1 (ONNX - OPTIMIZED)",
            "6": "Small model v2 (ONNX - OPTIMIZED)",
            "7": "Medium model v1 (ONNX - OPTIMIZED)",
            "8": "Medium model v2 (ONNX - OPTIMIZED)"
        },
        "recommendation": "Use indices 5-8 for maximum performance with ONNX optimizations",
        "endpoints": {
            "/docs": "API documentation",
            "/detect": "Single/Multi image detection",
            "/compare": "Compare before/after images (6 pairs) with DINOv2 ReID",
            "/uploads/{filename}": "Access saved visualization images",
            "/health": "Health check",
            "/model-info": "Get current model information",
            "/performance-info": "Get optimization details"
        }
    }


@app.get("/health")
async def health_check():
    """Enhanced health check with optimization status"""
    health_info = {
        "status": "healthy",
        "model": "YOLOv11",
        "reid": "DINOv2",
        "backend": "ONNX/PyTorch",
        "memory_optimization": "ThreadPoolExecutor with global models"
    }

    if detector and hasattr(detector, 'get_performance_info'):
        perf_info = detector.get_performance_info()
        health_info.update({
            "model_type": perf_info.get("model_type", "Unknown"),
            "optimization_status": "Optimized" if perf_info.get("model_type") == "ONNX" else "Standard"
        })

    return health_info


@app.get("/model-info")
async def get_model_info():
    """Get comprehensive information about currently loaded model"""
    if detector is None:
        return {"error": "No model loaded"}

    model_path = detector.model_path
    model_type = "ONNX" if model_path.endswith('.onnx') else "PyTorch"

    info = {
        "model_path": model_path,
        "model_type": model_type,
        "confidence_threshold": detector.confidence,
        "iou_threshold": detector.iou_threshold,
        "classes": detector.classes,
        "reid_model": "DINOv2",
        "optimization_status": "Optimized" if model_type == "ONNX" else "Standard"
    }

    if hasattr(detector, 'get_performance_info'):
        perf_info = detector.get_performance_info()
        info.update(perf_info)

    return info


@app.post("/detect")
async def detect_single_image(
        file: UploadFile = File(None),
        files: List[UploadFile] = File(None),
        select_models: int = Form(2),
        prefer_onnx: bool = Form(True)
):
    """Multi-view detection with ONNX Runtime optimizations and DINOv2 ReID"""
    try:
        # Validate select_models
        if select_models not in list(range(0, 12)):
            raise HTTPException(status_code=400,
                                detail="select_models must be 0-8 (0-4=PyTorch, 5-8=ONNX optimized)")

        # Load appropriate detector (if different from current)
        current_detector = load_detector(select_models, prefer_onnx)

        # Case 1: Single image (backward compatibility)
        if file is not None:
            contents = await file.read()
            image = Image.open(io.BytesIO(contents)).convert("RGB")
            image_np = np.array(image)
            image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)

            # Perform detection
            detections = current_detector.detect(image_bgr)

            # Create visualization
            visualized = visualizer.draw_detections(image_bgr, detections, 'new_damage')

            # Save and return
            filename = f"detection_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}.jpg"
            output_path = UPLOADS_DIR / filename
            cv2.imwrite(str(output_path), visualized)

            model_type = "ONNX" if current_detector.model_path.endswith('.onnx') else "PyTorch"
            optimization_status = "πŸš€ OPTIMIZED" if model_type == "ONNX" else "πŸ“¦ Standard"

            return JSONResponse({
                "status": "success",
                "model_type": model_type,
                "optimization_status": optimization_status,
                "detections": detections,
                "statistics": {
                    "total_damages": len(detections['boxes']),
                    "damage_types": list(set(detections['classes']))
                },
                "visualized_image_path": f"uploads/{filename}",
                "visualized_image_url": f"http://localhost:8000/uploads/{filename}",
                "filename": filename
            })

        # Case 2: Multiple images with DINOv2 deduplication
        elif files is not None:
            detections_list = []
            images = []

            for idx, f in enumerate(files):
                contents = await f.read()
                image = Image.open(io.BytesIO(contents)).convert("RGB")
                image_np = np.array(image)
                image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
                images.append(image_bgr)
                detections = current_detector.detect(image_bgr)
                detections_list.append(detections)

            # Deduplicate across views using DINOv2
            unique_damages = comparator.deduplicate_detections_across_views(detections_list, images)

            # Create combined visualization
            heights = [img.shape[0] for img in images]
            widths = [img.shape[1] for img in images]
            max_height = max(heights)
            total_width = sum(widths)
            combined_img = np.zeros((max_height, total_width, 3), dtype=np.uint8)
            x_offset = 0

            for img_idx, image in enumerate(images):
                h, w = image.shape[:2]
                if h != max_height:
                    image = cv2.resize(image, (w, max_height))
                detections = detections_list[img_idx]
                combined_img[:, x_offset:x_offset + w] = image

                # Draw detections with unique IDs
                for det_idx, bbox in enumerate(detections['boxes']):
                    # Find unique damage ID for this detection
                    damage_id = None
                    for uid, damage_info in unique_damages.items():
                        for d in damage_info['detections']:
                            if d['view_idx'] == img_idx and d['bbox'] == bbox:
                                damage_id = uid
                                break

                    # Draw with unique ID
                    x1, y1, x2, y2 = bbox
                    x1 += x_offset
                    x2 += x_offset

                    # Color based on unique ID
                    if damage_id:
                        color_hash = int(damage_id[-6:], 16)
                        color = ((color_hash >> 16) & 255, (color_hash >> 8) & 255, color_hash & 255)
                    else:
                        color = (0, 0, 255)

                    cv2.rectangle(combined_img, (x1, y1), (x2, y2), color, 2)

                    # Label
                    label = f"{damage_id[:8] if damage_id else 'Unknown'}"
                    cv2.putText(combined_img, label, (x1, y1 - 5),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

                x_offset += w

            # Save combined visualization
            filename = f"multiview_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}.jpg"
            output_path = UPLOADS_DIR / filename
            cv2.imwrite(str(output_path), combined_img)

            # Return results
            total_detections = sum(len(d['boxes']) for d in detections_list)
            model_type = "ONNX" if current_detector.model_path.endswith('.onnx') else "PyTorch"
            optimization_status = "πŸš€ OPTIMIZED" if model_type == "ONNX" else "πŸ“¦ Standard"

            return JSONResponse({
                "status": "success",
                "mode": "multi_view_with_dinov2_reid",
                "model_type": model_type,
                "optimization_status": optimization_status,
                "reid_model": "DINOv2",
                "total_detections_all_views": total_detections,
                "unique_damages_count": len(unique_damages),
                "unique_damages": {
                    damage_id: {
                        "appears_in_views": info['views'],
                        "class": info['class'],
                        "avg_confidence": float(info['avg_confidence']),
                        "detection_count": len(info['detections'])
                    }
                    for damage_id, info in unique_damages.items()
                },
                "reduction_rate": f"{(1 - len(unique_damages) / total_detections) * 100:.1f}%" if total_detections > 0 else "0%",
                "visualized_image_path": f"uploads/{filename}",
                "visualized_image_url": f"http://localhost:8000/uploads/{filename}",
                "message": f"Detected {total_detections} damages across {len(files)} views, "
                           f"identified {len(unique_damages)} unique damages using DINOv2 ReID"
            })

        else:
            raise HTTPException(status_code=400, detail="No image provided")

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Detection failed: {str(e)}")


def process_single_position_threaded(
        i: int,
        before_contents: bytes,
        after_contents: bytes,
        timestamp_str: str,
        session_id: str
) -> Dict:
    """
    Process single position comparison using GLOBAL models (thread-safe)
    No model loading here - uses global instances
    """
    # Use GLOBAL instances - no loading
    global detector, comparator, visualizer

    # Preprocess images
    before_img = Image.open(io.BytesIO(before_contents)).convert("RGB")
    after_img = Image.open(io.BytesIO(after_contents)).convert("RGB")
    before_np = np.array(before_img)
    after_np = np.array(after_img)
    before_bgr = cv2.cvtColor(before_np, cv2.COLOR_RGB2BGR)
    after_bgr = cv2.cvtColor(after_np, cv2.COLOR_RGB2BGR)

    # Detect using global detector
    before_detections = detector.detect(before_bgr)
    after_detections = detector.detect(after_bgr)

    # Compare using global comparator with DINOv2 ReID
    comparison = comparator.analyze_damage_status(
        before_detections, after_detections,
        before_bgr, after_bgr
    )

    # Visualize using global visualizer
    vis_img = visualizer.create_comparison_visualization(
        before_bgr, after_bgr,
        before_detections, after_detections,
        comparison
    )

    vis_filename = f"comparison_{timestamp_str}_{session_id}_pos{i + 1}.jpg"
    vis_path = UPLOADS_DIR / vis_filename
    cv2.imwrite(str(vis_path), vis_img)
    vis_url = f"http://localhost:8000/uploads/{vis_filename}"

    # Clear any GPU memory if used
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # Return result
    return {
        f"position_{i + 1}": {
            "case": comparison['case'],
            "message": comparison['message'],
            "statistics": comparison['statistics'],
            "new_damages": comparison['new_damages'],
            "matched_damages": comparison['matched_damages'],
            "repaired_damages": comparison['repaired_damages'],
            "using_reid": comparison['statistics'].get('using_reid', True),
            "reid_model": comparison['statistics'].get('reid_model', 'DINOv2'),
            "visualization_path": f"uploads/{vis_filename}",
            "visualization_url": vis_url,
            "filename": vis_filename
        },
        "before_detections": {
            "boxes": [np.array(box).tolist() for box in before_detections['boxes']],
            "confidences": [float(c) for c in before_detections['confidences']],
            "classes": before_detections['classes']
        },
        "after_detections": {
            "boxes": [np.array(box).tolist() for box in after_detections['boxes']],
            "confidences": [float(c) for c in after_detections['confidences']],
            "classes": after_detections['classes']
        },
        "_before_bgr": before_bgr,  # chỉ dΓΉng nα»™i bα»™
        "_after_bgr": after_bgr  # chỉ dΓΉng nα»™i bα»™
    }


@app.post("/compare")
async def compare_vehicle_damages(
        # Before delivery images (6 positions)
        before_1: UploadFile = File(...),
        before_2: UploadFile = File(...),
        before_3: UploadFile = File(...),
        before_4: UploadFile = File(...),
        before_5: UploadFile = File(...),
        before_6: UploadFile = File(...),
        # After delivery images (6 positions)
        after_1: UploadFile = File(...),
        after_2: UploadFile = File(...),
        after_3: UploadFile = File(...),
        after_4: UploadFile = File(...),
        after_5: UploadFile = File(...),
        after_6: UploadFile = File(...),
        # Model selection
        select_models: int = Form(2),
        prefer_onnx: bool = Form(True)
):
    """
    Enhanced comparison with DINOv2 ReID and Memory Optimization
    Uses ThreadPoolExecutor with global models to avoid OOM
    """
    try:
        # Validate select_models
        if select_models not in list(range(0, 12)):
            raise HTTPException(status_code=400,
                                detail="select_models must be 0-10 (0-5=PyTorch, 6-11=ONNX optimized)")

        # Load appropriate detector if different from current
        current_detector = load_detector(select_models, prefer_onnx)

        before_images = [before_1, before_2, before_3, before_4, before_5, before_6]
        after_images = [after_1, after_2, after_3, after_4, after_5, after_6]

        # Read contents first
        before_contents_list = [await img.read() for img in before_images]
        after_contents_list = [await img.read() for img in after_images]

        # Overall statistics
        total_new_damages = 0
        total_existing_damages = 0
        total_matched_damages = 0

        session_id = str(uuid.uuid4())[:8]
        timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")

        position_results = []
        all_visualizations = []
        image_pairs = []
        all_before_images = []
        all_after_images = []
        all_before_detections = []
        all_after_detections = []

        # Use ThreadPoolExecutor to share memory (avoid OOM)
        print(f"πŸ”„ Processing {len(before_images)} image pairs using ThreadPoolExecutor...")

        with ThreadPoolExecutor(max_workers=3) as executor:  # Limit workers to avoid memory issues
            futures = [
                executor.submit(
                    process_single_position_threaded,
                    i,
                    before_contents_list[i],
                    after_contents_list[i],
                    timestamp_str,
                    session_id
                )
                for i in range(6)
            ]

            for future in as_completed(futures):
                result = future.result()
                pos_key = list(result.keys())[0]  # e.g., 'position_1'
                position_results.append(result)
                all_visualizations.append(result[pos_key]["visualization_url"])

                # Collect for deduplication
                image_pairs.append((result["_before_bgr"], result["_after_bgr"]))
                all_before_images.append(result["_before_bgr"])
                all_after_images.append(result["_after_bgr"])
                result.pop("_before_bgr", None)
                result.pop("_after_bgr", None)

                all_before_detections.append(result["before_detections"])
                all_after_detections.append(result["after_detections"])

                # Update statistics
                comparison = result[pos_key]
                total_new_damages += len(comparison["new_damages"])
                total_existing_damages += len(comparison["repaired_damages"])
                total_matched_damages += len(comparison["matched_damages"])

        # Sort position_results by position number
        position_results.sort(key=lambda x: int(list(x.keys())[0].split('_')[1]))

        # Deduplicate BEFORE damages across all 6 views using DINOv2
        print("πŸ” Deduplicating damages across views using DINOv2...")
        unique_before = comparator.deduplicate_detections_across_views(
            all_before_detections, all_before_images
        )

        # Deduplicate AFTER damages across all 6 views using DINOv2
        unique_after = comparator.deduplicate_detections_across_views(
            all_after_detections, all_after_images
        )

        print(
            f"βœ… Before: {sum(len(d['boxes']) for d in all_before_detections)} detections β†’ {len(unique_before)} unique")
        print(f"βœ… After: {sum(len(d['boxes']) for d in all_after_detections)} detections β†’ {len(unique_after)} unique")

        # Determine overall case with deduplication
        actual_new_damages = max(0, len(unique_after) - len(unique_before))

        overall_case = "CASE_3_SUCCESS"
        overall_message = "Successful delivery - No damage detected"

        if actual_new_damages > 0:
            overall_case = "CASE_2_NEW_DAMAGE"
            overall_message = f"Error during delivery - {actual_new_damages} new unique damage(s) detected"
        elif len(unique_before) > 0 and actual_new_damages <= 0:
            overall_case = "CASE_1_EXISTING"
            overall_message = "Existing damages from beginning β†’ Delivery completed"

        # Create summary grid
        grid_results = [res[list(res.keys())[0]] for res in position_results]
        grid_img = visualizer.create_summary_grid(grid_results, image_pairs)

        grid_filename = f"summary_grid_{timestamp_str}_{session_id}.jpg"
        grid_path = UPLOADS_DIR / grid_filename
        cv2.imwrite(str(grid_path), grid_img)

        grid_url = f"http://localhost:8000/uploads/{grid_filename}"

        timestamp = datetime.now().isoformat()

        # Clean up memory
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        # Enhanced response
        model_type = "ONNX" if current_detector.model_path.endswith('.onnx') else "PyTorch"
        optimization_status = "πŸš€ OPTIMIZED" if model_type == "ONNX" else "πŸ“¦ Standard"

        return JSONResponse({
            "status": "success",
            "session_id": session_id,
            "timestamp": timestamp,
            "model_type": model_type,
            "optimization_status": optimization_status,
            "reid_enabled": True,
            "reid_model": "DINOv2",
            "memory_optimization": "ThreadPoolExecutor with global models",
            "overall_result": {
                "case": overall_case,
                "message": overall_message,
                "statistics": {
                    "total_new_damages": int(total_new_damages),
                    "total_matched_damages": int(total_matched_damages),
                    "total_repaired_damages": int(total_existing_damages),
                    "unique_damages_before": int(len(unique_before)),
                    "unique_damages_after": int(len(unique_after)),
                    "actual_new_unique_damages": int(actual_new_damages)
                }
            },
            "deduplication_info": {
                "model": "DINOv2",
                "before_total_detections": int(sum(len(d['boxes']) for d in all_before_detections)),
                "before_unique_damages": int(len(unique_before)),
                "after_total_detections": int(sum(len(d['boxes']) for d in all_after_detections)),
                "after_unique_damages": int(len(unique_after)),
                "duplicate_reduction_rate": f"{(1 - len(unique_after) / sum(len(d['boxes']) for d in all_after_detections)) * 100:.1f}%"
                if sum(len(d['boxes']) for d in all_after_detections) > 0 else "0%"
            },
            "position_results": position_results,
            "summary_visualization_path": f"uploads/{grid_filename}",
            "summary_visualization_url": grid_url,
            "all_visualizations": all_visualizations,
            "recommendations": {
                "action_required": bool(actual_new_damages > 0),
                "suggested_action": "Investigate delivery process" if actual_new_damages > 0
                else "Proceed with delivery completion"
            },
            "performance_note": f"Using {model_type} + DINOv2 ReID with memory optimization"
        })

    except Exception as e:
        # Clean up on error
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        raise HTTPException(status_code=500, detail=f"Comparison failed: {str(e)}")


if __name__ == "__main__":
    import os
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=int(os.environ.get("PORT", 7860)),  
        reload=False,  
        log_level="info"
    )