Update app.py

app.py

@@ -6,14 +6,15 @@ from PIL import Image, ImageDraw
 import io
 import random
 
-# YOLOv10 import (installed via requirements.txt)
+# Attempt to import YOLOv10 from THU-MIG/yolov10, which installs 'ultralytics'.
 try:
     from ultralytics import YOLOv10
 except ImportError:
-    st.error("Could not import YOLOv10. Please confirm the library installation.")
+    st.error("Could not import YOLOv10. Please confirm the library installation from THU-MIG/yolov10.")
+    st.stop()
 
 # ------------------------
-# 1. Chaotic Encryption Utilities
+# 1. Chaotic Logistic Map Encryption
 # ------------------------
 
 def logistic_map(r, x):
@@ -24,7 +25,7 @@ def generate_key(seed, n):
     x = seed
     for _ in range(n):
         x = logistic_map(3.9, x)
-        key.append(int(x * 255) % 256)  # Map float to [0..255]
+        key.append(int(x * 255) % 256)  # map float to 0-255
     return np.array(key, dtype=np.uint8)
 
 def shuffle_pixels(img_array, seed):
@@ -34,20 +35,20 @@ def shuffle_pixels(img_array, seed):
     indices = np.arange(num_pixels)
     
     random.seed(seed)
-    random.shuffle(indices)  # Shuffle indices
-    
+    random.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)."""
+    """Encrypt the given image array using a two-layer XOR + shuffle approach."""
     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_flat_1 = [p ^ chaotic_key_1[i] for i, p in enumerate(flat_image)]
     encrypted_array_1 = np.array(encrypted_flat_1, dtype=np.uint8).reshape(h, w, c)
 
     # Shuffle
@@ -56,50 +57,46 @@ def encrypt_image(img_array, 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)]
+    encrypted_flat_2 = [p ^ chaotic_key_2[i] for i, p 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
+# 2. YOLOv10 License Plate Detection
 # ------------------------
 
-def load_model(model_path):
-    """Load YOLOv10 model from local .pt weights."""
-    model = YOLOv10(model_path)
+@st.cache_data(show_spinner=False)  # <-- Removed allow_output_mutation here
+def load_model(weights_path):
+    """Loads the YOLOv10 model from local .pt weights."""
+    model = YOLOv10(weights_path)  # from ultralytics
     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
+      - image_with_boxes: PIL image with bounding boxes drawn
+      - bboxes: list of (x1, y1, x2, y2) for detected license plates
     """
-    # 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]
-    
+    results = model.predict(np_image)  # returns an object with .xyxy, etc.
+
+    # results.xyxy[0] -> [x1, y1, x2, y2, conf, class]
+    detections = results.xyxy[0] if len(results.xyxy) > 0 else []
+
     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'
+
+    for *box, conf, cls_id in detections:
+        cls_id = int(cls_id)
+        # If your model has a single class (license plate) as class 0:
         if cls_id == 0:
             x1, y1, x2, y2 = map(int, box)
             bboxes.append((x1, y1, x2, y2))
-            # Draw bounding box
+            # Draw bounding box for visualization
             draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
-    
+
     return pil_image, bboxes
 
 # ------------------------
@@ -117,77 +114,67 @@ def main():
         """
     )
 
-    # 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)
+    # Model weights path
+    default_path = "best.pt"  # Must exist in your repository if custom
+    model_path = st.sidebar.text_input("YOLOv10 Weights (.pt)", 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)
+        st.warning(f"Model file '{model_path}' not found in this directory. Upload or provide a correct path.")
+        st.stop()
 
     with st.spinner("Loading YOLOv10 model..."):
-        model = cached_load_model(model_path)
+        model = load_model(model_path)
     st.success("Model loaded successfully!")
 
-    # Option 1: URL input
+    # Image input
+    st.subheader("Image 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"])
+    uploaded_file = st.file_uploader("Or upload an image file", type=["jpg","jpeg","png"])
 
     # Encryption seed slider
-    key_seed = st.slider("Encryption Key Seed (0 < seed < 1)", 0.001, 0.999, 0.5, 0.001)
+    key_seed = st.slider("Encryption Key Seed (0 < seed < 1)", 0.001, 0.999, 0.5, step=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
+        # 1. Load the image from URL or file
+        if image_url and not uploaded_file:
             try:
                 response = requests.get(image_url, timeout=10)
-                pil_image = Image.open(io.BytesIO(response.content)).convert("RGB")
+                image_bytes = io.BytesIO(response.content)
+                pil_image = Image.open(image_bytes).convert("RGB")
             except Exception as e:
-                st.error(f"Failed to load image from URL. Error: {e}")
+                st.error(f"Failed to load image from URL. Error: {str(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.")
+            st.warning("Please either paste a valid URL or upload an image.")
             return
 
         st.image(pil_image, caption="Original Image", use_container_width=True)
 
-        # 2) Detect license plates
+        # 2. Detect 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)
 
+        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
+        # 3. Encrypt bounding box regions
         with st.spinner("Encrypting license plates..."):
-            np_original = np.array(pil_image)
-            encrypted_np = np_original.copy()
-
+            np_img = np.array(pil_image)
+            encrypted_np = np_img.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_region = encrypt_image(plate_region, key_seed)
+                encrypted_np[y1:y2, x1:x2] = encrypted_region
 
             encrypted_image = Image.fromarray(encrypted_np)
         
         st.image(encrypted_image, caption="Encrypted Image", use_container_width=True)
 
-        # 4) Provide download button
+        # 4. Download link
         buf = io.BytesIO()
         encrypted_image.save(buf, format="PNG")
         buf.seek(0)
@@ -195,7 +182,7 @@ def main():
             label="Download Encrypted Image",
             data=buf,
             file_name="encrypted_plate.png",
-            mime="image/png",
+            mime="image/png"
         )
 
 if __name__ == "__main__":