Spaces:
Build error
Build error
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() | |