teaycorp/Examination
0
1import streamlit as st2import os3from pathlib import Path4import numpy as np5import torch6import pandas as pd7from datetime import datetime8from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2, FasterRCNN_ResNet50_FPN_V2_Weights9from PIL import Image10from streamlit_webrtc import webrtc_streamer, WebRtcMode, RTCConfiguration11import logging12import time13 14# Configure logging15logging.basicConfig(level=logging.INFO)16logger = logging.getLogger(__name__)17 18# Hugging Face Spaces compatible paths19CACHE_DIR = Path("/data/cache")20TEMP_DIR = Path("/tmp") 21LOG_DIR = Path("/data/logs")22 23# Create directories24for dir_path in [CACHE_DIR, TEMP_DIR, LOG_DIR]:25 try:26 os.makedirs(dir_path, exist_ok=True, mode=0o777)27 except Exception as e:28 logger.warning(f"Could not create {dir_path}: {e}")29 30os.environ["TRANSFORMERS_CACHE"] = str(CACHE_DIR)31os.environ["TORCH_HOME"] = str(CACHE_DIR)32os.environ["TMPDIR"] = str(TEMP_DIR)33 34LOG_FILE = LOG_DIR / "object_logs.txt"35 36# Initialize session state37if 'label_stats' not in st.session_state:38 st.session_state.label_stats = {}39if 'camera_mode' not in st.session_state:40 st.session_state.camera_mode = "upload" # Default to upload mode41 42@st.cache_resource43def load_model():44 try:45 weights = FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT46 model = fasterrcnn_resnet50_fpn_v2(weights=weights, box_score_thresh=0.5)47 model.eval()48 return model, weights49 except Exception as e:50 st.error(f"Failed to load model: {str(e)}")51 return None, None52 53model, weights = load_model()54if model is None:55 st.error("Model loading failed. Please refresh the page.")56 st.stop()57 58categories = weights.meta["categories"]59img_preprocess = weights.transforms()60 61def detect_objects_in_image(image):62 """Detect objects in a single image"""63 try:64 # Preprocess image65 img_processed = img_preprocess(image)66 67 # Run inference68 with torch.no_grad():69 prediction = model(img_processed.unsqueeze(0))[0]70 71 labels = [categories[label] for label in prediction["labels"]]72 scores = prediction["scores"].detach().cpu().numpy()73 boxes = prediction["boxes"].detach().cpu().numpy()74 75 return {76 "labels": labels,77 "scores": scores,78 "boxes": boxes79 }80 except Exception as e:81 logger.error(f"Detection error: {str(e)}")82 return {"labels": [], "scores": [], "boxes": []}83 84def draw_boxes_on_image(image, prediction):85 """Draw bounding boxes on image"""86 from PIL import ImageDraw, ImageFont87 88 img_draw = image.copy()89 draw = ImageDraw.Draw(img_draw)90 91 try:92 font = ImageFont.load_default()93 except:94 font = None95 96 colors = ['red', 'blue', 'green', 'yellow', 'purple', 'orange', 'pink', 'brown']97 98 for i, (label, score, box) in enumerate(zip(99 prediction["labels"], 100 prediction["scores"], 101 prediction["boxes"]102 )):103 if score > 0.5: # Only show high confidence detections104 color = colors[i % len(colors)]105 x1, y1, x2, y2 = map(int, box)106 107 # Draw bounding box108 draw.rectangle([x1, y1, x2, y2], outline=color, width=3)109 110 # Draw label111 text = f"{label}: {score:.2f}"112 if font:113 draw.text((x1, y1-20), text, fill=color, font=font)114 else:115 draw.text((x1, y1-20), text, fill=color)116 117 return img_draw118 119def update_statistics(prediction):120 """Update detection statistics"""121 for label, score in zip(prediction["labels"], prediction["scores"]):122 if float(score) > 0.5: # Only count high-confidence detections123 if label not in st.session_state.label_stats:124 st.session_state.label_stats[label] = {125 "count": 0,126 "total_score": 0.0127 }128 st.session_state.label_stats[label]["count"] += 1129 st.session_state.label_stats[label]["total_score"] += float(score)130 131def log_predictions():132 """Log prediction statistics to file"""133 try:134 if not st.session_state.label_stats:135 st.warning("No statistics to log")136 return137 138 timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")139 total_score = sum(v["total_score"] for v in st.session_state.label_stats.values()) or 1140 141 with open(LOG_FILE, "a") as f:142 f.write(f"\n[{timestamp}]\n")143 f.write("Label\t\tCount\tConfidence %\n")144 f.write("-"*40 + "\n")145 for label, data in st.session_state.label_stats.items():146 percentage = (data["total_score"] / total_score * 100)147 f.write(f"{label.title()}\t\t{data['count']}\t{percentage:.1f}%\n")148 149 st.success("Statistics logged successfully!")150 logger.info("Statistics logged successfully")151 except Exception as e:152 error_msg = f"Logging failed: {str(e)}"153 st.error(error_msg)154 logger.error(error_msg)155 156# Enhanced RTC Configuration with TURN servers157RTC_CONFIGURATION = RTCConfiguration({158 "iceServers": [159 {"urls": ["stun:stun.l.google.com:19302"]},160 {"urls": ["stun:stun.cloudflare.com:3478"]},161 {"urls": ["stun:openrelay.metered.ca:80"]},162 {163 "urls": ["turn:openrelay.metered.ca:80"],164 "username": "openrelayproject",165 "credential": "openrelayproject"166 }167 ],168 "iceCandidatePoolSize": 10,169 "iceTransportPolicy": "all"170})171 172def video_frame_callback(frame):173 """Process video frames"""174 try:175 img = frame.to_image()176 prediction = detect_objects_in_image(img)177 update_statistics(prediction)178 except Exception as e:179 logger.error(f"Frame processing error: {str(e)}")180 return frame181 182# Main UI183st.title("🎥 Real-Time Object Detector")184 185# Mode selection186mode = st.radio(187 "Select Detection Mode:",188 ["📤 Upload Images", "📹 Live Camera (Beta)"],189 horizontal=True190)191 192if mode == "📤 Upload Images":193 st.header("Upload Image Detection")194 195 uploaded_files = st.file_uploader(196 "Choose image files", 197 type=['png', 'jpg', 'jpeg'],198 accept_multiple_files=True199 )200 201 if uploaded_files:202 for uploaded_file in uploaded_files:203 st.subheader(f"Results for: {uploaded_file.name}")204 205 # Load and display original image206 image = Image.open(uploaded_file).convert('RGB')207 208 col1, col2 = st.columns(2)209 210 with col1:211 st.write("**Original Image**")212 st.image(image, use_column_width=True)213 214 # Detect objects215 prediction = detect_objects_in_image(image)216 217 # Update statistics218 update_statistics(prediction)219 220 with col2:221 st.write("**Detection Results**")222 if prediction["labels"]:223 # Draw bounding boxes224 img_with_boxes = draw_boxes_on_image(image, prediction)225 st.image(img_with_boxes, use_column_width=True)226 227 # Show detection details228 st.write("**Detected Objects:**")229 for label, score in zip(prediction["labels"], prediction["scores"]):230 if score > 0.5:231 st.write(f"• {label.title()}: {score:.1%} confidence")232 else:233 st.write("No objects detected with high confidence")234 235elif mode == "📹 Live Camera (Beta)":236 st.header("Live Camera Detection")237 st.warning("⚠️ Camera mode may not work on all platforms. If it fails, use Upload mode instead.")238 239 try:240 ctx = webrtc_streamer(241 key="object-detector-v3",242 mode=WebRtcMode.SENDRECV,243 rtc_configuration=RTC_CONFIGURATION,244 video_frame_callback=video_frame_callback,245 media_stream_constraints={246 "video": {247 "width": {"min": 320, "ideal": 640, "max": 1280},248 "height": {"min": 240, "ideal": 480, "max": 720},249 "frameRate": {"min": 5, "ideal": 10, "max": 15}250 }, 251 "audio": False252 },253 async_processing=True254 )255 256 if ctx.state.playing:257 st.success("✅ Camera active - Detection running")258 elif ctx.state.signalling:259 st.info("🔄 Connecting to camera...")260 else:261 st.info("📷 Click START to begin camera detection")262 263 except Exception as e:264 st.error(f"Camera failed to initialize: {str(e)}")265 st.info("💡 **Tip:** Try the 'Upload Images' mode instead!")266 267# Statistics Display268if st.session_state.label_stats:269 st.header("📊 Detection Statistics")270 271 total_score = sum(v["total_score"] for v in st.session_state.label_stats.values()) or 1272 stats = []273 274 for label, data in st.session_state.label_stats.items():275 percentage = (data["total_score"] / total_score * 100)276 stats.append({277 "Object": label.title(),278 "Count": data["count"],279 "Avg Confidence": f"{(data['total_score']/data['count']):.1%}",280 "Total Confidence": f"{percentage:.1f}%"281 })282 283 stats_df = pd.DataFrame(stats).sort_values("Count", ascending=False)284 st.dataframe(stats_df, use_container_width=True)285else:286 st.info("📈 Statistics will appear here after detecting objects")287 288# Control buttons289st.header("🎛️ Controls")290col1, col2, col3 = st.columns(3)291 292with col1:293 if st.button("💾 Save Statistics", use_container_width=True):294 log_predictions()295 296with col2:297 if st.button("🗑️ Clear Statistics", use_container_width=True):298 st.session_state.label_stats = {}299 st.rerun()300 301with col3:302 if st.button("📜 View Log History", use_container_width=True):303 try:304 if LOG_FILE.exists():305 with open(LOG_FILE, "r") as f:306 log_content = f.read()307 if log_content.strip():308 st.text_area("Log History", log_content, height=200)309 else:310 st.info("Log file is empty")311 else:312 st.info("No log file found yet")313 except Exception as e:314 st.error(f"Could not read log file: {str(e)}")315 316# System info317with st.expander("ℹ️ System Information"):318 st.write(f"**Model:** Faster R-CNN ResNet50 FPN V2")319 st.write(f"**Confidence Threshold:** 50%")320 st.write(f"**Supported Objects:** {len(categories)} categories")321 st.write(f"**Cache Directory:** {CACHE_DIR}")322 st.write(f"**Log Directory:** {LOG_DIR}")323 324st.markdown("---")325st.markdown("🚀 **Powered by PyTorch & Streamlit** | 📧 Object detection with confidence scoring")