File size: 7,021 Bytes
48bc30e 0da364f 48bc30e |
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 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 |
import os
import numpy as np
import cv2
from ultralytics import SAM
from ultralytics.utils.downloads import safe_download
from PIL import Image
from flask import render_template, request, redirect, url_for, send_from_directory
from werkzeug.utils import secure_filename
from app import app
# Configurations
UPLOAD_FOLDER = '/tmp/uploads'
MODEL_DIR = os.path.join(app.root_path,'models')
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
os.makedirs(UPLOAD_FOLDER,exist_ok=True)
os.makedirs(MODEL_DIR,exist_ok=True)
MODEL_PATH = os.path.join(MODEL_DIR,'mobile_sam.pt')
if not os.path.exists(MODEL_PATH):
safe_download(url="https://github.com/ultralytics/assets/releases/download/v8.2.0/mobile_sam.pt", dir=MODEL_DIR)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
model = SAM(MODEL_PATH)
# File handling functions
def clear_uploads_folder():
"""
Clear all files in the uploads folder
"""
if os.path.exists(UPLOAD_FOLDER):
for filename in os.listdir(UPLOAD_FOLDER):
file_path = os.path.join(UPLOAD_FOLDER, filename)
try:
if os.path.isfile(file_path) or os.path.islink(file_path):
os.unlink(file_path)
except Exception as e:
print(f"Error deleting file {file_path}: {e}")
def is_allowed_file(filename):
"""
Checking if a file has an allowed extension
"""
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
def save_file(file, filename):
"""
Save the uploaded file to the uploads folder and return the file path.
"""
file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
file.save(file_path)
return file_path
def generate_preview(x, y, filename):
"""
Generate a preview of the selected region of interest using SAM.
Returns the preview image path or None if unsuccessful.
"""
image_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
image = cv2.imread(image_path)
if image is None:
return None
x_int = int(round(x))
y_int = int(round(y))
results = model.predict(source=image_path, points=[x_int, y_int], labels=[1], verbose=False)
if not results or not results[0].masks or len(results[0].masks) == 0:
return None
try:
mask = results[0].masks.data[0].cpu().numpy()
except (IndexError, AttributeError):
return None
mask = (mask < 0.5).astype(np.uint8) * 255
h, w = image.shape[:2]
if mask.shape != (h, w):
mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
overlay = np.zeros_like(image, dtype=np.uint8)
overlay[mask == 255] = (0, 0, 255) # Red color
alpha = 0.5
preview_image = cv2.addWeighted(image, 1 - alpha, overlay, alpha, 0.0)
preview_image_path = "preview_" + filename
preview_path = os.path.join(app.config['UPLOAD_FOLDER'], preview_image_path)
cv2.imwrite(preview_path, preview_image)
return preview_image_path
def generate_segment(x, y, filename):
"""
Segment the original image for the final result using SAM.
Returns the original filename and segmented filename with alpha channel, or (filename, None) if unsuccessful.
"""
image_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
image = cv2.imread(image_path)
image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
if image is None:
return filename, None
x_int = int(round(x))
y_int = int(round(y))
results = model.predict(source=image_path, points=[x_int, y_int], labels=[1], verbose=False)
if not results or not results[0].masks or len(results[0].masks) == 0:
return filename, None
try:
mask = results[0].masks.data[0].cpu().numpy()
except (IndexError, AttributeError):
return filename, None
mask = (mask < 0.5).astype(np.uint8) * 255
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_rgba = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2RGBA)
image_rgba[:, :, 3] = mask # Set alpha channel
segmented_filename = "segmented_" + filename
segmented_path = os.path.join(app.config['UPLOAD_FOLDER'], segmented_filename)
cv2.imwrite(segmented_path, image_rgba)
return filename, segmented_filename
# Flask Routes
@app.route('/', methods=['GET', 'POST'])
def index():
# Clear the uploads folder when showing the upload card
if request.method == 'GET':
clear_uploads_folder()
if request.method == 'POST':
if 'file' not in request.files:
return render_template('main.html', error="No file part")
file = request.files['file']
if file.filename == '':
return render_template('main.html', error="No selected file")
if file and is_allowed_file(file.filename):
filename = secure_filename(file.filename)
save_file(file, filename)
return render_template('main.html', uploaded_file=filename)
return render_template('main.html')
@app.route('/preview', methods=['POST'])
def preview():
x = request.form.get('x', '')
y = request.form.get('y', '')
filename = request.form.get('filename', '')
if not x or not y or not filename:
return render_template('main.html', uploaded_file=filename, error="Please select a region on the image")
try:
x = float(x)
y = float(y)
except ValueError:
return render_template('main.html', uploaded_file=filename, error="Invalid coordinates")
preview_file = generate_preview(x, y, filename)
if preview_file is None:
return render_template('main.html',uploaded_file=filename,error="There is no object in the selected region")
return render_template('main.html', uploaded_file=filename, preview_file=preview_file, x_coord=x, y_coord=y)
@app.route('/segment', methods=['POST'])
def segment():
x = request.form.get('x', '')
y = request.form.get('y', '')
filename = request.form.get('filename', '')
if not x or not y or not filename:
return render_template('main.html', uploaded_file=filename, error="Please select a region on the image")
try:
x = float(x)
y = float(y)
except ValueError:
return render_template('main.html', uploaded_file=filename, error="Invalid coordinates")
uploaded_file, result_file = generate_segment(x, y, filename)
return render_template('main.html', uploaded_file=uploaded_file, result_file=result_file)
@app.route('/reset', methods=['GET'])
def reset():
# Clear the uploads folder before returning to the upload phase
clear_uploads_folder()
return redirect(url_for('index'))
@app.route('/static/uploads/<filename>')
def uploaded_file(filename):
return send_from_directory(app.config['UPLOAD_FOLDER'], filename)
@app.route('/download/<filename>')
def download_file(filename):
return send_from_directory(app.config['UPLOAD_FOLDER'], filename, as_attachment=True)
@app.route('/about')
def about_page():
return render_template('about.html') |