File size: 11,738 Bytes
d790e98 |
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 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 |
from flask import Flask, render_template, request, jsonify, send_file, Response, stream_with_context
from werkzeug.utils import secure_filename
import os
from pathlib import Path
import shutil
import io
import json
import logging
from backend.pipeline import classify
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
app = Flask(__name__)
# Configuration
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
UPLOAD_FOLDER_SINGLE = 'static/uploads/single'
UPLOAD_FOLDER_MULTIPLE = 'static/uploads/multiple'
app.config['UPLOAD_FOLDER_SINGLE'] = UPLOAD_FOLDER_SINGLE
app.config['UPLOAD_FOLDER_MULTIPLE'] = UPLOAD_FOLDER_MULTIPLE
app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024
app.config['MAX_CONTENT_LENGTH_REPORT'] = 500 * 1024 * 1024
# Ensure upload directories exist
os.makedirs(UPLOAD_FOLDER_SINGLE, exist_ok=True)
os.makedirs(UPLOAD_FOLDER_MULTIPLE, exist_ok=True)
def allowed_file(filename):
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
def clear_uploads(folder):
"""Helper function to clear upload directories."""
if os.path.exists(folder):
for filename in os.listdir(folder):
file_path = os.path.join(folder, filename)
try:
if os.path.isfile(file_path) or os.path.islink(file_path):
os.unlink(file_path)
elif os.path.isdir(file_path):
shutil.rmtree(file_path)
except Exception as e:
logger.error(f'Failed to delete {file_path}. Reason: {e}')
@app.route('/', defaults={'path': ''})
@app.route('/<path:path>')
def index(path):
if path.startswith('static/'):
return send_file(path)
return send_file('static/react/index.html')
@app.route('/upload_single', methods=['POST'])
def upload_single():
if 'file' not in request.files:
return jsonify({'error': 'No file uploaded'}), 400
file = request.files['file']
if file.filename == '':
return jsonify({'error': 'No file selected'}), 400
if file and allowed_file(file.filename):
filename = secure_filename(file.filename)
filepath = os.path.join(app.config['UPLOAD_FOLDER_SINGLE'], filename)
file.save(filepath)
return jsonify({'filename': filename})
return jsonify({'error': 'Invalid file type'}), 400
@app.route('/classify_single', methods=['POST'])
def classify_single():
data = request.get_json()
filename = data.get('filename')
if not filename:
return jsonify({'error': 'No filename provided'}), 400
filepath = os.path.join(app.config['UPLOAD_FOLDER_SINGLE'], filename)
if not os.path.exists(filepath):
return jsonify({'error': 'File not found'}), 404
try:
classification_result, detailed_results, failure_labels = classify(filepath)
return jsonify({
'classification': classification_result,
'detailed_results': detailed_results
})
except Exception as e:
logger.error(f"Error in classify_single: {e}")
return jsonify({'error': str(e)}), 500
@app.route('/upload_multiple', methods=['POST'])
def upload_multiple():
logger.info("=== UPLOAD_MULTIPLE CALLED ===")
if 'file' not in request.files:
logger.warning("No 'file' in request.files")
return jsonify({'error': 'No file uploaded'}), 400
files = request.files.getlist('file')
logger.info(f"Received {len(files)} files in request")
if not files:
return jsonify({'error': 'No files selected'}), 400
try:
# Ensure temp directory exists (DON'T wipe it)
temp_dir = os.path.join(app.config['UPLOAD_FOLDER_MULTIPLE'], 'temp')
os.makedirs(temp_dir, exist_ok=True)
saved_count = 0
filename_map = {}
# Load existing map if present
map_path = os.path.join(temp_dir, 'filename_map.json')
if os.path.exists(map_path):
try:
with open(map_path, 'r') as f:
filename_map = json.load(f)
logger.info(f"Loaded existing filename map with {len(filename_map)} entries")
except Exception as e:
logger.error(f"Error loading existing filename map: {e}")
for file in files:
if file and allowed_file(file.filename):
original_filename = file.filename
filename = secure_filename(file.filename)
filepath = os.path.join(temp_dir, filename)
file.save(filepath)
filename_map[filename] = original_filename
saved_count += 1
logger.info(f"Saved: '{original_filename}' -> '{filename}'")
# Save updated filename map
with open(map_path, 'w') as f:
json.dump(filename_map, f)
# Count actual files in directory
actual_files = [f for f in os.listdir(temp_dir) if allowed_file(f)]
logger.info(f"Total files in temp directory after upload: {len(actual_files)}")
logger.info(f"Files: {actual_files}")
return jsonify({
'message': f'Successfully uploaded {saved_count} files',
'count': saved_count,
'status': 'Ready'
})
except Exception as e:
logger.error(f"Error in upload_multiple: {e}", exc_info=True)
return jsonify({'error': str(e)}), 500
@app.route('/classify_multiple', methods=['POST'])
def classify_multiple():
logger.info("=== CLASSIFY_MULTIPLE CALLED ===")
def generate():
temp_dir = os.path.join(app.config['UPLOAD_FOLDER_MULTIPLE'], 'temp')
if not os.path.exists(temp_dir):
logger.warning("Temp directory does not exist")
yield json.dumps({'error': 'No files to classify'}) + '\n'
return
# Load filename map
filename_map = {}
map_path = os.path.join(temp_dir, 'filename_map.json')
if os.path.exists(map_path):
try:
with open(map_path, 'r') as f:
filename_map = json.load(f)
logger.info(f"Loaded filename map: {filename_map}")
except Exception as e:
logger.error(f"Error loading filename map: {e}")
files = [f for f in os.listdir(temp_dir) if allowed_file(f)]
logger.info(f"Processing {len(files)} files from temp directory")
logger.info(f"Files to process: {files}")
for filename in files:
filepath = os.path.join(temp_dir, filename)
logger.info(f"Classifying: {filename}")
try:
classification_result, _, failure_labels = classify(filepath)
logger.info(f"Result for {filename}: {classification_result} with labels: {failure_labels}")
# Move file
dest_dir = os.path.join(app.config['UPLOAD_FOLDER_MULTIPLE'], classification_result.lower())
os.makedirs(dest_dir, exist_ok=True)
dest_path = os.path.join(dest_dir, filename)
shutil.move(filepath, dest_path)
# Get original filename if available
original_filename = filename_map.get(filename, filename)
logger.info(f"Sending result for original filename: {original_filename}")
result = {
'filename': original_filename,
'status': 'pass' if classification_result == 'Pass' else 'fail',
'labels': failure_labels,
'score': 0
}
yield json.dumps(result) + '\n'
except Exception as e:
logger.error(f"Error processing {filename}: {e}", exc_info=True)
# Use original filename for error reporting
original_filename = filename_map.get(filename, filename)
yield json.dumps({'filename': original_filename, 'status': 'error', 'error': str(e)}) + '\n'
return Response(stream_with_context(generate()), mimetype='application/x-ndjson')
@app.route('/clear_uploads', methods=['POST'])
def clear_uploads_route():
logger.info("=== CLEAR_UPLOADS CALLED ===")
try:
clear_uploads(app.config['UPLOAD_FOLDER_SINGLE'])
clear_uploads(app.config['UPLOAD_FOLDER_MULTIPLE'])
return jsonify({'success': True})
except Exception as e:
logger.error(f"Error in clear_uploads_route: {e}")
return jsonify({'error': str(e)}), 500
@app.route('/api/use_sample', methods=['POST'])
def use_sample():
try:
data = request.get_json()
filename = data.get('filename')
destination = data.get('destination') # 'single' or 'multiple'
if not filename or not destination:
return jsonify({'error': 'Missing filename or destination'}), 400
# Validate filename (security)
if not allowed_file(filename):
return jsonify({'error': 'Invalid filename'}), 400
# Source path
src_path = os.path.join('static', 'samples', filename)
if not os.path.exists(src_path):
return jsonify({'error': 'Sample not found'}), 404
# Destination path
if destination == 'single':
dest_folder = app.config['UPLOAD_FOLDER_SINGLE']
elif destination == 'multiple':
dest_folder = os.path.join(app.config['UPLOAD_FOLDER_MULTIPLE'], 'temp')
os.makedirs(dest_folder, exist_ok=True)
else:
return jsonify({'error': 'Invalid destination'}), 400
dest_path = os.path.join(dest_folder, filename)
# Copy file
shutil.copy2(src_path, dest_path)
# For multiple, we need to update the filename map
if destination == 'multiple':
map_path = os.path.join(dest_folder, 'filename_map.json')
filename_map = {}
if os.path.exists(map_path):
try:
with open(map_path, 'r') as f:
filename_map = json.load(f)
except:
pass
filename_map[filename] = filename # Map to itself for samples
with open(map_path, 'w') as f:
json.dump(filename_map, f)
return jsonify({'success': True, 'filename': filename})
except Exception as e:
logger.error(f"Error in use_sample: {e}")
return jsonify({'error': str(e)}), 500
@app.route('/api/samples', methods=['GET'])
def get_samples():
try:
samples_dir = os.path.join('static', 'samples')
if not os.path.exists(samples_dir):
return jsonify([])
files = [f for f in os.listdir(samples_dir) if allowed_file(f)]
# Sort files for consistent order
files.sort()
samples = []
for i, filename in enumerate(files):
samples.append({
'id': i + 1,
'url': f'/static/samples/{filename}',
'filename': filename
})
return jsonify(samples)
except Exception as e:
logger.error(f"Error in get_samples: {e}")
return jsonify({'error': str(e)}), 500
if __name__ == '__main__':
logger.info("SERVER STARTING ON PORT 7860")
# Disable reloader to prevent loading models twice (saves memory)
app.run(debug=False, use_reloader=False, host='0.0.0.0', port=7860)
|