File size: 35,853 Bytes
5c783e4
e2914ad
 
 
a028174
 
 
 
 
 
 
817736f
 
a028174
 
5c783e4
 
514a0e2
5c783e4
817736f
 
 
 
 
 
 
a028174
46b21a8
514a0e2
 
 
 
 
46b21a8
514a0e2
 
623cecd
a028174
 
275f52e
a028174
 
 
 
 
 
79a26c5
b9ac4ff
 
 
79a26c5
a028174
79a26c5
 
 
 
 
 
 
 
a028174
79a26c5
 
 
 
 
a028174
 
 
 
 
 
 
 
 
 
 
 
79a26c5
a028174
79a26c5
 
 
 
 
 
 
5c783e4
a028174
514a0e2
 
 
 
 
 
 
 
5c783e4
514a0e2
5c783e4
 
514a0e2
5c783e4
 
46b21a8
a028174
514a0e2
 
 
79a26c5
a028174
04f2345
 
 
a028174
79a26c5
 
514a0e2
79a26c5
a028174
79a26c5
 
 
 
a028174
04f2345
 
a028174
 
 
514a0e2
17aa4c9
 
79a26c5
 
46b21a8
817736f
 
514a0e2
817736f
 
46b21a8
514a0e2
817736f
 
46b21a8
817736f
 
 
46b21a8
514a0e2
817736f
514a0e2
46b21a8
817736f
 
 
514a0e2
 
 
 
 
 
 
 
817736f
514a0e2
817736f
514a0e2
817736f
85f8ca4
817736f
 
514a0e2
a028174
514a0e2
 
 
 
 
 
46b21a8
817736f
 
514a0e2
a028174
514a0e2
 
817736f
514a0e2
 
a028174
514a0e2
 
 
 
817736f
514a0e2
 
 
 
 
 
 
 
46b21a8
a028174
514a0e2
 
 
 
 
 
 
 
46b21a8
514a0e2
 
46b21a8
a028174
514a0e2
 
 
 
 
 
 
 
 
 
 
 
 
46b21a8
a028174
514a0e2
 
 
 
 
 
 
46b21a8
514a0e2
 
 
79a26c5
a028174
 
514a0e2
 
 
 
 
 
 
 
 
 
817736f
 
514a0e2
 
 
 
 
 
a028174
514a0e2
817736f
46b21a8
514a0e2
 
46b21a8
 
 
 
817736f
 
 
 
46b21a8
817736f
 
 
 
 
46b21a8
514a0e2
 
 
 
 
 
 
 
a028174
514a0e2
a028174
514a0e2
85f8ca4
514a0e2
817736f
514a0e2
 
85f8ca4
a028174
514a0e2
 
 
817736f
 
 
 
514a0e2
817736f
 
 
 
 
 
514a0e2
 
 
 
 
 
 
 
a028174
514a0e2
 
 
a028174
514a0e2
 
 
 
 
 
 
a028174
514a0e2
 
 
 
a028174
514a0e2
 
 
 
a028174
514a0e2
 
 
a028174
514a0e2
 
a028174
 
514a0e2
a028174
514a0e2
 
 
 
 
 
 
 
 
 
 
a028174
514a0e2
 
 
 
 
 
a028174
514a0e2
 
 
a028174
514a0e2
a028174
5c783e4
817736f
b9ac4ff
a028174
b9ac4ff
 
 
514a0e2
b9ac4ff
514a0e2
a028174
79a26c5
a028174
 
79a26c5
a028174
817736f
79a26c5
a028174
79a26c5
 
 
 
 
 
 
 
 
 
b9ac4ff
79a26c5
 
b9ac4ff
79a26c5
b9ac4ff
514a0e2
 
a028174
817736f
514a0e2
 
 
 
 
a028174
514a0e2
817736f
a028174
817736f
 
 
 
 
 
 
 
 
 
 
 
 
8ecb82d
a028174
817736f
46b21a8
a028174
 
514a0e2
817736f
 
 
 
 
 
514a0e2
 
817736f
514a0e2
817736f
 
a028174
 
 
 
817736f
514a0e2
a028174
 
 
 
 
 
817736f
a028174
 
 
817736f
a028174
 
514a0e2
817736f
8ecb82d
 
 
 
 
817736f
a028174
 
5c783e4
a028174
 
 
 
b9ac4ff
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
514a0e2
a028174
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9c3f062
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c10d9c6
a028174
 
 
 
 
 
b9ac4ff
 
 
a028174
 
