Spaces:
Running
Running
File size: 4,342 Bytes
f284af1 cd5b039 f284af1 3e4904b f284af1 |
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 |
from flask import Flask, render_template, request, redirect, url_for, jsonify
import cv2
import numpy as np
from tensorflow.lite.python.interpreter import Interpreter
import os
# Define paths to your model and label files
MODEL_PATH = "detect.tflite"
LABEL_PATH = "labelmap.txt"
# Function to load the TFLite model and labels
def load_model():
interpreter = Interpreter(model_path=MODEL_PATH)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
height = input_details[0]['shape'][1]
width = input_details[0]['shape'][2]
with open(LABEL_PATH, 'r') as f:
labels = [line.strip() for line in f.readlines()]
print(f"Model loaded. Input shape: {input_details[0]['shape']}")
return interpreter, input_details, output_details, height, width, labels
# Function to preprocess the image for the model
def preprocess_image(image, input_details, height, width):
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_resized = cv2.resize(image_rgb, (width, height))
input_data = np.expand_dims(image_resized, axis=0)
if input_details[0]['dtype'] == np.float32:
input_data = (np.float32(input_data) - 127.5) / 127.5
print(f"Image preprocessed: shape {input_data.shape}, dtype {input_data.dtype}")
return input_data
# Function to perform object detection and draw bounding boxes
def detect_objects(image, interpreter, input_details, output_details, labels):
input_data = preprocess_image(image, input_details, height, width)
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
boxes = interpreter.get_tensor(output_details[1]['index'])[0] # bounding box coordinates
classes = interpreter.get_tensor(output_details[3]['index'])[0] # class index
scores = interpreter.get_tensor(output_details[0]['index'])[0] # confidence scores
print(f"Detections: {len(scores)} objects detected")
for i in range(len(scores)):
if scores[i] > 0.1: # confidence threshold
ymin, xmin, ymax, xmax = boxes[i]
ymin = int(max(1, ymin * image.shape[0]))
xmin = int(max(1, xmin * image.shape[1]))
ymax = int(min(image.shape[0], ymax * image.shape[0]))
xmax = int(min(image.shape[1], xmax * image.shape[1]))
cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
label = f'{labels[int(classes[i])]}: {scores[i] * 100:.2f}%'
cv2.putText(image, label, (xmin, ymin - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
print(f"Object {i}: {label} at [{xmin}, {ymin}, {xmax}, {ymax}]")
return image
# Initialize the Flask app
app = Flask(__name__, static_folder='static')
# Load the TFLite model and labels
interpreter, input_details, output_details, height, width, labels = load_model()
@app.route('/', methods=['GET', 'POST'])
def upload_and_detect():
if request.method == 'POST':
if 'file' not in request.files:
print("No file part in the request")
return redirect(request.url)
file = request.files['file']
if file.filename == '':
print("No selected file")
return redirect(request.url)
# Read the image file
image = cv2.imdecode(np.frombuffer(file.read(), np.uint8), cv2.IMREAD_COLOR)
if image is None:
print("Failed to read image")
return redirect(request.url)
print(f"Image uploaded: {file.filename}, shape: {image.shape}")
# Perform object detection
processed_image = detect_objects(image, interpreter, input_details, output_details, labels)
# Ensure the static directory exists
if not os.path.exists(app.static_folder):
os.makedirs(app.static_folder)
# Save processed image
save_path = os.path.join(app.static_folder, 'detected.jpg')
cv2.imwrite(save_path, processed_image)
print(f"Processed image saved at: {save_path}")
# Send back the path to the processed image
return jsonify({'image_url': url_for('static', filename='detected.jpg')})
return render_template('index.html')
if __name__ == '__main__':
app.run(host='0.0.0.0', port=8000)
|