Update FastAPI_app.py

FastAPI_app.py

@@ -1,9 +1,11 @@
 # FastAPI application for Fridge2Dish
 
+# import libraries
 import os
 import io
 import time
 import traceback
+import threading
 
 import uvicorn
 import numpy as np
@@ -18,186 +20,195 @@ from fastapi.middleware.cors import CORSMiddleware
 import tensorflow as tf
 import google.generativeai as genai
 
-# Transformers libraries (for fallback)
+# Transformers libraries (Gemma local fallback)
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
-import threading
 
 
-# create presistent storage for GPT-2
-LOCAL_GPT2_DIR = "/data/gpt2"      # HF Spaces persistent folder
-REMOTE_GPT2_NAME = "gpt2-medium"
+# create presistent storage for Gemma-2-2b-it model
+LOCAL_GEMMA_DIR = "/data/gemma-2-2b-it"
+REMOTE_GEMMA_NAME = "google/gemma-2-2b-it"
+
+# Load ingredients model
+MODEL_PATH = "models/ingredient_model.h5"
 
-_local_generator = None
-_local_lock = threading.Lock()
+# Protect loading the large local Gemma model by locking.
+_gemma_lock = threading.Lock()
+_gemma_pipeline = None
 
 
-def load_or_download_gpt2():
+# load or download (as applicable) the Gemma model
+def load_or_download_gemma():
     """
-    This function downloads GPT-2-medium into `/data/gpt2` on first run.
-    And on subsequent runs, it loads the saved local version.
+    Loads a local Gemma-2-2b-it pipeline from LOCAL_GEMMA_DIR if present,
+    otherwise downloads from Hugging Face and saves into LOCAL_GEMMA_DIR.
+    Returns a transformers text-generation pipeline.
     """
+    global _gemma_pipeline
+    if _gemma_pipeline is not None:
+        return _gemma_pipeline
 
-    global _local_generator
-    if _local_generator is not None:
-        return _local_generator
+    with _gemma_lock:
+        if _gemma_pipeline is not None:
+            return _gemma_pipeline
 
-    with _local_lock:
-        if _local_generator is not None:
-            return _local_generator
-
-        os.makedirs(LOCAL_GPT2_DIR, exist_ok=True)
-
-        # Load from cache
-        if os.listdir(LOCAL_GPT2_DIR):
-            print("\n🔵 Loading GPT-2 from local cache...")
-            tokenizer = AutoTokenizer.from_pretrained(LOCAL_GPT2_DIR)
-            model = AutoModelForCausalLM.from_pretrained(LOCAL_GPT2_DIR)
+        os.makedirs(LOCAL_GEMMA_DIR, exist_ok=True)
 
+        # If local folder already populated, load from there
+        if os.listdir(LOCAL_GEMMA_DIR):
+            print("\n🔵 Loading Gemma-2-2b-it from local cache:", LOCAL_GEMMA_DIR)
+            tokenizer = AutoTokenizer.from_pretrained(LOCAL_GEMMA_DIR, trust_remote_code=True)
+            model = AutoModelForCausalLM.from_pretrained(LOCAL_GEMMA_DIR, trust_remote_code=True)
         else:
-            # First-time download
-            print("\n🟡 Downloading GPT-2-medium...")
-            tokenizer = AutoTokenizer.from_pretrained(REMOTE_GPT2_NAME)
-            model = AutoModelForCausalLM.from_pretrained(REMOTE_GPT2_NAME)
-
-            print("\n🟢 Saving GPT-2-medium to persistent storage...")
-            tokenizer.save_pretrained(LOCAL_GPT2_DIR)
-            model.save_pretrained(LOCAL_GPT2_DIR)
-
+            # Download and save locally
+            print("\n🟡 Downloading Gemma-2-2b-it from Hugging Face (first run)...")
+            tokenizer = AutoTokenizer.from_pretrained(REMOTE_GEMMA_NAME, trust_remote_code=True)
+            model = AutoModelForCausalLM.from_pretrained(REMOTE_GEMMA_NAME, trust_remote_code=True)
+            print("\n🟢 Saving Gemma to local persistent directory:", LOCAL_GEMMA_DIR)
+            tokenizer.save_pretrained(LOCAL_GEMMA_DIR)
+            model.save_pretrained(LOCAL_GEMMA_DIR)
+
+        # Choose device: GPU if available, otherwise CPU
         device = 0 if torch.cuda.is_available() else -1
