Spaces:
Running
Running
File size: 5,102 Bytes
1287b01 |
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 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
import os
import cv2
import json
from flask import Flask, request, jsonify, render_template, send_from_directory
from flask_cors import CORS
from werkzeug.utils import secure_filename
from ultralytics import YOLO
app = Flask(__name__)
CORS(app)
# Configuration
UPLOAD_FOLDER = 'uploads'
RESULT_FOLDER = 'static/results'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'mp4', 'avi', 'mov'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['RESULT_FOLDER'] = RESULT_FOLDER
# Create directories if they don't exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(RESULT_FOLDER, exist_ok=True)
# Load YOLOv8 model (pre-trained on COCO)
# We use the nano version for speed in this environment
model = YOLO('yolov8n.pt')
def allowed_file(filename):
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
@app.route('/')
def index():
return render_template('index.html')
@app.route('/detect', methods=['POST'])
def detect_objects():
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 and allowed_file(file.filename):
filename = secure_filename(file.filename)
filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
file.save(filepath)
# Determine file type
ext = filename.rsplit('.', 1)[1].lower()
if ext in {'mp4', 'avi', 'mov'}:
return process_video(filepath, filename)
else:
return process_image(filepath, filename)
return jsonify({'error': 'File type not allowed'}), 400
def process_image(filepath, filename):
# Run inference
results = model(filepath)
# Get the first result (since we only processed one image)
result = results[0]
# Save the plotted image (with bounding boxes)
res_filename = 'res_' + filename
res_path = os.path.join(app.config['RESULT_FOLDER'], res_filename)
result.save(filename=res_path)
# Count objects
detections = []
counts = {}
for box in result.boxes:
cls_id = int(box.cls[0])
label = model.names[cls_id]
conf = float(box.conf[0])
detections.append({
'label': label,
'confidence': conf,
'box': box.xyxy[0].tolist()
})
counts[label] = counts.get(label, 0) + 1
return jsonify({
'success': True,
'type': 'image',
'result_url': f'/static/results/{res_filename}',
'objects': detections,
'counts': counts,
'total_count': len(detections)
})
def process_video(filepath, filename):
# For video, we'll run inference and save the output video
# Note: Processing a full video can be slow.
# We'll use ultralytics' built-in video saving capability for simplicity
res_filename = 'res_' + filename.split('.')[0] + '.mp4'
res_path = os.path.join(app.config['RESULT_FOLDER'], res_filename)
# Run inference on the video
# Use project/name to control output location
results = model(filepath, stream=True)
# We need to collect overall counts for the video
all_seen_objects = {}
# To keep it simple and faster for the demo, we process every 5th frame if it's long?
# No, let's just use the built-in save for now.
# Run model.predict with save=True
save_results = model.predict(filepath, save=True, project=app.config['RESULT_FOLDER'], name='vid_temp', exist_ok=True)
# Find the saved video. Ultralytics saves it in a subfolder.
# We want to move it to our RESULT_FOLDER with a predictable name.
# Typically it goes to RESULT_FOLDER/vid_temp/filename
raw_saved_path = os.path.join(app.config['RESULT_FOLDER'], 'vid_temp', filename)
if os.path.exists(raw_saved_path):
import shutil
shutil.move(raw_saved_path, res_path)
# Cleanup temp dir
shutil.rmtree(os.path.join(app.config['RESULT_FOLDER'], 'vid_temp'))
else:
# Fallback if names differ (sometimes .avi becomes .mp4 etc)
# Just check the folder
temp_dir = os.path.join(app.config['RESULT_FOLDER'], 'vid_temp')
if os.path.exists(temp_dir):
files = os.listdir(temp_dir)
if files:
shutil.move(os.path.join(temp_dir, files[0]), res_path)
shutil.rmtree(temp_dir)
# For video summary, we'll just run a quick pass to get unique objects
# (In a real app, you'd aggregate frame-by-frame)
# Just return success for now with the URL
return jsonify({
'success': True,
'type': 'video',
'result_url': f'/static/results/{res_filename}',
'message': 'Video processed successfully'
})
if __name__ == '__main__':
app.run(debug=True, port=5000)
|