Update app.py

app.py

@@ -1,16 +1,16 @@
-import io
-import json
+import streamlit as st
+from transformers import DetrImageProcessor, DetrForObjectDetection
 import torch
 from PIL import Image
-from transformers import DetrImageProcessor, DetrForObjectDetection
+import requests
 
-# Initialize the DETR model and processor
-processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
-model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
+st.set_page_config(page_title="SnapSpot", page_icon="📸", layout="wide", initial_sidebar_state="collapsed")
 
-def predict(inputs):
-    # Load the image from the provided inputs
-    image = Image.open(io.BytesIO(inputs["image"]))
+# Function to perform object detection
+def detect_objects(image):
+    # Load DETR model and processor
+    processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
+    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
 
     # Preprocess the image
     inputs = processor(images=image, return_tensors="pt")
@@ -22,16 +22,48 @@ def predict(inputs):
     target_sizes = torch.tensor([image.size[::-1]])
     results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
 
-    # Prepare the results in a dictionary format
-    detections = [{"label": model.config.id2label[label.item()], "confidence": score.item(), "box": box.tolist()} 
-                  for score, label, box in zip(results["scores"], results["labels"], results["boxes"])]
+    return results
+
+# Main Streamlit app
+def main():
+    st.title("SnapSpot")
+    st.markdown(
+        """
+        <style>
+            .reportview-container {
+                background: #0e1117;
+                color: #f0f6fc;
+            }
+            .st-bq {
+                background-color: #0e1117;
+            }
+            .st-bm {
+                padding-top: 2rem;
+            }
+        </style>
+        """,
+        unsafe_allow_html=True,
+    )
+
+    # Upload image
+    uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+
+    if uploaded_image is not None:
+        # Display uploaded image
+        image = Image.open(uploaded_image)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
 
-    return detections
+        # Perform object detection
+        results = detect_objects(image)
 
-# Define the API endpoint for Hugging Face Spaces
-def huggingface_spaces_endpoint(inputs):
-    # Call the predict function with the provided inputs
-    detections = predict(inputs)
+        # Display detection results
+        st.subheader("Detection Results:")
+        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+            box = [round(i, 2) for i in box.tolist()]
+            st.write(
+                f"Detected {model.config.id2label[label.item()]} with confidence "
+                f"{round(score.item(), 3)} at location {box}"
+            )
 
-    # Return the detections as a JSON object
-    return json.dumps(detections)
+if __name__ == "__main__":
+    main()