Update app.py

app.py

@@ -6,21 +6,24 @@ from PIL import Image, ImageDraw
 import io
 import random
 
-# Attempt to import YOLOv10 from THU-MIG/yolov10, which installs 'ultralytics'.
+# Attempt to import YOLOv10 from the ultralytics package provided by THU-MIG/yolov10
 try:
     from ultralytics import YOLOv10
 except ImportError:
-    st.error("Could not import YOLOv10. Please confirm the library installation from THU-MIG/yolov10.")
+    st.error("Could not import YOLOv10. Please confirm the THU-MIG/yolov10 installation in requirements.txt.")
     st.stop()
 
-# ------------------------
-# 1. Chaotic Logistic Map Encryption
-# ------------------------
+# -----------------------------------------------------------------------------
+# 1. Chaotic Logistic Map Encryption Functions
+# -----------------------------------------------------------------------------
 
 def logistic_map(r, x):
     return r * x * (1 - x)
 
 def generate_key(seed, n):
+    """
+    Generate a chaotic key (array of size n) using a logistic map and the given seed.
+    """
     key = []
     x = seed
     for _ in range(n):
@@ -29,11 +32,14 @@ def generate_key(seed, n):
     return np.array(key, dtype=np.uint8)
 
 def shuffle_pixels(img_array, seed):
+    """
+    Shuffle the pixels in img_array based on a random sequence seeded by '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)
 
@@ -41,7 +47,9 @@ def shuffle_pixels(img_array, seed):
     return shuffled.reshape(h, w, c), indices
 
 def encrypt_image(img_array, seed):
-    """Encrypt the given image array using a two-layer XOR + shuffle approach."""
+    """
+    Encrypt the given image array using a two-layer XOR + pixel shuffling approach.
+    """
     h, w, c = img_array.shape
     flat_image = img_array.flatten()
 
@@ -62,46 +70,58 @@ def encrypt_image(img_array, seed):
 
     return doubly_encrypted_array
 
-# ------------------------
+# -----------------------------------------------------------------------------
 # 2. YOLOv10 License Plate Detection
-# ------------------------
+# -----------------------------------------------------------------------------
 
-@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
+@st.cache_data(show_spinner=False)
+def load_model(weights_path: str):
+    """
+    Loads the YOLOv10 model from local .pt weights.
+    Make sure the path is correct and the .pt file is in your Space if it's custom.
+    """
+    model = YOLOv10(weights_path)
     return model
 
 def detect_license_plates(model, pil_image):
     """
     Runs YOLOv10 detection on the PIL image.
+    
+    According to the THU-MIG/yolov10 usage, 'model.predict(np_image)'
+    typically returns a list of bounding boxes in the format:
+      [ [x1, y1, x2, y2, conf, class_id], [x1, y1, x2, y2, conf, class_id], ... ]
+    
+    We assume class_id == 0 corresponds to a license plate.
+    
     Returns:
       - image_with_boxes: PIL image with bounding boxes drawn
-      - bboxes: list of (x1, y1, x2, y2) for detected license plates
+      - bboxes: list of (x1, y1, x2, y2) for any detected license plates
     """
     np_image = np.array(pil_image)
-    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 []
+    results = model.predict(np_image)  # Usually a list of detections for each image
+    # If there's only one image, results is something like [[x1, y1, x2, y2, conf, cls], ...]
+    # So we grab the first list:
+    detections = results[0] if len(results) > 0 else []
 
     bboxes = []
     draw = ImageDraw.Draw(pil_image)
 
-    for *box, conf, cls_id in detections:
+    for det in detections:
+        # Unpack [x1, y1, x2, y2, conf, cls_id]
+        x1, y1, x2, y2, conf, cls_id = det
         cls_id = int(cls_id)
-        # If your model has a single class (license plate) as class 0:
+        # e.g., if your license plate is class 0
         if cls_id == 0:
-            x1, y1, x2, y2 = map(int, box)
+            x1, y1, x2, y2 = map(int, (x1, y1, x2, y2))
             bboxes.append((x1, y1, x2, y2))
-            # Draw bounding box for visualization
+            # Draw bounding box (optional, for visualization)
             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")
@@ -115,11 +135,11 @@ def main():
     )
 
     # 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)
+    default_model_path = "best.pt"  # Adjust if your model file has a different name
+    model_path = st.sidebar.text_input("YOLOv10 Weights (.pt)", value=default_model_path)
 
     if not os.path.isfile(model_path):
-        st.warning(f"Model file '{model_path}' not found in this directory. Upload or provide a correct path.")
+        st.warning(f"Model file '{model_path}' not found. Upload or provide a correct path.")
         st.stop()
 
     with st.spinner("Loading YOLOv10 model..."):
@@ -129,7 +149,7 @@ def main():
     # Image input
     st.subheader("Image Input")
     image_url = st.text_input("Image URL (optional)")
-    uploaded_file = st.file_uploader("Or upload an image file", type=["jpg","jpeg","png"])
+    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, step=0.001)
@@ -171,7 +191,7 @@ def main():
                 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. Download link