Spaces:
Runtime error
Runtime error
| #!/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 |