LovnishVerma's picture
Update app.py
a34751a verified
raw
history blame
3.12 kB
import streamlit as st
import os
from PIL import Image
from huggingface_hub import HfApi
# Constants
KNOWN_FACES_DIR = "known_faces" # Directory to save user images
REPO_NAME = "face_and_emotion_detection"
REPO_ID = f"LovnishVerma/{REPO_NAME}"
# Ensure the directory exists
os.makedirs(KNOWN_FACES_DIR, exist_ok=True)
# Retrieve Hugging Face token from environment variable
hf_token = os.getenv("upload") # Replace with your actual Hugging Face token
if not hf_token:
st.error("Hugging Face token not found. Please set the environment variable.")
st.stop()
# Initialize Hugging Face API
api = HfApi()
def save_image_to_hugging_face(image, name, roll_no):
""" Saves the image locally to the KNOWN_FACES_DIR and uploads it to Hugging Face. """
# Ensure the directory exists
if not os.path.exists(KNOWN_FACES_DIR):
os.makedirs(KNOWN_FACES_DIR)
# Construct the local file path
filename = f"{name}_{roll_no}.jpg"
local_path = os.path.join(KNOWN_FACES_DIR, filename)
try:
# Convert image to RGB if necessary
if image.mode != "RGB":
image = image.convert("RGB")
# Save the image to the known_faces directory
image.save(local_path)
# Upload the saved file to Hugging Face
api.upload_file(
path_or_fileobj=local_path,
path_in_repo=filename,
repo_id=REPO_ID,
repo_type="space",
token=hf_token,
)
st.success(f"Image saved to {KNOWN_FACES_DIR} and uploaded to Hugging Face as {filename}.")
except Exception as e:
st.error(f"Error saving or uploading image: {e}")
return local_path
# Streamlit UI
st.title("Student Registration with Hugging Face Image Upload")
# Input fields for student details
name = st.text_input("Enter your name")
roll_no = st.text_input("Enter your roll number")
# Choose input method for the image (webcam or file upload)
capture_mode = st.radio("Choose an option to upload your image", ["Use Webcam", "Upload File"])
# Handle webcam capture or file upload
if capture_mode == "Use Webcam":
picture = st.camera_input("Take a picture") # Capture image using webcam
elif capture_mode == "Upload File":
picture = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"]) # Upload image from file system
# Save data and process image on button click
if st.button("Register"):
if not name or not roll_no:
st.error("Please fill in both name and roll number.")
elif not picture:
st.error("Please upload or capture an image.")
else:
try:
# Open the image based on capture mode
if capture_mode == "Use Webcam" and picture:
image = Image.open(picture)
elif capture_mode == "Upload File" and picture:
image = Image.open(picture)
# Save the image locally and upload it to Hugging Face
image_path = save_image_to_hugging_face(image, name, roll_no)
except Exception as e:
st.error(f"An error occurred: {e}")