munnae commited on
Commit
0d3324b
·
verified ·
1 Parent(s): f2e83e4

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +38 -0
app.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from transformers import AutoModelForCausalLM, AutoProcessor
3
+ from PIL import Image
4
+ import torch
5
+
6
+ # Load the model and processor
7
+ @st.cache_resource
8
+ def load_model():
9
+ st.text("Loading model...")
10
+ processor = AutoProcessor.from_pretrained("vikhyatk/moondream2", trust_remote_code=True)
11
+ model = AutoModelForCausalLM.from_pretrained(
12
+ "vikhyatk/moondream2",
13
+ revision="2025-03-27",
14
+ trust_remote_code=True,
15
+ # device_map="auto" # Enable if running on GPU
16
+ )
17
+ st.text("Model loaded successfully!")
18
+ return model, processor
19
+
20
+ model, processor = load_model()
21
+
22
+ # File uploader
23
+ uploaded_file = st.file_uploader("Upload an image of a dish", type=["jpg", "jpeg", "png"])
24
+
25
+ if uploaded_file:
26
+ image = Image.open(uploaded_file).convert("RGB")
27
+ st.image(image, caption="Uploaded Image", use_column_width=True)
28
+
29
+ # Auto-ask the food question
30
+ question = "What food is in this image?" # or try: "What is the name of the dish?"
31
+
32
+ if st.button("Classify Dish"):
33
+ inputs = processor(image, question, return_tensors="pt").to(model.device)
34
+ with torch.no_grad():
35
+ output = model.generate(**inputs, max_new_tokens=64)
36
+ answer = processor.batch_decode(output, skip_special_tokens=True)[0].strip()
37
+
38
+ st.success(f"🍽️ Predicted Dish: **{answer}**")