Update app.py

app.py

@@ -1,13 +1,9 @@
 import os
 import subprocess
-import logging
-from flask import Flask, request, send_file, jsonify
-from werkzeug.utils import secure_filename
-from flask_cors import CORS
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import FileResponse
+from fastapi.middleware.cors import CORSMiddleware
+import uvicorn
 
 # Define directories for uploads and outputs
 UPLOAD_FOLDER = 'uploads_gradio'
@@ -17,22 +13,22 @@ OUTPUT_FOLDER = 'outputs_gradio'
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 os.makedirs(OUTPUT_FOLDER, exist_ok=True)
 
-# Initialize Flask app
-app = Flask(__name__)
-app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-
-# Enable CORS for all routes (optional, adjust origins as needed)
-CORS(app)
-
-# Allowed file extensions
-ALLOWED_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.bmp'}
-
-def allowed_file(filename):
-    """Check if the file has an allowed extension."""
-    _, ext = os.path.splitext(filename)
-    return ext.lower() in ALLOWED_EXTENSIONS
-
-def animate_image(file_path):
+app = FastAPI(
+    title="Animated Drawings API",
+    description="Upload your drawing and receive an animated GIF.",
+    version="1.0.0"
+)
+
+# Enable CORS (optional, adjust origins as needed)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Update with specific origins in production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+def animate_image(file_path: str) -> str:
     """
     Process the uploaded image and generate an animated GIF.
 
@@ -53,12 +49,11 @@ def animate_image(file_path):
     char_anno_dir = os.path.join(OUTPUT_FOLDER, f"{base}_out")
     os.makedirs(char_anno_dir, exist_ok=True)
 
-    # Validate file extension
-    if ext.lower() not in ALLOWED_EXTENSIONS:
-        raise ValueError("Unsupported file type. Please upload an image file (png, jpg, jpeg, bmp).")
-
     try:
-        logger.info(f"Starting animation for {input_path}")
+        # Validate file extension
+        allowed_extensions = ['.png', '.jpg', '.jpeg', '.bmp']
+        if ext.lower() not in allowed_extensions:
+            raise ValueError("Unsupported file type. Please upload an image file (png, jpg, jpeg, bmp).")
 
         # Run the image_to_animation.py script with required arguments
         subprocess.run([
@@ -70,79 +65,67 @@ def animate_image(file_path):
         gif_path = os.path.join(char_anno_dir, "video.gif")
 
         if os.path.exists(gif_path):
-            logger.info(f"Animation successful: {gif_path}")
             return gif_path
         else:
             raise FileNotFoundError("Animation failed to generate. Please ensure the input image contains clear humanoid drawings.")
 
     except subprocess.CalledProcessError as e:
-        logger.error(f"Error during processing: {e}")
         raise RuntimeError(f"Error during processing: {e}")
     except Exception as e:
-        logger.error(f"Unexpected error: {e}")
         raise RuntimeError(f"Unexpected error: {e}")
 
-@app.route('/animate', methods=['POST'])
-def animate():
+@app.post("/animate", summary="Generate Animated GIF from Image")
+async def generate_gif(image: UploadFile = File(...)):
     """
-    Endpoint to receive an image and return the animated GIF.
+    Endpoint to upload an image and receive an animated GIF.
+
+    Args:
+        image (UploadFile): The image file to be animated.
+
+    Returns:
+        FileResponse: The generated animated GIF.
     """
-    logger.info("Received request to /animate")
-    if 'file' not in request.files:
-        logger.warning("No file part in the request.")
-        return jsonify({"error": "No file part in the request."}), 400
-
-    file = request.files['file']
-
-    if file.filename == '':
-        logger.warning("No file selected for uploading.")
-        return jsonify({"error": "No file selected for uploading."}), 400
-
-    if file and allowed_file(file.filename):
-        filename = secure_filename(file.filename)
-        input_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
-        file.save(input_path)
-        logger.info(f"File saved to {input_path}")
-
-        try:
-            gif_path = animate_image(input_path)
-            logger.info(f"Sending GIF: {gif_path}")
-            return send_file(
-                gif_path,
-                mimetype='image/gif',
-                as_attachment=True,
-                download_name=f"{os.path.splitext(filename)[0]}.gif"
-            )
-        except Exception as e:
-            logger.error(f"Error in /animate: {e}")
-            return jsonify({"error": str(e)}), 500
+    # Validate the uploaded file
+    if not image:
+        raise HTTPException(status_code=400, detail="No file uploaded.")
+
+    filename = image.filename
+    base, ext = os.path.splitext(filename)
+    allowed_extensions = ['.png', '.jpg', '.jpeg', '.bmp']
+    if ext.lower() not in allowed_extensions:
+        raise HTTPException(status_code=400, detail="Unsupported file type. Please upload an image file (png, jpg, jpeg, bmp).")
+
+    # Save the uploaded file to the upload directory
+    upload_path = os.path.join(UPLOAD_FOLDER, filename)
+    try:
+        with open(upload_path, "wb") as f:
+            f.write(await image.read())
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to save uploaded file: {e}")
+
+    # Process the image to generate GIF
+    try:
+        gif_path = animate_image(upload_path)
+    except ValueError as ve:
+        raise HTTPException(status_code=400, detail=str(ve))
+    except FileNotFoundError as fnfe:
+        raise HTTPException(status_code=500, detail=str(fnfe))
+    except RuntimeError as re:
+        raise HTTPException(status_code=500, detail=str(re))
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")
+
+    # Return the generated GIF as a response
+    if os.path.exists(gif_path):
+        return FileResponse(path=gif_path, media_type="image/gif", filename="animated.gif")
     else:
-        logger.warning("Allowed file types are png, jpg, jpeg, bmp.")
-        return jsonify({"error": "Allowed file types are png, jpg, jpeg, bmp."}), 400
+        raise HTTPException(status_code=500, detail="Failed to generate GIF.")
 
-@app.route('/', methods=['GET'])
-def index():
-    """
-    Root endpoint to provide basic information.
-    """
-    logger.info("Received request to /")
-    return jsonify({
-        "message": "Animated Drawings API",
-        "endpoints": {
-            "/animate": "POST an image to receive an animated GIF."
-        }
-    })
-
-@app.route('/health', methods=['GET'])
-def health():
-    """
-    Health check endpoint.
-    """
-    logger.info("Received request to /health")
-    return jsonify({"status": "healthy"}), 200
+@app.get("/", summary="Root Endpoint")
+def read_root():
+    return {"message": "Welcome to the Animated Drawings API. Use the /animate endpoint to upload images and receive animated GIFs."}
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     # Use the PORT environment variable provided by Hugging Face Spaces or default to 7860
     port = int(os.getenv("PORT", "7860"))
-    logger.info(f"Starting Flask app on port {port}")
-    app.run(host='0.0.0.0', port=port)
+    uvicorn.run(app, host="0.0.0.0", port=port)