added exception handling

.ipynb_checkpoints/app-checkpoint.py

@@ -1,6 +1,6 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 from transformers import AutoProcessor, Blip2ForConditionalGeneration
 import torch
 import io
@@ -9,26 +9,41 @@ app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Adjust this as needed for security
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
 # Load the model and processor
-model = Blip2ForConditionalGeneration.from_pretrained("ybelkada/blip2-opt-2.7b-fp16-sharded")
-model.load_adapter('blip-cpu-model')
-processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+try:
+    model = Blip2ForConditionalGeneration.from_pretrained("ybelkada/blip2-opt-2.7b-fp16-sharded")
+    model.load_adapter('blip-cpu-model')
+    processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+except Exception as e:
+    raise RuntimeError(f"Failed to load the model or processor: {str(e)}")
 
 @app.post("/generate-caption/")
 async def generate_caption(file: UploadFile = File(...)):
-    image = Image.open(io.BytesIO(await file.read()))
-    inputs = processor(images=image, return_tensors="pt").to(device, torch.float16)
+    try:
+        image = Image.open(io.BytesIO(await file.read()))
+    except UnidentifiedImageError:
+        # Raise a 400 error if the file is not a valid image
+        raise HTTPException(status_code=400, detail="Uploaded file is not a valid image.")
+    except Exception as e:
+        # Catch any other unexpected errors related to image processing
+        raise HTTPException(status_code=500, detail=f"An unexpected error occurred while processing the image: {str(e)}")    
     
-    with torch.no_grad():
-        caption_ids = model.generate(**inputs, max_length=128)
-        caption = processor.decode(caption_ids[0], skip_special_tokens=True)
-    
-    return {"caption": caption}
+    try:
+        inputs = processor(images=image, return_tensors="pt").to(device, torch.float16)
+
+        with torch.no_grad():
+            caption_ids = model.generate(**inputs, max_length=128)
+            caption = processor.decode(caption_ids[0], skip_special_tokens=True)
+
+        return {"caption": caption}
+    except Exception as e:
+        # Catch any errors during the caption generation process
+        raise HTTPException(status_code=500, detail=f"An error occurred while generating the caption: {str(e)}")

.ipynb_checkpoints/streamlit_app-checkpoint.py

@@ -1,41 +1,38 @@
 import streamlit as st
 import requests
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 from transformers import AutoProcessor, Blip2ForConditionalGeneration
 import torch
 import io
 
 
-# @st.cache_resource
-# def load_model():
-#     model = Blip2ForConditionalGeneration.from_pretrained("ybelkada/blip2-opt-2.7b-fp16-sharded")
-#     model.load_adapter('blip-cpu-model')
-#     processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
-#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-#     model.to(device)
-#     return model, processor
-
-# model, processor = load_model()
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 st.title("Image Captioning with Fine-Tuned BLiPv2 Model")
 
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-    
-    files = {"file": uploaded_file.getvalue()}
-    print("Sending API request")
-    response = requests.post("http://0.0.0.0:8502/generate-caption/", files=files)
-    caption = response.json().get("caption")
-    
-#     inputs = processor(images=image, return_tensors="pt").to(device, torch.float16)
-    
-#     with torch.no_grad():
-#         caption_ids = model.generate(**inputs, max_length=128)
-#         caption = processor.decode(caption_ids[0], skip_special_tokens=True)
-    
-    st.write("Generated Caption:")
-    st.write(f"**{caption}**")
+    try:
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+        
+        try:
+            files = {"file": uploaded_file.getvalue()}
+            print("Sending API request")
+            response = requests.post("http://0.0.0.0:8502/generate-caption/", files=files)
+            caption = response.json().get("caption")
+        except requests.exceptions.RequestException as e:
+            st.error(f"An error occurred while making the API request: {str(e)}")
+            caption = "Error generating caption"
+        except ValueError as e:
+            st.error(f"An error occurred while parsing the API response: {str(e)}")
+            caption = "Error generating caption"
+
+        st.write("Generated Caption:")
+        st.write(f"**{caption}**")
+        
+    except UnidentifiedImageError:
+        st.error("The uploaded file is not a valid image. Please upload a JPG, JPEG, or PNG file.")
+    except Exception as e:
+        st.error(f"An unexpected error occurred: {str(e)}")
+else:
+    st.write("Please upload an image to generate a caption.")