-        _local_generator = pipeline(
+        print(f"\n[Gemma] creating pipeline (device={device}) -- this may take a moment")
+        _gemma_pipeline = pipeline(
             "text-generation",
             model=model,
             tokenizer=tokenizer,
             device=device,
+            # reduce returned tokens to keep small responses
+            max_new_tokens=300,
+            do_sample=True,
+            top_p=0.95,
+            temperature=0.7
         )
+        print("[Gemma] loaded and ready")
+        return _gemma_pipeline
 
-        print("\n\n✅ GPT-2 ready for generation.")
-        return _local_generator
 
 
-# improve GPT-2 recipe generation
-def clean_output(text: str) -> str:
+# improve LM output by cleaning
+def _clean_generated_text(text: str) -> str:
     """
-    Remove garbage, repeated sentences, disclaimers, and anything before the recipe.
+    Basic cleaning of the LM output:
+    - remove obvious leading garbage,
+    - remove repeated lines,
+    - trim long tails after a natural stopping point.
     """
+    if not text:
+        return ""
+
+    # If model echoes prompt, try to cut at 'Recipe' or '### Ingredients' or similar markers
+    markers = ["### Ingredients", "### Steps", "Ingredients:", "Steps:", "Recipe"]
+    for m in markers:
+        if m in text:
+            # keep starting at the marker if there is garbage before
+            try:
+                idx = text.index(m)
+                text = text[idx:]
+                break
+            except ValueError:
+                pass
+
+    # Deduplicate repeated consecutive lines
+    out_lines = []
+    prev = None
+    for line in text.splitlines():
+        s = line.rstrip()
+        if s and s == prev:
+            continue
+        out_lines.append(line)
+        prev = s
+
+    cleaned = "\n".join(out_lines).strip()
+    # Trim at a long trailing repeated token if present
+    if len(cleaned) > 2000:
+        cleaned = cleaned[:2000].rsplit("\n", 1)[0]
 
-    # Remove leading garbage before a recognizable title
-    for key in ["Ingredients", "Recipe", "###", "Steps"]:
-        if key in text:
-            text = text.split(key, 1)[1]
-            text = key + text
-            break
-
-    # Remove repeated lines
-    cleaned = []
-    seen = set()
-
-    for line in text.split("\n"):
-        l = line.strip()
-        if l not in seen:
-            seen.add(l)
-            cleaned.append(line)
-
-    return "\n".join(cleaned).strip()
-
+    return cleaned
 
 
-def generate_recipe_local(ingredient_names):
+# generate recipe using local Gemma
+def generate_recipe_local_gemma(ingredient_names):
     """
-    Improved GPT-2 recipe generation with strict formatting.
+    Use local Gemma pipeline to generate a well-formatted recipe in markdown.
     """
-    generator = load_or_download_gpt2()
-
-    prompt = f"""
-You are an AI chef. Create a clean, structured recipe using ONLY these ingredients:
-{', '.join(ingredient_names)}.
-
-STRICT RULES:
-1. Start with a short recipe title (one line).
-2. Then one-sentence description.
-3. Then a section titled "### Ingredients" with bullet points.
-4. Use quantities (approximate is okay).
-5. Then "### Steps" with 6–10 numbered steps.
-6. Keep it short, clear, and well formatted.
-7. No rambling. No repeating. No intros. No disclaimers.
-8. End after the steps.
-
-FORMAT EXAMPLE:
-
-Title
-Short description.
-
-### Ingredients
-- item
-- item
-
-### Steps
-1. step
-2. step
-3. step
-
-Generate the recipe:
-"""
-
-    output = generator(
-        prompt,
-        max_new_tokens=180,
-        temperature=0.7,
-        do_sample=True,
-        top_p=0.95,
-        num_return_sequences=1
-    )[0]["generated_text"]
-
-    cleaned = clean_output(output)
-
+    gen = load_or_download_gemma()
+
+    prompt = (
+        "You are a professional chef and recipe writer. Create a concise, well-formatted recipe in Markdown "
+        f"using ONLY the following ingredients: {', '.join(ingredient_names)}.\n\n"
+        "Requirements:\n"
+        "- Start with the recipe title on one line.\n"
+        "- One-sentence description.\n"
+        "- Then a '### Ingredients' section with bullet points and approximate quantities.\n"
+        "- Then a '### Steps' section with 6-8 numbered steps.\n"
+        "- Keep it concise, no filler, no disclaimers, and end after the steps.\n\n"
+        "Output only the recipe in Markdown.\n\nRecipe:\n"
+    )
+
+    out = gen(prompt, do_sample=True, temperature=0.7, top_p=0.95, max_new_tokens=300, num_return_sequences=1)
+    generated = out[0].get("generated_text", "")
+    # If the model reprints the prompt, remove the leading prompt part:
+    if "Recipe:" in generated:
+        generated = generated.split("Recipe:", 1)[1].strip()
+    cleaned = _clean_generated_text(generated)
     return cleaned
 
 
 
-# Load ingredients model once startup.
-MODEL_PATH = "models/ingredient_model.h5"
+# Ingredient detection model loading
 MODEL = tf.keras.models.load_model(MODEL_PATH)
 
-CLASS_NAMES = [
-    'apple', 'banana', 'beetroot', 'bell pepper', 'cabbage', 'capsicum', 'carrot', 'cauliflower', 'chilli pepper', 
-     'corn', 'cucumber', 'eggplant', 'garlic', 'ginger', 'grapes', 'jalepeno', 'kiwi', 'lemon', 'lettuce', 'mango',
-     'onion', 'orange', 'paprika', 'pear', 'peas', 'pineapple', 'pomegranate', 'potato', 'raddish', 'soy beans', 
-     'spinach', 'sweetcorn', 'sweetpotato', 'tomato', 'turnip', 'watermelon']
+
+# Class names from folder or manual.
+if os.path.isdir("dataset/dataset_2/train"):
+    CLASS_NAMES = sorted(os.listdir("dataset/dataset_2/train"))
+    
+else:
+    CLASS_NAMES = [
+        'apple', 'banana', 'beetroot', 'bell pepper', 'cabbage', 'capsicum', 'carrot', 'cauliflower', 'chilli pepper', 
+        'corn', 'cucumber', 'eggplant', 'garlic', 'ginger', 'grapes', 'jalepeno', 'kiwi', 'lemon', 'lettuce', 'mango',
+        'onion', 'orange', 'paprika', 'pear', 'peas', 'pineapple', 'pomegranate', 'potato', 'raddish', 'soy beans', 
+        'spinach', 'sweetcorn', 'sweetpotato', 'tomato', 'turnip', 'watermelon']
+
 
 # Infer uploaded image function
 def infer_image(pil_image):
+    """
+    Returns a list of dicts: [{ "name": CapitalizedName, "confidence": 0.xx }, ...]
+    """
     img = pil_image.resize((224, 224))
     arr = np.expand_dims(np.array(img) / 255.0, axis=0)
-
     preds = MODEL.predict(arr)[0]
+    # Top 3 predictions
     top_idxs = np.argsort(preds)[::-1][:3]
-
-    ingredients = [
-        {"name": CLASS_NAMES[i].capitalize(), "confidence": float(preds[i])}
-        for i in top_idxs
-    ]
-
-    return ingredients or [{"name": "Unknown", "confidence": 0.0}]
-
+    ingredients = []
+    for i in top_idxs:
+        ingredients.append({"name": CLASS_NAMES[i].capitalize(), "confidence": float(preds[i])})
+    if not ingredients:
+        return [{"name": "Unknown", "confidence": 0.0}]
+    return ingredients
 
 
 # initialize FastAPI app
 app = FastAPI(
     title="Fridge2Dish",
-    description="Upload image → Detect ingredients → Generate recipes",
-    version="2.0.0"
+    description="Upload an image → Detect ingredients → Generate recipes",
+    version="3.0.0"
 )
 
-# Serve static files
+# static/templates
 app.mount("/static", StaticFiles(directory="static"), name="static")
 templates = Jinja2Templates(directory="templates")
 
 # CORS
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"], allow_credentials=True,
-    allow_methods=["*"], allow_headers=["*"]
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
 
@@ -215,24 +226,32 @@ def home(request: Request):
 async def upload_image(
     file: UploadFile = File(...),
     user_api_key: str = Form(alias="api_key", default="")
-):
-
+    ):
+         
     try:
         if not file.filename.lower().endswith((".jpg", ".jpeg", ".png")):
             raise HTTPException(status_code=400, detail="Invalid image format.")
