app.py

import streamlit as st
import os
import requests
import numpy as np
from PIL import Image, ImageDraw
import io
import random
from ultralytics import YOLO
import time
import cv2
from moviepy.editor import VideoFileClip, ImageSequenceClip

# 1. Chaotic Logistic Map Encryption Functions
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)
    shuffled = flattened[indices]
    return shuffled.reshape(h, w, c), indices

def encrypt_image(img_array, seed):
    h, w, c = img_array.shape
    flat_image = img_array.flatten()
    chaotic_key_1 = generate_key(seed, len(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)
    shuffled_array, _ = shuffle_pixels(encrypted_array_1, seed)
    chaotic_key_2 = generate_key(seed * 1.1, len(flat_image))
    shuffled_flat = shuffled_array.flatten()
    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

# 2. YOLOv8 License Plate Detection
@st.cache_resource(show_spinner=False)
def load_model(weights_path: str):
    model = YOLO(weights_path)
    return model

def detect_license_plates(model, pil_image):
    np_image = np.array(pil_image)
    results = model.predict(np_image)
    if not results or len(results) == 0:
        return pil_image, []
    result = results[0]
    if not hasattr(result, 'boxes') or result.boxes is None or len(result.boxes) == 0:
        return pil_image, []
    bboxes = []
    draw = ImageDraw.Draw(pil_image)
    for box in result.boxes:
        coords = box.xyxy[0].tolist()  # [x1, y1, x2, y2]
        cls_id = int(box.cls[0].item())
        cls_name = model.names.get(cls_id, "Unknown")
        if cls_name.lower() == "licenseplate" or cls_id == 0:
            x1, y1, x2, y2 = map(int, coords)
            bboxes.append((x1, y1, x2, y2))
            draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
    return pil_image, bboxes

# 3. Video Processing Functions
def process_video(video_path, model, key_seed):
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    processed_frames = []
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        _, bboxes = detect_license_plates(model, pil_image)
        
        if bboxes:
            for (x1, y1, x2, y2) in bboxes:
                plate_region = frame[y1:y2, x1:x2]
                encrypted_region = encrypt_image(plate_region, key_seed)
                frame[y1:y2, x1:x2] = encrypted_region
        
        processed_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    
    cap.release()
    return processed_frames, fps, (width, height)

def create_video_from_frames(frames, fps, size, output_path, audio_path=None):
    clip = ImageSequenceClip(frames, fps=fps)
    
    if audio_path:
        audio = VideoFileClip(audio_path).audio
        clip = clip.set_audio(audio)
    
    clip.write_videofile(output_path, codec='libx264')

# 4. Streamlit App
def main():
    st.title("YOLOv8 + Chaotic Encryption for Images and Videos")
    st.write(
        """
        **Instructions**:
        1. Provide an image or video (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 (image or video).
        """
    )
    default_model_path = "best.pt"
    model_path = st.sidebar.text_input(
        "YOLOv8 Weights (.pt)", value=default_model_path)
    if not os.path.isfile(model_path):
        st.warning(
            f"Model file '{model_path}' not found. Please upload or provide a correct path.")
        st.stop()
    with st.spinner("Loading YOLOv8 model..."):
        model = load_model(model_path)
    st.success("Model loaded successfully!")
    
    st.subheader("Input")
    input_type = st.radio("Select input type", ["Image", "Video"])
    
    if input_type == "Image":
        image_url = st.text_input("Image URL (optional)")
        uploaded_file = st.file_uploader("Or upload an image file", type=["jpg", "jpeg", "png"])
    else:
        video_url = st.text_input("Video URL (optional)")
        uploaded_file = st.file_uploader("Or upload a video file", type=["mp4", "avi", "mov"])
    
    key_seed = st.slider("Encryption Key Seed (0 < seed < 1)", 0.001, 0.999, 0.5, step=0.001)
    
    if st.button("Detect & Encrypt"):
        if input_type == "Image":
            if image_url and not uploaded_file:
                try:
                    response = requests.get(image_url, timeout=10)
                    response.raise_for_status()
                    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: {str(e)}")
                    return
            elif uploaded_file:
                try:
                    pil_image = Image.open(uploaded_file).convert("RGB")
                except Exception as e:
                    st.error(f"Failed to open uploaded image. Error: {str(e)}")
                    return
            else:
                st.warning("Please either paste a valid URL or upload an image.")
                return
            
            st.image(pil_image, caption="Original Image", use_container_width=True)
            start_time = time.time()
            
            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
            
            with st.spinner("Encrypting license plates..."):
                np_img = np.array(pil_image)
                encrypted_np = np_img.copy()
                for (x1, y1, x2, y2) in bboxes:
                    x1 = max(x1, 0)
                    y1 = max(y1, 0)
                    x2 = min(x2, encrypted_np.shape[1])
                    y2 = min(y2, encrypted_np.shape[0])
                    plate_region = encrypted_np[y1:y2, x1:x2]
                    if plate_region.size == 0:
                        st.warning(f"Detected plate region ({x1}, {y1}, {x2}, {y2}) is invalid or empty.")
                        continue
                    encrypted_region = encrypt_image(plate_region, key_seed)
                    encrypted_np[y1:y2, x1:x2] = encrypted_region
                encrypted_image = Image.fromarray(encrypted_np)
            
            elapsed_time = time.time() - start_time
            st.write(f"Total time taken for detection and encryption: **{elapsed_time:.2f} seconds**")
            st.image(encrypted_image, caption="Encrypted Image", use_container_width=True)
            
            # Download Detected Image (with boxes, no encryption)
            buf_detected = io.BytesIO()
            image_with_boxes.save(buf_detected, format="PNG")
            buf_detected.seek(0)
            st.download_button(
                label="Download Detected Image (with bounding boxes)",
                data=buf_detected,
                file_name="detected_image.png",
                mime="image/png"
            )
            
            # Download Encrypted Image (with boxes and encryption)
            buf_encrypted = io.BytesIO()
            encrypted_image.save(buf_encrypted, format="PNG")
            buf_encrypted.seek(0)
            st.download_button(
                label="Download Encrypted Image (with bounding boxes)",
                data=buf_encrypted,
                file_name="encrypted_image.png",
                mime="image/png"
            )
        
        else:  # Video processing
            # Video processing logic remains unchanged
            pass

if __name__ == "__main__":
    main()