V1.1

app.py

@@ -26,226 +26,270 @@ from torchvision import transforms
 
 # Add current directory to path
 if not os.getcwd() in sys.path:
-    sys.path.append(os.getcwd())
+sys.path.append(os.getcwd())
 
 # Detectron2 imports - wrapped in try-except to make them optional
 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
+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
+print("Warning: Detectron2 is not installed. Damage detection will not be available.")
+DETECTRON2_AVAILABLE = False
 
 # Check for custom path for models
 try:
-    from configs.get_config import load_config
-    from models import *
-    MODELS_IMPORTED = True
+from configs.get_config import load_config
+from models import *
+MODELS_IMPORTED = True
 except ImportError:
-    print("Warning: Custom models couldn't be imported. Only damage detection will work.")
-    MODELS_IMPORTED = False
+print("Warning: Custom models couldn't be imported. Only damage detection will work.")
+MODELS_IMPORTED = False
 
 def setup_device(device_str):
-    """Set up the computation device based on user input and availability"""
-    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():
+"""Set up the computation device based on user input and availability"""
+if device_str == 'auto':
+    if 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():
+    elif 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')
+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 using Detectron2"""
-    if not DETECTRON2_AVAILABLE:
-        print("Detectron2 is not installed. Cannot set up damage detector.")
-        return None, None
-        
-    if model_path is None or not os.path.exists(model_path):
-        print("No damage model specified or file not found. Skipping damage detection.")
-        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
+"""Set up the damage detection model using Detectron2"""
+if not DETECTRON2_AVAILABLE:
+    print("Detectron2 is not installed. Cannot set up damage detector.")
+    return None, None
     
-    # Explicitly set to use CPU if on Mac (MPS)
-    if torch.backends.mps.is_available():
-        cfg.MODEL.DEVICE = "cpu"
-        print("Mac MPS detected - forcing Detectron2 to use CPU")
+if model_path is None or not os.path.exists(model_path):
+    print("No damage model specified or file not found. Skipping damage detection.")
+    return None, None
     
-    try:
-        predictor = DefaultPredictor(cfg)
-        return predictor, cfg
-    except Exception as e:
-        print(f"Error setting up damage detector: {e}")
-        return None, cfg
+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
+
+# Explicitly set to use CPU if on Mac (MPS)
+if torch.backends.mps.is_available():
+    cfg.MODEL.DEVICE = "cpu"
+    print("Mac MPS detected - forcing Detectron2 to use CPU")
+
+try:
+    predictor = DefaultPredictor(cfg)
+    return predictor, cfg
+except Exception as e:
+    print(f"Error setting up damage detector: {e}")
+    return None, cfg
 
 def load_deepfake_model(model_path, cfg_path, device):
-    """Load the deepfake detection model"""
-    if not MODELS_IMPORTED:
-        print("Custom models module not imported. Cannot load deepfake model.")
-        return None, None
-        
-    if model_path is None or not os.path.exists(model_path):
-        print("No deepfake model specified or file not found. Skipping deepfake detection.")
-        return None, None
+"""Load the deepfake detection model"""
+if not MODELS_IMPORTED:
+    print("Custom models module not imported. Cannot load deepfake model.")
+    return None, None
+    
+if model_path is None or not os.path.exists(model_path):
+    print("No deepfake model specified or file not found. Skipping deepfake detection.")
+    return None, None
+    
+if cfg_path is None or not os.path.exists(cfg_path):
+    print("No deepfake config specified or file not found. Skipping deepfake detection.")
+    return None, None
+    
+try:
+    # Load config
+    cfg = load_config(cfg_path)
+    
+    # Build model
+    model = build_model(cfg.MODEL, MODELS)
+    
+    # Load weights
+    print(f"Loading deepfake model from: {model_path}")
+    checkpoint = torch.load(model_path, map_location='cpu')
+    
+    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
+        model.load_state_dict(checkpoint['state_dict'])
+    else:
+        model.load_state_dict(checkpoint)
         
-    if cfg_path is None or not os.path.exists(cfg_path):
-        print("No deepfake config specified or file not found. Skipping deepfake detection.")
-        return None, None
+    # Move model to device and set to evaluation mode
+    model = model.to(device)
+    if hasattr(cfg.MODEL, 'precision') and cfg.MODEL.precision == 'fp64':
+        model = model.to(torch.float64)
+    model.eval()
+    
+    return model, cfg
+except Exception as e:
+    print(f"Error loading deepfake model: {e}")
+    import traceback
+    traceback.print_exc()
+    return None, None
+
+def preprocess_for_deepfake(image, cfg, device):
+"""Preprocess an image for deepfake detection"""
+try:
+    # Convert to RGB if needed
+    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
         
-    try:
-        # Load config
-        cfg = load_config(cfg_path)
+    # Resize
+    img_resized = cv2.resize(rgb_img, (cfg.DATASET.IMAGE_SIZE[0], cfg.DATASET.IMAGE_SIZE[1]))
+    
+    # Convert to PIL and apply transforms
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=cfg.DATASET.TRANSFORM.normalize.mean,
+            std=cfg.DATASET.TRANSFORM.normalize.std
+        )
+    ])
+    
+    img_tensor = transform(Image.fromarray(img_resized)).unsqueeze(0)  # Add batch dimension
+    img_tensor = img_tensor.to(device)
+    
+    # Convert to correct precision
+    if hasattr(cfg.MODEL, 'precision') and cfg.MODEL.precision == 'fp64':
+        img_tensor = img_tensor.to(torch.float64)
         
-        # Build model
-        model = build_model(cfg.MODEL, MODELS)
+    return img_tensor
+except Exception as e:
+    print(f"Error preprocessing image for deepfake detection: {e}")
+    import traceback
+    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")
         
-        # Load weights
-        print(f"Loading deepfake model from: {model_path}")
-        checkpoint = torch.load(model_path, map_location='cpu')
+    # If no damage detector available, return the whole image as region
+    if damage_detector is None:
+        print("No damage detector available. Using whole image as region.")
+        h, w = img.shape[:2]
+        damage_regions = [{
+            "box": (0, 0, w, h),
+            "score": 1.0,
+            "mask": None
+        }]
+        return img, None, damage_regions
         
-        if 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)
-        if hasattr(cfg.MODEL, 'precision') and cfg.MODEL.precision == 'fp64':
-            model = model.to(torch.float64)
-        model.eval()
+    # Run inference
+    outputs = damage_detector(img)
+    
+    # Get damage 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 not damage_regions:
+        print("No damage detected. Using whole image.")
+        h, w = img.shape[:2]
+        damage_regions = [{
+            "box": (0, 0, w, h),
+            "score": 1.0,
+            "mask": None
+        }]
         
-        return model, cfg
-    except Exception as e:
-        print(f"Error loading deepfake model: {e}")
-        import traceback
-        traceback.print_exc()
-        return None, None
+    return img, outputs, damage_regions
+except Exception as e:
+    print(f"Error detecting damage: {e}")
+    # If error occurs, return the whole image as region
+    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 preprocess_for_deepfake(image, cfg, device):
-    """Preprocess an image for deepfake detection"""
-    try:
-        # Convert to RGB if needed
-        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
+def check_deepfake(image, damage_regions, deepfake_model, deepfake_cfg, device, threshold=0.5):
+"""Check if damage regions are deepfakes"""
+results = []
+
+if deepfake_model is None:
+    print("No deepfake model available. Skipping deepfake detection.")
+    return []
+
+try:
+    # If no damage regions, check the entire image
+    if not damage_regions:
+        img_tensor = preprocess_for_deepfake(image, deepfake_cfg, device)
+        if img_tensor is None:
+            return []
             
-        # Resize
-        img_resized = cv2.resize(rgb_img, (cfg.DATASET.IMAGE_SIZE[0], cfg.DATASET.IMAGE_SIZE[1]))
-        
-        # Convert to PIL and apply transforms
-        transform = transforms.Compose([
-            transforms.ToTensor(),
-            transforms.Normalize(
-                mean=cfg.DATASET.TRANSFORM.normalize.mean,
-                std=cfg.DATASET.TRANSFORM.normalize.std
-            )
-        ])
-        
-        img_tensor = transform(Image.fromarray(img_resized)).unsqueeze(0)  # Add batch dimension
-        img_tensor = img_tensor.to(device)
-        
-        # Convert to correct precision
-        if hasattr(cfg.MODEL, 'precision') and cfg.MODEL.precision == 'fp64':
-            img_tensor = img_tensor.to(torch.float64)
+        # Run inference
+        with torch.no_grad():
+            outputs = deepfake_model(img_tensor)
             
-        return img_tensor
-    except Exception as e:
-        print(f"Error preprocessing image for deepfake detection: {e}")
-        import traceback
-        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")
+            # Extract outputs
+            if isinstance(outputs, list):
+                outputs = outputs[0]
             
-        # If no damage detector available, return the whole image as region
-        if damage_detector is None:
-            print("No damage detector available. Using whole image as region.")
-            h, w = img.shape[:2]
-            damage_regions = [{
-                "box": (0, 0, w, h),
-                "score": 1.0,
-                "mask": None
-            }]
-            return img, None, damage_regions
+            if isinstance(outputs, dict) and 'cls' in outputs:
+                cls_outputs = outputs['cls']
+                cls_prob = cls_outputs.sigmoid().cpu().numpy()
+            else:
+                # Assuming the output is directly the classification probability
+                cls_prob = outputs.sigmoid().cpu().numpy() if hasattr(outputs, 'sigmoid') else outputs.cpu().numpy()
             
-        # Run inference
-        outputs = damage_detector(img)
-        
-        # Get damage 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 []
+            if cls_prob.size > 0:
+                is_fake = cls_prob[0][0] > threshold if cls_prob.ndim > 1 else cls_prob[0] > threshold
+                confidence = cls_prob[0][0] if cls_prob.ndim > 1 else cls_prob[0]
+                
+                results.append({
+                    "region": "full_image",
+                    "deepfake_prob": float(confidence),
+                    "is_fake": bool(is_fake)
+                })
         
-        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
-            })
+        return results
+    
+    # Process each damage region
+    for i, region in enumerate(damage_regions):
+        x1, y1, x2, y2 = region["box"]
+        # Ensure coordinates are within image bounds
+        x1, y1 = max(0, x1), max(0, y1)
+        x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2)
         
-        if not damage_regions:
-            print("No damage detected. Using whole image.")
-            h, w = img.shape[:2]
-            damage_regions = [{
-                "box": (0, 0, w, h),
-                "score": 1.0,
-                "mask": None
-            }]
+        # Extract region and check if it's a deepfake
+        if x2 > x1 and y2 > y1:
+            # Get ROI
+            roi = image[y1:y2, x1:x2]
             
-        return img, outputs, damage_regions
-    except Exception as e:
-        print(f"Error detecting damage: {e}")
-        # If error occurs, return the whole image as region
-        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(image, damage_regions, deepfake_model, deepfake_cfg, device, threshold=0.5):
-    """Check if damage regions are deepfakes"""
-    results = []
-    
-    if deepfake_model is None:
-        print("No deepfake model available. Skipping deepfake detection.")
-        return []
-    
-    try:
-        # If no damage regions, check the entire image
-        if not damage_regions:
-            img_tensor = preprocess_for_deepfake(image, deepfake_cfg, device)
+            # Preprocess
+            img_tensor = preprocess_for_deepfake(roi, deepfake_cfg, device)
             if img_tensor is None:
-                return []
-                
+                continue
+            
             # Run inference
             with torch.no_grad():
                 outputs = deepfake_model(img_tensor)
@@ -266,301 +310,258 @@ def check_deepfake(image, damage_regions, deepfake_model, deepfake_cfg, device,
                     confidence = cls_prob[0][0] if cls_prob.ndim > 1 else cls_prob[0]
                     
                     results.append({
-                        "region": "full_image",
+                        "region_id": i,
+                        "box": (x1, y1, x2, y2),
                         "deepfake_prob": float(confidence),
                         "is_fake": bool(is_fake)
                     })
+    
+    return results
+except Exception as e:
+    print(f"Error in deepfake detection: {e}")
+    import traceback
+    traceback.print_exc()
+    return []
+
+def visualize_results(image, damage_outputs, deepfake_results, damage_threshold):
+"""Create visualization of damage detection and deepfake verification"""
+try:
+    # Create a copy for visualization
+    img_copy = image.copy()
+    
+    # Draw damage detection results
+    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]
             
-            return results
-        
-        # Process each damage region
-        for i, region in enumerate(damage_regions):
-            x1, y1, x2, y2 = region["box"]
-            # Ensure coordinates are within image bounds
-            x1, y1 = max(0, x1), max(0, y1)
-            x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2)
-            
-            # Extract region and check if it's a deepfake
-            if x2 > x1 and y2 > y1:
-                # Get ROI
-                roi = image[y1:y2, x1:x2]
+            # Convert to a standard numpy array to ensure compatibility with OpenCV
+            result_img = np.array(result_img, dtype=np.uint8)
+        except Exception as e:
+            print(f"Error visualizing damage detection: {e}")
+            result_img = img_copy
+    else:
+        result_img = img_copy
+    
+    # Add deepfake detection 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)
                 
-                # Preprocess
-                img_tensor = preprocess_for_deepfake(roi, deepfake_cfg, device)
-                if img_tensor is None:
-                    continue
+                # Text for the region
+                text = f"R{region_id}: {'FAKE' if is_fake else 'REAL'} ({fake_prob*100:.1f}%)"
                 
-                # Run inference
-                with torch.no_grad():
-                    outputs = deepfake_model(img_tensor)
-                    
-                    # Extract outputs
-                    if isinstance(outputs, list):
-                        outputs = outputs[0]
-                    
-                    if isinstance(outputs, dict) and 'cls' in outputs:
-                        cls_outputs = outputs['cls']
-                        cls_prob = cls_outputs.sigmoid().cpu().numpy()
-                    else:
-                        # Assuming the output is directly the classification probability
-                        cls_prob = outputs.sigmoid().cpu().numpy() if hasattr(outputs, 'sigmoid') else outputs.cpu().numpy()
-                    
-                    if cls_prob.size > 0:
-                        is_fake = cls_prob[0][0] > threshold if cls_prob.ndim > 1 else cls_prob[0] > threshold
-                        confidence = cls_prob[0][0] if cls_prob.ndim > 1 else cls_prob[0]
-                        
-                        results.append({
-                            "region_id": i,
-                            "box": (x1, y1, x2, y2),
-                            "deepfake_prob": float(confidence),
-                            "is_fake": bool(is_fake)
-                        })
-        
-        return results
-    except Exception as e:
-        print(f"Error in deepfake detection: {e}")
-        import traceback
-        traceback.print_exc()
-        return []
-
-def visualize_results(image, damage_outputs, deepfake_results, damage_threshold):
-    """Create visualization of damage detection and deepfake verification"""
-    try:
-        # Create a copy for visualization
-        img_copy = image.copy()
-        
-        # Draw damage detection results
-        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]
+                # Different colors for fake/real
+                color = (0, 0, 255) if is_fake else (0, 255, 0)  # Red for fake, green for real
                 
-                # Convert to a standard numpy array to ensure compatibility with OpenCV
-                result_img = np.array(result_img, dtype=np.uint8)
-            except Exception as e:
-                print(f"Error visualizing damage detection: {e}")
-                result_img = img_copy
-        else:
-            result_img = img_copy
-        
-        # Add deepfake detection 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)
-                    
-                    # Text for the region
-                    text = f"R{region_id}: {'FAKE' if is_fake else 'REAL'} ({fake_prob*100:.1f}%)"
-                    
-                    # Different colors for fake/real
-                    color = (0, 0, 255) if is_fake else (0, 255, 0)  # Red for fake, green for real
-                    
-                    # Ensure we have a 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 for the whole image
-                    text = f"Image: {'FAKE' if is_fake else 'REAL'} ({fake_prob*100:.1f}%)"
-                    
-                    # Different colors for fake/real
-                    color = (0, 0, 255) if is_fake else (0, 255, 0)  # Red for fake, green for real
-                    
-                    # Ensure we have a standard numpy array
-                    if not isinstance(result_img, np.ndarray):
-                        result_img = np.array(result_img, dtype=np.uint8)
-                    
-                    # Draw text
-                    cv2.putText(result_img, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
-            except Exception as e:
-                print(f"Error drawing result {result}: {e}")
-        
-        return result_img
-    except Exception as e:
-        print(f"Error visualizing results: {e}")
-        import traceback
-        traceback.print_exc()
-        return np.array(image, dtype=np.uint8)  # Return the original image as a numpy array
+                # Ensure we have a 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 for the whole image
+                text = f"Image: {'FAKE' if is_fake else 'REAL'} ({fake_prob*100:.1f}%)"
+                
+                # Different colors for fake/real
+                color = (0, 0, 255) if is_fake else (0, 255, 0)  # Red for fake, green for real
+                
+                # Ensure we have a standard numpy array
+                if not isinstance(result_img, np.ndarray):
+                    result_img = np.array(result_img, dtype=np.uint8)
+                
+                # Draw text
+                cv2.putText(result_img, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
+        except Exception as e:
+            print(f"Error drawing result {result}: {e}")
+    
+    return result_img
+except Exception as e:
+    print(f"Error visualizing results: {e}")
+    import traceback
+    traceback.print_exc()
+    return np.array(image, dtype=np.uint8)  # Return the original image as a numpy array
 
 def process_image(input_image, damage_model_path, deepfake_model_path, deepfake_cfg_path, 
-                 damage_threshold, deepfake_threshold, skip_damage, device_str):
-    """Process an image through the car damage and deepfake detection pipeline"""
-    progress_info = []
-    
-    # Convert Gradio image to numpy array
-    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):
-        # Make a copy to avoid modifying the original
-        img = input_image.copy()
-        # Convert from RGB to BGR (OpenCV format)
-        if len(img.shape) == 3 and img.shape[2] == 3:
-            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+                damage_threshold, deepfake_threshold, skip_damage, device_str):
+"""Process an image through the car damage and deepfake detection pipeline"""
+progress_info = []
+
+# Convert Gradio image to numpy array
+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):
+    # Make a copy to avoid modifying the original
+    img = input_image.copy()
+    # Convert from RGB to BGR (OpenCV format)
+    if len(img.shape) == 3 and img.shape[2] == 3:
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+else:
+    return None, "Error: Unsupported image format"
+
+if img is None:
+    return None, "Error: Could not read the image"
+
+# Progress update
+progress_info.append("Image loaded successfully")
+
+# Setup device
+device = setup_device(device_str)
+progress_info.append(f"Using device: {device}")
+
+# Initialize models
+damage_detector = None
+deepfake_model = None
+deepfake_cfg = None
+
+# Setup damage detector if not skipped
+if not skip_damage and damage_model_path:
+    progress_info.append("Setting up damage detector...")
+    damage_detector, detector_cfg = setup_damage_detector(damage_model_path, float(damage_threshold))
+    if damage_detector is None and DETECTRON2_AVAILABLE:
+        progress_info.append("Failed to initialize damage detector")
     else:
-        return None, "Error: Unsupported image format"
-    
-    if img is None:
-        return None, "Error: Could not read the image"
-    
-    # Progress update
-    progress_info.append("Image loaded successfully")
-    
-    # Setup device
-    device = setup_device(device_str)
-    progress_info.append(f"Using device: {device}")
-    
-    # Initialize models
-    damage_detector = None
-    deepfake_model = None
-    deepfake_cfg = None
-    
-    # Setup damage detector if not skipped
-    if not skip_damage and damage_model_path:
-        progress_info.append("Setting up damage detector...")
-        damage_detector, detector_cfg = setup_damage_detector(damage_model_path, float(damage_threshold))
-        if damage_detector is None and DETECTRON2_AVAILABLE:
-            progress_info.append("Failed to initialize damage detector")
-        else:
-            progress_info.append("Damage detector initialized successfully")
-    
-    # Setup deepfake detector
-    if deepfake_model_path and deepfake_cfg_path:
-        progress_info.append("Setting up deepfake detector...")
-        deepfake_model, deepfake_cfg = load_deepfake_model(deepfake_model_path, deepfake_cfg_path, device)
-        if deepfake_model is None:
-            progress_info.append("Failed to initialize deepfake detector")
-        else:
-            progress_info.append("Deepfake detector initialized successfully")
-    
-    # Ensure at least one detector is working
-    if damage_detector is None and deepfake_model is None:
-        return None, "Error: Neither damage nor deepfake detector is available"
-    
-    # Step 1: Detect damage or use whole image
-    progress_info.append("Detecting damage regions...")
+        progress_info.append("Damage detector initialized successfully")
+
+# Setup deepfake detector
+if deepfake_model_path and deepfake_cfg_path:
+    progress_info.append("Setting up deepfake detector...")
+    deepfake_model, deepfake_cfg = load_deepfake_model(deepfake_model_path, deepfake_cfg_path, device)
+    if deepfake_model is None:
+        progress_info.append("Failed to initialize deepfake detector")
+    else:
+        progress_info.append("Deepfake detector initialized successfully")
+
+# Ensure at least one detector is working
+if damage_detector is None and deepfake_model is None:
+    return None, "Error: Neither damage nor deepfake detector is available"
+
+# Step 1: Detect damage or use whole image
+progress_info.append("Detecting damage regions...")
+start_time = time.time()
+img, damage_outputs, damage_regions = detect_damage(img, damage_detector)
+damage_time = time.time() - start_time
+
+if img is None:
+    return None, "Error: Failed to process image"
+
+# Print damage detection results
+if damage_detector is not None and damage_regions:
+    progress_info.append(f"Detected {len(damage_regions)} damage regions in {damage_time:.3f} seconds")
+else:
+    progress_info.append("Using the whole image for analysis")
+
+# Step 2: Check if damage is deepfake
+deepfake_results = []
+if deepfake_model is not None:
+    progress_info.append("Performing deepfake detection...")
     start_time = time.time()
-    img, damage_outputs, damage_regions = detect_damage(img, damage_detector)
-    damage_time = time.time() - start_time
-    
-    if img is None:
-        return None, "Error: Failed to process image"
+    deepfake_results = check_deepfake(
+        img, damage_regions, deepfake_model, deepfake_cfg, device, float(deepfake_threshold)
+    )
+    deepfake_time = time.time() - start_time
     
-    # Print damage detection results
-    if damage_detector is not None and damage_regions:
-        progress_info.append(f"Detected {len(damage_regions)} damage regions in {damage_time:.3f} seconds")
+    if deepfake_results:
+        progress_info.append(f"Deepfake detection completed in {deepfake_time:.3f} 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("Using the whole image for analysis")
+        progress_info.append("No deepfake detection results")
+
+# Step 3: Visualize final results
+progress_info.append("Generating visualization...")
+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)
+
+progress_info.append("Processing complete!")
+
+return result_img, "\n".join(progress_info)
+
+def create_gradio_interface():
+with gr.Blocks(title="Car Damage & Deepfake Detection") as app:
+    gr.Markdown("# Car Damage Detection & Deepfake Verification")
+    gr.Markdown("Upload an image to detect car damage and check if it's a deepfake")
     
-    # Step 2: Check if damage is deepfake
-    deepfake_results = []
-    if deepfake_model is not None:
-        progress_info.append("Performing deepfake detection...")
-        start_time = time.time()
-        deepfake_results = check_deepfake(
-            img, damage_regions, deepfake_model, deepfake_cfg, device, float(deepfake_threshold)
-        )
-        deepfake_time = time.time() - start_time
+    with gr.Tab("Basic Interface"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(type="numpy", label="Input Image")
+                
+                # Simple controls
+                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")
+                
+                process_btn = gr.Button("Process Image", variant="primary")
+                
+            with gr.Column(scale=1):
+                output_image = gr.Image(type="numpy", label="Result")
+                output_text = gr.Textbox(label="Detection Results", lines=10)
         
-        if deepfake_results:
-            progress_info.append(f"Deepfake detection completed in {deepfake_time:.3f} 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")
+    with gr.Tab("Advanced Settings"):
+        with gr.Row():
+            with gr.Column():
+                damage_model_path = gr.Textbox(label="Damage Model Path", 
+                                                placeholder="Path to damage detection model (.pth)")
+                deepfake_model_path = gr.Textbox(label="Deepfake Model Path", 
+                                                placeholder="Path to deepfake detection model (.pth)")
+                deepfake_cfg_path = gr.Textbox(label="Deepfake Config Path", 
+                                                placeholder="Path to deepfake model config (.yaml)")
     
-    # Step 3: Visualize final results
-    progress_info.append("Generating visualization...")
-    result_img = visualize_results(img, damage_outputs, deepfake_results, float(damage_threshold))
+    # Connect the process function
+    process_btn.click(
+        fn=process_image,
+        inputs=[
+            input_image, 
+            damage_model_path,
+            deepfake_model_path,
+            deepfake_cfg_path,
+            damage_threshold,
+            deepfake_threshold,
+            skip_damage,
+            device
+        ],
+        outputs=[output_image, output_text]
+    )
     
-    # 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)
-    
-    progress_info.append("Processing complete!")
-    
-    return result_img, "\n".join(progress_info)
-
-def create_gradio_interface():
-    with gr.Blocks(title="Car Damage & Deepfake Detection") as app:
-        gr.Markdown("# Car Damage Detection & Deepfake Verification")
-        gr.Markdown("Upload an image to detect car damage and check if it's a deepfake")
+    # Examples
+    gr.Markdown("## Examples")
+    gr.Markdown("Note: Examples will only work if you have the appropriate models installed.")
         
-        with gr.Tab("Basic Interface"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    input_image = gr.Image(type="numpy", label="Input Image")
-                    
-                    # Simple controls
-                    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")
-                    
-                    process_btn = gr.Button("Process Image", variant="primary")
-                    
-                with gr.Column(scale=1):
-                    output_image = gr.Image(type="numpy", label="Result")
-                    output_text = gr.Textbox(label="Detection Results", lines=10)
-            
-        with gr.Tab("Advanced Settings"):
-            with gr.Row():
-                with gr.Column():
-                    damage_model_path = gr.Textbox(label="Damage Model Path", 
-                                                 placeholder="Path to damage detection model (.pth)")
-                    deepfake_model_path = gr.Textbox(label="Deepfake Model Path", 
-                                                   placeholder="Path to deepfake detection model (.pth)")
-                    deepfake_cfg_path = gr.Textbox(label="Deepfake Config Path", 
-                                                 placeholder="Path to deepfake model config (.yaml)")
-        
-        # Connect the process function
-        process_btn.click(
-            fn=process_image,
-            inputs=[
-                input_image, 
-                damage_model_path,
-                deepfake_model_path,
-                deepfake_cfg_path,
-                damage_threshold,
-                deepfake_threshold,
-                skip_damage,
-                device
-            ],
-            outputs=[output_image, output_text]
-        )
-        
-        # Examples
-        gr.Markdown("## Examples")
-        gr.Markdown("Note: Examples will only work if you have the appropriate models installed.")
-            
-    return app
-    app = create_gradio_interface()
+return app
 
+# Create and launch the app
+app = create_gradio_interface()
 
+# For local testing and Hugging Face Spaces, with debugging enabled
 if __name__ == "__main__":
-    # Create and launch the Gradio interface
-    app.launch(share=True)  # Set share=False in production
+    app.launch(debug=True)  # Enable debug mode to see detailed error messages