alanvww/p5js-object-detection-webgpu

1// Alan Ren @NYU ITP2// Spring 20243 4// Global variables for the detector, image elements, and status paragraph5let detector;6let imgElement,7	defaultImg1,8	defaultImg1DataURL,9	defaultImg2,10	defaultImg2DataURL;11let statusP;12 13// Preload function to load the default image and convert it to a base64 data URL14function preload() {15	defaultImg1 = loadAndConvertImage('https://cors-anywhere-ajr.up.railway.app/https://hdwallpaperim.com/wp-content/uploads/2017/08/25/452511-street-street_view-cityscape-city.jpg', (dataURL) => {16		defaultImg1DataURL = dataURL;17	});18 19	defaultImg2 = loadAndConvertImage('https://cors-anywhere-ajr.up.railway.app/https://external-content.duckduckgo.com/iu/?u=https%3A%2F%2Fdirectoffice.com%2Fwp-content%2Fuploads%2F2024%2F06%2FOffice-Storage-Cabinets-2048x1366.jpg&f=1&nofb=1&ipt=f6b591aded06e6451a47bac25c3b6fe04aa285fb7796c7bc35b881e063a07a5f&ipo=images', (dataURL) => {20		defaultImg2DataURL = dataURL;21	});22}23 24// Function to load an image and convert it to a base64 data URL25function loadAndConvertImage(url, callback) {26	return loadImage(url, (img) => {27		let offscreenGraphics = createGraphics(img.width, img.height);28		offscreenGraphics.image(img, 0, 0);29		let dataURL = offscreenGraphics.elt.toDataURL();30		callback(dataURL);31	});32}33 34// Setup function initializes the p5 sketch35function setup() {36	// Create a paragraph element to show status messages37	statusP = createP('Loading model...').style('color', 'blue');38 39	// Check if the transformersPipeline is loaded and available40	if (typeof window.transformersPipeline === 'undefined') {41		console.error('Transformers pipeline not available!');42		statusP.html('Pipeline loading failed.');43		return;44	}45 46	// Initialize the object detection model using the transformersPipeline47	window48		.transformersPipeline('object-detection', 'Xenova/yolos-tiny', {49			device: 'webgpu',50		})51		.then((d) => {52			detector = d;53			statusP.html('Model ready. Upload an image.');54		})55		.catch((error) => {56			console.error('Error loading the model: ', error);57			statusP.html('Model loading failed.');58		});59 60	// Create a file input for image upload61	createFileInput(imageUploaded).attribute('accept', 'image/*');62	createButton('Example 1').mousePressed(() => loadExampleImage(defaultImg1DataURL));63	createButton('Example 2').mousePressed(() => loadExampleImage(defaultImg2DataURL));64}65 66// Function to load an example image67function loadExampleImage(dataURL) {68	if (imgElement) {69		imgElement.remove();70	}71	imgElement = createImg(dataURL, '').hide();72	detect(imgElement);73}74 75// Function to handle uploaded images76function imageUploaded(file) {77	if (file.type === 'image') {78		if (imgElement) {79			imgElement.remove(); // Remove the previous image if exists80		}81		imgElement = createImg(file.data, '').hide();82		detect(imgElement);83	} else {84		statusP.html('Please upload an image file.');85	}86}87 88// Function to perform object detection89async function detect(image) {90	statusP.html('Analysing...');91	const results = await detector(image.elt.src, {92		threshold: 0.5,93		percentage: true,94	});95 96	displayImageAndBoxes(image, results);97	statusP.html('Image Processed');98}99 100// Function to display the image and bounding boxes on the canvas101function displayImageAndBoxes(img, results) {102	let scaleX = windowWidth / img.width;103	let scaleY = scaleX;104 105	let scaledWidth = img.width * scaleX;106	let scaledHeight = img.height * scaleY;107 108	resizeCanvas(windowWidth, scaledHeight + 200);109	image(img, 0, 100, scaledWidth, scaledHeight);110 111	for (const result of results) {112		const { box, label } = result;113		const { xmax, xmin, ymax, ymin } = box;114 115		let rectX = xmin * scaledWidth;116		let rectY = ymin * scaledHeight + 100;117		let rectWidth = (xmax - xmin) * scaledWidth;118		let rectHeight = (ymax - ymin) * scaledHeight;119 120		stroke(255, 0, 0);121		noFill();122		rect(rectX, rectY, rectWidth, rectHeight);123 124		fill(255);125		strokeWeight(1);126		textSize(16);127		text(label, rectX, rectY - 10);128	}129}130 131// The draw function is left empty because updates only occur during detection132function draw() {133	// Leave this function empty if updates only occur during detection134}135