Spaces:
Runtime error
Runtime error
| from fastapi import FastAPI, HTTPException, Request | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from transformers import pipeline | |
| from PIL import Image | |
| import requests | |
| from io import BytesIO | |
| import uvicorn | |
| from sentence_transformers import SentenceTransformer, util | |
| app = FastAPI() | |
| # CORS setup | |
| origins = ["*"] | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=origins, | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # Initialize the BLIP and GIT models for image captioning | |
| blip_model_name = "Salesforce/blip-image-captioning-large" | |
| git_model_name = "microsoft/git-large-coco" | |
| blip_pipeline = pipeline("image-to-text", model=blip_model_name ) | |
| git_coco_pipeline = pipeline("image-to-text", model=git_model_name) | |
| # Initialize the model for semantic similarity | |
| model = SentenceTransformer('all-MiniLM-L6-v2') | |
| classification_server_url = 'https://raviipandey-classification.hf.space/classify_text' | |
| def calculate_similarity(caption1, caption2): | |
| embeddings = model.encode([caption1, caption2], convert_to_tensor=True) | |
| similarity = util.pytorch_cos_sim(embeddings[0], embeddings[1]) | |
| return similarity.item() | |
| def select_best_caption(caption1, caption2): | |
| similarity_threshold = 0.6486401639 | |
| similarity_score = calculate_similarity(caption1, caption2) | |
| if similarity_score > similarity_threshold: | |
| print("Similarity score:",similarity_score) | |
| return caption1 if len(caption1) > len(caption2) else caption2 | |
| else: | |
| print(f"Git is selected with score :{similarity_score}") | |
| return caption1 | |
| async def generate_caption_and_classify(request: Request): | |
| try: | |
| body = await request.json() | |
| image_urls = body.get('image_urls', []) | |
| results = [] | |
| for index, image_url in enumerate(image_urls): | |
| response = requests.get(image_url) | |
| image = Image.open(BytesIO(response.content)) | |
| # Generate captions from both models | |
| caption_git = git_coco_pipeline(image, max_new_tokens=50)[0]['generated_text'] | |
| print("Caption generated by GIT:" , caption_git) | |
| caption_blip = blip_pipeline(image, max_new_tokens=50)[0]['generated_text'] | |
| print("Caption generated by BLIP:" , caption_blip) | |
| # Select the best caption | |
| best_caption = select_best_caption(caption_git, caption_blip) | |
| # Classify the best caption | |
| classification_response = requests.post( | |
| classification_server_url, | |
| json={"texts": [best_caption]} | |
| ) | |
| classification_result = classification_response.json()[0] | |
| results.append({ | |
| "image_url": image_url, | |
| "caption": best_caption, | |
| "classification": classification_result | |
| }) | |
| return results | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=7860) | |