Update app.py

app.py

@@ -17,6 +17,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 IMAGE_SHAPE = 640
 data_path = 'employees'
 webcam_path = 'captured_image.jpg'
+STUDENT_DB = "students.db"
 
 # Set Streamlit title
 st.title("AIML-Student Attendance System")
@@ -28,6 +29,26 @@ image_paths = glob(os.path.join(data_path, '*.jpg'))
 app = FaceAnalysis(name="buffalo_l")  # ArcFace model
 app.prepare(ctx_id=0 if torch.cuda.is_available() else -1, det_size=(IMAGE_SHAPE, IMAGE_SHAPE))
 
+# Check if roll exists
+def roll_exists(rno):
+    with sqlite3.connect(STUDENT_DB) as conn:
+        cur = conn.execute("SELECT * FROM students WHERE rno=?", (rno,))
+        return cur.fetchone()
+
+# Register new student
+def register_student(rno, sname, sclass, image):
+    image_path = os.path.join(data_path, f"{rno}.jpg")
+    image.save(image_path)
+    with sqlite3.connect(STUDENT_DB) as conn:
+        conn.execute("INSERT INTO students (rno, sname, sclass, image_path) VALUES (?, ?, ?, ?)",
+                     (rno, sname, sclass, image_path))
+
+
+# Load all registered images
+def load_registered_images():
+    with sqlite3.connect(STUDENT_DB) as conn:
+        return [row[3] for row in conn.execute("SELECT * FROM students")]
+        
 # Define function to match face embeddings
 def prod_function(app, prod_path, webcam_img_pil):
     np_webcam = np.array(webcam_img_pil)
@@ -126,3 +147,29 @@ with app_tab:
                 # Convert back to RGB for display
                 image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
                 st.image(image_rgb, caption="Detected Faces", use_container_width=True)
+
+with register_tab:
+    upload_mode = st.radio("Choose Method", ["📤 Upload Image", "📷 Take Photo"])
+    if upload_mode == "📤 Upload Image":
+        upload = st.file_uploader("Upload Student Image with plain background(.jpg)", type="jpg")
+        if upload:
+            image_pil = Image.open(upload)
+    else:
+        camera_capture = st.camera_input("Capture Image")
+        if camera_capture:
+            image_pil = Image.open(camera_capture)
+
+    if 'image_pil' in locals():
+        st.image(image_pil, caption="Student Image", use_container_width=True)
+        rno = st.text_input("Roll Number")
+        sname = st.text_input("Student Name")
+        sclass = st.text_input("Class")
+
+        if st.button("Register"):
+            if not (rno and sname and sclass):
+                st.warning("Fill all details.")
+            elif roll_exists(rno):
+                st.error("Student with this roll_no registered already.")
+            else:
+                register_student(rno, sname, sclass, image_pil)
+                st.success("Student registered successfully.")