Update app.py

app.py

@@ -1,36 +1,101 @@
-
 import streamlit as st
 import torch
-from torchvision import transforms
+import torchvision.transforms as transforms
+from torchvision import models
 from PIL import Image
-import torchvision.models as models
+import numpy as np
+import cv2
+import tempfile
+import pyttsx3
+import os
+from datetime import datetime
 
-# ⚠️ Only do this if you trust the file!
-model = torch.load("pkr_currency_classifier.pt", map_location=torch.device('cpu'), weights_only=False)
-model.eval()
+# ========== TTS ENGINE ==========
+def speak(text):
+    engine = pyttsx3.init()
+    engine.say(text)
+    engine.runAndWait()
 
+# ========== CURRENCY TYPE ==========
+currency_type = st.selectbox("Select currency type:", ["PKR (Pakistani Rupees)", "USD (US Dollars)", "INR (Indian Rupees)"])
 
-# Define transform
+# 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 ==========
+model = models.mobilenet_v2(pretrained=False)
+model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, 2)
+model.load_state_dict(torch.load("pkr_currency_classifier.pt", map_location=torch.device('cpu')))
+model.eval()
+
+# ========== TRANSFORMS ==========
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406],
-                         [0.229, 0.224, 0.225])
+    transforms.ToTensor()
 ])
 
-st.title("💵 PKR Currency Fake vs Real Detector")
-st.write("Upload an image of a Pakistani Rupee note to detect if it's **Real or Fake**.")
+# ========== PREDICT ==========
+def predict(image):
+    image_tensor = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        output = model(image_tensor)
+        _, pred = torch.max(output, 1)
+    return "Real Currency" if pred.item() == 1 else "Fake Currency"
 
-uploaded_file = st.file_uploader("Choose a currency note image...", type=["jpg", "jpeg", "png"])
+# ========== 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")
 
-if uploaded_file:
-    image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-    st.write("Classifying...")
+# ========== MAIN UI ==========
+st.title("💵 Currency Authenticity Detector")
+st.subheader("Check if your currency is real or fake!")
 
-    img = transform(image).unsqueeze(0)
-    with torch.no_grad():
-        outputs = model(img)
-        _, predicted = torch.max(outputs, 1)
-        label = "Real" if predicted.item() == 0 else "Fake"
-        st.success(f"This currency note is **{label}**.")
+option = st.radio("Choose method:", ["Upload Image", "Scan via Camera"])
+
+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(f"This is a {prediction}")
+        save_history(image, prediction)
+
+elif option == "Scan via Camera":
+    st.write("Hold currency in front of your webcam and press 'Start Camera'.")
+    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:
+            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)
+
+        cap.release()
+        stframe.empty()
+        result_box.empty()
+
+# ========== VIEW 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)