Update app.py

app.py

@@ -15,91 +15,44 @@ from PIL import Image as PILImage
 from io import BytesIO
 
 # Streamlit title
-st.title("🚗 Vehicle Information Extraction from Images 🚙")
-
-# Custom CSS Styling for the app
-st.markdown("""
-    <style>
-        body {
-            background-color: #f0f8ff;
-            font-family: 'Arial', sans-serif;
-            color: #333;
-        }
-
-        .stButton button {
-            background-color: #1E90FF;
-            color: white;
-            border-radius: 5px;
-            font-size: 16px;
-            padding: 10px 20px;
-            transition: transform 0.2s ease-in-out;
-        }
-
-        .stButton button:hover {
-            background-color: #4682b4;
-            transform: scale(1.05);
-        }
-
-        .stTitle {
-            font-size: 36px;
-            font-weight: bold;
-            color: #1E90FF;
-        }
-
-        .stSubheader {
-            color: #4682b4;
-            font-size: 24px;
-            font-weight: 600;
-            text-align: center;
-        }
-
-        .stImage {
-            border-radius: 10px;
-            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
-        }
-
-        .stJson {
-            background-color: #f8f9fa;
-            padding: 15px;
-            border-radius: 5px;
-            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
-            font-size: 16px;
-        }
-
-        .grid-container {
-            display: grid;
-            grid-template-columns: repeat(3, 1fr);
-            gap: 10px;
-            padding: 10px;
-        }
-
-        .grid-container img {
-            border-radius: 10px;
-            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
-            transition: transform 0.2s ease-in-out;
-        }
-
-        .grid-container img:hover {
-            transform: scale(1.05);
-        }
-
-        @keyframes fadeIn {
-            from {
-                opacity: 0;
-            }
-            to {
-                opacity: 1;
-            }
-        }
-
-        .fade-in {
-            animation: fadeIn 1s ease-in-out;
-        }
-
-    </style>
-""", unsafe_allow_html=True)
-
-# Image display function with animation
+st.title("Vehicle Information Extraction from Images")
+
+translateimg = PILImage.open("car.JPG")  # Ensure the file is in the correct directory
+st.image(translateimg, use_container_width=True)  # Adjust the size as per preference
+
+# Prompt user for OpenAI API key
+openai_api_key = st.text_input("Enter your OpenAI API Key:", type="password")
+
+# Set the OpenAI API key if provided
+if openai_api_key:
+    os.environ["OPENAI_API_KEY"] = openai_api_key
+
+# Vehicle class (same as in the original code)
+class Vehicle(BaseModel):
+    Type: str = Field(..., examples=["Car", "Truck", "Motorcycle", 'Bus', 'Van'], description="The type of the vehicle.")
+    License: str = Field(..., description="The license plate number of the vehicle.")
+    Make: str = Field(..., examples=["Toyota", "Honda", "Ford", "Suzuki"], description="The Make of the vehicle.")
+    Model: str = Field(..., examples=["Corolla", "Civic", "F-150"], description="The Model of the vehicle.")
+    Color: str = Field(..., example=["Red", "Blue", "Black", "White"], description="Return the color of the vehicle.")
+    Year: str = Field(None, description="The year of the vehicle.")
+    Condition: str = Field(None, description="The condition of the vehicle.")
+    Logo: str = Field(None, description="The visible logo of the vehicle, if applicable.")
+    Damage: str = Field(None, description="Any visible damage or wear and tear on the vehicle.")
+    Region: str = Field(None, description="Region or country based on the license plate or clues from the image.")
+    PlateType: str = Field(None, description="Type of license plate, e.g., government, personal.")
+
+# Parser for vehicle details
+parser = JsonOutputParser(pydantic_object=Vehicle)
+instructions = parser.get_format_instructions()
+
+# Image encoding function (for base64 encoding)
+def image_encoding(inputs):
+    """Load and convert image to base64 encoding"""
+    with open(inputs["image_path"], "rb") as image_file:
+        image_base64 = base64.b64encode(image_file.read()).decode("utf-8")
+    return {"image": image_base64}
+
+# Image display in grid (for multiple images)
 def display_image_grid(image_paths, rows=2, cols=3, figsize=(10, 7)):
     fig = plt.figure(figsize=figsize)
     max_images = rows * cols
@@ -116,15 +69,41 @@ def display_image_grid(image_paths, rows=2, cols=3, figsize=(10, 7)):
     plt.tight_layout()
     st.pyplot(fig)
 
