Update app.py

app.py

@@ -5,112 +5,113 @@ from torchvision import models
 from PIL import Image
 import numpy as np
 import cv2
-import tempfile
-import pyttsx3
-import os
 from datetime import datetime
-import torch
-from torchvision.models.mobilenetv2 import MobileNetV2 
+from torchvision.models.mobilenetv2 import MobileNetV2
 from gtts import gTTS
 import os
 
-# Create history folder if it doesn't exist
+# ====== Setup ======
 os.makedirs("history", exist_ok=True)
-
-# ========== TTS ENGINE ==========
-
-def speak_streamlit(text):
-    tts = gTTS(text=text, lang='en')
-    tts.save("temp.mp3")
-    st.audio("temp.mp3", format="audio/mp3")
-
-
-# ========== CURRENCY TYPE ==========
-currency_type = st.selectbox("Select currency type:", ["PKR (Pakistani Rupees)", "USD (US Dollars)", "INR (Indian Rupees)"])
-
-# For now, only PKR model is supported
-if "PKR" not in currency_type:
-    st.warning("Currently only Pakistani Rupees (PKR) is supported. Other currencies coming soon!")
-    st.stop()
-
-# ========== LOAD MODEL ==========
-
 torch.serialization.add_safe_globals({'MobileNetV2': MobileNetV2})
 
+# Load model
 model = torch.load("pkr_currency_classifier.pt", map_location='cpu', weights_only=False)
 model.eval()
+class_names = ['Fake', 'Not Currency', 'Real']
 
-
-# ========== TRANSFORMS ==========
+# Transforms
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor()
 ])
 
-# ========== PREDICT ==========
-prediction = predict(image)
+# Text-to-speech
+def speak_streamlit(text):
+    tts = gTTS(text=text, lang='en')
+    tts.save("temp.mp3")
+    st.audio("temp.mp3", format="audio/mp3")
 
-if prediction.lower() == "not currency":
-    st.warning("This image does not appear to be a currency note.")
-    speak("This is not a currency note.")
-else:
-    st.success(f"Prediction: **{prediction}**")
-    speak(f"This is a {prediction} currency note.")
-    save_history(image, prediction)  # Only save if it's currency
+# Prediction function
+def predict(image):
+    img = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        outputs = model(img)
+        _, predicted = torch.max(outputs, 1)
+    return class_names[predicted.item()]
 
-# ========== SAVE HISTORY ==========
+# Save history
 def save_history(image, result):
-    os.makedirs("history", exist_ok=True)
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     image.save(f"history/{timestamp}_{result}.png")
 
-# ========== MAIN UI ==========
+# ====== UI ======
+st.set_page_config(page_title="Currency Detector", page_icon="💵")
 st.title("💵 Currency Authenticity Detector")
 st.subheader("Check if your currency is real or fake!")
 
+currency_type = st.selectbox("Select currency type:", ["PKR (Pakistani Rupees)", "USD (US Dollars)", "INR (Indian Rupees)"])
+if "PKR" not in currency_type:
+    st.warning("Currently only Pakistani Rupees (PKR) is supported. Other currencies coming soon!")
+    st.stop()
+
+# Choose input method
 option = st.radio("Choose method:", ["Upload Image", "Scan via Camera"])
 
+# ===== Upload Mode =====
 if option == "Upload Image":
     uploaded_file = st.file_uploader("Upload a currency image", type=["jpg", "jpeg", "png"])
     if uploaded_file:
         image = Image.open(uploaded_file).convert("RGB")
         st.image(image, caption="Uploaded Image", use_column_width=True)
         prediction = predict(image)
-        st.success(f"Prediction: **{prediction}**")
-        speak_streamlit(f"This is a {prediction}")
-        save_history(image, prediction)
 
+        if prediction == "Not Currency":
+            st.warning("⚠️ This image does not appear to be a currency note.")
+            speak_streamlit("This is not a currency note.")
+        else:
+            st.success(f"Prediction: **{prediction}**")
+            speak_streamlit(f"This is a {prediction} currency note.")
+            save_history(image, prediction)
+
+# ===== Camera Mode =====
 elif option == "Scan via Camera":
-    st.write("Hold currency in front of your webcam and press 'Start Camera'.")
+    st.write("Hold the currency in front of your webcam and press Start.")
     if st.button("Start Camera"):
         cap = cv2.VideoCapture(0)
         stframe = st.empty()
         result_box = st.empty()
-        stop = st.button("Stop Camera")
 
-        while cap.isOpened() and not stop:
+        stop_button = st.button("Stop Camera")
+
+        while cap.isOpened():
             ret, frame = cap.read()
             if not ret:
                 break
             frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
             pil_image = Image.fromarray(frame_rgb)
-
-            # Show frame
             stframe.image(pil_image, channels="RGB", use_column_width=True)
 
-            # Predict
             prediction = predict(pil_image)
             result_box.markdown(f"### Prediction: **{prediction}**")
-            speak(f"This is a {prediction}")
-            save_history(pil_image, prediction)
+
+            if prediction != "Not Currency":
+                speak_streamlit(f"This is a {prediction} currency note.")
+                save_history(pil_image, prediction)
+            else:
+                speak_streamlit("This is not a currency note.")
+
+            if stop_button:
+                break
 
         cap.release()
         stframe.empty()
         result_box.empty()
 
-# ========== VIEW HISTORY ==========
+# ===== History =====
 if st.checkbox("📁 Show Scan History"):
     st.write("Previously scanned images:")
     for img_file in sorted(os.listdir("history"))[::-1][:5]:
-        st.image(f"history/{img_file}", caption=img_file)
-st.write("Design by MERAJ GRAPHICS")
+        st.image(f"history/{img_file}", caption=img_file, width=200)
+
+st.markdown("---")
+st.markdown("👨‍💻 Designed by MERAJ GRAPHICS")