Spaces:
Paused
Paused
| 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 | |
| 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() | |