-# Image encoding function
-def image_encoding(inputs):
-    """Load and convert image to base64 encoding"""
-    with open(inputs["image_path"], "rb") as image_file:
-        image_base64 = base64.b64encode(image_file.read()).decode("utf-8")
-    return {"image": image_base64}
+# Create the prompt for the AI model
+@chain
+def prompt(inputs):
+    prompt = [
+        SystemMessage(content="""You are an AI assistant tasked with extracting detailed information from a vehicle image. Please extract the following details:
+        - Vehicle type (e.g., Car, Truck, Bus)
+        - License plate number and type (if identifiable, such as personal, commercial, government)
+        - Vehicle make, model, and year (e.g., 2020 Toyota Corolla)
+        - Vehicle color and condition (e.g., Red, well-maintained, damaged)
+        - Any visible brand logos or distinguishing marks (e.g., Tesla logo)
+        - Details of any visible damage (e.g., scratches, dents)
+        - Vehicle’s region or country (based on the license plate or other clues)
+        If some details are unclear or not visible, return `None` for those fields. Do not guess or provide inaccurate information."""),
+        HumanMessage(
+            content=[ 
+                {"type": "text", "text": "Analyze the vehicle in the image and extract as many details as possible, including type, license plate, make, model, year, condition, damage, etc."},
+                {"type": "text", "text": instructions},  # include any other format instructions here
+                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{inputs['image']}", "detail": "low"}}
+            ]
+        )
+    ]
+    return prompt
+
+# Invoke the model for extracting vehicle details
+@chain
+def MLLM_response(inputs):
+    model: ChatOpenAI = ChatOpenAI(model="gpt-4o-2024-08-06", temperature=0.0, max_tokens=1024)
+    output = model.invoke(inputs)
+    return output.content
+
+# The complete pipeline for extracting vehicle details
+pipeline = image_encoding | prompt | MLLM_response | parser
 
 # Streamlit Interface for uploading images and showing results
-st.header("🚗 Upload Vehicle Images for Information Extraction 🚙")
+st.header("Upload Vehicle Images for Information Extraction")
 
 # Option to select either single or batch image upload
 upload_option = st.radio("Select Upload Type", ["Single Image Upload", "Batch Images Upload"])
@@ -135,23 +114,23 @@ if upload_option == "Single Image Upload":
     uploaded_image = st.file_uploader("Choose an Image (JPEG, PNG, GIF, BMP, etc.)", type=["jpeg", "png", "gif", "bmp", "jpg"])
 
     if uploaded_image is not None:
-        # Display the uploaded image with animation
+        # Display the uploaded image
         image = PILImage.open(uploaded_image)
-        st.image(image, caption="Uploaded Image", use_container_width=True, class_="fade-in")
+        st.image(image, caption="Uploaded Image", use_container_width=True)
 
         # Convert the uploaded image to base64
         image_path = "/tmp/uploaded_image" + os.path.splitext(uploaded_image.name)[1]
         with open(image_path, "wb") as f:
             f.write(uploaded_image.getbuffer())
 
-        # Add button to trigger information extraction with animated hover effect
+        # Add button to trigger information extraction
         if st.button("Extract Vehicle Information"):
             # Process the image through the pipeline
             output = pipeline.invoke({"image_path": image_path})
 
             # Show the results in a user-friendly format
-            st.subheader("Extracted Vehicle Information", class_="fade-in")
-            st.json(output, class_="fade-in")
+            st.subheader("Extracted Vehicle Information")
+            st.json(output)
 
 # Batch Images Upload
 elif upload_option == "Batch Images Upload":
@@ -169,12 +148,8 @@ elif upload_option == "Batch Images Upload":
             # Process the batch and display the results in a DataFrame
             batch_output = pipeline.batch(batch_input)
             df = pd.DataFrame(batch_output)
-            st.dataframe(df, use_container_width=True)
+            st.dataframe(df)
 
-        # Display batch images in a grid with animation
-        st.subheader("Images in Grid", class_="fade-in")
+        # Show images in a grid
         image_paths = [f"/tmp/{file.name}" for file in batch_images]
-        st.markdown('<div class="grid-container">', unsafe_allow_html=True)
-        for image_path in image_paths:
-            st.markdown(f'<img src="data:image/jpeg;base64,{base64.b64encode(open(image_path, "rb").read()).decode()}" class="fade-in"/>', unsafe_allow_html=True)
-        st.markdown('</div>', unsafe_allow_html=True)
+        display_image_grid(image_paths)