Update app.py

app.py

@@ -1,49 +1,27 @@
+from fastapi import FastAPI, UploadFile, File
+import shutil
 import os
-
-os.environ["HF_HOME"] = "/tmp/huggingface"
-os.makedirs("/tmp/huggingface", exist_ok=True)
-from fastapi import FastAPI
-from pydantic import BaseModel
-import torch
-import numpy as np
-from transformers import AutoTokenizer, AutoModel
-from sklearn.linear_model import LogisticRegression
-import uvicorn
+from image_caption import generate_caption
 
 app = FastAPI()
 
-# Load Hugging Face model
-model_name = "bert-base-uncased"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name)
-
-# Function to get text embeddings
-def get_embedding(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    return outputs.last_hidden_state[:, 0, :].numpy()
-
-# Sample dataset
-texts = ["I love this!", "This is terrible.", "Fantastic experience!", "I hate it.", "Absolutely wonderful!", "Worst ever!"]
-labels = [1, 0, 1, 0, 1, 0]  # 1 = Positive, 0 = Negative
-X = np.vstack([get_embedding(text) for text in texts])
-y = np.array(labels)
+UPLOAD_DIR = "uploads"
+os.makedirs(UPLOAD_DIR, exist_ok=True)
 
-# Train model
-clf = LogisticRegression()
-clf.fit(X, y)
+@app.post("/upload-image/")
+async def upload_image(file: UploadFile = File(...)):
+    """Endpoint to accept image and return generated caption"""
+    file_path = f"{UPLOAD_DIR}/{file.filename}"
+    
+    # Save the uploaded file
+    with open(file_path, "wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
 
-# Define request format
-class InputText(BaseModel):
-    text: str
+    # Generate caption
+    caption = generate_caption(file_path)
 
-@app.post("/predict")
-def predict_sentiment(data: InputText):
-    user_embedding = get_embedding(data.text)
-    prediction = clf.predict(user_embedding)
-    sentiment = "Positive 😊" if prediction[0] == 1 else "Negative 😡"
-    return {"sentiment": sentiment}
+    return {"filename": file.filename, "caption": caption}
 
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+@app.get("/")
+async def root():
+    return {"message": "Image-to-Text API is running. Use /upload-image to send an image."}