Update app.py

app.py

@@ -16,38 +16,39 @@ import base64
 # Set Gemini API Key
 client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
 
+# System prompt
+SYSTEM_PROMPT = (
+    "You are part of an intelligent driver distraction detection system. "
+    "Your role is to analyze documents, images, and videos to help identify and classify driver states, "
+    "including: drowsy, alert, yawning, microsleep, eyes/lips open or closed, and other distractions. "
+    "The user may ask for timestamps of events in video, technical explanations from code/docs, "
+    "or structured summaries. Prioritize your answers for this use case and act as a helpful assistant "
+    "within this domain. "
+    "Additionally, your classification may rely on the State Farm Distracted Driver Detection dataset, "
+    "which defines 10 distraction categories such as texting, phone use, reaching behind, adjusting the radio, "
+    "and safe driving. Use this labeling scheme when interpreting relevant images or video frames."
+)
+
 # --------- File Extractors ---------
 def extract_text_from_docx(docx_file):
-    document = Document(docx_file)
-    return "\n".join([para.text for para in document.paragraphs])
+    return "\n".join([para.text for para in Document(docx_file).paragraphs])
 
 def extract_text_from_csv(csv_file):
-    df = pd.read_csv(csv_file)
-    return df.to_string(index=False)
+    return pd.read_csv(csv_file).to_string(index=False)
 
 def extract_text_from_xlsx(xlsx_file):
-    df = pd.read_excel(xlsx_file)
-    return df.to_string(index=False)
+    return pd.read_excel(xlsx_file).to_string(index=False)
 
 def extract_text_from_pptx(pptx_file):
     prs = Presentation(pptx_file)
