smart_mirror / app.py
Aliashraf's picture
Update app.py
b8d6bec verified
raw
history blame
1.13 kB
import cv2
import numpy as np
import torch
import streamlit as st
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
# Load AI Model from Hugging Face
model_name = "microsoft/resnet-50"
model = AutoModelForImageClassification.from_pretrained(model_name)
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
def process_frame(frame):
# Convert frame to tensor
inputs = feature_extractor(images=[frame], return_tensors="pt")
outputs = model(**inputs)
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
return predictions
def main():
st.title("Smart Mirror AI")
run = st.checkbox("Run Camera")
cap = cv2.VideoCapture(0)
frame_placeholder = st.empty()
while run:
success, frame = cap.read()
if not success:
st.write("Failed to capture video")
break
else:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
predictions = process_frame(frame)
frame_placeholder.image(frame, channels="RGB")
cap.release()
if __name__ == "__main__":
main()