Spaces:
Running
Running
import React, { useState, useEffect } from 'react'; | |
import { Card } from '@/components/ui/card'; | |
import { Button } from '@/components/ui/button'; | |
import { Upload, Image as ImageIcon } from 'lucide-react'; | |
const ObjectDetector = () => { | |
const [status, setStatus] = useState('Loading model...'); | |
const [detections, setDetections] = useState([]); | |
const [currentImage, setCurrentImage] = useState(null); | |
const [detector, setDetector] = useState(null); | |
// Initialize the model | |
useEffect(() => { | |
const initModel = async () => { | |
try { | |
const { pipeline, env } = await import('https://cdn.jsdelivr.net/npm/@xenova/[email protected]'); | |
env.allowLocalModels = false; | |
const model = await pipeline('object-detection', 'Xenova/detr-resnet-50'); | |
setDetector(model); | |
setStatus('Ready'); | |
} catch (error) { | |
console.error('Error loading model:', error); | |
setStatus('Error loading model'); | |
} | |
}; | |
initModel(); | |
}, []); | |
const handleFileUpload = (event) => { | |
const file = event.target.files[0]; | |
if (!file) return; | |
const reader = new FileReader(); | |
reader.onload = (e) => { | |
setCurrentImage(e.target.result); | |
detectObjects(e.target.result); | |
}; | |
reader.readAsDataURL(file); | |
}; | |
const handleExampleImage = () => { | |
const EXAMPLE_URL = 'https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/city-streets.jpg'; | |
setCurrentImage(EXAMPLE_URL); | |
detectObjects(EXAMPLE_URL); | |
}; | |
const detectObjects = async (img) => { | |
if (!detector) return; | |
setStatus('Analyzing...'); | |
try { | |
const output = await detector(img, { | |
threshold: 0.5, | |
percentage: true, | |
}); | |
setDetections(output.map((det, index) => ({ | |
...det, | |
id: index + 1 | |
}))); | |
setStatus(''); | |
} catch (error) { | |
console.error('Error detecting objects:', error); | |
setStatus('Error analyzing image'); | |
} | |
}; | |
return ( | |
<Card className="p-6 max-w-4xl mx-auto"> | |
<div className="space-y-6"> | |
<div className="flex gap-4 items-center justify-center"> | |
<Button | |
onClick={() => document.getElementById('fileUpload').click()} | |
className="flex items-center gap-2" | |
> | |
<Upload className="w-4 h-4" /> | |
Upload Image | |
</Button> | |
<Button | |
onClick={handleExampleImage} | |
className="flex items-center gap-2" | |
> | |
<ImageIcon className="w-4 h-4" /> | |
Try Example | |
</Button> | |
<input | |
id="fileUpload" | |
type="file" | |
accept="image/*" | |
className="hidden" | |
onChange={handleFileUpload} | |
/> | |
</div> | |
{status && ( | |
<div className="text-center text-sm text-gray-600"> | |
{status} | |
</div> | |
)} | |
{currentImage && ( | |
<div className="relative w-full h-96 bg-gray-100 rounded overflow-hidden"> | |
<img | |
src={currentImage} | |
alt="Uploaded" | |
className="w-full h-full object-contain" | |
/> | |
{detections.map((detection) => ( | |
<div | |
key={detection.id} | |
className="absolute border-2 border-blue-500" | |
style={{ | |
left: `${detection.box.xmin}%`, | |
top: `${detection.box.ymin}%`, | |
width: `${detection.box.xmax - detection.box.xmin}%`, | |
height: `${detection.box.ymax - detection.box.ymin}%`, | |
}} | |
> | |
<span className="absolute top-0 left-0 transform -translate-y-full bg-blue-500 text-white px-1 py-0.5 text-xs rounded"> | |
{detection.label} ({(detection.score * 100).toFixed(1)}%) | |
</span> | |
</div> | |
))} | |
</div> | |
)} | |
{detections.length > 0 && ( | |
<div className="overflow-x-auto"> | |
<table className="w-full"> | |
<thead> | |
<tr className="border-b"> | |
<th className="p-2 text-left font-semibold">Object</th> | |
<th className="p-2 text-left font-semibold">Confidence</th> | |
<th className="p-2 text-left font-semibold">Location</th> | |
</tr> | |
</thead> | |
<tbody> | |
{detections.map((detection) => ( | |
<tr key={detection.id} className="border-b"> | |
<td className="p-2"> | |
<span className="inline-block px-2 py-1 bg-blue-100 text-blue-800 rounded"> | |
{detection.label} | |
</span> | |
</td> | |
<td className="p-2"> | |
{(detection.score * 100).toFixed(1)}% | |
</td> | |
<td className="p-2"> | |
<div className="text-sm text-gray-600"> | |
<div>x: {detection.box.xmin.toFixed(1)}% - {detection.box.xmax.toFixed(1)}%</div> | |
<div>y: {detection.box.ymin.toFixed(1)}% - {detection.box.ymax.toFixed(1)}%</div> | |
</div> | |
</td> | |
</tr> | |
))} | |
</tbody> | |
</table> | |
</div> | |
)} | |
</div> | |
</Card> | |
); | |
}; | |
export default ObjectDetector; |