-        
-        
-        image_bytes = await file.read()
-        pil_img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-        
-        # Load image
+
+        # read image
+        img_bytes = await file.read()
+        pil_img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+
+        # detect ingredients
+        start = time.time()
         ingredients = infer_image(pil_img)
-        ingredient_names = [i["name"] for i in ingredients]
+        dur = time.time() - start
+        print(f"Detected ingredients: {ingredients} (took {dur:.2f}s)")
 
-        # Try Gemini if user provided a key
-        if user_api_key.strip():
+        ingredient_names = [it["name"] for it in ingredients]
+
+        recipe_text = None
+        api_key = user_api_key.strip()
+
+        # Try server Gemini if api_key provided
+        if api_key:
             try:
-                genai.configure(api_key=user_api_key.strip())
+                # Try Gemini first...
+                genai.configure(api_key=api_key)
                 model = genai.GenerativeModel("gemini-2.5-flash")
 
                 prompt = f"""
@@ -248,17 +267,23 @@ async def upload_image(
 
                 response = model.generate_content(prompt)
                 recipe_text = response.text.strip()
-
-            except Exception as e1:
-                print("\n⚠ Gemini failed. Switching to GPT-2 fallback.", e1)
-                recipe_text = generate_recipe_local(ingredient_names)
+                print("\nGemini succeeded.")
+                
+            except Exception as e_gem:
+                # Log and fallback to local Gemma
+                print("Gemini failed or threw exception; falling back to local Gemma:", e_gem)
+                recipe_text = generate_recipe_local_gemma(ingredient_names)
 
         else:
-            print("\nNo API key → Using GPT-2 fallback")
-            recipe_text = generate_recipe_local(ingredient_names)
+            # No API key -> local Gemma
+            print("\nNo API key provided -> Using local Gemma fallback.")
+            recipe_text = generate_recipe_local_gemma(ingredient_names)
 
+        # Return structured response (ingredients keep confidence)
         return {"ingredients": ingredients, "recipe": recipe_text}
 
+    except HTTPException:
+        raise
     except Exception as e:
         traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"Server Error: {str(e)}")