Update app.py

app.py

@@ -1,63 +1,135 @@
+import streamlit as st
+import uuid
+import cv2
+import numpy as np
+from PIL import Image
+import tempfile
+import pickle
+import os
+from datetime import datetime
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaFileUpload
+from google.auth.transport.requests import Request
+from simple_salesforce import Salesforce
+
+# === Google Drive Setup ===
+REG_FOLDER_ID = '1qkcR7nQTEtiMH9OFUv2bGxVn08E3dKjF'  # Registered images folder
+INTRUDER_FOLDER_ID = '1PPAUWU-wMx7fek73p-hqPqYQypYtG8Ob'  # Intruder folder
+
+# === Salesforce Setup ===
+sf = Salesforce(
+    username='suriraja822@agentforce.com',
+    password='Sati@1010',
+    security_token='B9nS0HEyBE7YmWllqXCyiOJpY',
+    domain='login'
+)
+
+def upload_to_drive(image, filename, folder_id):
+    creds = None
+    if os.path.exists("token.pickle"):
+        with open("token.pickle", "rb") as token:
+            creds = pickle.load(token)
+    if creds and creds.expired and creds.refresh_token:
+        creds.refresh(Request())
+    service = build('drive', 'v3', credentials=creds)
+
+    image_pil = Image.fromarray(image)
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg")
+    image_pil.save(temp_file.name)
+
+    file_metadata = {'name': filename, 'parents': [folder_id]}
+    media = MediaFileUpload(temp_file.name, mimetype='image/jpeg')
+    file = service.files().create(body=file_metadata, media_body=media, fields='id').execute()
+
+    file_id = file.get('id')
+    public_link = f"https://drive.google.com/uc?id={file_id}"
+    os.remove(temp_file.name)
+
+    # Make file public
+    try:
+        service.permissions().create(
+            fileId=file_id,
+            body={'role': 'reader', 'type': 'anyone'}
+        ).execute()
+    except Exception as e:
+        st.warning(f"Failed to set permission: {e}")
+
+    return public_link
+
+# === Streamlit UI ===
+st.set_page_config(page_title="Hotel Registration", layout="wide")
+st.title("📸 Hotel Guest & Staff Registration")
+
+tab1, tab2 = st.tabs(["🛎️ Guest Registration", "👮 Staff Registration"])
+
+# === Guest Tab ===
 with tab1:
     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:
-            st.write("🧪 Starting image conversion...")
-            image_np = np.array(Image.open(guest_cam))
-            try:
-                img_url = upload_image_to_drive(image_np, f"GUEST")
-                st.write(f"✅ Guest photo uploaded: {img_url}")
-            except Exception as e:
-                st.error(f"❌ Upload Guest Image Failed: {e}")
-                st.stop()
 
-            try:
-                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")
-                st.write(f"✅ ID photo uploaded: {id_url}")
-            except Exception as e:
-                st.error(f"❌ Upload ID Photo Failed: {e}")
-                st.stop()
+    guest_name = st.text_input("Guest Name")
+    guest_aadhar = st.text_input("Aadhar Number")
+    guest_phone = st.text_input("Phone Number")
+    guest_room = st.text_input("Room Number")
+
+    id_photo = st.file_uploader("Upload ID Card Photo", type=["jpg", "jpeg", "png"])
+    guest_image = st.camera_input("Capture Guest Photo")
+
+    if st.button("✅ Register Guest"):
+        if guest_image and id_photo and all([guest_name, guest_aadhar, guest_phone, guest_room]):
+            guest_np = np.array(Image.open(guest_image).convert('RGB'))
 
+            filename = f"G_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
             try:
-                # 1. Update or Create Room
-                st.write(f"🔍 Searching for room: {room_number}")
-                rooms = sf.query_all(f"SELECT Id FROM Room__c WHERE Name = '{room_number}'")
-                if rooms['totalSize'] == 0:
-                    st.write("🏗️ Creating new Room record...")
-                    room_rec = sf.Room__c.create({'Name': room_number})
-                    room_id = room_rec['id']
-                    st.write(f"✅ Room created with ID: {room_id}")
-                else:
-                    room_id = rooms['records'][0]['Id']
-                    st.write(f"🏨 Found existing Room ID: {room_id}")
+                image_link = upload_to_drive(guest_np, filename, REG_FOLDER_ID)
+                st.image(guest_np, caption="Guest Registered")
+
+                result = sf.Guest__c.create({
+                    "Name": guest_name,
+                    "Aadhar_Number__c": guest_aadhar,
+                    "Phone_Number__c": guest_phone,
+                    "Room_Number__c": guest_room,
+                    "Check_In__c": datetime.utcnow().isoformat(),
+                    "Image_Link__c": image_link
+                })
+
+                st.success("🎉 Guest Registered in Salesforce!")
+                st.json(result)
             except Exception as e:
-                st.error(f"❌ Room handling failed: {e}")
-                st.stop()
+                st.error(f"Salesforce Error: {e}")
+        else:
+            st.warning("📋 Please fill all fields and upload both photos.")
 
+# === Staff Tab ===
+with tab2:
+    st.subheader("Staff Registration")
+
+    staff_name = st.text_input("Staff Name")
+    staff_id = st.text_input("Staff ID")
+    staff_dept = st.selectbox("Department", ["Security", "Reception", "Housekeeping", "Maintenance", "Kitchen"])
+
+    id_card = st.file_uploader("Upload ID Card (Staff)", type=["jpg", "jpeg", "png"], key="staff")
+    staff_image = st.camera_input("Capture Staff Photo")
+
+    if st.button("✅ Register Staff"):
+        if staff_image and id_card and all([staff_name, staff_id, staff_dept]):
+            staff_np = np.array(Image.open(staff_image).convert('RGB'))
+
+            filename = f"S_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
             try:
-                # 2. Create Guest record
-                st.write("📤 Attempting to create Guest__c 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
-                }
-                result = sf.Guest__c.create(guest_data)
-                st.write(f"✅ Guest creation result: {result}")
-                st.success("🎉 Guest registered successfully!")
+                image_link = upload_to_drive(staff_np, filename, REG_FOLDER_ID)
+                st.image(staff_np, caption="Staff Registered")
+
+                result = sf.Staff__c.create({
+                    "Name": staff_name,
+                    "Staff_ID__c": staff_id,
+                    "Department__c": staff_dept,
+                    "Registered_On__c": datetime.utcnow().isoformat(),
+                    "Image_Link__c": image_link
+                })
+
+                st.success("✅ Staff Registered in Salesforce!")
+                st.json(result)
             except Exception as e:
-                st.error(f"❌ Guest creation failed: {e}")
+                st.error(f"Salesforce Error: {e}")
         else:
-            st.warning("⚠️ Please complete all fields and capture/upload images.")
+            st.warning("📋 Please fill all fields and upload both photos.")