Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import numpy as np | |
| import gradio as gr | |
| import json | |
| import tempfile | |
| from ultralytics import YOLO | |
| from insightface.app import FaceAnalysis | |
| from huggingface_hub import hf_hub_download | |
| # ========================================== | |
| # CONFIGURATION & STATE MANAGEMENT | |
| # ========================================== | |
| DB_FILE = "face_db.json" | |
| EMBEDDINGS_FILE = "face_embeddings.npy" | |
| class FaceSystem: | |
| def __init__(self): | |
| print("🚀 Initializing AI Models...") | |
| # 1. Load YOLOv8-Face | |
| model_path = hf_hub_download(repo_id="arnabdhar/YOLOv8-Face-Detection", filename="model.pt") | |
| self.detector = YOLO(model_path) | |
| # 2. Load InsightFace | |
| self.recognizer = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider']) | |
| self.recognizer.prepare(ctx_id=0, det_size=(640, 640)) | |
| # 3. Load Database | |
| self.known_names = [] | |
| self.known_embeddings = np.empty((0, 512)) | |
| self.load_db() | |
| print("✅ System Ready.") | |
| def load_db(self): | |
| if os.path.exists(DB_FILE) and os.path.exists(EMBEDDINGS_FILE): | |
| with open(DB_FILE, 'r') as f: | |
| self.known_names = json.load(f) | |
| self.known_embeddings = np.load(EMBEDDINGS_FILE) | |
| print(f"📂 Loaded {len(self.known_names)} identities.") | |
| else: | |
| print("📂 Database empty. Starting fresh.") | |
| def save_db(self): | |
| with open(DB_FILE, 'w') as f: | |
| json.dump(self.known_names, f) | |
| np.save(EMBEDDINGS_FILE, self.known_embeddings) | |
| def enroll_user(self, name, image): | |
| if image is None or name.strip() == "": | |
| return "⚠️ Error: Missing name or photo." | |
| img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) | |
| faces = self.recognizer.get(img_bgr) | |
| if len(faces) == 0: | |
| return "⚠️ Error: No face detected." | |
| # Get largest face | |
| face = sorted(faces, key=lambda x: (x.bbox[2]-x.bbox[0]) * (x.bbox[3]-x.bbox[1]))[-1] | |
| embedding = face.normed_embedding.reshape(1, -1) | |
| if self.known_embeddings.shape[0] == 0: | |
| self.known_embeddings = embedding | |
| else: | |
| self.known_embeddings = np.vstack([self.known_embeddings, embedding]) | |
| self.known_names.append(name) | |
| self.save_db() | |
| return f"✅ Success: '{name}' added to database." | |
| def recognize_and_process(self, frame, blur_intensity=20, threshold=0.5): | |
| """Core processing logic for a single frame""" | |
| if frame is None: return None | |
| img_vis = frame.copy() | |
| h, w = img_vis.shape[:2] | |
| # Detect | |
| results = self.detector(img_vis, conf=0.5, verbose=False) | |
| for result in results: | |
| boxes = result.boxes | |
| for box in boxes: | |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) | |
| # Add context margin | |
| margin = 0 | |
| cx1 = max(0, x1 - margin); cy1 = max(0, y1 - margin) | |
| cx2 = min(w, x2 + margin); cy2 = min(h, y2 + margin) | |
| face_crop = img_vis[cy1:cy2, cx1:cx2] | |
| # Identify | |
| name = "Unknown" | |
| color = (200, 0, 0) # Red | |
| if self.known_embeddings.shape[0] > 0 and face_crop.size > 0: | |
| face_crop_bgr = cv2.cvtColor(face_crop, cv2.COLOR_RGB2BGR) | |
| analysis = self.recognizer.get(face_crop_bgr) | |
| if len(analysis) > 0: | |
| target_emb = analysis[0].normed_embedding | |
| similarities = np.dot(self.known_embeddings, target_emb) | |
| best_idx = np.argmax(similarities) | |
| if similarities[best_idx] > threshold: | |
| name = self.known_names[best_idx] | |
| color = (0, 255, 0) # Green | |
| # Privacy Blur (Pixelate) | |
| roi = img_vis[y1:y2, x1:x2] | |
| if roi.size > 0: | |
| block_size = max(3, int(30 - (blur_intensity / 4))) | |
| h_roi, w_roi = roi.shape[:2] | |
| small = cv2.resize(roi, (max(1, w_roi // block_size), max(1, h_roi // block_size)), interpolation=cv2.INTER_LINEAR) | |
| pixelated = cv2.resize(small, (w_roi, h_roi), interpolation=cv2.INTER_NEAREST) | |
| img_vis[y1:y2, x1:x2] = pixelated | |
| # Overlay | |
| cv2.rectangle(img_vis, (x1, y1), (x2, y2), color, 2) | |
| label_size, _ = cv2.getTextSize(name, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2) | |
| cv2.rectangle(img_vis, (x1, y1 - label_size[1] - 10), (x1 + label_size[0] + 10, y1), color, -1) | |
| cv2.putText(img_vis, name, (x1 + 5, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) | |
| return img_vis | |
| # Initialize | |
| system = FaceSystem() | |
| # ========================================== | |
| # VIDEO PROCESSING HELPER | |
| # ========================================== | |
| def process_video_file(video_path, blur_intensity, threshold): | |
| """Reads a video file, processes every frame, saves it, and returns path.""" | |
| if video_path is None: return None | |
| cap = cv2.VideoCapture(video_path) | |
| fps = int(cap.get(cv2.CAP_PROP_FPS)) | |
| width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| # Create temp output file | |
| temp_out = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) | |
| output_path = temp_out.name | |
| # Setup writer (mp4v is usually safe for CPU) | |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
| writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: break | |
| # Process frame using our existing core function | |
| # We convert BGR (OpenCV) to RGB (needed for our function) then back to BGR | |
| frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| processed_rgb = system.recognize_and_process(frame_rgb, blur_intensity, threshold) | |
| processed_bgr = cv2.cvtColor(processed_rgb, cv2.COLOR_RGB2BGR) | |
| writer.write(processed_bgr) | |
| cap.release() | |
| writer.release() | |
| return output_path | |
| # ========================================== | |
| # GRADIO INTERFACE (Fixed UI) | |
| # ========================================== | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 👁️ SecureVision Pro") | |
| with gr.Tabs(): | |
| # --- SURVEILLANCE TAB --- | |
| with gr.Tab("📹 Surveillance"): | |
| # Global Settings for this tab | |
| with gr.Accordion("⚙️ Settings", open=False): | |
| blur_slider = gr.Slider(1, 100, value=50, label="Privacy Level") | |
| conf_slider = gr.Slider(0.1, 0.9, value=0.5, label="Recognition Threshold") | |
| with gr.Tabs(): | |
| # Sub-Tab 1: Live Webcam | |
| with gr.TabItem("Live Webcam"): | |
| with gr.Row(): | |
| web_in = gr.Image(sources=["webcam"], streaming=True, label="Live Feed") | |
| web_out = gr.Image(label="Live Output") | |
| web_in.stream(system.recognize_and_process, [web_in, blur_slider, conf_slider], web_out) | |
| # Sub-Tab 2: Upload Image | |
| with gr.TabItem("Upload Image"): | |
| with gr.Row(): | |
| img_in = gr.Image(sources=["upload", "clipboard"], label="Upload Image") | |
| img_out = gr.Image(label="Processed Image") | |
| img_btn = gr.Button("Analyze Image", variant="primary") | |
| img_btn.click(system.recognize_and_process, [img_in, blur_slider, conf_slider], img_out) | |
| # Sub-Tab 3: Upload Video | |
| with gr.TabItem("Upload Video"): | |
| with gr.Row(): | |
| vid_in = gr.Video(label="Upload Video") | |
| vid_out = gr.Video(label="Processed Output") | |
| vid_btn = gr.Button("Process Video", variant="primary") | |
| vid_btn.click(process_video_file, [vid_in, blur_slider, conf_slider], vid_out) | |
| # --- DATABASE TAB --- | |
| with gr.Tab("👤 Database"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown("### Enroll New Personnel") | |
| new_name = gr.Textbox(label="Full Name / ID") | |
| new_photo = gr.Image(sources=["upload", "webcam"], label="Reference Photo") | |
| add_btn = gr.Button("Enroll User", variant="primary") | |
| status_msg = gr.Markdown("") | |
| with gr.Column(): | |
| gr.Markdown("### Database Status") | |
| def get_user_list(): | |
| if not system.known_names: return "No users enrolled." | |
| return "\n".join([f"• {n}" for n in system.known_names]) | |
| user_list = gr.Markdown(get_user_list) | |
| refresh_btn = gr.Button("Refresh List") | |
| add_btn.click(system.enroll_user, [new_name, new_photo], status_msg) | |
| refresh_btn.click(get_user_list, outputs=user_list) | |
| add_btn.click(get_user_list, outputs=user_list) | |
| if __name__ == "__main__": | |
| demo.launch() |