Spaces:
Sleeping
Sleeping
| document.addEventListener('DOMContentLoaded', () => { | |
| // State management | |
| let currentFile = null; | |
| let isDrawing = false; | |
| let startX, startY; | |
| let roi = { x1: 0, y1: 0, x2: 100, y2: 100 }; | |
| let previewImage = new Image(); | |
| // Elements | |
| const modelDropZone = document.getElementById('model-drop-zone'); | |
| const modelInput = document.getElementById('model-input'); | |
| const modelStatus = document.getElementById('model-status'); | |
| const statusText = document.getElementById('status-text'); | |
| const statusIcon = modelStatus.querySelector('i'); | |
| const mediaDropZone = document.getElementById('media-drop-zone'); | |
| const mediaInput = document.getElementById('media-input'); | |
| const previewSection = document.getElementById('preview-section'); | |
| const roiCanvas = document.getElementById('roi-canvas'); | |
| const ctx = roiCanvas.getContext('2d'); | |
| const thresholdInput = document.getElementById('threshold-input'); | |
| const confMaxInput = document.getElementById('conf-max-input'); | |
| const confRangeVal = document.getElementById('conf-range-val'); | |
| const roiX1 = document.getElementById('roi-x1'); | |
| const roiY1 = document.getElementById('roi-y1'); | |
| const roiX2 = document.getElementById('roi-x2'); | |
| const roiY2 = document.getElementById('roi-y2'); | |
| const resetRoiBtn = document.getElementById('reset-roi-btn'); | |
| const progressCard = document.getElementById('progress-card'); | |
| const loading = document.getElementById('loading'); | |
| const videoProgressContainer = document.getElementById('video-progress-container'); | |
| const videoProgressBar = document.getElementById('video-progress-bar'); | |
| const videoStatusMsg = document.getElementById('video-status-msg'); | |
| const videoPercentage = document.getElementById('video-percentage'); | |
| const analyzeBtn = document.getElementById('analyze-btn'); | |
| const resultSection = document.getElementById('result-section'); | |
| const resultImage = document.getElementById('result-image'); | |
| const resultCount = document.getElementById('result-count'); | |
| const downloadBtn = document.getElementById('download-btn'); | |
| const videoResultSection = document.getElementById('video-result-section'); | |
| const resultVideo = document.getElementById('result-video'); | |
| const videoDownloadBtn = document.getElementById('video-download-btn'); | |
| // Drag and Drop Setup | |
| [modelDropZone, mediaDropZone].forEach(zone => { | |
| ['dragenter', 'dragover', 'dragleave', 'drop'].forEach(eventName => { | |
| zone.addEventListener(eventName, e => { | |
| e.preventDefault(); | |
| e.stopPropagation(); | |
| }); | |
| }); | |
| ['dragenter', 'dragover'].forEach(eventName => { | |
| zone.addEventListener(eventName, () => zone.classList.add('dragover')); | |
| }); | |
| ['dragleave', 'drop'].forEach(eventName => { | |
| zone.addEventListener(eventName, () => zone.classList.remove('dragover')); | |
| }); | |
| }); | |
| // Clicks | |
| modelDropZone.addEventListener('click', () => modelInput.click()); | |
| mediaDropZone.addEventListener('click', () => mediaInput.click()); | |
| // File Handlers | |
| modelInput.addEventListener('change', e => handleModelUpload(e.target.files[0])); | |
| modelDropZone.addEventListener('drop', e => handleModelUpload(e.dataTransfer.files[0])); | |
| mediaInput.addEventListener('change', e => handleMediaSelection(e.target.files[0])); | |
| mediaDropZone.addEventListener('drop', e => handleMediaSelection(e.dataTransfer.files[0])); | |
| // Settings | |
| const updateConfLabel = () => { | |
| const min = Math.round(thresholdInput.value * 100); | |
| const max = Math.round(confMaxInput.value * 100); | |
| confRangeVal.innerText = `${min}% - ${max}%`; | |
| }; | |
| thresholdInput.addEventListener('input', updateConfLabel); | |
| confMaxInput.addEventListener('input', updateConfLabel); | |
| [roiX1, roiY1, roiX2, roiY2].forEach(input => { | |
| input.addEventListener('change', updateROIFromInputs); | |
| }); | |
| resetRoiBtn.addEventListener('click', () => { | |
| roi = { x1: 0, y1: 0, x2: 100, y2: 100 }; | |
| updateInputsFromROI(); | |
| drawROI(); | |
| }); | |
| analyzeBtn.addEventListener('click', startInference); | |
| // --- Functions --- | |
| async function handleModelUpload(file) { | |
| if (!file || !file.name.endsWith('.pt')) { | |
| showToast('Please upload a valid YOLO .pt model.', 'error'); | |
| return; | |
| } | |
| const formData = new FormData(); | |
| formData.append('file', file); | |
| statusText.innerText = 'Uploading model...'; | |
| modelStatus.classList.remove('loaded'); | |
| statusIcon.className = 'fas fa-spinner fa-spin'; | |
| try { | |
| const resp = await fetch('/upload-model', { method: 'POST', body: formData }); | |
| const data = await resp.json(); | |
| if (data.status === 'success') { | |
| statusText.innerText = `Model: ${file.name}`; | |
| modelStatus.classList.add('loaded'); | |
| statusIcon.className = 'fas fa-check-circle'; | |
| showToast('Model loaded successfully!', 'success'); | |
| } else { | |
| throw new Error(data.detail); | |
| } | |
| } catch (err) { | |
| statusText.innerText = 'Error loading model'; | |
| statusIcon.className = 'fas fa-exclamation-circle'; | |
| showToast(err.message, 'error'); | |
| } | |
| } | |
| async function handleMediaSelection(file) { | |
| if (!file) return; | |
| currentFile = file; | |
| // Reset state | |
| resultSection.classList.add('hidden'); | |
| videoResultSection.classList.add('hidden'); | |
| progressCard.classList.add('hidden'); | |
| if (file.type.startsWith('image/')) { | |
| const reader = new FileReader(); | |
| reader.onload = e => { | |
| previewImage.onload = () => initCanvas(); | |
| previewImage.src = e.target.result; | |
| }; | |
| reader.readAsDataURL(file); | |
| } else if (file.type.startsWith('video/')) { | |
| extractVideoFrame(file); | |
| } else { | |
| showToast('Unsupported file type.', 'error'); | |
| } | |
| } | |
| function extractVideoFrame(file) { | |
| const video = document.createElement('video'); | |
| video.preload = 'metadata'; | |
| video.src = URL.createObjectURL(file); | |
| video.onloadedmetadata = () => { | |
| video.currentTime = 0.1; // Seek a bit in to avoid black frames | |
| }; | |
| video.onseeked = () => { | |
| const tempCanvas = document.createElement('canvas'); | |
| tempCanvas.width = video.videoWidth; | |
| tempCanvas.height = video.videoHeight; | |
| const tempCtx = tempCanvas.getContext('2d'); | |
| tempCtx.drawImage(video, 0, 0); | |
| previewImage.onload = () => initCanvas(); | |
| previewImage.src = tempCanvas.toDataURL('image/jpeg'); | |
| URL.revokeObjectURL(video.src); | |
| }; | |
| } | |
| function initCanvas() { | |
| previewSection.classList.remove('hidden'); | |
| previewSection.scrollIntoView({ behavior: 'smooth' }); | |
| // Scale canvas to fit container but keep aspect ratio | |
| const containerWidth = roiCanvas.parentElement.clientWidth; | |
| const scale = containerWidth / previewImage.width; | |
| roiCanvas.width = previewImage.width * scale; | |
| roiCanvas.height = previewImage.height * scale; | |
| drawROI(); | |
| } | |
| function drawROI() { | |
| ctx.clearRect(0, 0, roiCanvas.width, roiCanvas.height); | |
| ctx.drawImage(previewImage, 0, 0, roiCanvas.width, roiCanvas.height); | |
| // Darken outside | |
| ctx.fillStyle = 'rgba(0, 0, 0, 0.4)'; | |
| const x1 = (roi.x1 / 100) * roiCanvas.width; | |
| const y1 = (roi.y1 / 100) * roiCanvas.height; | |
| const x2 = (roi.x2 / 100) * roiCanvas.width; | |
| const y2 = (roi.y2 / 100) * roiCanvas.height; | |
| const w = x2 - x1; | |
| const h = y2 - y1; | |
| // Draw overlay path with a "hole" for the ROI | |
| ctx.beginPath(); | |
| ctx.rect(0, 0, roiCanvas.width, roiCanvas.height); | |
| ctx.rect(x1, y1, w, h); | |
| ctx.fill('evenodd'); | |
| // Draw border | |
| ctx.strokeStyle = '#f59e0b'; | |
| ctx.lineWidth = 3; | |
| ctx.setLineDash([5, 5]); | |
| ctx.strokeRect(x1, y1, w, h); | |
| // Optional: corner handles design look | |
| ctx.fillStyle = '#f59e0b'; | |
| ctx.fillRect(x1-4, y1-4, 8, 8); | |
| ctx.fillRect(x2-4, y1-4, 8, 8); | |
| ctx.fillRect(x1-4, y2-4, 8, 8); | |
| ctx.fillRect(x2-4, y2-4, 8, 8); | |
| } | |
| // Canvas Events | |
| roiCanvas.addEventListener('mousedown', e => { | |
| isDrawing = true; | |
| const rect = roiCanvas.getBoundingClientRect(); | |
| startX = e.clientX - rect.left; | |
| startY = e.clientY - rect.top; | |
| roi.x1 = (startX / roiCanvas.width) * 100; | |
| roi.y1 = (startY / roiCanvas.height) * 100; | |
| }); | |
| roiCanvas.addEventListener('mousemove', e => { | |
| if (!isDrawing) return; | |
| const rect = roiCanvas.getBoundingClientRect(); | |
| const curX = e.clientX - rect.left; | |
| const curY = e.clientY - rect.top; | |
| roi.x2 = (curX / roiCanvas.width) * 100; | |
| roi.y2 = (curY / roiCanvas.height) * 100; | |
| updateInputsFromROI(); | |
| drawROI(); | |
| }); | |
| roiCanvas.addEventListener('mouseup', () => { | |
| isDrawing = false; | |
| // Normalize coordinates (ensure x1 < x2, y1 < y2) | |
| if (roi.x1 > roi.x2) [roi.x1, roi.x2] = [roi.x2, roi.x1]; | |
| if (roi.y1 > roi.y2) [roi.y1, roi.y2] = [roi.y2, roi.y1]; | |
| updateInputsFromROI(); | |
| drawROI(); | |
| }); | |
| function updateInputsFromROI() { | |
| roiX1.value = Math.round(roi.x1); | |
| roiY1.value = Math.round(roi.y1); | |
| roiX2.value = Math.round(roi.x2); | |
| roiY2.value = Math.round(roi.y2); | |
| } | |
| function updateROIFromInputs() { | |
| roi.x1 = parseInt(roiX1.value); | |
| roi.y1 = parseInt(roiY1.value); | |
| roi.x2 = parseInt(roiX2.value); | |
| roi.y2 = parseInt(roiY2.value); | |
| drawROI(); | |
| } | |
| async function startInference() { | |
| if (!currentFile) return; | |
| progressCard.classList.remove('hidden'); | |
| progressCard.scrollIntoView({ behavior: 'smooth' }); | |
| const isVideo = currentFile.type.startsWith('video/'); | |
| const formData = new FormData(); | |
| formData.append('file', currentFile); | |
| formData.append('conf_min', thresholdInput.value); | |
| formData.append('conf_max', confMaxInput.value); | |
| formData.append('roi', JSON.stringify(roi)); | |
| if (isVideo) { | |
| handleVideoInference(formData); | |
| } else { | |
| handleImageInference(formData); | |
| } | |
| } | |
| async function handleImageInference(formData) { | |
| loading.classList.remove('hidden'); | |
| videoProgressContainer.classList.add('hidden'); | |
| try { | |
| const resp = await fetch('/inference', { method: 'POST', body: formData }); | |
| const data = await resp.json(); | |
| if (data.status === 'success') { | |
| resultImage.src = data.image; | |
| resultCount.innerText = `${data.count} Detections`; | |
| resultSection.classList.remove('hidden'); | |
| resultSection.scrollIntoView({ behavior: 'smooth' }); | |
| } else { | |
| throw new Error(data.detail); | |
| } | |
| } catch (err) { | |
| showToast(err.message, 'error'); | |
| } finally { | |
| progressCard.classList.add('hidden'); | |
| } | |
| } | |
| async function handleVideoInference(formData) { | |
| loading.classList.add('hidden'); | |
| videoProgressContainer.classList.remove('hidden'); | |
| videoProgressBar.style.width = '0%'; | |
| videoPercentage.innerText = '0%'; | |
| videoStatusMsg.innerText = 'Uploading video...'; | |
| try { | |
| const resp = await fetch('/inference-video', { method: 'POST', body: formData }); | |
| const data = await resp.json(); | |
| if (data.status === 'success') { | |
| videoStatusMsg.innerText = 'Processing frames...'; | |
| pollVideoProgress(data.task_id); | |
| } else { | |
| throw new Error(data.detail); | |
| } | |
| } catch (err) { | |
| showToast(err.message, 'error'); | |
| progressCard.classList.add('hidden'); | |
| } | |
| } | |
| function pollVideoProgress(taskId) { | |
| const interval = setInterval(async () => { | |
| try { | |
| const resp = await fetch(`/video-progress/${taskId}`); | |
| const data = await resp.json(); | |
| if (data.status === 'processing') { | |
| videoProgressBar.style.width = `${data.progress}%`; | |
| videoPercentage.innerText = `${data.progress}%`; | |
| } else if (data.status === 'completed') { | |
| clearInterval(interval); | |
| videoProgressBar.style.width = '100%'; | |
| videoPercentage.innerText = '100%'; | |
| videoStatusMsg.innerText = 'Processing complete!'; | |
| showVideoResult(taskId); | |
| } else if (data.status === 'error') { | |
| clearInterval(interval); | |
| showToast(data.message, 'error'); | |
| progressCard.classList.add('hidden'); | |
| } | |
| } catch (err) { | |
| console.error(err); | |
| } | |
| }, 1000); | |
| } | |
| function showVideoResult(taskId) { | |
| const url = `/video-result/${taskId}`; | |
| resultVideo.src = url; | |
| videoDownloadBtn.href = url; | |
| videoResultSection.classList.remove('hidden'); | |
| videoResultSection.scrollIntoView({ behavior: 'smooth' }); | |
| progressCard.classList.add('hidden'); | |
| } | |
| function showToast(message, type = 'info') { | |
| const toast = document.createElement('div'); | |
| toast.className = `toast ${type}`; | |
| Object.assign(toast.style, { | |
| position: 'fixed', bottom: '20px', right: '20px', padding: '1rem 1.5rem', | |
| borderRadius: '10px', color: 'white', zIndex: '1000', | |
| background: type === 'error' ? '#ef4444' : '#10b981', | |
| boxShadow: '0 4px 15px rgba(0,0,0,0.3)', animation: 'slideIn 0.3s ease forwards' | |
| }); | |
| toast.innerText = message; | |
| document.body.appendChild(toast); | |
| setTimeout(() => { | |
| toast.style.animation = 'slideOut 0.3s ease forwards'; | |
| setTimeout(() => toast.remove(), 300); | |
| }, 3000); | |
| } | |
| }); | |