Spaces:
Runtime error
Runtime error
Ravi Pandey
commited on
Commit
•
0045e3d
1
Parent(s):
bab727f
Docker deploy
Browse files- Dockerfile +1 -1
- caption_server.py +89 -0
- classification_server.py +41 -0
Dockerfile
CHANGED
@@ -14,4 +14,4 @@ ENV HF_HOME=/tmp/transformers_cache
|
|
14 |
COPY . .
|
15 |
|
16 |
# Command to run your application
|
17 |
-
CMD ["uvicorn", "
|
|
|
14 |
COPY . .
|
15 |
|
16 |
# Command to run your application
|
17 |
+
CMD ["uvicorn", "caption_server:app", "--reload", "--host", "0.0.0.0", "--port", "8000"]
|
caption_server.py
ADDED
@@ -0,0 +1,89 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from fastapi import FastAPI, HTTPException, Request
|
2 |
+
from fastapi.middleware.cors import CORSMiddleware
|
3 |
+
from transformers import pipeline
|
4 |
+
from PIL import Image
|
5 |
+
import requests
|
6 |
+
from io import BytesIO
|
7 |
+
import uvicorn
|
8 |
+
from sentence_transformers import SentenceTransformer, util
|
9 |
+
|
10 |
+
app = FastAPI()
|
11 |
+
|
12 |
+
# CORS setup
|
13 |
+
origins = ["*"]
|
14 |
+
app.add_middleware(
|
15 |
+
CORSMiddleware,
|
16 |
+
allow_origins=origins,
|
17 |
+
allow_credentials=True,
|
18 |
+
allow_methods=["*"],
|
19 |
+
allow_headers=["*"],
|
20 |
+
)
|
21 |
+
|
22 |
+
|
23 |
+
|
24 |
+
# Initialize the BLIP and GIT models for image captioning
|
25 |
+
blip_model_name = "Salesforce/blip-image-captioning-large"
|
26 |
+
git_model_name = "microsoft/git-large-coco"
|
27 |
+
blip_pipeline = pipeline("image-to-text", model=blip_model_name )
|
28 |
+
git_coco_pipeline = pipeline("image-to-text", model=git_model_name)
|
29 |
+
|
30 |
+
# Initialize the model for semantic similarity
|
31 |
+
model = SentenceTransformer('all-MiniLM-L6-v2')
|
32 |
+
|
33 |
+
classification_server_url = 'http://localhost:8001/classify_text'
|
34 |
+
|
35 |
+
def calculate_similarity(caption1, caption2):
|
36 |
+
embeddings = model.encode([caption1, caption2], convert_to_tensor=True)
|
37 |
+
similarity = util.pytorch_cos_sim(embeddings[0], embeddings[1])
|
38 |
+
return similarity.item()
|
39 |
+
|
40 |
+
def select_best_caption(caption1, caption2):
|
41 |
+
similarity_threshold = 0.6486401639
|
42 |
+
similarity_score = calculate_similarity(caption1, caption2)
|
43 |
+
|
44 |
+
if similarity_score > similarity_threshold:
|
45 |
+
print("Similarity score:",similarity_score)
|
46 |
+
return caption1 if len(caption1) > len(caption2) else caption2
|
47 |
+
|
48 |
+
else:
|
49 |
+
print(f"Git is selected with score :{similarity_score}")
|
50 |
+
return caption1
|
51 |
+
|
52 |
+
@app.post("/generate_caption_and_classify")
|
53 |
+
async def generate_caption_and_classify(request: Request):
|
54 |
+
try:
|
55 |
+
body = await request.json()
|
56 |
+
image_urls = body.get('image_urls', [])
|
57 |
+
results = []
|
58 |
+
|
59 |
+
for index, image_url in enumerate(image_urls):
|
60 |
+
response = requests.get(image_url)
|
61 |
+
image = Image.open(BytesIO(response.content))
|
62 |
+
|
63 |
+
# Generate captions from both models
|
64 |
+
caption_git = git_coco_pipeline(image, max_new_tokens=50)[0]['generated_text']
|
65 |
+
print("Caption generated by GIT:" , caption_git)
|
66 |
+
caption_blip = blip_pipeline(image, max_new_tokens=50)[0]['generated_text']
|
67 |
+
print("Caption generated by BLIP:" , caption_blip)
|
68 |
+
# Select the best caption
|
69 |
+
best_caption = select_best_caption(caption_git, caption_blip)
|
70 |
+
|
71 |
+
# Classify the best caption
|
72 |
+
classification_response = requests.post(
|
73 |
+
classification_server_url,
|
74 |
+
json={"texts": [best_caption]}
|
75 |
+
)
|
76 |
+
classification_result = classification_response.json()[0]
|
77 |
+
|
78 |
+
results.append({
|
79 |
+
"image_url": image_url,
|
80 |
+
"caption": best_caption,
|
81 |
+
"classification": classification_result
|
82 |
+
})
|
83 |
+
|
84 |
+
return results
|
85 |
+
except Exception as e:
|
86 |
+
raise HTTPException(status_code=500, detail=str(e))
|
87 |
+
|
88 |
+
if __name__ == "__main__":
|
89 |
+
uvicorn.run(app, host="0.0.0.0", port=8000)
|
classification_server.py
ADDED
@@ -0,0 +1,41 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from fastapi import FastAPI, HTTPException
|
2 |
+
from fastapi.middleware.cors import CORSMiddleware
|
3 |
+
from transformers import pipeline
|
4 |
+
from fastapi import FastAPI, Request
|
5 |
+
import uvicorn
|
6 |
+
|
7 |
+
app = FastAPI()
|
8 |
+
|
9 |
+
app.add_middleware(
|
10 |
+
CORSMiddleware,
|
11 |
+
allow_origins=["https://www.google.com"], # or specify other origins
|
12 |
+
allow_credentials=True,
|
13 |
+
allow_methods=["*"],
|
14 |
+
allow_headers=["*"],
|
15 |
+
)
|
16 |
+
|
17 |
+
classifier = pipeline("zero-shot-classification", model="MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli")
|
18 |
+
|
19 |
+
@app.post("/classify_text")
|
20 |
+
async def classify_text(request: Request):
|
21 |
+
try:
|
22 |
+
body = await request.json()
|
23 |
+
texts = body.get('texts', [])
|
24 |
+
labels = body.get('labels', ["Violent", "Neutral", "Sexually explicit"])
|
25 |
+
results = []
|
26 |
+
|
27 |
+
# Classify the selected caption and send response to server
|
28 |
+
for text in texts:
|
29 |
+
result = classifier(text, labels , max_new_tokens=50)
|
30 |
+
results.append({
|
31 |
+
"text": text,
|
32 |
+
"category": result["labels"][0],
|
33 |
+
"scores": result["scores"][0]
|
34 |
+
})
|
35 |
+
|
36 |
+
return results
|
37 |
+
except Exception as e:
|
38 |
+
raise HTTPException(status_code=500, detail=str(e))
|
39 |
+
|
40 |
+
if __name__ == "__main__":
|
41 |
+
uvicorn.run(app, host="0.0.0.0", port=8001)
|