File size: 4,627 Bytes
b338c44 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | import os
import torch
import torchvision.transforms as transforms
import torchvision.models.detection as detection
from flask import Flask, request, render_template, send_from_directory, url_for, Response
from PIL import Image
import cv2
import numpy as np
# Initialize Flask app
app = Flask(__name__)
UPLOAD_FOLDER = os.path.abspath(os.path.join("deployment", "static", "uploads"))
OUTPUT_FOLDER = os.path.abspath(os.path.join("deployment", "outputs"))
# Ensure directories exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)
# Define class names
CLASS_NAMES = {
0: "vehicle", 1: "bicycle", 2: "bus", 3: "car", 4: "lorry"
}
NUM_CLASSES = len(CLASS_NAMES) + 1 # Include background class
# Load Faster R-CNN model
def load_model(model_path, num_classes):
model = detection.fasterrcnn_resnet50_fpn(weights=None, num_classes=num_classes)
model.load_state_dict(torch.load(model_path, map_location='cpu'))
model.eval()
return model
model_path = os.path.join("models", "fasterrcnn_model.pth")
model = load_model(model_path, NUM_CLASSES)
# Image preprocessing
def preprocess_image(image):
transform = transforms.Compose([transforms.ToTensor()])
return transform(image).unsqueeze(0)
# Draw predictions
def draw_predictions(image, prediction):
boxes, scores, labels = prediction['boxes'], prediction['scores'], prediction['labels']
for i in range(len(boxes)):
if scores[i] > 0.5: # Confidence threshold
x1, y1, x2, y2 = map(int, boxes[i].tolist())
class_name = CLASS_NAMES.get(labels[i].item(), f"Class {labels[i].item()}")
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(image, f"{class_name}: {scores[i]:.2f}", (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
return image
# Home route (upload and detect)
@app.route('/', methods=['GET', 'POST'])
def upload_and_predict():
if request.method == 'POST':
file = request.files['file']
if file:
filename = file.filename
upload_path = os.path.join(UPLOAD_FOLDER, filename)
result_path = os.path.join(OUTPUT_FOLDER, filename)
file.save(upload_path)
print(f"✅ Uploaded file saved at: {upload_path}")
# Run inference
img = Image.open(upload_path).convert("RGB")
img_tensor = preprocess_image(img)
with torch.no_grad():
prediction = model(img_tensor)[0]
# Draw results
image_cv = cv2.imread(upload_path)
if image_cv is None:
print(f"❌ ERROR: OpenCV could not read {upload_path}")
else:
result_image = draw_predictions(image_cv, prediction)
cv2.imwrite(result_path, result_image)
print(f"✅ Result image saved at: {result_path}")
return render_template('index.html',
uploaded_image=url_for('uploaded_file', filename=filename),
result_image=url_for('output_file', filename=filename))
return render_template('index.html')
# Serve uploaded images correctly
@app.route('/uploads/<filename>')
def uploaded_file(filename):
return send_from_directory(os.path.abspath(UPLOAD_FOLDER), filename)
# Serve output images correctly
@app.route('/outputs/<filename>')
def output_file(filename):
return send_from_directory(os.path.abspath(OUTPUT_FOLDER), filename)
# Real-time webcam detection
def generate_frames():
cap = cv2.VideoCapture(0)
while cap.isOpened():
success, frame = cap.read()
if not success:
break
# Convert frame and preprocess
img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
img_tensor = preprocess_image(img)
# Run inference
with torch.no_grad():
prediction = model(img_tensor)[0]
# Draw results
frame = draw_predictions(frame, prediction)
# Convert to JPEG
_, buffer = cv2.imencode('.jpg', frame)
frame_bytes = buffer.tobytes()
yield (b'--frame\r\n'
b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')
cap.release()
@app.route('/video_feed')
def video_feed():
return Response(generate_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')
if __name__ == '__main__':
app.run(debug=True)
|