126 lines
4.8 KiB
Python
126 lines
4.8 KiB
Python
|
import numpy as np
|
||
|
|
import cv2
|
||
|
|
from typing import List, Dict, Any, Tuple
|
||
|
|
|
||
|
|
class Model:
|
||
|
|
"""
|
||
|
|
Smoke detection model implementation
|
||
|
|
|
||
|
|
This is a simple example that could be replaced with an actual
|
||
|
|
TensorFlow, PyTorch, or other ML framework implementation.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
"""Initialize smoke detection model"""
|
||
|
|
# In a real implementation, you would load your model here
|
||
|
|
print("Smoke detection model initialized")
|
||
|
|
|
||
|
|
# Define smoke class IDs
|
||
|
|
self.smoke_classes = {
|
||
|
|
0: "smoke",
|
||
|
|
1: "fire"
|
||
|
|
}
|
||
|
|
|
||
|
|
def preprocess(self, image: np.ndarray) -> np.ndarray:
|
||
|
|
"""Preprocess image for model input"""
|
||
|
|
# Convert BGR to grayscale for smoke detection
|
||
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
|
|
# Convert back to 3 channels to match model expected input shape
|
||
|
|
gray_3ch = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
|
||
|
|
|
||
|
|
# In a real implementation, you would do normalization, etc.
|
||
|
|
return gray_3ch
|
||
|
|
|
||
|
|
def predict(self, image: np.ndarray) -> List[Dict[str, Any]]:
|
||
|
|
"""
|
||
|
|
Run smoke detection on the image
|
||
|
|
|
||
|
|
This is a simplified example that uses basic image processing
|
||
|
|
In a real implementation, you would use your ML model
|
||
|
|
"""
|
||
|
|
# Convert to grayscale for processing
|
||
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
|
|
|
||
|
|
# Apply Gaussian blur to reduce noise
|
||
|
|
blurred = cv2.GaussianBlur(gray, (15, 15), 0)
|
||
|
|
|
||
|
|
# Simple thresholding to find potential smoke regions
|
||
|
|
# In a real implementation, you'd use a trained model
|
||
|
|
_, thresh = cv2.threshold(blurred, 100, 255, cv2.THRESH_BINARY)
|
||
|
|
|
||
|
|
# Find contours in the thresholded image
|
||
|
|
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
|
|
||
|
|
# Process contours to find potential smoke regions
|
||
|
|
detections = []
|
||
|
|
height, width = image.shape[:2]
|
||
|
|
|
||
|
|
for contour in contours:
|
||
|
|
# Get bounding box
|
||
|
|
x, y, w, h = cv2.boundingRect(contour)
|
||
|
|
|
||
|
|
# Filter small regions
|
||
|
|
if w > width * 0.05 and h > height * 0.05:
|
||
|
|
# Calculate area ratio
|
||
|
|
area = cv2.contourArea(contour)
|
||
|
|
rect_area = w * h
|
||
|
|
fill_ratio = area / rect_area if rect_area > 0 else 0
|
||
|
|
|
||
|
|
# Smoke tends to have irregular shapes
|
||
|
|
# This is just for demonstration purposes
|
||
|
|
if fill_ratio > 0.2 and fill_ratio < 0.8:
|
||
|
|
# Normalize coordinates
|
||
|
|
x_norm = x / width
|
||
|
|
y_norm = y / height
|
||
|
|
w_norm = w / width
|
||
|
|
h_norm = h / height
|
||
|
|
|
||
|
|
# Determine if it's smoke or fire (just a simple heuristic for demo)
|
||
|
|
# In a real model, this would be determined by the model prediction
|
||
|
|
class_id = 0 # Default to smoke
|
||
|
|
|
||
|
|
# Check if the region has high red values (fire)
|
||
|
|
roi = image[y:y+h, x:x+w]
|
||
|
|
if roi.size > 0: # Make sure ROI is not empty
|
||
|
|
avg_color = np.mean(roi, axis=(0, 1))
|
||
|
|
if avg_color[2] > 150 and avg_color[2] > avg_color[0] * 1.5: # High red, indicating fire
|
||
|
|
class_id = 1 # Fire
|
||
|
|
|
||
|
|
# Calculate confidence based on fill ratio
|
||
|
|
# This is just for demonstration
|
||
|
|
confidence = 0.5 + fill_ratio * 0.3
|
||
|
|
|
||
|
|
# Add to detections
|
||
|
|
detections.append({
|
||
|
|
'bbox': (x_norm, y_norm, w_norm, h_norm),
|
||
|
|
'class_id': class_id,
|
||
|
|
'confidence': confidence
|
||
|
|
})
|
||
|
|
|
||
|
|
# For demo purposes, if no smoke detected by algorithm,
|
||
|
|
# add a small chance of random detection
|
||
|
|
if not detections and np.random.random() < 0.1: # 10% chance
|
||
|
|
# Random smoke detection
|
||
|
|
x = np.random.random() * 0.7
|
||
|
|
y = np.random.random() * 0.7
|
||
|
|
w = 0.1 + np.random.random() * 0.2
|
||
|
|
h = 0.1 + np.random.random() * 0.2
|
||
|
|
confidence = 0.5 + np.random.random() * 0.3
|
||
|
|
|
||
|
|
detections.append({
|
||
|
|
'bbox': (x, y, w, h),
|
||
|
|
'class_id': 0, # Smoke
|
||
|
|
'confidence': confidence
|
||
|
|
})
|
||
|
|
|
||
|
|
return detections
|
||
|
|
|
||
|
|
@property
|
||
|
|
def applies_nms(self) -> bool:
|
||
|
|
"""Model does not apply NMS internally"""
|
||
|
|
return False
|
||
|
|
|
||
|
|
def close(self):
|
||
|
|
"""Release resources"""
|
||
|
|
# In a real implementation, you would release model resources here
|
||
|
|
pass
|