817736f
a028174
 
 
 
 
817736f
b9ac4ff
 
a028174
 
 
 
 
 
 
 
 
 
 
 
b9ac4ff
 
 
a028174
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9bd1724
 
a028174
79a26c5
a028174
 
 
 
79a26c5
a028174
79a26c5
a028174
79a26c5
a028174
79a26c5
a028174
 
 
 
 
 
9bd1724
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a028174
79a26c5
 
 
817736f
79a26c5
817736f
a028174
 
 
 
 
79a26c5
a028174
 
 
 
79a26c5
a028174
 
 
 
 
 
79a26c5
 
a028174
 
 
817736f
79a26c5
a028174
 
 
 
 
79a26c5
a028174
79a26c5
a028174
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9bd1724
 
a028174
 
83ba4a0
 
 
a028174
 
 
 
 
79a26c5
a028174
817736f
a028174
817736f
79a26c5
a028174
 
 
 
 
817736f
a028174
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
817736f
 
 
 
 
 
a028174
 
 
 
 
 
 
 
 
 
 
817736f
a028174
 
 
 
 
 
 
 
9bd1724
a028174
 
64df5a3
a028174
 
 
 
 
83ba4a0
79a26c5
83ba4a0
 
a028174
 
83ba4a0
 
 
64df5a3
 
83ba4a0
64df5a3
83ba4a0
 
a028174
 
64df5a3
a028174
64df5a3
a028174
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
817736f
a028174
 
 
64df5a3
 
 
 
 
 
 
79a26c5
623cecd
5c783e4
9bd1724
 
a028174
 
8acd934
9bd1724
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
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
#!/usr/bin/env python3
import importlib.util
import os
import sys
import time
import cv2
import torch
import numpy as np
import gradio as gr
from PIL import Image
from torchvision import transforms
from torchvision.models import vit_b_16
import torch.nn as nn
import traceback

# Add current directory to path
if not os.getcwd() in sys.path:
    sys.path.append(os.getcwd())

# Check if detectron2 is installed
if importlib.util.find_spec("detectron2") is None:
    print("Installing PyTorch and Detectron2...")
    os.system("pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu")
    os.system("pip install git+https://github.com/facebookresearch/detectron2.git")
    print("Installation complete!")

# Check for detectron2
try:
    from detectron2.engine import DefaultPredictor
    from detectron2.config import get_cfg
    from detectron2.utils.visualizer import Visualizer, ColorMode
    from detectron2 import model_zoo
    DETECTRON2_AVAILABLE = True
except ImportError:
    print("Warning: Detectron2 is not installed. Damage detection will not be available.")
    DETECTRON2_AVAILABLE = False

# Define model paths
DEFAULT_DAMAGE_MODEL_PATH = "./model_final.pth"
DEFAULT_DEEPFAKE_MODEL_PATH = "./vit_deepfake_best.pth"

# Sample images for demo (add your own paths)
SAMPLE_IMAGES = [
    "./test3.png",
    "./test5.png",
]

# Maximum number of tries allowed
MAX_TRIES = 3

def verify_detectron2_installation():
    """Verify that Detectron2 is properly installed"""
    results = {
        "detectron2_installed": False,
        "model_zoo_accessible": False,
        "can_create_cfg": False,
        "error_messages": []
    }
    
    try:
        import importlib.util
        if importlib.util.find_spec("detectron2") is not None:
            results["detectron2_installed"] = True
            
            try:
                import detectron2
                from detectron2 import model_zoo
                config_file = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
                config_path = model_zoo.get_config_file(config_file)
                if os.path.exists(config_path):
                    results["model_zoo_accessible"] = True
            except Exception as e:
                results["error_messages"].append(f"Error accessing model zoo: {str(e)}")
            
            try:
                from detectron2.config import get_cfg
                cfg = get_cfg()
                results["can_create_cfg"] = True
            except Exception as e:
                results["error_messages"].append(f"Error creating Detectron2 config: {str(e)}")
        else:
            results["error_messages"].append("Detectron2 is not installed")
    except Exception as e:
        results["error_messages"].append(f"Error checking Detectron2 installation: {str(e)}")
    
    return results

def setup_device(device_str):
    """Set up the computation device"""
    if device_str == 'auto':
        if torch.cuda.is_available():
            return torch.device('cuda:0')
        elif hasattr(torch, 'backends') and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            return torch.device('mps')
        else:
            return torch.device('cpu')
    elif device_str == 'cuda' and torch.cuda.is_available():
        return torch.device('cuda:0')
    elif device_str == 'mps' and hasattr(torch, 'backends') and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        return torch.device('mps')
    else:
        print(f"Warning: Device {device_str} not available, using CPU instead.")
        return torch.device('cpu')

