Spaces:
Sleeping
Sleeping
from flask import Flask, request, jsonify | |
import cv2 | |
import numpy as np | |
import torch | |
import yolov5 | |
from yolov5 import YOLOv5 | |
app = Flask(__name__) | |
# Load YOLOv5 model | |
model = YOLOv5('yolov5s.pt') # Replace with your model path | |
def detect_number_plate(frame): | |
results = model(frame) | |
detections = results.pandas().xyxy[0] | |
plates = [] | |
for _, row in detections.iterrows(): | |
if row['name'] == 'number_plate': # Adjust class name | |
plates.append({ | |
'class': row['name'], | |
'confidence': row['confidence'], | |
'x_min': row['xmin'], | |
'y_min': row['ymin'], | |
'x_max': row['xmax'], | |
'y_max': row['ymax'] | |
}) | |
return plates | |
def detect_smoke(frame): | |
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) | |
blur = cv2.GaussianBlur(gray, (21, 21), 0) | |
_, thresh = cv2.threshold(blur, 200, 255, cv2.THRESH_BINARY) | |
smoke_intensity = np.sum(thresh) / (thresh.shape[0] * thresh.shape[1]) | |
smoke_detected = smoke_intensity > 0.1 # Adjust this threshold | |
return smoke_detected, smoke_intensity | |
def process_frame(frame): | |
plates = detect_number_plate(frame) | |
smoke_detected, smoke_intensity = detect_smoke(frame) | |
return { | |
'smoke_detected': smoke_detected, | |
'smoke_intensity': smoke_intensity, | |
'number_plates': plates | |
} | |
def upload_file(): | |
if 'file' not in request.files: | |
return jsonify({'error': 'No file part'}), 400 | |
file = request.files['file'] | |
if file.filename == '': | |
return jsonify({'error': 'No selected file'}), 400 | |
if file: | |
in_memory_file = file.read() | |
np_arr = np.frombuffer(in_memory_file, np.uint8) | |
frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR) | |
results = process_frame(frame) | |
return jsonify(results) | |
if __name__ == '__main__': | |
app.run(port=5000, debug=True) | |