Update app.py

app.py

@@ -1,6 +1,6 @@
-from flask import Flask, render_template, request, jsonify, send_from_directory
+from flask import Flask, request, jsonify, send_from_directory, render_template
 import os
-import shutil
+import subprocess
 
 app = Flask(__name__)
 
@@ -14,7 +14,6 @@ os.makedirs(OUTPUT_FOLDER, exist_ok=True)
 def index():
     return render_template('index.html')
 
-
 @app.route('/process_image', methods=['POST'])
 def process_image():
     if 'file' not in request.files:
@@ -22,24 +21,27 @@ def process_image():
 
     file = request.files['file']
     input_path = os.path.join(UPLOAD_FOLDER, file.filename)
-    output_path = os.path.join(OUTPUT_FOLDER, 'enhanced_' + file.filename)
+    output_path = os.path.join(OUTPUT_FOLDER, 'output.png')
 
-    # Save original upload
+    # Save the uploaded image
     file.save(input_path)
 
-    # Dummy enhancement (copy image as demo)
-    shutil.copy(input_path, output_path)
-
-    return jsonify({
-        'status': 'success',
-        'output_path': f'/outputs/enhanced_{file.filename}'
-    })
+    try:
+        # Run NAFNet model using subprocess
+        subprocess.run([
+            'python', 'NAFNet/demo.py',
+            '-opt', 'NAFNet/options/test/REDS/NAFNet-width64.yml',
+            '--input_path', input_path,
+            '--output_path', output_path
+        ], check=True)
 
+        return jsonify({'status': 'success', 'output_path': f'/outputs/output.png'})
+    except subprocess.CalledProcessError as e:
+        return jsonify({'status': 'error', 'message': f'Failed to run model: {str(e)}'})
 
 @app.route('/outputs/<filename>')
 def get_output_image(filename):
     return send_from_directory(OUTPUT_FOLDER, filename)
 
-
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)