app.py

@@ -1,6 +1,6 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 from transformers import AutoProcessor, Blip2ForConditionalGeneration
 import torch
 import io
@@ -9,26 +9,41 @@ app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Adjust this as needed for security
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
 # Load the model and processor
-model = Blip2ForConditionalGeneration.from_pretrained("ybelkada/blip2-opt-2.7b-fp16-sharded")
-model.load_adapter('blip-cpu-model')
-processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+try:
+    model = Blip2ForConditionalGeneration.from_pretrained("ybelkada/blip2-opt-2.7b-fp16-sharded")
+    model.load_adapter('blip-cpu-model')
+    processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+except Exception as e:
+    raise RuntimeError(f"Failed to load the model or processor: {str(e)}")
 
 @app.post("/generate-caption/")
 async def generate_caption(file: UploadFile = File(...)):
-    image = Image.open(io.BytesIO(await file.read()))
-    inputs = processor(images=image, return_tensors="pt").to(device, torch.float16)
+    try:
+        image = Image.open(io.BytesIO(await file.read()))
+    except UnidentifiedImageError:
+        # Raise a 400 error if the file is not a valid image
+        raise HTTPException(status_code=400, detail="Uploaded file is not a valid image.")
+    except Exception as e:
+        # Catch any other unexpected errors related to image processing
+        raise HTTPException(status_code=500, detail=f"An unexpected error occurred while processing the image: {str(e)}")    
     
-    with torch.no_grad():
-        caption_ids = model.generate(**inputs, max_length=128)
-        caption = processor.decode(caption_ids[0], skip_special_tokens=True)
-    
-    return {"caption": caption}
+    try:
+        inputs = processor(images=image, return_tensors="pt").to(device, torch.float16)
+
+        with torch.no_grad():
+            caption_ids = model.generate(**inputs, max_length=128)
+            caption = processor.decode(caption_ids[0], skip_special_tokens=True)
+
+        return {"caption": caption}
+    except Exception as e:
+        # Catch any errors during the caption generation process
+        raise HTTPException(status_code=500, detail=f"An error occurred while generating the caption: {str(e)}")

streamlit_app.py

@@ -1,41 +1,38 @@
 import streamlit as st
 import requests
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 from transformers import AutoProcessor, Blip2ForConditionalGeneration
 import torch
 import io
 
 
-# @st.cache_resource
-# def load_model():
-#     model = Blip2ForConditionalGeneration.from_pretrained("ybelkada/blip2-opt-2.7b-fp16-sharded")
-#     model.load_adapter('blip-cpu-model')
-#     processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
-#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-#     model.to(device)
-#     return model, processor
-
-# model, processor = load_model()
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 st.title("Image Captioning with Fine-Tuned BLiPv2 Model")
 
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-    
-    files = {"file": uploaded_file.getvalue()}
-    print("Sending API request")
-    response = requests.post("http://0.0.0.0:8502/generate-caption/", files=files)
-    caption = response.json().get("caption")
-    
-#     inputs = processor(images=image, return_tensors="pt").to(device, torch.float16)
-    
-#     with torch.no_grad():
-#         caption_ids = model.generate(**inputs, max_length=128)
-#         caption = processor.decode(caption_ids[0], skip_special_tokens=True)
-    
-    st.write("Generated Caption:")
-    st.write(f"**{caption}**")
+    try:
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+        
+        try:
+            files = {"file": uploaded_file.getvalue()}
+            print("Sending API request")
+            response = requests.post("http://0.0.0.0:8502/generate-caption/", files=files)
+            caption = response.json().get("caption")
+        except requests.exceptions.RequestException as e:
+            st.error(f"An error occurred while making the API request: {str(e)}")
+            caption = "Error generating caption"
+        except ValueError as e:
+            st.error(f"An error occurred while parsing the API response: {str(e)}")
+            caption = "Error generating caption"
+
+        st.write("Generated Caption:")
+        st.write(f"**{caption}**")
+        
+    except UnidentifiedImageError:
+        st.error("The uploaded file is not a valid image. Please upload a JPG, JPEG, or PNG file.")
+    except Exception as e:
+        st.error(f"An unexpected error occurred: {str(e)}")
+else:
+    st.write("Please upload an image to generate a caption.")