import torch import requests from PIL import Image, ImageDraw, ImageFont import matplotlib.pyplot as plt from transformers import ( DetrImageProcessor, DetrForObjectDetection, ViTImageProcessor, ViTForImageClassification, AutoModelForCausalLM, AutoTokenizer, pipeline ) from ultralytics import YOLO import cv2 import numpy as np from datetime import datetime from fpdf import FPDF import os import io import traceback # Import traceback for better error logging import torch import joblib import numpy as np from transformers import CLIPProcessor, CLIPModel from config import DEVICE, MODEL_SAVE_PATH, LABEL_ENCODER_PATH import google.generativeai as genai # ====================== # 1. INITIALIZATION # ====================== class CrimeSceneAnalyzer: def __init__(self): self.yolo_model = YOLO('yolov8x.pt') # Automatically downloads if missing # Load CLIP model and processor self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(DEVICE) self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") # Load crime scene classification model and label encoder self.crime_scene_model = joblib.load(MODEL_SAVE_PATH) self.label_encoder = joblib.load(LABEL_ENCODER_PATH) # Object Detection (DETR) self.detr_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") self.detr_model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") self.detr_model.config.auxiliary_loss = True self.detr_model.config.num_queries = 150 # Evidence Classification (ViT) self.vit_processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224') self.vit_model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224') # Report Generation (GPT-2 fine-tuned) gemini_api_key = os.environ.get("GEMINI_API_KEY", "AIzaSyBG6e9rH2tPinqRQay2QrXTkMphkYEjyeY") if not gemini_api_key: print("⚠️ GEMINI_API_KEY environment variable not set. Report/Summary generation will be affected.") self.gemini_model = None else: genai.configure(api_key=gemini_api_key) # Choose the model best suited for report generation. # gemini-1.5-flash is fast and cost-effective. # gemini-1.0-pro or gemini-1.5-pro for potentially more detailed/nuanced reports. self.gemini_model = genai.GenerativeModel('gemini-2.5-flash-lite') # Evidence Mapping (your existing map) self.EVIDENCE_MAP = { "knife": {"type": "weapon", "priority": 1}, "gun": {"type": "weapon", "priority": 1}, "pistol": {"type": "weapon", "priority": 1}, "revolver": {"type": "weapon", "priority": 1}, "rifle": {"type": "weapon", "priority": 1}, "shotgun": {"type": "weapon", "priority": 1}, "sword": {"type": "weapon", "priority": 1}, "machete": {"type": "weapon", "priority": 1}, "brass knuckles": {"type": "weapon", "priority": 1}, "taser": {"type": "weapon", "priority": 1}, "pepper spray": {"type": "weapon", "priority": 1}, "crossbow": {"type": "weapon", "priority": 1}, "axe": {"type": "weapon", "priority": 1}, "hammer": {"type": "weapon", "priority": 1}, "scissors": {"type": "weapon", "priority": 1}, "phone": {"type": "electronic", "priority": 2}, "laptop": {"type": "electronic", "priority": 2}, "tablet": {"type": "electronic", "priority": 2}, "camera": {"type": "electronic", "priority": 2}, "usb": {"type": "electronic", "priority": 2}, "hard drive": {"type": "electronic", "priority": 2}, "sd card": {"type": "electronic", "priority": 2}, "dvr": {"type": "electronic", "priority": 2}, "router": {"type": "electronic", "priority": 2}, "sim card": {"type": "electronic", "priority": 2}, "blood": {"type": "biological", "priority": 1}, "hair": {"type": "biological", "priority": 1}, "fingerprint": {"type": "biological", "priority": 1}, "dna": {"type": "biological", "priority": 1}, "saliva": {"type": "biological", "priority": 1}, "semen": {"type": "biological", "priority": 1}, "tissue": {"type": "biological", "priority": 1}, "bone": {"type": "biological", "priority": 1}, "tooth": {"type": "biological", "priority": 1}, "bottle": {"type": "container", "priority": 3}, "syringe": {"type": "container", "priority": 1}, "needle": {"type": "drug", "priority": 1}, "pill": {"type": "drug", "priority": 1}, "powder": {"type": "drug", "priority": 1}, "marijuana": {"type": "drug", "priority": 1}, "cocaine": {"type": "drug", "priority": 1}, "heroin": {"type": "drug", "priority": 1}, "meth": {"type": "drug", "priority": 1}, "pipe": {"type": "drug", "priority": 1}, "scale": {"type": "drug", "priority": 1}, "paper": {"type": "document", "priority": 2}, "key": {"type": "tool", "priority": 2}, "id": {"type": "document", "priority": 2}, "passport": {"type": "document", "priority": 2}, "license": {"type": "document", "priority": 2}, "credit card": {"type": "document", "priority": 2}, "money": {"type": "document", "priority": 2}, "note": {"type": "document", "priority": 2}, "letter": {"type": "document", "priority": 2}, "diary": {"type": "document", "priority": 2}, "map": {"type": "document", "priority": 2}, "blueprint": {"type": "document", "priority": 2}, "shoe": {"type": "clothing", "priority": 3}, "glove": {"type": "clothing", "priority": 3}, "mask": {"type": "clothing", "priority": 3}, "hat": {"type": "clothing", "priority": 3}, "jacket": {"type": "clothing", "priority": 3}, "backpack": {"type": "clothing", "priority": 3}, "watch": {"type": "clothing", "priority": 3}, "jewelry": {"type": "clothing", "priority": 3}, "eyeglasses": {"type": "clothing", "priority": 3}, "crowbar": {"type": "tool", "priority": 2}, "screwdriver": {"type": "tool", "priority": 2}, "wrench": {"type": "tool", "priority": 2}, "pliers": {"type": "tool", "priority": 2}, "lockpick": {"type": "tool", "priority": 2}, "shovel": {"type": "tool", "priority": 2}, "rope": {"type": "tool", "priority": 2}, "duct tape": {"type": "tool", "priority": 2}, "wire": {"type": "tool", "priority": 2}, "car": {"type": "vehicle", "priority": 3}, "bicycle": {"type": "vehicle", "priority": 3}, "motorcycle": {"type": "vehicle", "priority": 3}, "license plate": {"type": "vehicle", "priority": 2}, "key": {"type": "vehicle", "priority": 3}, "person": {"type": "person", "priority": 3} } # Visualization try: self.font = ImageFont.truetype("arial.ttf", 12) except: try: self.font = ImageFont.truetype("LiberationSans-Regular.ttf", 12) except: self.font = ImageFont.load_default() self.colors = { "weapon": "red", "electronic": "blue", "biological": "green", "drug":"orange", "person": "purple", "document":"black", "tool":"brown", "clothing":"pink", "vehicle":"gold", "default": "yellow" } # ====================== # 2. CORE FUNCTIONALITY # ====================== def classify_scene(self, image: Image.Image): """ Classifies the crime scene image using the CLIP model and the trained crime scene classification model. """ try: inputs = self.clip_processor(images=image, return_tensors="pt").to(DEVICE) with torch.no_grad(): image_features = self.clip_model.get_image_features(**inputs) # Newer transformers may return BaseModelOutputWithPooling if hasattr(image_features, 'pooler_output'): image_features = image_features.pooler_output image_features = image_features / image_features.norm(p=2, dim=-1, keepdim=True) features = image_features.cpu().numpy() pred = self.crime_scene_model.predict(features) label = self.label_encoder.inverse_transform(pred) return label[0] except Exception as e: print(f"Error classifying scene: {e}") traceback.print_exc() return "Undetermined" def analyze_scene(self, image_pil: Image.Image): """ Analyzes the crime scene image, generates analysis data, and a PDF report. Returns a dictionary containing both for API consumption. """ try: image = image_pil # Classify the scene crime_type = self.classify_scene(image) # Detection and classification detections = self._detect_objects(image) evidence = self._classify_evidence(image, detections) # Generate report text (for GPT-2 and PDF) if evidence: report_text = self._generate_report(evidence, image.size) else: report_text = "No significant forensic evidence detected. Detected objects:\n" + \ "\n".join(f"- {d['label']} (confidence: {d['score']:.2f})" for d in detections if d['score'] > 0.4) # Generate visualization (for PDF embedding) visualization = self._visualize_results(image.copy(), evidence if evidence else detections) # Prepare data for React Native Frontend (AnalysisResult type) scene_summary = "AI analysis complete. Refer to the 'Digital Forensic Analysis' section or the full PDF report for detailed findings." frontend_evidence = [] if evidence: for i, item in enumerate(evidence): # For `item.image` in React Native, you would typically save # a cropped thumbnail or provide a URL to a pre-generated visualization. # For this example, we'll leave it empty. frontend_evidence.append({ "id": i, "type": item['type'].capitalize(), "confidence": item['score'], # Approximate location string from bounding box "location": f"X:{int(item['box'][0])}-Y:{int(item['box'][1])} (W:{int(item['box'][2]-item['box'][0])}, H:{int(item['box'][3]-item['box'][1])})", "image": "" # Placeholder: No cropped image sent directly in this example }) else: scene_summary = "No primary forensic evidence detected. General object detection was performed. Detailed analysis is in the report." # Determine crime type and recommendations based on evidence and scene classification crime_type_primary = crime_type crime_type_confidence = 0.8 # Adjust confidence based on model performance crime_type_secondary = ["General Investigation", "Scene Documentation"] if any(item.get("type") == "weapon" for item in evidence): crime_type_primary = "Assault/Homicide Related" crime_type_confidence = 0.9 crime_type_secondary.insert(0, "Weapon Related Incident") if any(item.get("type") == "drug" for item in evidence): crime_type_primary = "Drug Related Incident" crime_type_confidence = max(crime_type_confidence, 0.8) crime_type_secondary.insert(0, "Substance Abuse Investigation") if any(item.get("type") == "electronic" for item in evidence): if "Digital Forensics" not in crime_type_secondary: crime_type_secondary.append("Digital Forensics Required") recommendations = [ "Ensure scene is fully secured and all personnel wear appropriate PPE.", "Prioritize collection of high-priority evidence (weapons, biological).", "Carefully document and photograph all evidence 'in situ' before collection.", "Establish a strict chain of custody for every item collected.", "Consider requesting specialized forensic units for biological or digital evidence.", "Review immediate vicinity for additional evidence or potential witnesses/CCTV." ] analysis_data_for_frontend = { "sceneSummary": report_text, # Using the generated report text as scene summary "evidenceDetected": frontend_evidence, "crimeType": { "primary": crime_type_primary, "confidence": crime_type_confidence, "secondary": list(set(crime_type_secondary)) # Remove duplicates }, "recommendations": recommendations, } # Generate PDF bytes in memory pdf_bytes = self._generate_pdf_report_in_memory( evidence=evidence or [], report_text=report_text, visualization_img=visualization ) return { "analysis_data": analysis_data_for_frontend, "pdf_bytes": pdf_bytes } except Exception as e: print(f"Analysis failed in CrimeSceneAnalyzer: {str(e)}") traceback.print_exc() return {"error": f"Analysis failed: {str(e)}", "analysis_data": None, "pdf_bytes": None} # ====================== # 3. HELPER METHODS (These remain largely unchanged) # ====================== def _load_image(self, source): # Not directly used by the API endpoint, but good to keep if source.startswith(('http:', 'https:')): return Image.open(requests.get(source, stream=True).raw) return Image.open(source) def _detect_objects(self, image): img_array = np.array(image) results = self.yolo_model(img_array) detections = [] for result in results: for detection in result.boxes.data.tolist(): x1, y1, x2, y2, confidence, class_id = detection label = self.yolo_model.names[int(class_id)] detections.append({ "label": label, "score": confidence, "box": [x1, y1, x2, y2] }) print(f"\nDetection results (showing >40% confidence):") for obj in detections: if obj["score"] > 0.4: print(f"- {obj['label']}: {obj['score']:.2f} at {[round(x,1) for x in obj['box']]}") return detections def _classify_evidence(self, image, detections): # Using the instance's EVIDENCE_MAP evidence = [] for obj in detections: obj_name = obj["label"].lower() matched = False for key in self.EVIDENCE_MAP: # Use self.EVIDENCE_MAP if key in obj_name: evidence.append({ **obj, **self.EVIDENCE_MAP[key], # Use self.EVIDENCE_MAP "exact_match": key == obj_name }) matched = True break if not matched and obj["score"] > 0.1: print(f"⚠️ Unclassified object: {obj_name} (score: {obj['score']:.2f})") return sorted(evidence, key=lambda x: (-x["priority"], -x["score"])) def _generate_report(self, evidence, image_size): # Renamed for clarity if you keep separate summary if not self.gemini_model: return "Report generation skipped: Gemini API key not configured." if not evidence: return "No evidence provided to generate a report." evidence_text = "\n".join( f"- Object: {e.get('label', 'Unknown')}\n Type: {e.get('type', 'N/A').upper()}\n Confidence: {e.get('score', 0):.0%}\n Assessed Priority: {e.get('priority', 'N/A')}" for e in evidence ) # You can add image_size or other context to the prompt if needed # image_context = f"The analysis was performed on an image of size: {image_size[0]}x{image_size[1]} pixels." prompt = f""" You are a professional detective providing a detailed forensic analysis report. Based *only* on the evidence items listed below, generate a plausible narrative of what could have happened at the scene. Structure your report clearly. Be objective and stick to the facts presented by the evidence. Do not invent evidence not listed. EVIDENCE FOUND: {evidence_text} ANALYSIS OF WHAT COULD HAVE HAPPENED: """ try: print("Generating detailed report with Gemini...") response = self.gemini_model.generate_content(prompt, request_options={'timeout': 15}) report_text = "" if response.parts: report_text = "".join(part.text for part in response.parts if hasattr(part, 'text')) elif response.candidates and response.candidates[0].content and response.candidates[0].content.parts: report_text = "".join(part.text for part in response.candidates[0].content.parts if hasattr(part, 'text')) if not report_text.strip(): return "Gemini generated an empty report for the provided evidence." # The prompt is already part of the generated text by Gemini, # so we usually don't need to prepend it. # If Gemini's response *only* contains the analysis part, and not the "EVIDENCE FOUND" preamble, # then you might want to structure the final return differently. # For now, let's assume Gemini's response is the full text you need. return report_text except Exception as e: print(f"‼️ Error calling Gemini API for report generation: {str(e)}") return f"Detailed report generation failed due to an API error: {str(e)}" def _visualize_results(self, image, items): draw = ImageDraw.Draw(image) for item in items: box = [int(b) for b in item["box"]] # Ensure integers for drawing if len(box) != 4: continue # Skip malformed boxes if "type" in item: color = self.colors.get(item["type"], self.colors["default"]) label = f"{item['label']} ({item['score']:.0%})" else: color = "gray" label = f"{item['label']} ({item['score']:.2f})" draw.rectangle(box, outline=color, width=3) try: # Use textbbox for modern Pillow bbox = draw.textbbox((0, 0), label, font=self.font) text_width = bbox[2] - bbox[0] text_height = bbox[3] - bbox[1] except AttributeError: # Fallback for older versions text_width, text_height = draw.textsize(label, font=self.font) draw.rectangle( [box[0], box[1], box[0] + text_width + 4, box[1] + text_height], fill=color ) draw.text( (box[0] + 2, box[1]), label, fill="white", font=self.font ) return image def _generate_pdf_report_in_memory(self, evidence, report_text, visualization_img): try: pdf = FPDF(format='A4', unit='mm') pdf.set_auto_page_break(auto=True, margin=15) pdf.add_page() # Header logo_path = "SceneX_logo.png" # Make sure this file exists in the same directory or provide full path if os.path.exists(logo_path): pdf.image(logo_path, x=10, y=8, w=25, h=25) pdf.set_font("Courier", 'B', 18) pdf.set_text_color(0, 51, 102) pdf.cell(0, 20, "OFFICIAL CRIME SCENE REPORT", ln=1, align='C') pdf.set_font("Courier", '', 12) pdf.set_text_color(0, 0, 0) pdf.cell(0, 6, f"Case Reference: CSR-{datetime.now().strftime('%Y%m%d-%H%M%S')}", ln=1) pdf.cell(0, 6, f"Report Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=1) pdf.ln(15) # Visualization img = visualization_img.convert('RGB') img_buffer = io.BytesIO() img.save(img_buffer, format="JPEG", quality=95) img_buffer.seek(0) pdf.image(img_buffer, x=12, y=pdf.get_y()+2, w=186) pdf.ln(125) # Evidence Catalog pdf.set_font("Courier", 'B', 14) pdf.set_fill_color(220, 230, 242) pdf.cell(0, 10, "EVIDENCE CATALOG", ln=1, fill=True) pdf.ln(5) if evidence: pdf.set_font("Courier", 'B', 12) col_widths = [70, 40, 40, 40] headers = ["Item Description", "Evidence Type", "Confidence", "Priority"] for i, header in enumerate(headers): pdf.cell(col_widths[i], 8, header, border=1, fill=True) pdf.ln() pdf.set_font("Courier", '', 10) row_fill = False for item in evidence: pdf.set_fill_color(255, 255, 255) if row_fill else pdf.set_fill_color(240, 240, 240) pdf.cell(col_widths[0], 8, item['label'], border=1, fill=True) pdf.cell(col_widths[1], 8, item['type'].upper(), border=1, fill=True) pdf.cell(col_widths[2], 8, f"{item['score']:.0%}", border=1, fill=True) pdf.cell(col_widths[3], 8, str(item['priority']), border=1, ln=1, fill=True) row_fill = not row_fill else: pdf.set_font("Courier", 'I', 12) pdf.set_text_color(150, 150, 150) pdf.cell(0, 8, "No specific forensic evidence categorized by the AI model.", ln=1, align='C') pdf.set_text_color(0, 0, 0) pdf.ln(10) pdf.set_font("Courier", 'B', 14) pdf.set_text_color(0, 51, 102) pdf.cell(0, 10, "DIGITAL FORENSIC ANALYSIS", ln=1) pdf.set_draw_color(200, 200, 200) pdf.line(10, pdf.get_y(), 200, pdf.get_y()) pdf.ln(5) pdf.set_font("Courier", '', 12) cleaned_report_text = report_text.encode('latin-1', 'replace').decode('latin-1') paragraphs = cleaned_report_text.split('\n') for para in paragraphs: if para.strip(): pdf.multi_cell(0, 6, para) pdf.ln(3) # Footer pdf.set_y(-20) pdf.set_font("Courier", 'I', 8) pdf.set_text_color(100, 100, 100) pdf.cell(0, 5, "CONFIDENTIAL - Law Enforcement Use Only", 0, 0, 'L') pdf.cell(0, 5, f"Page {pdf.page_no()}", 0, 0, 'R') print("✓ Professional report generated in-memory.") return pdf.output() except Exception as e: print(f"!! Critical PDF generation error: {str(e)}") traceback.print_exc() return None def _add_watermark(self, img): try: draw = ImageDraw.Draw(img) font = ImageFont.load_default() text = f"SceneX Analysis {datetime.now().strftime('%Y-%m-%d')}" for i in range(0, img.width, 200): for j in range(0, img.height, 200): draw.text((i, j), text, fill=(200, 200, 200, 128), font=font) return img except: return img