Update app.py

app.py

@@ -1,33 +1,29 @@
 import streamlit as st
+import os
 import cv2
-import numpy as np
 import tempfile
-from PIL import Image
-import requests
-import os
 import pickle
 import io
+import numpy as np
+from PIL import Image
 from datetime import datetime
-import uuid
 from googleapiclient.discovery import build
 from googleapiclient.http import MediaFileUpload
 from google.auth.transport.requests import Request
-
 from simple_salesforce import Salesforce
 
-# === SALESFORCE SETUP ===
+# ==== Salesforce Login ====
 sf = Salesforce(
-    username='suriraja822@agentforce.com',
-    password='Sati@1010',
-    security_token='B9nS0HEyBE7YmWllqXCyiOJpY',
-    domain='login'
+    username="suriraja822@agentforce.com",
+    password="Sati@1010",
+    security_token="B9nS0HEyBE7YmWllqXCyiOJpY",
+    domain="login"
 )
 
-# === GOOGLE DRIVE SETUP ===
-GUEST_FOLDER_ID = '1qkcR7nQTEtiMH9OFUv2bGxVn08E3dKjF'
-STAFF_FOLDER_ID = '1PPAUWU-wMx7fek73p-hqPqYQypYtG8Ob'
+# ==== Google Drive Setup ====
+REG_FOLDER_ID = '1qkcR7nQTEtiMH9OFUv2bGxVn08E3dKjF'
 
-def upload_image_to_drive(image, folder_id, filename):
+def upload_image_to_drive(image_np, prefix):
     creds = None
     if os.path.exists("token.pickle"):
         with open("token.pickle", "rb") as token:
@@ -37,92 +33,96 @@ def upload_image_to_drive(image, folder_id, filename):
             creds.refresh(Request())
     service = build('drive', 'v3', credentials=creds)
 
+    filename = f"{prefix}_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}.jpg"
+    image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
+    buf = io.BytesIO()
+    image_pil.save(buf, format="JPEG")
+    buf.seek(0)
     with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:
-        image.save(tmp.name)
-        file_metadata = {'name': filename, 'parents': [folder_id]}
+        tmp.write(buf.read())
+        tmp.flush()
+        file_metadata = {'name': filename, 'parents': [REG_FOLDER_ID]}
         media = MediaFileUpload(tmp.name, mimetype='image/jpeg')
-        uploaded = service.files().create(body=file_metadata, media_body=media, fields='id').execute()
+        uploaded_file = service.files().create(body=file_metadata, media_body=media, fields='id').execute()
+        file_id = uploaded_file.get("id")
+        service.permissions().create(
+            fileId=file_id,
+            body={'role': 'reader', 'type': 'anyone'}
+        ).execute()
         os.remove(tmp.name)
-        return f"https://drive.google.com/uc?id={uploaded.get('id')}"
-
-# === CAMERA CAPTURE ===
-def capture_image():
-    cap = cv2.VideoCapture(0)
-    st.info("📸 Press 's' to take a snapshot, then close the camera window.")
-    img = None
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        cv2.imshow("Press 's' to save, 'q' to quit", frame)
-        k = cv2.waitKey(1)
-        if k % 256 == ord('s'):
-            img = frame
-            break
-        elif k % 256 == ord('q'):
-            break
-    cap.release()
-    cv2.destroyAllWindows()
-    if img is not None:
-        return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
-    return None
-
-# === STREAMLIT UI ===
-st.title("🏨 Staff & Guest Registration (With Camera & Salesforce)")
+        return f"https://drive.google.com/uc?id={file_id}&export=download"
 
-tab1, tab2 = st.tabs(["🧳 Guest", "👮 Staff"])
+# ==== UI ====
+st.title("🏨 Hotel Guest & Staff Registration")
+tab1, tab2 = st.tabs(["👤 Guest Check-In", "👔 Staff Registration"])
 
 with tab1:
