Update app.py

app.py

@@ -1,30 +1,21 @@
 from flask import Flask, request, render_template, jsonify
-from flask_limiter import Limiter
-from flask_limiter.util import get_remote_address
-from flask_limiter.errors import RateLimitExceeded
-from flask_wtf.csrf import CSRFProtect
-from flask_wtf import FlaskForm
-from wtforms import FileField
-from wtforms.validators import DataRequired
 from werkzeug.utils import secure_filename
-from functools import wraps
 from PIL import Image
 import numpy as np
 import torch
 import torchvision.transforms as T
-import random
 import os
 import time
 import hashlib
 import logging
-import secrets
 
 # ===============================
-# Deterministic inference
+# Logging
 # ===============================
-torch.manual_seed(42)
-np.random.seed(42)
-random.seed(42)
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
 
 # ===============================
 # App initialization
@@ -32,29 +23,11 @@ random.seed(42)
 app = Flask(__name__)
 
 app.config.update(
-    SECRET_KEY=secrets.token_hex(32),
+    SECRET_KEY=os.environ.get('SECRET_KEY', 'dev-secret-key-change-in-production'),
     MAX_CONTENT_LENGTH=30 * 1024 * 1024,  # 30 MB
-    WTF_CSRF_TIME_LIMIT=None,
     UPLOAD_FOLDER="uploads"
 )
 
-csrf = CSRFProtect(app)
-
-limiter = Limiter(
-    key_func=get_remote_address,
-    default_limits=["100 per hour", "20 per minute"]
-)
-limiter.init_app(app)
-
-# ===============================
-# Logging
-# ===============================
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s",
-    handlers=[logging.FileHandler("security.log"), logging.StreamHandler()]
-)
-
 # ===============================
 # Load Model
 # ===============================
@@ -69,15 +42,16 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 try:
     from transformers import SiglipForImageClassification
     
+    logging.info(f"Loading model from {MODEL_NAME}...")
     model = SiglipForImageClassification.from_pretrained(
         MODEL_NAME,
         torch_dtype=torch.float32
     ).to(device)
     model.eval()
     
-    logging.info(f"Model loaded successfully on {device}")
+    logging.info(f"✓ Model loaded successfully on {device}")
 except Exception as e:
-    logging.error(f"Model load failed: {e}")
+    logging.error(f"✗ Model load failed: {e}")
     model = None
 
 # ===============================
@@ -98,18 +72,8 @@ MAX_FILE_SIZE = 30 * 1024 * 1024
 os.makedirs(app.config["UPLOAD_FOLDER"], exist_ok=True)
 
 # ===============================
-# Helpers
+# Helper Functions
 # ===============================
-def security_validate(f):
-    @wraps(f)
-    def wrapper(*args, **kwargs):
-        logging.info(f"{request.remote_addr} → {request.endpoint}")
-        return f(*args, **kwargs)
-    return wrapper
-
-class UploadForm(FlaskForm):
-    image = FileField("Image", validators=[DataRequired()])
-
 def allowed_file(filename):
     return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS
 
@@ -181,15 +145,15 @@ def predict_image(path):
 # Routes
 # ===============================
 @app.route("/", methods=["GET"])
-@limiter.limit("30 per minute")
-@security_validate
 def index():
-    return render_template("index.html", form=UploadForm())
+    return render_template("index.html")
+
+@app.route("/health", methods=["GET"])
+def health():
+    """Health check endpoint for deployment platforms"""
+    return jsonify({"status": "healthy", "model_loaded": model is not None}), 200
 
 @app.route("/predict", methods=["POST"])
-@csrf.exempt
-@limiter.limit("60 per minute")
-@security_validate
 def predict():
     if "image" not in request.files:
         return jsonify({"error": "No image uploaded"}), 400
@@ -223,7 +187,10 @@ def predict():
 
     finally:
         if os.path.exists(path):
-            os.remove(path)
+            try:
+                os.remove(path)
+            except:
+                pass
 
 # ===============================
 # Security headers
@@ -231,30 +198,16 @@ def predict():
 @app.after_request
 def security_headers(res):
     res.headers["X-Content-Type-Options"] = "nosniff"
-    res.headers["X-Frame-Options"] = "DENY"
-    res.headers["X-XSS-Protection"] = "1; mode=block"
-    res.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains"
-    res.headers["Content-Security-Policy"] = (
-        "default-src 'self' https://cdnjs.cloudflare.com; "
-        "style-src 'self' 'unsafe-inline' https://cdnjs.cloudflare.com; "
-        "script-src 'self' 'unsafe-inline'; "
-        "img-src 'self' data:; "
-        "font-src 'self' https://cdnjs.cloudflare.com; "
-        "frame-src https://www.youtube.com https://www.youtube-nocookie.com;"
-    )
+    res.headers["X-Frame-Options"] = "SAMEORIGIN"
     return res
 
 # ===============================
-# JSON Error Handlers
+# Error Handlers
 # ===============================
 @app.errorhandler(413)
 def request_entity_too_large(error):
     return jsonify({"error": "File too large. Maximum allowed size is 30 MB."}), 413
 
-@app.errorhandler(RateLimitExceeded)
-def rate_limit_handler(e):
-    return jsonify({"error": "Too many requests. Please slow down."}), 429
-
 @app.errorhandler(400)
 def bad_request(error):
     return jsonify({"error": "Invalid request or image file."}), 400
@@ -267,8 +220,8 @@ def internal_error(error):
 # Run
 # ===============================
 if __name__ == "__main__":
-    # Try multiple port options for different platforms
-    port = int(os.environ.get("PORT", os.environ.get("GRADIO_SERVER_PORT", 7860)))
+    port = int(os.environ.get("PORT", 7860))
+    logging.info(f"Starting Flask app on port {port}...")
     app.run(
         debug=False,
         host="0.0.0.0",