CrimeSceneAnalyzer.py

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