def setup_damage_detector(model_path, threshold=0.7):
    """Set up the damage detection model"""
    if not DETECTRON2_AVAILABLE:
        print("Detectron2 is not installed. Cannot set up damage detector.")
        return None, None
    
    try:
        print(f"Checking model path: {model_path}")
        print(f"Model exists: {os.path.exists(model_path)}")
        
        if model_path is None or not os.path.exists(model_path):
            print(f"Error: Damage model file not found at {model_path}")
            return None, None
        
        cfg = get_cfg()
        cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
        cfg.MODEL.WEIGHTS = model_path
        cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # Only one class (damage)
        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = threshold
        
        # Use CPU if on Mac (MPS)
        cfg.MODEL.DEVICE = "cpu"
        print("Forcing Detectron2 to use CPU")        
        
        predictor = DefaultPredictor(cfg)
        return predictor, cfg
    except Exception as e:
        print(f"Detailed error: {str(e)}")
        import traceback
        traceback.print_exc()
        return None, None

def load_vit_deepfake_model(model_path, device):
    """Load the Vision Transformer (ViT) model for deepfake detection"""
    if model_path is None or not os.path.exists(model_path):
        print(f"Error: ViT deepfake model file not found at {model_path}")
        return None
        
    try:
        # Create ViT model with binary classification head
        model = vit_b_16(weights=None)
        
        # Modify the classifier head for binary classification (real vs fake)
        in_features = model.heads.head.in_features
        model.heads.head = nn.Linear(in_features, 2)
        
        # Load weights
        print(f"Loading ViT deepfake model from: {model_path}")
        checkpoint = torch.load(model_path, map_location='cpu')
        
        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'])
        elif isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
            model.load_state_dict(checkpoint['state_dict'])
        else:
            model.load_state_dict(checkpoint)
            
        # Move model to device and set to evaluation mode
        model = model.to(device)
        model.eval()
        
        return model
    except Exception as e:
        print(f"Error loading ViT deepfake model: {e}")
        traceback.print_exc()
        return None

def preprocess_for_vit(image, device):
    """Preprocess an image for Vision Transformer deepfake detection"""
    try:
        # Ensure image is RGB
        if len(image.shape) == 3 and image.shape[2] == 3:
            if image.dtype != np.uint8:
                image = (image * 255).astype(np.uint8)
            rgb_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        else:
            rgb_img = image
            
        # Resize to standard ViT input size (224x224)
        img_resized = cv2.resize(rgb_img, (224, 224))
        
        # Apply transforms
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
        img_tensor = transform(Image.fromarray(img_resized)).unsqueeze(0)
        img_tensor = img_tensor.to(device)
        
        return img_tensor
    except Exception as e:
        print(f"Error preprocessing image for ViT: {e}")
        traceback.print_exc()
        return None

def detect_damage(img, damage_detector):
    """Detect damage in an image"""
    try:
        if img is None:
            raise ValueError("Invalid image")
            
        # If no detector, use whole image
        if damage_detector is None:
            h, w = img.shape[:2]
            damage_regions = [{
                "box": (0, 0, w, h),
                "score": 1.0,
                "mask": None
            }]
            return img, None, damage_regions
            
        # Run inference
        outputs = damage_detector(img)
        
        # Get regions
        instances = outputs["instances"].to("cpu")
        boxes = instances.pred_boxes.tensor.numpy() if instances.has("pred_boxes") else []
        scores = instances.scores.numpy() if instances.has("scores") else []
        masks = instances.pred_masks.numpy() if instances.has("pred_masks") else []
        
        damage_regions = []
        for i in range(len(boxes)):
            x1, y1, x2, y2 = map(int, boxes[i])
            damage_regions.append({
                "box": (x1, y1, x2, y2),
                "score": float(scores[i]),
                "mask": masks[i] if len(masks) > i else None
            })
        
        # If no regions found, use whole image
        if not damage_regions:
            h, w = img.shape[:2]
            damage_regions = [{
                "box": (0, 0, w, h),
                "score": 1.0,
                "mask": None
            }]
            
        return img, outputs, damage_regions
    except Exception as e:
        print(f"Error detecting damage: {e}")
        traceback.print_exc()
        
        # Return whole image if error
        if 'img' in locals() and img is not None:
            h, w = img.shape[:2]
            damage_regions = [{
                "box": (0, 0, w, h),
                "score": 1.0,
                "mask": None
            }]
            return img, None, damage_regions
        return None, None, []

