yolo_model / app.py
Harika12323's picture
Update app.py
78bdbf3 verified
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
}
@app.route('/upload', methods=['POST'])
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)