Update app.py

app.py

@@ -6,17 +6,18 @@ from PIL import Image, ImageDraw
 import io
 import random
 
-# Attempt to import YOLOv10 from the ultralytics package provided by THU-MIG/yolov10
+##############################################################################
+# 1. Attempt to import YOLOv10 from the ultralytics package (THU-MIG/yolov10)
+##############################################################################
 try:
     from ultralytics import YOLOv10
 except ImportError:
-    st.error("Could not import YOLOv10. Please confirm the THU-MIG/yolov10 installation in requirements.txt.")
+    st.error("Could not import YOLOv10. Please confirm THU-MIG/yolov10 installation in requirements.txt.")
     st.stop()
 
-# -----------------------------------------------------------------------------
-# 1. Chaotic Logistic Map Encryption Functions
-# -----------------------------------------------------------------------------
-
+##############################################################################
+# 2. Chaotic Logistic Map Encryption Functions
+##############################################################################
 def logistic_map(r, x):
     return r * x * (1 - x)
 
@@ -25,7 +26,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)
     return np.array(key, dtype=np.uint8)
 
 def shuffle_pixels(img_array, seed):
@@ -64,66 +65,72 @@ def encrypt_image(img_array, seed):
 
     return doubly_encrypted_array
 
-# -----------------------------------------------------------------------------
-# 2. YOLOv10 License Plate Detection
-# -----------------------------------------------------------------------------
-
+##############################################################################
+# 3. YOLOv10 License Plate Detection
+##############################################################################
 @st.cache_data(show_spinner=False)
 def load_model(weights_path: str):
     """
-    Loads the YOLOv10 model from local .pt weights.
+    Loads the YOLOv10 model from local .pt weights (Ultralytics style).
     """
-    model = YOLOv10(weights_path)
+    model = YOLOv10(weights_path)  # e.g., 'best.pt'
     return model
 
 def detect_license_plates(model, pil_image):
     """
-    Runs YOLOv10 detection on the PIL image.
+    Runs YOLOv10 detection on the PIL image using ultralytics-style output:
+      results -> list of ultralytics.engine.results.Results
+        each Results has .boxes, .masks, .names, etc.
 
-    After calling model.predict(), we log the raw output for debugging:
-    - This helps see how bounding boxes or classes are returned.
+    We'll handle only the first Results object (single image).
     """
     np_image = np.array(pil_image)
-    # Perform detection
     results = model.predict(np_image)
 
-    # *** DEBUG: Print the raw model output ***
-    print("Raw model output:", results)
+    # 1) Check how many Results objects we have
+    if not results:
+        print("No results returned by model.")
+        return pil_image, []
+
+    # 2) Take the first Results object
+    r = results[0]
 
-    # If the model output is expected to be a list of bounding boxes
-    # for a single image, typically we'd do something like:
-    if len(results) > 0:
-        detections = results[0]
-    else:
-        detections = []
+    # Debug: print the entire Results object
+    print("Raw model output (first Results object):", r)
 
+    # 3) If r.boxes is None or empty, we have no detections
+    if not hasattr(r, 'boxes') or r.boxes is None or len(r.boxes) == 0:
+        print("No boxes found in results[0].")
+        return pil_image, []
+
+    # 4) Parse bounding boxes
     bboxes = []
     draw = ImageDraw.Draw(pil_image)
 
-    # Example of how you'd parse each detection if it's [x1, y1, x2, y2, conf, cls_id]
-    for det in detections:
-        # If the model returns more or fewer elements, you'll need to adjust this line
-        # e.g., if len(det) == 6, then you can unpack:
-        #   x1, y1, x2, y2, conf, cls_id = det
-        # Otherwise, print(det) or log it to see how to parse properly
-        if len(det) == 6:
-            x1, y1, x2, y2, conf, cls_id = det
-            cls_id = int(cls_id)
-            # If your model’s license plate class is 0:
-            if cls_id == 0:
-                x1, y1, x2, y2 = map(int, (x1, y1, x2, y2))
-                bboxes.append((x1, y1, x2, y2))
-                draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
-        else:
-            # Debug print if you run into mismatch
-            print("Unexpected detection format:", det)
+    # r.boxes is an ultralytics.engine.results.Boxes object
+    # We can iterate over each box in r.boxes:
+    for box in r.boxes:
+        # box has .xyxy, .conf, .cls as 1D tensors
+        # e.g. box.xyxy[0] is [x1, y1, x2, y2]
+        # box.conf[0] is confidence
+        # box.cls[0] is class ID
+        coords = box.xyxy[0].tolist()  # [x1, y1, x2, y2]
+        conf = float(box.conf[0])
+        cls_id = int(box.cls[0])
+
+        # If your license plate class is 0:
+        if cls_id == 0:
+            x1, y1, x2, y2 = map(int, coords)
+            bboxes.append((x1, y1, x2, y2))
+
+            # Optional: draw bounding box for visualization
+            draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
 
     return pil_image, bboxes
 
-# -----------------------------------------------------------------------------
-# 3. Streamlit App
-# -----------------------------------------------------------------------------
-
+##############################################################################
+# 4. Streamlit App
+##############################################################################
 def main():
     st.title("YOLOv10 + Chaotic Encryption Demo")
     st.write(
@@ -135,8 +142,8 @@ def main():
         """
     )
 
-    # Model weights path
-    default_model_path = "best.pt"  # Adjust if your model file is named differently
+    # A. Model weights path
+    default_model_path = "best.pt"
     model_path = st.sidebar.text_input("YOLOv10 Weights (.pt)", value=default_model_path)
 
     if not os.path.isfile(model_path):
@@ -147,16 +154,16 @@ def main():
         model = load_model(model_path)
     st.success("Model loaded successfully!")
 
-    # Image input
+    # B. 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"])
 
-    # Encryption seed slider
+    # C. Encryption seed slider
     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 file
+        # 1) Load the image from URL or file
         if image_url and not uploaded_file:
             try:
                 response = requests.get(image_url, timeout=10)
@@ -173,7 +180,7 @@ def main():
 
         st.image(pil_image, caption="Original Image", use_container_width=True)
 
-        # 2. Detect plates (includes debug print to logs)
+        # 2) Detect plates
         with st.spinner("Detecting license plates..."):
             image_with_boxes, bboxes = detect_license_plates(model, pil_image.copy())
 
@@ -182,7 +189,7 @@ def main():
             st.warning("No license plates detected.")
             return
 
-        # 3. Encrypt bounding box regions
+        # 3) Encrypt bounding box regions
         with st.spinner("Encrypting license plates..."):
             np_img = np.array(pil_image)
             encrypted_np = np_img.copy()
@@ -195,7 +202,7 @@ def main():
 
         st.image(encrypted_image, caption="Encrypted Image", use_container_width=True)
 
-        # 4. Download link
+        # 4) Download link
         buf = io.BytesIO()
         encrypted_image.save(buf, format="PNG")
         buf.seek(0)