def check_deepfake_vit(image, damage_regions, deepfake_model, device, threshold=0.5):
    """Check if damage regions are deepfakes using ViT model"""
    results = []
    
    if deepfake_model is None:
        return []
    
    try:
        # If no damage regions, check entire image
        if not damage_regions:
            img_tensor = preprocess_for_vit(image, device)
            if img_tensor is None:
                return []
                
            # Run inference
            with torch.no_grad():
                outputs = deepfake_model(img_tensor)
                
                # Get predictions
                probabilities = torch.nn.functional.softmax(outputs, dim=1)
                fake_prob = probabilities[0, 1].item()  # Probability of being fake (class 1)
                is_fake = fake_prob > threshold
                
                results.append({
                    "region": "full_image",
                    "deepfake_prob": float(fake_prob),
                    "is_fake": bool(is_fake)
                })
            
            return results
        
        # Process each damage region
        for i, region in enumerate(damage_regions):
            x1, y1, x2, y2 = region["box"]
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2)
            
            # Only process valid regions
            if x2 > x1 and y2 > y1:
                # Extract region
                roi = image[y1:y2, x1:x2]
                
                # Preprocess
                img_tensor = preprocess_for_vit(roi, device)
                if img_tensor is None:
                    continue
                
                # Inference
                with torch.no_grad():
                    outputs = deepfake_model(img_tensor)
                    
                    # Get predictions
                    probabilities = torch.nn.functional.softmax(outputs, dim=1)
                    fake_prob = probabilities[0, 1].item()  # Probability of being fake (class 1)
                    is_fake = fake_prob > threshold
                    
                    results.append({
                        "region_id": i,
                        "box": (x1, y1, x2, y2),
                        "deepfake_prob": float(fake_prob),
                        "is_fake": bool(is_fake)
                    })
        
        return results
    except Exception as e:
        print(f"Error in deepfake detection: {e}")
        traceback.print_exc()
        return []