-    text_runs = []
-    for slide in prs.slides:
-        for shape in slide.shapes:
-            if hasattr(shape, "text"):
-                text_runs.append(shape.text)
-    return "\n".join(text_runs)
+    return "\n".join([shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text")])
 
 def extract_text_from_pdf(pdf_file):
     doc = fitz.open(stream=pdf_file.read(), filetype="pdf")
-    text = ""
-    for page in doc:
-        text += page.get_text()
-    return text
+    return "".join([page.get_text() for page in doc])
 
 def extract_text_from_html(html_file):
-    soup = BeautifulSoup(html_file.read(), "html.parser")
-    return soup.get_text()
+    return BeautifulSoup(html_file.read(), "html.parser").get_text()
 
 def extract_text_from_tex(tex_file):
     content = tex_file.read().decode("utf-8")
@@ -56,183 +57,131 @@ def extract_text_from_tex(tex_file):
     return content
 
 def process_image(image_file):
-    image_bytes = image_file.read()
-    encoded_image = base64.b64encode(image_bytes).decode("utf-8")
-    return {
-        "inline_data": {
-            "mime_type": image_file.type,
-            "data": encoded_image
-        }
-    }
+    return types.Part.from_data(
+        data=image_file.read(),
+        mime_type=image_file.type
+    )
 
 def process_video(video_file):
-    video_bytes = video_file.read()
     filetype = video_file.name.split(".")[-1].lower()
     mime_map = {
-        "avi": "video/x-msvideo",
-        "mkv": "video/x-matroska",
-        "flv": "video/x-flv",
-        "wmv": "video/x-ms-wmv",
+        "avi": "video/x-msvideo", "mkv": "video/x-matroska",
+        "flv": "video/x-flv", "wmv": "video/x-ms-wmv",
         "mpeg": "video/mpeg"
     }
     mime_type = mime_map.get(filetype, video_file.type)
-    encoded_video = base64.b64encode(video_bytes).decode("utf-8")
-    return {
-        "inline_data": {
-            "mime_type": mime_type,
-            "data": encoded_video
-        }
-    }
+    return types.Part.from_data(data=video_file.read(), mime_type=mime_type)
 
 def export_conversation_to_pdf(conversation_history):
     pdf_path = "Conversation.pdf"
     doc = SimpleDocTemplate(pdf_path, pagesize=A4)
     styles = getSampleStyleSheet()
     elements = []
-
     for i, (user_q, gemini_a) in enumerate(conversation_history):
         elements.append(Paragraph(f"<b>Q{i+1}:</b> {user_q}", styles["Normal"]))
         elements.append(Spacer(1, 8))
         elements.append(Paragraph(f"<b>A{i+1}:</b> {gemini_a.replace(chr(10), '<br/>')}", styles["Normal"]))
         elements.append(Spacer(1, 16))
-
     doc.build(elements)
     return pdf_path
 
 # --------- Main App ---------
 def main():
-    st.set_page_config(page_title="Gemini 2.0 Flash Chat Q&A", layout="wide")
-    st.title("\U0001F4C4 Chat with Your Documents leveraging Internet Using Gemini 2.0 Flash (Team 18)")
-
-    if "conversation" not in st.session_state:
-        st.session_state.conversation = []
-    if "documents_text" not in st.session_state:
-        st.session_state.documents_text = []
-    if "chat_active" not in st.session_state:
-        st.session_state.chat_active = True
-    if "chat_history" not in st.session_state:
-        st.session_state.chat_history = []
+    st.set_page_config(page_title="Driver State Chat with Gemini", layout="wide")
+    st.title("🚗 Chat with Driver Distraction Files using Gemini Flash (Team 18)")
+
+    for key in ["conversation", "documents_text", "chat_active", "chat_history", "image_data", "video_data"]:
+        if key not in st.session_state:
+            st.session_state[key] = [] if key == "conversation" else ""
 
     uploaded_files = st.file_uploader(
-        "Upload files (.txt, .docx, .csv, .xlsx, .pptx, .pdf, .html, .htm, .tex, .jpg, .jpeg, .png, video formats):",
-        type=[
-            "txt", "docx", "csv", "xlsx", "pptx", "pdf", "html", "htm", "tex",
-            "jpg", "jpeg", "png",
-            "mp4", "webm", "mov", "avi", "mkv", "flv", "wmv", "mpeg"
-        ],
+        "Upload documents, images, or videos:",
+        type=["txt", "docx", "csv", "xlsx", "pptx", "pdf", "html", "htm", "tex",
+              "jpg", "jpeg", "png", "mp4", "webm", "mov", "avi", "mkv", "flv", "wmv", "mpeg"],
         accept_multiple_files=True
     )
 
     if uploaded_files:
-        all_text_content = ""
-        for uploaded_file in uploaded_files:
-            filetype = uploaded_file.name.split(".")[-1].lower()
+        all_text = ""
+        for f in uploaded_files:
+            ext = f.name.split(".")[-1].lower()
             try:
-                if filetype == "txt":
-                    all_text_content += uploaded_file.read().decode("utf-8") + "\n"
-                elif filetype == "docx":
-                    all_text_content += extract_text_from_docx(uploaded_file) + "\n"
-                elif filetype == "csv":
-                    all_text_content += extract_text_from_csv(uploaded_file) + "\n"
-                elif filetype == "xlsx":
-                    all_text_content += extract_text_from_xlsx(uploaded_file) + "\n"
-                elif filetype == "pptx":
-                    all_text_content += extract_text_from_pptx(uploaded_file) + "\n"
-                elif filetype == "pdf":
-                    all_text_content += extract_text_from_pdf(uploaded_file) + "\n"
-                elif filetype in ["html", "htm"]:
-                    all_text_content += extract_text_from_html(uploaded_file) + "\n"
-                elif filetype == "tex":
-                    all_text_content += extract_text_from_tex(uploaded_file) + "\n"
-                elif filetype in ["jpg", "jpeg", "png"]:
-                    st.session_state.image_data = process_image(uploaded_file)
-                    all_text_content += "Image uploaded and processed.\n"
-                    st.image(uploaded_file, caption=uploaded_file.name, use_container_width=True)
-                elif filetype in ["mp4", "webm", "mov", "avi", "mkv", "flv", "wmv", "mpeg"]:
-                    st.session_state.video_data = process_video(uploaded_file)
-                    all_text_content += f"Video file '{uploaded_file.name}' uploaded and processed.\n"
-                    st.video(uploaded_file)
-                else:
-                    st.error(f"Unsupported file format: {uploaded_file.name}")
+                if ext == "txt":
+                    all_text += f.read().decode("utf-8") + "\n"
+                elif ext == "docx":
+                    all_text += extract_text_from_docx(f) + "\n"
+                elif ext == "csv":
+                    all_text += extract_text_from_csv(f) + "\n"
+                elif ext == "xlsx":
+                    all_text += extract_text_from_xlsx(f) + "\n"
+                elif ext == "pptx":
+                    all_text += extract_text_from_pptx(f) + "\n"
+                elif ext == "pdf":
+                    all_text += extract_text_from_pdf(f) + "\n"
+                elif ext in ["html", "htm"]:
+                    all_text += extract_text_from_html(f) + "\n"
+                elif ext == "tex":
+                    all_text += extract_text_from_tex(f) + "\n"
+                elif ext in ["jpg", "jpeg", "png"]:
+                    st.session_state.image_data = process_image(f)
+                    st.image(f, caption=f.name, use_container_width=True)
+                    all_text += f"Image {f.name} uploaded.\n"
+                elif ext in ["mp4", "webm", "mov", "avi", "mkv", "flv", "wmv", "mpeg"]:
+                    st.session_state.video_data = process_video(f)
+                    st.video(f)
+                    all_text += f"Video {f.name} uploaded.\n"
             except Exception as e:
-                st.error(f"Failed to extract/process {uploaded_file.name}: {e}")
-
-        st.session_state.documents_text = all_text_content
+                st.error(f"Error processing {f.name}: {e}")
+        st.session_state.documents_text = all_text
 
     if st.session_state.documents_text:
-        st.markdown("### \U0001F4AC Conversation")
-        for user_q, gemini_a in st.session_state.conversation:
-            st.markdown(f"**You:** {user_q}")
-            st.markdown(f"**Gemini:**\n\n{gemini_a}")
+        st.markdown("### 💬 Chat")
+        for q, a in st.session_state.conversation:
+            st.markdown(f"**You:** {q}")
+            st.markdown(f"**Gemini:** {a}")
 
         if st.session_state.chat_active:
-            with st.form(key="chat_form", clear_on_submit=True):
-                user_input = st.text_input("Ask a question about the documents (type 'exit' to stop):")
-                submit = st.form_submit_button("Send")
-
-                if submit and user_input:
+            with st.form("chat_form", clear_on_submit=True):
+                user_input = st.text_input("Ask a question about the uploaded files:")
+                submitted = st.form_submit_button("Send")
+                if submitted and user_input:
                     if user_input.strip().lower() == "exit":
                         st.session_state.chat_active = False
-                        st.success("Chat ended. Reload to start again.")
+                        st.success("Chat ended.")
                     else:
                         with st.spinner("Gemini is thinking..."):
-                            # -------- System context for Gemini --------
-                            system_context = {
-                                "role": "user",
-                                "text": (
-                                    "You are part of an intelligent driver distraction detection system. "
-                                    "Your role is to process images and videos to help identify and classify driver states, analyze documents, and answer the queries. "
-                                    "including: drowsy, alert, yawning, microsleep, eyes/lips open or closed, and other distractions like safe driving, texting - right, talking on the phone - right,texting - left, talking on the phone - left, operating the radio, drinking, reaching behind,hair and makeup,talking to passenger."
-                                    "The user may ask for timestamps of events in video, technical explanations from code/docs, "
-                                    "or structured summaries. Prioritize your answers for this use case and be precise."
-                                    "within this domain."
-                                )
-                            }
-
-                            content_blocks = [system_context]
+                            messages = [
+                                types.Content(role="user", parts=[types.Part(text=SYSTEM_PROMPT)])
+                            ]
 
                             if st.session_state.conversation:
-                                history_text = "\n\n".join(
-                                    f"Q: {entry[0]}\nA: {entry[1]}"
-                                    for entry in st.session_state.conversation
-                                )
-                                content_blocks.append({"text": f"Previous conversation:\n{history_text}"})
+                                conv = "\n".join([f"Q: {q}\nA: {a}" for q, a in st.session_state.conversation])
+                                messages.append(types.Content(role="user", parts=[types.Part(text=conv)]))
 
-                            if st.session_state.documents_text.strip():
-                                content_blocks.append({"text": f"Context:\n{st.session_state.documents_text}"})
+                            if st.session_state.documents_text:
+                                messages.append(types.Content(role="user", parts=[types.Part(text=st.session_state.documents_text)]))
 
-                            if "image_data" in st.session_state:
-                                content_blocks.append(st.session_state.image_data)
+                            if st.session_state.image_data:
+                                messages.append(types.Content(role="user", parts=[st.session_state.image_data]))
 
-                            if "video_data" in st.session_state:
-                                content_blocks.append(st.session_state.video_data)
+                            if st.session_state.video_data:
+                                messages.append(types.Content(role="user", parts=[st.session_state.video_data]))
 
-                            content_blocks.append({"text": f"Question: {user_input}"})
+                            messages.append(types.Content(role="user", parts=[types.Part(text=user_input)]))
 
-                            response = client.models.generate_content(
+                            response = client.generate_content(
                                 model="gemini-2.0-flash",
-                                contents=content_blocks,
-                                config={"tools": [{"google_search": {}}]}
+                                contents=messages
                             )
 
                             st.session_state.conversation.append((user_input, response.text))
-                            st.success("\U0001F4A1 Answer:")
+                            st.success("💡 Answer:")
                             st.write(response.text)
 
-                            st.session_state.chat_history.append({
-                                "question": user_input,
-                                "answer": response.text
-                            })
-
-        if st.session_state.conversation:
-            pdf_path = export_conversation_to_pdf(st.session_state.conversation)
-            with open(pdf_path, "rb") as f:
-                st.download_button(
-                    label="\U0001F4E5 Export Conversation as PDF",
-                    data=f,
-                    file_name="Conversation.pdf",
-                    mime="application/pdf"
-                )
+    if st.session_state.conversation:
+        pdf_path = export_conversation_to_pdf(st.session_state.conversation)
+        with open(pdf_path, "rb") as f:
+            st.download_button("⬇️ Export Conversation as PDF", f, "Conversation.pdf", "application/pdf")
 
 if __name__ == "__main__":
-    main()
+    main()