-    st.header("📋 Register Guest")
-    first_name = st.text_input("First Name")
-    last_name = st.text_input("Last Name")
-    room_number = st.text_input("Room Number")
-    st.subheader("📸 Capture ID Photo")
-    if st.button("Capture Guest ID Photo"):
-        guest_img = capture_image()
-        if guest_img:
-            st.image(guest_img, caption="Captured Guest ID", width=300)
-            filename = f"guest_{first_name}_{uuid.uuid4().hex[:6]}.jpg"
-            img_url = upload_image_to_drive(guest_img, GUEST_FOLDER_ID, filename)
-            st.success("✅ Uploaded to Drive")
+    st.subheader("Guest Registration")
+    guest_name = st.text_input("Full Name")
+    guest_mobile = st.text_input("Mobile Number")
+    guest_email = st.text_input("Email")
+    guest_id_photo = st.file_uploader("Upload ID Proof", type=["jpg", "jpeg", "png"])
+    room_number = st.text_input("Assign Room Number")
+
+    guest_cam = st.camera_input("Capture Guest Photo")
+    if st.button("📥 Register Guest"):
+        if guest_cam and guest_id_photo and guest_name and room_number:
+            image_np = np.array(Image.open(guest_cam))
+            img_url = upload_image_to_drive(image_np, f"GUEST")
+            id_bytes = guest_id_photo.read()
+            id_np = np.array(Image.open(io.BytesIO(id_bytes)))
+            id_url = upload_image_to_drive(id_np, f"GUEST_ID")
 
             try:
-                sf.Hotel_Guest__c.create({
-                    'First_Name__c': first_name,
-                    'Last_Name__c': last_name,
-                    'Room_Number__c': room_number,
-                    'HF_Registered__c': True,
-                    'Photo_Uploaded__c': True,
-                    'Guest_Status__c': 'Checked In',
-                    'Check_In_Date__c': datetime.utcnow().isoformat(),
-                })
-                st.success("✅ Guest Registered in Salesforce")
+                # 1. Update Room
+                rooms = sf.query_all(f"SELECT Id FROM Room__c WHERE Name = '{room_number}'")
+                if rooms['totalSize'] == 0:
+                    room_rec = sf.Room__c.create({'Name': room_number})
+                    room_id = room_rec['id']
+                else:
+                    room_id = rooms['records'][0]['Id']
+
+                # 2. Create Guest record
+                guest_data = {
+                    'Name': guest_name,
+                    'Mobile__c': guest_mobile,
+                    'Email__c': guest_email,
+                    'Room__c': room_id,
+                    'Image_Link__c': img_url,
+                    'ID_Proof_Link__c': id_url
+                }
+                sf.Guest__c.create(guest_data)
+                st.success("✅ Guest registered successfully!")
             except Exception as e:
-                st.error(f"❌ Salesforce Error: {e}")
+                st.error(f"❌ Error: {e}")
+        else:
+            st.warning("⚠️ Please complete all fields and capture/upload images.")
 
 with tab2:
-    st.header("📋 Register Staff")
-    sf_first = st.text_input("First Name", key="staff_fn")
-    sf_last = st.text_input("Last Name", key="staff_ln")
-    sf_dept = st.selectbox("Department", ["Security", "Housekeeping", "Reception", "Engineering"])
+    st.subheader("Staff Registration")
+    staff_name = st.text_input("Staff Full Name")
+    staff_role = st.text_input("Role / Designation")
+    staff_id_photo = st.file_uploader("Upload ID Card", type=["jpg", "jpeg", "png"])
+    staff_cam = st.camera_input("Capture Staff Photo")
 
-    st.subheader("📸 Capture Staff ID Photo")
-    if st.button("Capture Staff ID Photo"):
-        staff_img = capture_image()
-        if staff_img:
-            st.image(staff_img, caption="Captured Staff ID", width=300)
-            filename = f"staff_{sf_first}_{uuid.uuid4().hex[:6]}.jpg"
-            img_url = upload_image_to_drive(staff_img, STAFF_FOLDER_ID, filename)
-            st.success("✅ Uploaded to Drive")
+    if st.button("📥 Register Staff"):
+        if staff_cam and staff_id_photo and staff_name and staff_role:
+            image_np = np.array(Image.open(staff_cam))
+            img_url = upload_image_to_drive(image_np, f"STAFF")
+            id_bytes = staff_id_photo.read()
+            id_np = np.array(Image.open(io.BytesIO(id_bytes)))
+            id_url = upload_image_to_drive(id_np, f"STAFF_ID")
 
             try:
-                sf.Hotel_Staff__c.create({
-                    'First_Name__c': sf_first,
-                    'Last_Name__c': sf_last,
-                    'Department__c': sf_dept,
-                    'HF_Registered__c': True,
-                    'Photo_Uploaded__c': True,
-                })
-                st.success("✅ Staff Registered in Salesforce")
+                staff_data = {
+                    'Name': staff_name,
+                    'Role__c': staff_role,
+                    'Image_Link__c': img_url,
+                    'ID_Card_Link__c': id_url
+                }
+                sf.Staff__c.create(staff_data)
+                st.success("✅ Staff registered successfully!")
             except Exception as e:
-                st.error(f"❌ Salesforce Error: {e}")
+                st.error(f"❌ Error: {e}")
+        else:
+            st.warning("⚠️ Please complete all fields and capture/upload images.")