58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
|
from flask import Blueprint, request, jsonify, send_file, render_template
|
||
|
|
from app.utils.detector import MemoryDetector
|
||
|
|
import os
|
||
|
|
from PIL import Image
|
||
|
|
import io
|
||
|
|
|
||
|
|
main_bp = Blueprint('main', __name__)
|
||
|
|
detector = MemoryDetector()
|
||
|
|
|
||
|
|
@main_bp.route('/')
|
||
|
|
def index():
|
||
|
|
return render_template('index.html')
|
||
|
|
|
||
|
|
@main_bp.route('/detect', methods=['POST'])
|
||
|
|
def detect_memory():
|
||
|
|
if 'image' not in request.files:
|
||
|
|
return jsonify({'error': 'No image provided'}), 400
|
||
|
|
|
||
|
|
file = request.files['image']
|
||
|
|
|
||
|
|
# Read the image
|
||
|
|
img = Image.open(file.stream)
|
||
|
|
|
||
|
|
# Process the image and get annotated image and detections
|
||
|
|
annotated_img, detections = detector.detect(img)
|
||
|
|
|
||
|
|
# Convert PIL image to bytes
|
||
|
|
img_byte_arr = io.BytesIO()
|
||
|
|
annotated_img.save(img_byte_arr, format='PNG')
|
||
|
|
img_byte_arr.seek(0)
|
||
|
|
|
||
|
|
return send_file(
|
||
|
|
img_byte_arr,
|
||
|
|
mimetype='image/png'
|
||
|
|
)
|
||
|
|
|
||
|
|
@main_bp.route('/detect/test', methods=['GET'])
|
||
|
|
def detect_test():
|
||
|
|
"""Endpoint for testing with a hardcoded image"""
|
||
|
|
# Using an existing image from the validation set
|
||
|
|
test_image_path = os.path.join('training', 'val', 'images', 'memory_out19.png')
|
||
|
|
|
||
|
|
if not os.path.exists(test_image_path):
|
||
|
|
return jsonify({'error': f'Test image not found at {test_image_path}'}), 404
|
||
|
|
|
||
|
|
img = Image.open(test_image_path)
|
||
|
|
# Get both the annotated image and detections
|
||
|
|
annotated_img, detections = detector.detect(img)
|
||
|
|
|
||
|
|
img_byte_arr = io.BytesIO()
|
||
|
|
annotated_img.save(img_byte_arr, format='PNG')
|
||
|
|
img_byte_arr.seek(0)
|
||
|
|
|
||
|
|
return send_file(
|
||
|
|
img_byte_arr,
|
||
|
|
mimetype='image/png'
|
||
|
|
)
|