Update app.py

app.py

@@ -7,6 +7,8 @@ from tensorflow.keras.models import load_model
 from keras.preprocessing.image import img_to_array
 from keras.applications.inception_v3 import preprocess_input
 import os
+from PIL import Image
+import io
 
 openai.api_key = os.getenv('OPENAI_API_KEY')
 
@@ -18,14 +20,34 @@ class_labels = [
     "Peeling",
 ]
 
-def encode_image(image_path):
-    with open(image_path, "rb") as image_file:
-        return base64.b64encode(image_file.read()).decode("utf-8")
-
 @st.cache_resource
 def load_trained_model():
     return load_model('my_new_model12.h5')
 
+def compress_image(image_bytes, max_size_kb=500):
+    # Open the image
+    img = Image.open(io.BytesIO(image_bytes))
+    
+    # Initialize quality
+    quality = 95
+    output_bytes = io.BytesIO()
+    
+    # Compress until size is under max_size_kb
+    while True:
+        output_bytes.seek(0)
+        output_bytes.truncate()
+        img.save(output_bytes, format='JPEG', quality=quality)
+        if len(output_bytes.getvalue()) <= max_size_kb * 1024 or quality <= 5:
+            break
+        quality -= 5
+    
+    return output_bytes.getvalue()
+
+def process_image_for_openai(image_bytes):
+    # Compress image to ensure it fits within token limits
+    compressed_image = compress_image(image_bytes)
+    return base64.b64encode(compressed_image).decode('utf-8')
+
 loaded_model = load_trained_model()
 
 st.title("Wall Defect Classification and AI Analysis")
@@ -34,46 +56,36 @@ st.write("Upload an image to classify wall defects and generate AI-based descrip
 uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
+    # Read file bytes once
+    file_bytes = uploaded_file.getvalue()
+    
     # Display the uploaded image
-    st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
-
-    # Read and preprocess the input image
-    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
-    input_img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
+    st.image(file_bytes, caption="Uploaded Image", use_column_width=True)
+    
+    # Process for model prediction
+    input_img = cv2.imdecode(np.frombuffer(file_bytes, np.uint8), cv2.IMREAD_COLOR)
     input_img_resized = cv2.resize(input_img, dsize=(256,256), interpolation=cv2.INTER_CUBIC)
-
     x = img_to_array(input_img_resized)
     x = np.expand_dims(x, axis=0)
     x = preprocess_input(x)
-
+    
     preds = loaded_model.predict(x)
-
-    # Get the index of the class with the maximum probability
     class_index = np.argmax(preds[0])
-    
-    # The corresponding maximum probability
     max_probability = preds[0][class_index]
-    
-    # Get the class name for the predicted index
     class_name = class_labels[class_index]
     
-    # Prepare the results text
     results_text = f"{class_name} (Class {class_index}): Probability {max_probability:.2f}\n"
-    
-    # Display classification results in a text box
     st.text_area("Classification Results:", value=results_text, height=200)
 
-    # Encode the uploaded image as Base64
-    base64_image = base64.b64encode(file_bytes).decode("utf-8")
-
-    # If probability < 0.59, show a warning and skip AI analysis
     if max_probability < 0.59:
         st.warning(
             "The confidence for this prediction is below 59%. "
             "Please do a manual review."
         )
     else:
-        # Generate AI-based descriptions using OpenAI API
+        # Compress and encode image for OpenAI
+        compressed_base64 = process_image_for_openai(file_bytes)
+        
         defects_string = class_name
         ai_prompt = (
             f"Our trained model predicts the following defect: {defects_string}. "
@@ -81,31 +93,30 @@ if uploaded_file is not None:
             f"for this defect? The output format should be:\n"
             f"Category ID: <Category_ID>\n"
             f"Title: <Title>\n"
-            f"Description: <description>"
-            # f"Please generate description in 150 words"
+            f"Description: <description in 100 words or less>"
         )
-
+        
         st.write("Analyzing image with AI...")
         try:
             response = openai.ChatCompletion.create(
-                model="gpt-4o",
+                model="gpt-4-vision-preview",  # Using vision model instead of gpt-4
                 messages=[
                     {
                         "role": "user",
-                        "content": ai_prompt
-                    },
-                    {
-                        # Some OpenAI endpoints may not support directly passing an image in a single message like this.
-                        # If necessary, provide a direct link or your own approach to referencing the image.
-                        "role": "user",
-                        "content": f"Base64 Image:\n{base64_image}"
+                        "content": [
+                            {"type": "text", "text": ai_prompt},
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": f"data:image/jpeg;base64,{compressed_base64}"
+                                }
+                            }
+                        ]
                     }
                 ],
                 max_tokens=300,
             )
-            # Extract AI-generated descriptions
             ai_description = response.choices[0].message.content
             st.text_area("AI-Generated Description:", value=ai_description, height=200)
-
         except Exception as e:
             st.error(f"An error occurred while generating AI-based descriptions: {str(e)}")