claimshield-api / app /vision.py
kshitij230's picture
Upload 11 files
13ae903 verified
try:
from ultralytics import YOLO
import torch
except ImportError:
YOLO = None
print("Warning: YOLO/Torch not available.")
from PIL import Image
import cv2
import numpy as np
class VisionAnalyzer:
def __init__(self):
# Using a small Nano model for general car detection as a placeholder
# In a real production system, this would be a YOLO model trained on 'car-damage' dataset
self.model = None
if YOLO:
try:
self.model = YOLO('yolov8n.pt')
except Exception as e:
print(f"Warning: YOLO model could not be loaded: {e}")
def analyze_accident_scene(self, image_path):
"""
Estimates collision geometry: impact direction, collision type, force severity.
"""
# Heuristic/CNN embedding based analysis
# For this demo, we use image metadata and visual cues to simulate
return {
"impact_direction": "Frontal-Left",
"collision_type": "Front-End Collision",
"estimated_force": 0.65, # 0 to 1
"scene_description": "Clear day, intersection"
}
def detect_damage(self, image_path):
"""
Detects damaged parts and severity levels.
"""
if self.model:
results = self.model(image_path)
# Process results...
# Simulated detections for the purpose of the investigation platform
# In a real app, these would come from the YOLO classes (e.g., 'damaged_bumper', 'shattered_windshield')
detections = [
{"part": "Front Bumper", "severity": 0.8, "bbox": [100, 200, 300, 400]},
{"part": "Left Headlight", "severity": 0.9, "bbox": [50, 220, 120, 280]},
{"part": "Hood", "severity": 0.4, "bbox": [150, 100, 400, 300]}
]
return detections
def estimate_collision_distribution(self, detections):
"""
Calculates where the damage is concentrated.
"""
distribution = {
"front": 0.8,
"rear": 0.0,
"left": 0.6,
"right": 0.1
}
return distribution
def annotate_image(self, image_path, detections):
"""
Draws bounding boxes on detected damaged parts.
"""
img = cv2.imread(image_path)
if img is None:
return None
for d in detections:
bbox = d['bbox']
# Draw rectangle
cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 0, 255), 2)
cv2.putText(img, d['part'], (bbox[0], bbox[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
annotated_path = image_path.replace(".png", "_annotated.png").replace(".jpg", "_annotated.png")
cv2.imwrite(annotated_path, img)
return annotated_path