Create app.py

app.py

@@ -0,0 +1,202 @@
+import streamlit as st
+import os
+import requests
+import numpy as np
+from PIL import Image, ImageDraw
+import io
+import random
+
+# YOLOv10 import (installed via requirements.txt)
+try:
+    from yolov10 import YOLOv10
+except ImportError:
+    st.error("Could not import YOLOv10. Please confirm the library installation.")
+
+# ------------------------
+# 1. Chaotic Encryption Utilities
+# ------------------------
+
+def logistic_map(r, x):
+    return r * x * (1 - x)
+
+def generate_key(seed, n):
+    key = []
+    x = seed
+    for _ in range(n):
+        x = logistic_map(3.9, x)
+        key.append(int(x * 255) % 256)  # Map float to [0..255]
+    return np.array(key, dtype=np.uint8)
+
+def shuffle_pixels(img_array, seed):
+    h, w, c = img_array.shape
+    num_pixels = h * w
+    flattened = img_array.reshape(-1, c)
+    indices = np.arange(num_pixels)
+    
+    random.seed(seed)
+    random.shuffle(indices)  # Shuffle indices
+    
+    shuffled = flattened[indices]
+    return shuffled.reshape(h, w, c), indices
+
+def encrypt_image(img_array, seed):
+    """Encrypt the given image array using a chaotic logistic map (two-layer XOR + shuffle)."""
+    h, w, c = img_array.shape
+    flat_image = img_array.flatten()
+
+    # First chaotic key
+    chaotic_key_1 = generate_key(seed, len(flat_image))
+    # XOR-based encryption (first layer)
+    encrypted_flat_1 = [pixel ^ chaotic_key_1[i] for i, pixel in enumerate(flat_image)]
+    encrypted_array_1 = np.array(encrypted_flat_1, dtype=np.uint8).reshape(h, w, c)
+
+    # Shuffle
+    shuffled_array, _ = shuffle_pixels(encrypted_array_1, seed)
+
+    # Second chaotic key
+    chaotic_key_2 = generate_key(seed * 1.1, len(flat_image))
+    shuffled_flat = shuffled_array.flatten()
+    encrypted_flat_2 = [pixel ^ chaotic_key_2[i] for i, pixel in enumerate(shuffled_flat)]
+    doubly_encrypted_array = np.array(encrypted_flat_2, dtype=np.uint8).reshape(h, w, c)
+
+    return doubly_encrypted_array
+
+# (A decrypt_image function could be implemented if needed)
+
+# ------------------------
+# 2. YOLOv10 Detection Logic
+# ------------------------
+
+def load_model(model_path):
+    """Load YOLOv10 model from local .pt weights."""
+    model = YOLOv10(model_path)
+    return model
+
+def detect_license_plates(model, pil_image):
+    """
+    Runs YOLOv10 detection on the PIL image.
+    Returns:
+      - image_with_boxes (PIL Image) with bounding boxes drawn
+      - bboxes (list of (x1, y1, x2, y2)) for each license plate
+    """
+    # Convert PIL image to np array
+    np_image = np.array(pil_image)
+    
+    # YOLO inference
+    results = model.predict(np_image)
+    
+    # Extract detections
+    detections = results.xyxy[0]  # shape: (N, 6) -> [x1, y1, x2, y2, conf, class]
+    
+    bboxes = []
+    draw = ImageDraw.Draw(pil_image)
+    for *box, conf, cls in detections:
+        cls_id = int(cls)
+        # Adjust class check if your model uses a different label for license plate
+        # e.g., if class 0 is 'license_plate'
+        if cls_id == 0:
+            x1, y1, x2, y2 = map(int, box)
+            bboxes.append((x1, y1, x2, y2))
+            # Draw bounding box
+            draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
+    
+    return pil_image, bboxes
+
+# ------------------------
+# 3. Streamlit App
+# ------------------------
+
+def main():
+    st.title("YOLOv10 + Chaotic Encryption Demo")
+    st.write(
+        """
+        **Instructions**:
+        1. Provide an image (URL or file upload).
+        2. If a license plate is detected, only that region will be **encrypted** using Chaotic Logistic Map.
+        3. Download the final result.
+        """
+    )
+
+    # Sidebar: Model path
+    st.sidebar.header("Model Config")
+    default_path = "best.pt"  # Make sure you've uploaded this file in your HF Space
+    model_path = st.sidebar.text_input("YOLOv10 Model Path", value=default_path)
+
+    if not os.path.isfile(model_path):
+        st.error(f"Model weights not found at {model_path}. Please upload best.pt!")
+        return
+
+    # Load model once
+    @st.cache_data(show_spinner=False, allow_output_mutation=True)
+    def cached_load_model(path):
+        return load_model(path)
+
+    with st.spinner("Loading YOLOv10 model..."):
+        model = cached_load_model(model_path)
+    st.success("Model loaded successfully!")
+
+    # Option 1: URL input
+    image_url = st.text_input("Image URL (optional)")
+
+    # Option 2: File upload
+    uploaded_file = st.file_uploader("OR Upload an Image", type=["png", "jpg", "jpeg"])
+
+    # Encryption seed slider
+    key_seed = st.slider("Encryption Key Seed (0 < seed < 1)", 0.001, 0.999, 0.5, 0.001)
+
+    if st.button("Detect & Encrypt"):
+        # 1) Load the image (from URL or uploaded file)
+        if image_url and (not uploaded_file):
+            # Download from URL
+            try:
+                response = requests.get(image_url, timeout=10)
+                pil_image = Image.open(io.BytesIO(response.content)).convert("RGB")
+            except Exception as e:
+                st.error(f"Failed to load image from URL. Error: {e}")
+                return
+        elif uploaded_file:
+            # Use uploaded file
+            pil_image = Image.open(uploaded_file).convert("RGB")
+        else:
+            st.warning("Please provide an image URL or upload an image.")
+            return
+
+        st.image(pil_image, caption="Original Image", use_container_width=True)
+
+        # 2) Detect license plates
+        with st.spinner("Detecting license plates..."):
+            image_with_boxes, bboxes = detect_license_plates(model, pil_image.copy())
+        
+        st.image(image_with_boxes, caption="Detected Plate(s)", use_container_width=True)
+
+        if not bboxes:
+            st.warning("No license plates detected.")
+            return
+
+        # 3) Encrypt only the bounding box regions
+        with st.spinner("Encrypting license plates..."):
+            np_original = np.array(pil_image)
+            encrypted_np = np_original.copy()
+
+            for (x1, y1, x2, y2) in bboxes:
+                plate_region = encrypted_np[y1:y2, x1:x2]
+                plate_encrypted = encrypt_image(plate_region, key_seed)
+                encrypted_np[y1:y2, x1:x2] = plate_encrypted
+
+            encrypted_image = Image.fromarray(encrypted_np)
+        
+        st.image(encrypted_image, caption="Encrypted Image", use_container_width=True)
+
+        # 4) Provide download button
+        buf = io.BytesIO()
+        encrypted_image.save(buf, format="PNG")
+        buf.seek(0)
+        st.download_button(
+            label="Download Encrypted Image",
+            data=buf,
+            file_name="encrypted_plate.png",
+            mime="image/png",
+        )
+
+if __name__ == "__main__":
+    main()