Update app.py

app.py

@@ -1,40 +1,89 @@
+import streamlit as st
 import torch
 from PIL import Image
-import gradio as gr
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
 
-# Load trained model
-model = torch.load("best.pt", map_location=torch.device("cpu"))
-model.eval()
+# Konfigurasi model
+MODEL_PATH = "model.pth"  # Pastikan model ada di direktori ini
+CLASS_NAMES = ["bag", "person-static-object"]  # Sesuaikan dengan kelas Anda
 
-# Function to process image
-def predict(image):
-    # Convert to PIL image if needed
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
-    
-    # Preprocess image (adjust as per your model's requirements)
-    results = model([image])  # Assuming YOLOv11 inference works like this
-    detections = results.xyxy[0].numpy()  # Extract bounding boxes, scores, etc.
+# Fungsi untuk memuat model
+@st.cache_resource
+def load_model():
+    model = torch.load(MODEL_PATH, map_location=torch.device('cpu'))
+    model.eval()  # Mode evaluasi
+    return model
 
-    # Draw boxes on image
-    for box in detections:
-        x1, y1, x2, y2, conf, cls = box
-        label = f"Class {int(cls)}: {conf:.2f}"
-        draw = ImageDraw.Draw(image)
-        draw.rectangle([(x1, y1), (x2, y2)], outline="red", width=3)
-        draw.text((x1, y1), label, fill="red")
-    
+# Preprocessing gambar
+def preprocess(image):
+    # Resize dan konversi ke tensor
+    image = np.array(image)
+    image = cv2.resize(image, (640, 640))  # Sesuaikan dengan input size model
+    image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
+    image = image.unsqueeze(0)  # Tambahkan batch dimension
     return image
 
-# Gradio interface
-interface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs="image",
-    title="YOLOv11 Object Detection",
-    description="Upload an image to detect objects using YOLOv11.",
-)
+# Postprocessing hasil prediksi
+def postprocess(prediction, confidence_threshold=0.5):
+    # Contoh untuk deteksi objek YOLO-style (sesuaikan dengan output model Anda)
+    boxes = prediction[0]["boxes"].detach().numpy()
+    scores = prediction[0]["scores"].detach().numpy()
+    labels = prediction[0]["labels"].detach().numpy()
+    
+    # Filter berdasarkan confidence threshold
+    keep = scores >= confidence_threshold
+    return boxes[keep], scores[keep], labels[keep]
+
+# Antarmuka Streamlit
+st.title("Deteksi Objek Tertinggal 👜👤")
+st.write("Upload gambar untuk mendeteksi objek 'bag' atau 'person-static-object'")
+
+# Upload gambar
+uploaded_file = st.file_uploader("Pilih gambar...", type=["jpg", "png", "jpeg"])
 
-# Launch app
-if __name__ == "__main__":
-    interface.launch()
+if uploaded_file is not None:
+    # Memuat gambar
+    image = Image.open(uploaded_file).convert("RGB")
+    st.image(image, caption="Gambar Input", use_column_width=True)
+    
+    # Proses deteksi
+    if st.button("Deteksi Objek"):
+        with st.spinner("Memproses..."):
+            try:
+                # Memuat model
+                model = load_model()
+                
+                # Preprocessing
+                input_tensor = preprocess(image)
+                
+                # Prediksi
+                with torch.no_grad():
+                    prediction = model(input_tensor)
+                
+                # Postprocessing
+                boxes, scores, labels = postprocess(prediction)
+                
+                # Visualisasi hasil
+                fig, ax = plt.subplots(1, figsize=(12, 6))
+                ax.imshow(image)
+                
+                for box, score, label in zip(boxes, scores, labels):
+                    x1, y1, x2, y2 = box
+                    rect = plt.Rectangle(
+                        (x1, y1), x2 - x1, y2 - y1,
+                        linewidth=2, edgecolor="lime", facecolor="none"
+                    )
+                    ax.add_patch(rect)
+                    ax.text(
+                        x1, y1 - 5,
+                        f"{CLASS_NAMES[label]}: {score:.2f}",
+                        color="lime", fontsize=10, backgroundcolor="black"
+                    )
+                
+                st.pyplot(fig)
+                st.success("Selesai!")
+            
+            except Exception as e:
+                st.error(f"Error: {str(e)}")