def visualize_results(image, damage_outputs, deepfake_results, damage_threshold):
    """Create visualization of results"""
    try:
        img_copy = image.copy()
        
        # Draw damage detection
        if damage_outputs is not None and DETECTRON2_AVAILABLE:
            try:
                v = Visualizer(img_copy[:, :, ::-1], scale=1.0, instance_mode=ColorMode.IMAGE_BW)
                v = v.draw_instance_predictions(damage_outputs["instances"].to("cpu"))
                result_img = v.get_image()[:, :, ::-1]
                result_img = np.array(result_img, dtype=np.uint8)
            except Exception as e:
                print(f"Error visualizing damage: {e}")
                result_img = img_copy
        else:
            result_img = img_copy
        
        # Add deepfake results
        for result in deepfake_results:
            try:
                if "box" in result:
                    x1, y1, x2, y2 = result["box"]
                    fake_prob = result["deepfake_prob"]
                    is_fake = result["is_fake"]
                    region_id = result.get("region_id", 0)
                    
                    # Status text
                    text = f"R{region_id}: {'FAKE' if is_fake else 'REAL'} ({fake_prob*100:.1f}%)"
                    
                    # Red for fake, green for real
                    color = (0, 0, 255) if is_fake else (0, 255, 0)
                    
                    # Ensure standard numpy array
                    if not isinstance(result_img, np.ndarray):
                        result_img = np.array(result_img, dtype=np.uint8)
                    
                    # Draw rectangle and text
                    cv2.rectangle(result_img, (x1, y1), (x2, y2), color, 2)
                    cv2.putText(result_img, text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
                elif "region" in result and result["region"] == "full_image":
                    fake_prob = result["deepfake_prob"]
                    is_fake = result["is_fake"]
                    
                    text = f"Image: {'FAKE' if is_fake else 'REAL'} ({fake_prob*100:.1f}%)"
                    color = (0, 0, 255) if is_fake else (0, 255, 0)
                    
                    if not isinstance(result_img, np.ndarray):
                        result_img = np.array(result_img, dtype=np.uint8)
                    
                    cv2.putText(result_img, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
            except Exception as e:
                print(f"Error drawing result: {e}")
        
        return result_img
    except Exception as e:
        print(f"Error in visualization: {e}")
        traceback.print_exc()
        return np.array(image, dtype=np.uint8)

def process_image(input_image, damage_model_path, deepfake_model_path, damage_threshold, 
                 deepfake_threshold, skip_damage, device_str, usage_count):
    """Process an image through the detection pipeline"""
    # Increment usage count and check if limit reached
    usage_count = usage_count + 1
    
    progress_info = []
    progress_info.append(f"Usage: {usage_count}/{MAX_TRIES}")
    
    # Check model files
    if not skip_damage and damage_model_path:
        if not os.path.exists(damage_model_path):
            progress_info.append(f"ERROR: Damage model not found at {damage_model_path}")
    
    if deepfake_model_path and not os.path.exists(deepfake_model_path):
        progress_info.append(f"ERROR: ViT deepfake model not found at {deepfake_model_path}")
    
    # Convert image to proper format
    try:
        if isinstance(input_image, dict) and "path" in input_image:
            img = cv2.imread(input_image["path"])
        elif isinstance(input_image, str):
            img = cv2.imread(input_image)
        elif isinstance(input_image, np.ndarray):
            img = input_image.copy()
            if len(img.shape) == 3 and img.shape[2] == 3:
                img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        else:
            return None, "Error: Unsupported image format", usage_count
        
        if img is None:
            return None, "Error: Could not read the image", usage_count
    except Exception as e:
        return None, f"Error loading image: {str(e)}", usage_count
    
    # Setup device
    device = setup_device(device_str)
    progress_info.append(f"Using device: {device}")
    
    # Initialize models
    damage_detector = None
    deepfake_model = None
    
    # Setup damage detector
    if not skip_damage and damage_model_path:
        progress_info.append("Setting up damage detector...")
        damage_detector, _ = setup_damage_detector(damage_model_path, float(damage_threshold))
        if damage_detector:
            progress_info.append("βœ… Damage detector initialized")
        else:
            progress_info.append("❌ Failed to initialize damage detector")
        
    # Setup ViT deepfake detector
    if deepfake_model_path:
        progress_info.append("Setting up ViT deepfake detector...")
        deepfake_model = load_vit_deepfake_model(deepfake_model_path, device)
        if deepfake_model:
            progress_info.append("βœ… ViT deepfake detector initialized")
        else:
            progress_info.append("❌ Failed to initialize ViT deepfake detector")
     
    # Step 1: Detect damage
    progress_info.append("Detecting damaged regions...")
    start_time = time.time()
    img, damage_outputs, damage_regions = detect_damage(img, damage_detector)
    damage_time = time.time() - start_time
    
    if damage_regions:
        progress_info.append(f"Found {len(damage_regions)} damage regions in {damage_time:.2f} seconds")
    else:
        progress_info.append("No damage regions detected")
    
    # Step 2: Check for deepfakes using ViT
    deepfake_results = []
    if deepfake_model is not None:
        progress_info.append("Analyzing regions for deepfakes using ViT model...")
        start_time = time.time()
        deepfake_results = check_deepfake_vit(
            img, damage_regions, deepfake_model, device, float(deepfake_threshold)
        )
        deepfake_time = time.time() - start_time
        
        if deepfake_results:
            progress_info.append(f"Deepfake analysis completed in {deepfake_time:.2f} seconds")
            
            # Generate report
            for result in deepfake_results:
                if "region_id" in result:
                    region_id = result["region_id"]
                    fake_prob = result["deepfake_prob"]
                    is_fake = result["is_fake"]
                    progress_info.append(f"Region {region_id}: {'FAKE' if is_fake else 'REAL'} (Probability: {fake_prob*100:.2f}%)")
                elif "region" in result and result["region"] == "full_image":
                    fake_prob = result["deepfake_prob"]
                    is_fake = result["is_fake"]
                    progress_info.append(f"Whole image: {'FAKE' if is_fake else 'REAL'} (Probability: {fake_prob*100:.2f}%)")
        else:
            progress_info.append("No deepfake detection results")
    
    # Final verdict
    fake_regions = [r for r in deepfake_results if r.get("is_fake", False)]
    if fake_regions:
        progress_info.append("\n🚨 VERDICT: This image contains FAKE damage 🚨")
    else:
        progress_info.append("\nβœ… VERDICT: All damage appears REAL")
        
    # Step 3: Visualize results
    result_img = visualize_results(img, damage_outputs, deepfake_results, float(damage_threshold))
    
    # Convert back to RGB for Gradio
    if len(result_img.shape) == 3 and result_img.shape[2] == 3:
        result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
    
    # Add usage information at the bottom of the image
    if usage_count >= MAX_TRIES:
        # Add a "Usage limit reached" message to the bottom of the image in red
        cv2.putText(result_img, f"USAGE LIMIT REACHED: {usage_count}/{MAX_TRIES}", 
                    (10, result_img.shape[0] - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
    else:
        # Add a usage counter to the bottom of the image
        cv2.putText(result_img, f"Usage: {usage_count}/{MAX_TRIES}", 
                    (10, result_img.shape[0] - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
    
    # Add usage info to the progress text
    if usage_count >= MAX_TRIES:
        progress_info.append("\n⚠️ You have reached the maximum number of tries allowed ⚠️")
    else:
        progress_info.append(f"\nRemaining tries: {MAX_TRIES - usage_count}")
    
    return result_img, "\n".join(progress_info), usage_count

def auto_install_dependencies():
    """Attempt to install dependencies if needed"""
    try:
        import importlib.util
        
        # Check for PyTorch
        if importlib.util.find_spec("torch") is None:
            print("Installing PyTorch...")
            os.system("pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu")
        
        # Check for Detectron2
        if importlib.util.find_spec("detectron2") is None:
            print("Installing Detectron2...")
            os.system("pip install git+https://github.com/facebookresearch/detectron2.git")
        
        # Check for Gradio
        if importlib.util.find_spec("gradio") is None:
            print("Installing Gradio...")
            os.system("pip install gradio")
        
        print("Dependencies installation complete!")
        return True
    except Exception as e:
        print(f"Error installing dependencies: {e}")
        return False

def check_model_paths(damage_path, deepfake_path):
    """Check if model paths are valid and exist"""
    output = ["## Path Verification Results\n"]
    
    # Check damage model
    if os.path.exists(damage_path):
        file_size = os.path.getsize(damage_path) / (1024 * 1024)  # Size in MB
        output.append(f"βœ… **Damage model:** Found at {damage_path} ({file_size:.2f} MB)")
    else:
        output.append(f"❌ **Damage model:** NOT found at {damage_path}")
    
    # Check deepfake model
    if os.path.exists(deepfake_path):
        file_size = os.path.getsize(deepfake_path) / (1024 * 1024)  # Size in MB
        output.append(f"βœ… **ViT deepfake model:** Found at {deepfake_path} ({file_size:.2f} MB)")
    else:
        output.append(f"❌ **ViT deepfake model:** NOT found at {deepfake_path}")
    
    return "\n".join(output)

def create_gradio_interface():
    # Define a theme
    theme = gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="orange",
    )
    
    # Initialize usage counter
    usage_counter = gr.State(0)
    
    with gr.Blocks(title="Car Damage & Deepfake Detector", theme=theme) as app:
        gr.Markdown("""
        # πŸš— Car Damage Fraud Detector with Vision Transformer
        
        Upload a car image to:
        1. Detect damaged areas
        2. Verify if the damage is real or artificially generated (deepfake)
        
        *This app uses a Vision Transformer (ViT) model for deepfake detection.*
        
        ⚠️ **Note: You have a maximum of 3 tries to analyze images.**
        """)
        
        # Main Interface Tab
        with gr.Tab("Analyze Image"):
            with gr.Row():
                with gr.Column(scale=1):
                    input_image = gr.Image(type="numpy", label="Upload Car Image")
                    
                    with gr.Row():
                        process_btn = gr.Button("Analyze Image", variant="primary")
                        clear_btn = gr.Button("Clear", variant="secondary")
                    
                    # Usage limit display
                    usage_display = gr.Markdown("**Usage: 0/3**")
                    
                    with gr.Accordion("Advanced Settings", open=False):
                        skip_damage = gr.Checkbox(label="Skip Damage Detection", value=False)
                        damage_threshold = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.05,
                                                   label="Damage Detection Threshold")
                        deepfake_threshold = gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.05,
                                                     label="Deepfake Detection Threshold")
                        device = gr.Dropdown(choices=["auto", "cuda", "cpu", "mps"], value="auto",
                                            label="Computation Device")
                        
                with gr.Column(scale=1):
                    output_image = gr.Image(type="numpy", label="Analysis Result")
                    
                    # Analysis info with nice formatting
                    with gr.Accordion("Analysis Details", open=True):
                        output_text = gr.Markdown(label="Detection Results")
        
 # This is the continuation of the code, starting from the System Diagnostics tab

        with gr.Tab("System Diagnostics"):
            with gr.Row():
                run_diagnostic_btn = gr.Button("Run Diagnostics", variant="primary")
                install_deps_btn = gr.Button("Install Dependencies", variant="secondary")
                
            diagnostic_output = gr.Markdown(label="Diagnostic Results")
            
            # Function to run system diagnostics
            def run_diagnostics():
                detectron2_results = verify_detectron2_installation()
                
                output = ["## System Diagnostics\n"]
                
                # Python & PyTorch versions
                import sys
                output.append(f"**Python:** {sys.version.split()[0]}")
                
                try:
                    import torch
                    output.append(f"**PyTorch:** {torch.__version__}")
                    output.append(f"**CUDA Available:** {'Yes βœ…' if torch.cuda.is_available() else 'No ❌'}")
                    if torch.cuda.is_available():
                        output.append(f"**GPU:** {torch.cuda.get_device_name(0)}")
                    output.append(f"**MPS Available:** {'Yes βœ…' if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() else 'No ❌'}")
                except ImportError:
                    output.append("**PyTorch:** Not installed ❌")
                
                output.append("\n## Model Files")
                
                # Check model files
                if os.path.exists(DEFAULT_DAMAGE_MODEL_PATH):
                    file_size = os.path.getsize(DEFAULT_DAMAGE_MODEL_PATH) / (1024 * 1024)  # Size in MB
                    output.append(f"**Damage Model:** Available βœ… ({file_size:.2f} MB)")
                else:
                    output.append(f"**Damage Model:** Not found ❌ at {DEFAULT_DAMAGE_MODEL_PATH}")
                
                if os.path.exists(DEFAULT_DEEPFAKE_MODEL_PATH):
                    file_size = os.path.getsize(DEFAULT_DEEPFAKE_MODEL_PATH) / (1024 * 1024)  # Size in MB
                    output.append(f"**ViT Deepfake Model:** Available βœ… ({file_size:.2f} MB)")
                else:
                    output.append(f"**ViT Deepfake Model:** Not found ❌ at {DEFAULT_DEEPFAKE_MODEL_PATH}")
                
                output.append("\n## Detectron2 Status")
                output.append(f"**Installed:** {'Yes βœ…' if detectron2_results['detectron2_installed'] else 'No ❌'}")
                output.append(f"**Model Zoo Access:** {'Yes βœ…' if detectron2_results['model_zoo_accessible'] else 'No ❌'}")
                output.append(f"**Config Creation:** {'Yes βœ…' if detectron2_results['can_create_cfg'] else 'No ❌'}")
                
                if detectron2_results['error_messages']:
                    output.append("\n## Error Messages")
                    for error in detectron2_results['error_messages']:
                        output.append(f"- {error}")
                
                output.append("\n## Recommendations")
                recommendations = []
                
                if not detectron2_results['detectron2_installed']:
                    recommendations.append("Install Detectron2: `pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu`")
                    recommendations.append("Then: `pip install git+https://github.com/facebookresearch/detectron2.git`")
                
                if not os.path.exists(DEFAULT_DAMAGE_MODEL_PATH):
                    recommendations.append(f"Place the damage model file at: {DEFAULT_DAMAGE_MODEL_PATH}")
                
                if not os.path.exists(DEFAULT_DEEPFAKE_MODEL_PATH):
                    recommendations.append(f"Place the ViT deepfake model file at: {DEFAULT_DEEPFAKE_MODEL_PATH}")
                
                if recommendations:
                    for rec in recommendations:
                        output.append(f"- {rec}")
                else:
                    output.append("- All systems are ready! πŸŽ‰")
                
                return "\n".join(output)
            
            # Function to install dependencies
            def install_dependencies():
                output = ["## Installing Dependencies\n"]
                
                # Try to install PyTorch
                output.append("Installing PyTorch...")
                try:
                    import importlib.util
                    if importlib.util.find_spec("torch") is None:
                        os.system("pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu")
                        output.append("βœ… PyTorch installed")
                    else:
                        output.append("βœ… PyTorch already installed")
                except Exception as e:
                    output.append(f"❌ Error installing PyTorch: {str(e)}")
                
                # Try to install Detectron2
                output.append("\nInstalling Detectron2...")
                try:
                    if importlib.util.find_spec("detectron2") is None:
                        os.system("pip install git+https://github.com/facebookresearch/detectron2.git")
                        output.append("βœ… Detectron2 installed")
                    else:
                        output.append("βœ… Detectron2 already installed")
                except Exception as e:
                    output.append(f"❌ Error installing Detectron2: {str(e)}")
                
                output.append("\n**Note:** You may need to restart the application for changes to take effect.")
                
                return "\n".join(output)
        
        # Settings Tab
        with gr.Tab("Settings"):
            gr.Markdown("## Model Settings")
            
            with gr.Row():
                with gr.Column():
                    custom_damage_model = gr.Textbox(
                        label="Custom Damage Model Path",
                        value=DEFAULT_DAMAGE_MODEL_PATH,
                        placeholder="Path to damage detection model (.pth)"
                    )
                    
                    custom_deepfake_model = gr.Textbox(
                        label="Custom ViT Deepfake Model Path",
                        value=DEFAULT_DEEPFAKE_MODEL_PATH,
                        placeholder="Path to ViT deepfake detection model (.pth)"
                    )
                    
                    check_paths_btn = gr.Button("Verify Paths", variant="primary")
                
                with gr.Column():
                    paths_result = gr.Markdown(label="Path Verification Results")

        with gr.Tab("Help"):
            gr.Markdown("""
            ## πŸ“‹ How to Use This Tool
            
            ### Basic Usage
            1. Upload a car image in the "Analyze Image" tab
            2. Click "Analyze Image" to process it
            3. View the results - damaged areas will be highlighted and classified as real or fake
            
            ### Understanding Results
            - **Green boxes/text:** Real damage
            - **Red boxes/text:** Potential deepfake damage
            - Percentage values show the confidence score
            
            ### Requirements
            - Damage detection model file (Detectron2 Mask R-CNN)
            - Vision Transformer (ViT) deepfake detection model file
            
            ### About the Vision Transformer (ViT)
            - This version uses a state-of-the-art Vision Transformer model for deepfake detection
            - The ViT model has been trained specifically to detect artificially generated car damage
            - ViT models are better at capturing global features in images compared to traditional CNNs
            
            ### Troubleshooting
            - If models aren't loading, check the "System Diagnostics" tab
            - Ensure all model files are in the expected locations
            - Try installing dependencies from the Diagnostics tab
            
            ### Advanced Settings
            - **Damage Threshold:** Higher values mean only high-confidence damage regions are shown
            - **Deepfake Threshold:** Higher values make the system more selective in flagging fakes
            - **Skip Damage Detection:** Analyze the entire image for deepfakes without damage detection
            """)
            
        # Examples
        if any(os.path.exists(img) for img in SAMPLE_IMAGES):
            gr.Markdown("## Example Images")
            with gr.Row():
                example_inputs = [img for img in SAMPLE_IMAGES if os.path.exists(img)]
                gr.Examples(
                    examples=example_inputs,
                    inputs=input_image,
                    outputs=[output_image, output_text, usage_counter],  # Add usage_counter to outputs
                    fn=lambda x: process_image(
                        x, DEFAULT_DAMAGE_MODEL_PATH, DEFAULT_DEEPFAKE_MODEL_PATH, 
                        0.7, 0.5, False, "auto", 0  # Added initial usage_count value of 0
                    ),
                    cache_examples=True
                )
        
        # Connect functions to the UI
        process_btn.click(
            fn=process_image,
            inputs=[
                input_image, 
                custom_damage_model,
                custom_deepfake_model,
                damage_threshold,
                deepfake_threshold,
                skip_damage,
                device,
                usage_counter  # Add the usage counter state here
            ],
            outputs=[output_image, output_text, usage_counter]  # Update usage counter in outputs
        )
        
        # Clear button functionality
        clear_btn.click(
            fn=lambda: [None, "", 0],  # Reset usage counter to 0 when clearing
            inputs=[],
            outputs=[output_image, output_text, usage_counter]
        )
        
        # System diagnostic buttons
        run_diagnostic_btn.click(
            fn=run_diagnostics,
            inputs=[],
            outputs=diagnostic_output
        )
        
        install_deps_btn.click(
            fn=install_dependencies,
            inputs=[],
            outputs=diagnostic_output
        )
        
        # Settings tab
        check_paths_btn.click(
            fn=check_model_paths,
            inputs=[custom_damage_model, custom_deepfake_model],
            outputs=paths_result
        )
        
        # Update usage display when counter changes
        usage_counter.change(
            fn=lambda count: f"**Usage: {count}/{MAX_TRIES}**",
            inputs=[usage_counter],
            outputs=[usage_display]
        )
        
    return app

if __name__ == "__main__":
    # Check if dependencies are installed
    auto_install_dependencies()
    
    # Create and launch the Gradio app
    app = create_gradio_interface()
    app.launch(share=False)  # Set share=True to create a public link