Update app.py

app.py

@@ -4,60 +4,67 @@ import os
 import tempfile
 import base64
 import json
+from werkzeug.utils import secure_filename
 
 app = Flask(__name__)
 
+# Configure maximum content length for file uploads (10MB)
+app.config['MAX_CONTENT_LENGTH'] = 10 * 1024 * 1024
+
 # Load OSMF model
-osmf = YOLO('best.pt', task="classify")
+model_path = os.getenv('MODEL_PATH', 'best.pt')
+osmf = YOLO(model_path, task="classify")
 
 @app.route('/')
 def home():
-    return render_template('index.html')
+return render_template('index.html')
 
 @app.route('/predict', methods=['POST'])
 def predict():
-    if 'image' not in request.files:
-        return jsonify({'error': 'No image provided'})
-    
-    file = request.files['image']
-    
-    try:
-        # Create temporary file
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
-            file.save(temp_file.name)
-            temp_path = temp_file.name
-        
-        # Process image
-        predict = osmf(temp_path)
-        results = predict[0].to_json()  # Using the new recommended method
-        predict_dict = json.loads(results)
-        
-        name = predict_dict[0]["name"]
-        confidence = float(predict_dict[0]["confidence"]) * 100
-        
-        # Close file handle and remove
-        try:
-            os.close(temp_file.fileno())
-        except:
-            pass
-        
-        try:
-            os.unlink(temp_path)
-        except:
-            pass
-            
-        return jsonify({
-            'class': name,
-            'confidence': f"{confidence:.2f}%",
-            'status': 'success'
-        })
-        
-    except Exception as e:
-        return jsonify({
-            'error': str(e),
-            'status': 'error'
-        }), 500
+if 'image' not in request.files:
+return jsonify({'error': 'No image provided', 'status': 'error'}), 400
+
+file = request.files['image']
+if not file.filename:
+return jsonify({'error': 'Empty file provided', 'status': 'error'}), 400
+
+try:
+# Secure the filename and create temporary file
+filename = secure_filename(file.filename)
+with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
+file.save(temp_file.name)
+temp_path = temp_file.name
+
+# Process image
+predict = osmf(temp_path)
+results = predict[0].to_json()
+predict_dict = json.loads(results)
+
+# Extract prediction results
+name = predict_dict[0]["name"]
+confidence = float(predict_dict[0]["confidence"]) * 100
+
+except Exception as e:
+return jsonify({
+'error': str(e),
+'status': 'error'
+}), 500
+
+finally:
+# Cleanup temporary files
+try:
+if 'temp_file' in locals():
+os.close(temp_file.fileno())
+os.unlink(temp_path)
+except Exception:
+pass
 
+return jsonify({
+'class': name,
+'confidence': f"{confidence:.2f}%",
+'status': 'success'
+})
 
 if __name__ == '__main__':
-    app.run(debug=True)
+port = int(os.environ.get('PORT', 7860))
+app.run(host='0.0.0.0', port=port)