Code Update
This commit is contained in:
@@ -6,11 +6,9 @@ from flask_cors import CORS
|
||||
from werkzeug.utils import secure_filename
|
||||
from processor import AudioImageProcessor
|
||||
|
||||
# Serve the build folder from the parent directory
|
||||
app = Flask(__name__, static_folder='../build', static_url_path='')
|
||||
CORS(app) # Allow Svelte to communicate
|
||||
CORS(app)
|
||||
|
||||
# Configuration
|
||||
UPLOAD_FOLDER = os.path.join(os.getcwd(), 'uploads')
|
||||
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
|
||||
processor = AudioImageProcessor(UPLOAD_FOLDER)
|
||||
@@ -21,29 +19,24 @@ def save_upload(file_obj):
|
||||
file_obj.save(path)
|
||||
return path
|
||||
|
||||
# --- Frontend Routes ---
|
||||
@app.route('/')
|
||||
def index():
|
||||
return send_from_directory(app.static_folder, 'index.html')
|
||||
|
||||
@app.errorhandler(404)
|
||||
def not_found(e):
|
||||
# If the path starts with /api, return actual 404
|
||||
if request.path.startswith('/api/'):
|
||||
return jsonify({"error": "Not found"}), 404
|
||||
# Otherwise return index.html for SPA routing
|
||||
return send_from_directory(app.static_folder, 'index.html')
|
||||
|
||||
@app.route('/health', methods=['GET'])
|
||||
def health():
|
||||
return jsonify({"status": "ok", "max_mb": 40})
|
||||
|
||||
# --- Background Cleanup ---
|
||||
import threading
|
||||
|
||||
def cleanup_task():
|
||||
"""Background thread to clean up old files."""
|
||||
expiration_seconds = 600 # 10 minutes
|
||||
expiration_seconds = 600
|
||||
while True:
|
||||
try:
|
||||
now = time.time()
|
||||
@@ -51,7 +44,6 @@ def cleanup_task():
|
||||
for filename in os.listdir(UPLOAD_FOLDER):
|
||||
filepath = os.path.join(UPLOAD_FOLDER, filename)
|
||||
if os.path.isfile(filepath):
|
||||
# check creation time
|
||||
if now - os.path.getctime(filepath) > expiration_seconds:
|
||||
try:
|
||||
os.remove(filepath)
|
||||
@@ -61,16 +53,12 @@ def cleanup_task():
|
||||
except Exception as e:
|
||||
print(f"Cleanup Error: {e}")
|
||||
|
||||
time.sleep(60) # Run every minute
|
||||
time.sleep(60)
|
||||
|
||||
# Start cleanup thread safely
|
||||
if os.environ.get('WERKZEUG_RUN_MAIN') == 'true' or not os.environ.get('WERKZEUG_RUN_MAIN'):
|
||||
# Simple check to try and avoid double threads in reloader, though not perfect
|
||||
t = threading.Thread(target=cleanup_task, daemon=True)
|
||||
t.start()
|
||||
|
||||
|
||||
# --- Endpoint 1: Create Art (Optional: Embed Audio in it) ---
|
||||
@app.route('/api/generate-art', methods=['POST'])
|
||||
def generate_art():
|
||||
if 'audio' not in request.files:
|
||||
@@ -83,27 +71,19 @@ def generate_art():
|
||||
art_path = None
|
||||
|
||||
try:
|
||||
# 1. Save Audio
|
||||
audio_path = save_upload(audio_file)
|
||||
|
||||
# 2. Generate Art
|
||||
min_pixels = 0
|
||||
if should_embed:
|
||||
# Calculate required pixels: File Bytes * 8 (bits) / 3 (channels)
|
||||
# Add 5% buffer for header and safety
|
||||
file_size = os.path.getsize(audio_path)
|
||||
min_pixels = int((file_size * 8 / 3) * 1.05)
|
||||
|
||||
art_path = processor.generate_spectrogram(audio_path, min_pixels=min_pixels)
|
||||
|
||||
# 3. If Embed requested, run Steganography immediately using the art as host
|
||||
final_path = art_path
|
||||
if should_embed:
|
||||
# art_path becomes the host, audio_path is the data
|
||||
final_path = processor.encode_stego(audio_path, art_path)
|
||||
|
||||
# If we created a new stego image, the pure art_path is intermediate (and audio is input)
|
||||
# We can delete art_path now if it's different (it is)
|
||||
if art_path != final_path:
|
||||
try: os.remove(art_path)
|
||||
except: pass
|
||||
@@ -115,14 +95,10 @@ def generate_art():
|
||||
except Exception as e:
|
||||
return jsonify({"error": str(e)}), 500
|
||||
finally:
|
||||
# Cleanup Inputs
|
||||
if audio_path and os.path.exists(audio_path):
|
||||
try: os.remove(audio_path)
|
||||
except: pass
|
||||
|
||||
|
||||
|
||||
# --- Endpoint 3: Steganography (Audio + Custom Host) ---
|
||||
@app.route('/api/hide', methods=['POST'])
|
||||
def hide_data():
|
||||
if 'data' not in request.files or 'host' not in request.files:
|
||||
@@ -148,7 +124,6 @@ def hide_data():
|
||||
try: os.remove(host_path)
|
||||
except: pass
|
||||
|
||||
# --- Endpoint 4: Decode (Universal) ---
|
||||
@app.route('/api/decode', methods=['POST'])
|
||||
def decode():
|
||||
if 'image' not in request.files:
|
||||
@@ -159,7 +134,6 @@ def decode():
|
||||
img_path = save_upload(request.files['image'])
|
||||
restored_path = processor.decode_image(img_path)
|
||||
|
||||
# Determine mimetype based on extension for browser friendliness
|
||||
filename = os.path.basename(restored_path)
|
||||
return send_file(restored_path, as_attachment=True, download_name=filename)
|
||||
except ValueError as e:
|
||||
@@ -171,13 +145,8 @@ def decode():
|
||||
try: os.remove(img_path)
|
||||
except: pass
|
||||
|
||||
# --- Endpoint 4: Visualizer SSE Stream ---
|
||||
@app.route('/api/visualize', methods=['POST'])
|
||||
def visualize():
|
||||
"""
|
||||
SSE endpoint that streams the spectrogram generation process.
|
||||
Returns step-by-step updates for visualization.
|
||||
"""
|
||||
if 'audio' not in request.files:
|
||||
return jsonify({"error": "No audio file provided"}), 400
|
||||
|
||||
@@ -196,56 +165,48 @@ def visualize():
|
||||
try:
|
||||
import base64
|
||||
|
||||
# Step 1: Audio loaded
|
||||
yield f"data: {json.dumps({'step': 1, 'status': 'loading', 'message': 'Loading audio file...', 'progress': 10})}\n\n"
|
||||
time.sleep(0.8)
|
||||
|
||||
yield f"data: {json.dumps({'step': 1, 'status': 'complete', 'message': f'Audio loaded: {audio_file.filename}', 'progress': 20, 'fileSize': file_size})}\n\n"
|
||||
time.sleep(0.5)
|
||||
|
||||
# Step 2: Analyzing audio
|
||||
yield f"data: {json.dumps({'step': 2, 'status': 'loading', 'message': 'Analyzing audio frequencies...', 'progress': 30})}\n\n"
|
||||
time.sleep(1.0)
|
||||
|
||||
yield f"data: {json.dumps({'step': 2, 'status': 'complete', 'message': 'Frequency analysis complete', 'progress': 40})}\n\n"
|
||||
time.sleep(0.5)
|
||||
|
||||
# Step 3: Generating spectrogram
|
||||
yield f"data: {json.dumps({'step': 3, 'status': 'loading', 'message': 'Generating spectrogram image...', 'progress': 50})}\n\n"
|
||||
|
||||
print(f"[VISUALIZE] Starting spectrogram generation for {audio_path}")
|
||||
art_path = processor.generate_spectrogram(audio_path, min_pixels=min_pixels)
|
||||
print(f"[VISUALIZE] Spectrogram generated at {art_path}")
|
||||
|
||||
# Read the spectrogram image and encode as base64
|
||||
with open(art_path, 'rb') as img_file:
|
||||
spectrogram_b64 = base64.b64encode(img_file.read()).decode('utf-8')
|
||||
print(f"[VISUALIZE] Spectrogram base64 length: {len(spectrogram_b64)}")
|
||||
|
||||
yield f"data: {json.dumps({'step': 3, 'status': 'complete', 'message': 'Spectrogram generated!', 'progress': 70, 'spectrogramImage': f'data:image/png;base64,{spectrogram_b64}'})}\n\n"
|
||||
print("[VISUALIZE] Sent spectrogram image")
|
||||
time.sleep(2.0) # Pause to let user see the spectrogram
|
||||
time.sleep(2.0)
|
||||
|
||||
# Step 4: Embedding audio
|
||||
yield f"data: {json.dumps({'step': 4, 'status': 'loading', 'message': 'Embedding audio into image (LSB steganography)...', 'progress': 80})}\n\n"
|
||||
|
||||
final_path = processor.encode_stego(audio_path, art_path)
|
||||
|
||||
# Read the final image and encode as base64
|
||||
with open(final_path, 'rb') as img_file:
|
||||
final_b64 = base64.b64encode(img_file.read()).decode('utf-8')
|
||||
|
||||
yield f"data: {json.dumps({'step': 4, 'status': 'complete', 'message': 'Audio embedded successfully!', 'progress': 95, 'finalImage': f'data:image/png;base64,{final_b64}'})}\n\n"
|
||||
time.sleep(2.0) # Pause to let user see the final image
|
||||
time.sleep(2.0)
|
||||
|
||||
# Step 5: Complete - send the result URL
|
||||
result_id = os.path.basename(final_path)
|
||||
yield f"data: {json.dumps({'step': 5, 'status': 'complete', 'message': 'Process complete!', 'progress': 100, 'resultId': result_id})}\n\n"
|
||||
|
||||
except Exception as e:
|
||||
yield f"data: {json.dumps({'step': 0, 'status': 'error', 'message': str(e), 'progress': 0})}\n\n"
|
||||
finally:
|
||||
# Clean up intermediate files (but keep final)
|
||||
if art_path and art_path != final_path and os.path.exists(art_path):
|
||||
try: os.remove(art_path)
|
||||
except: pass
|
||||
@@ -265,16 +226,12 @@ def visualize():
|
||||
except: pass
|
||||
return jsonify({"error": str(e)}), 500
|
||||
|
||||
|
||||
@app.route('/api/result/<result_id>', methods=['GET'])
|
||||
def get_result(result_id):
|
||||
"""Serve the result image by ID."""
|
||||
result_path = os.path.join(app.config['UPLOAD_FOLDER'], result_id)
|
||||
if os.path.exists(result_path):
|
||||
return send_file(result_path, mimetype='image/png', as_attachment=False)
|
||||
return jsonify({"error": "Result not found"}), 404
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# Threaded=True is important for processing images without blocking
|
||||
app.run(debug=True, port=5000, threaded=True)
|
||||
|
||||
@@ -6,12 +6,10 @@ import numpy as np
|
||||
import librosa
|
||||
import librosa.display
|
||||
import matplotlib
|
||||
# Set backend to Agg (Anti-Grain Geometry) to render without a GUI (essential for servers)
|
||||
matplotlib.use('Agg')
|
||||
import matplotlib.pyplot as plt
|
||||
from PIL import Image
|
||||
|
||||
# --- Constants ---
|
||||
MAX_MB = 40
|
||||
SIG_SHIFT = b'B2I!'
|
||||
SIG_STEGO = b'B2S!'
|
||||
@@ -25,7 +23,6 @@ class AudioImageProcessor:
|
||||
os.makedirs(upload_folder, exist_ok=True)
|
||||
|
||||
def _get_bytes(self, path):
|
||||
"""Helper to safely read bytes"""
|
||||
if os.path.getsize(path) > (MAX_MB * 1024 * 1024):
|
||||
raise ValueError("File too large (Max 40MB)")
|
||||
with open(path, 'rb') as f:
|
||||
@@ -36,9 +33,7 @@ class AudioImageProcessor:
|
||||
ext_bytes = ext.encode('utf-8')
|
||||
return struct.pack(HEADER_FMT, signature, file_size, len(ext_bytes)) + ext_bytes
|
||||
|
||||
# --- Feature 1: Spectrogram Art ---
|
||||
def generate_spectrogram(self, audio_path, min_pixels=0):
|
||||
"""Generates a visual spectrogram from audio."""
|
||||
try:
|
||||
import torch
|
||||
import torchaudio
|
||||
@@ -48,13 +43,10 @@ class AudioImageProcessor:
|
||||
|
||||
if has_torch and torch.cuda.is_available():
|
||||
try:
|
||||
# GPU Accelerated Path
|
||||
device = "cuda"
|
||||
waveform, sr = torchaudio.load(audio_path)
|
||||
waveform = waveform.to(device)
|
||||
|
||||
# Create transformation
|
||||
# Mimic librosa defaults roughly: n_fft=2048, hop_length=512
|
||||
n_fft = 2048
|
||||
win_length = n_fft
|
||||
hop_length = 512
|
||||
@@ -72,17 +64,13 @@ class AudioImageProcessor:
|
||||
S = mel_spectrogram(waveform)
|
||||
S_dB = torchaudio.transforms.AmplitudeToDB()(S)
|
||||
|
||||
# Back to CPU for plotting
|
||||
S_dB = S_dB.cpu().numpy()[0] # Take first channel
|
||||
# Librosa display expects numpy
|
||||
S_dB = S_dB.cpu().numpy()[0]
|
||||
except Exception as e:
|
||||
# Fallback to CPU/Librosa if any error occurs
|
||||
print(f"GPU processing failed, falling back to CPU: {e}")
|
||||
return self._generate_spectrogram_cpu(audio_path, min_pixels)
|
||||
else:
|
||||
return self._generate_spectrogram_cpu(audio_path, min_pixels)
|
||||
|
||||
# Plotting (Common)
|
||||
return self._plot_spectrogram(S_dB, sr, min_pixels)
|
||||
|
||||
def _generate_spectrogram_cpu(self, audio_path, min_pixels=0):
|
||||
@@ -92,44 +80,31 @@ class AudioImageProcessor:
|
||||
return self._plot_spectrogram(S_dB, sr, min_pixels)
|
||||
|
||||
def _plot_spectrogram(self, S_dB, sr, min_pixels=0):
|
||||
# Calculate DPI dynamically to ensure we have enough pixels for steganography
|
||||
dpi = 300
|
||||
if min_pixels > 0:
|
||||
# Figure is 12x6 inches. Area = 72 sq inches.
|
||||
# Total Pixels = 72 * dpi^2
|
||||
required_dpi = math.ceil((min_pixels / 72) ** 0.5)
|
||||
# Add a small buffer
|
||||
dpi = max(dpi, int(required_dpi * 1.05))
|
||||
|
||||
# Use exact dimensions without margins
|
||||
width_in = 12
|
||||
height_in = 6
|
||||
fig = plt.figure(figsize=(width_in, height_in))
|
||||
|
||||
# Add axes covering the entire figure [left, bottom, width, height]
|
||||
ax = plt.axes([0, 0, 1, 1], frameon=False)
|
||||
ax.set_axis_off()
|
||||
|
||||
# 'magma' is a nice default
|
||||
librosa.display.specshow(S_dB, sr=sr, fmax=8000, cmap='magma', ax=ax)
|
||||
|
||||
output_path = os.path.join(self.upload_folder, f"art_{int(time.time())}.png")
|
||||
|
||||
# specific DPI, no bbox_inches='tight' (which shrinks the image)
|
||||
plt.savefig(output_path, dpi=dpi)
|
||||
plt.close()
|
||||
return output_path
|
||||
|
||||
|
||||
|
||||
# --- Feature 3: Steganography (Embed in Host) ---
|
||||
def encode_stego(self, data_path, host_path):
|
||||
# 1. Prepare Data
|
||||
file_data = self._get_bytes(data_path)
|
||||
header = self._create_header(SIG_STEGO, len(file_data), data_path)
|
||||
payload_bits = np.unpackbits(np.frombuffer(header + file_data, dtype=np.uint8))
|
||||
|
||||
# 2. Prepare Host
|
||||
host = Image.open(host_path).convert('RGB')
|
||||
host_arr = np.array(host)
|
||||
flat_host = host_arr.flatten()
|
||||
@@ -137,7 +112,6 @@ class AudioImageProcessor:
|
||||
if len(payload_bits) > len(flat_host):
|
||||
raise ValueError(f"Host image too small. Need {len(payload_bits)/3/1e6:.2f} MP.")
|
||||
|
||||
# 3. Embed (LSB)
|
||||
padded_bits = np.pad(payload_bits, (0, len(flat_host) - len(payload_bits)), 'constant')
|
||||
embedded_flat = (flat_host & 0xFE) + padded_bits
|
||||
|
||||
@@ -147,19 +121,16 @@ class AudioImageProcessor:
|
||||
embedded_img.save(output_path, "PNG")
|
||||
return output_path
|
||||
|
||||
# --- Feature 4: Universal Decoder ---
|
||||
def decode_image(self, image_path):
|
||||
img = Image.open(image_path).convert('RGB')
|
||||
flat_bytes = np.array(img).flatten()
|
||||
|
||||
# Strategy A: Check for Shift Signature (Raw Bytes)
|
||||
try:
|
||||
sig = struct.unpack('>4s', flat_bytes[:4])[0]
|
||||
if sig == SIG_SHIFT:
|
||||
return self._extract(flat_bytes, image_path, is_bits=False)
|
||||
except: pass
|
||||
|
||||
# Strategy B: Check for Stego Signature (LSB)
|
||||
try:
|
||||
sample_bytes = np.packbits(flat_bytes[:300] & 1)
|
||||
sig = struct.unpack('>4s', sample_bytes[:4])[0]
|
||||
|
||||
Reference in New Issue
Block a user