sketch.js

// Alan Ren @NYU ITP
// Spring 2024

// Global variables for the detector, image elements, and status paragraph
let detector;
let imgElement,
	defaultImg1,
	defaultImg1DataURL,
	defaultImg2,
	defaultImg2DataURL;
let statusP;

// Preload function to load the default image and convert it to a base64 data URL
function preload() {
	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) => {
		defaultImg1DataURL = dataURL;
	});

	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) => {
		defaultImg2DataURL = dataURL;
	});
}

// Function to load an image and convert it to a base64 data URL
function loadAndConvertImage(url, callback) {
	return loadImage(url, (img) => {
		let offscreenGraphics = createGraphics(img.width, img.height);
		offscreenGraphics.image(img, 0, 0);
		let dataURL = offscreenGraphics.elt.toDataURL();
		callback(dataURL);
	});
}

// Setup function initializes the p5 sketch
function setup() {
	// Create a paragraph element to show status messages
	statusP = createP('Loading model...').style('color', 'blue');

	// Check if the transformersPipeline is loaded and available
	if (typeof window.transformersPipeline === 'undefined') {
		console.error('Transformers pipeline not available!');
		statusP.html('Pipeline loading failed.');
		return;
	}

	// Initialize the object detection model using the transformersPipeline
	window
		.transformersPipeline('object-detection', 'Xenova/yolos-tiny', {
			device: 'webgpu',
		})
		.then((d) => {
			detector = d;
			statusP.html('Model ready. Upload an image.');
		})
		.catch((error) => {
			console.error('Error loading the model: ', error);
			statusP.html('Model loading failed.');
		});

	// Create a file input for image upload
	createFileInput(imageUploaded).attribute('accept', 'image/*');
	createButton('Example 1').mousePressed(() => loadExampleImage(defaultImg1DataURL));
	createButton('Example 2').mousePressed(() => loadExampleImage(defaultImg2DataURL));
}

// Function to load an example image
function loadExampleImage(dataURL) {
	if (imgElement) {
		imgElement.remove();
	}
	imgElement = createImg(dataURL, '').hide();
	detect(imgElement);
}

// Function to handle uploaded images
function imageUploaded(file) {
	if (file.type === 'image') {
		if (imgElement) {
			imgElement.remove(); // Remove the previous image if exists
		}
		imgElement = createImg(file.data, '').hide();
		detect(imgElement);
	} else {
		statusP.html('Please upload an image file.');
	}
}

// Function to perform object detection
async function detect(image) {
	statusP.html('Analysing...');
	const results = await detector(image.elt.src, {
		threshold: 0.5,
		percentage: true,
	});

	displayImageAndBoxes(image, results);
	statusP.html('Image Processed');
}

// Function to display the image and bounding boxes on the canvas
function displayImageAndBoxes(img, results) {
	let scaleX = windowWidth / img.width;
	let scaleY = scaleX;

	let scaledWidth = img.width * scaleX;
	let scaledHeight = img.height * scaleY;

	resizeCanvas(windowWidth, scaledHeight + 200);
	image(img, 0, 100, scaledWidth, scaledHeight);

	for (const result of results) {
		const { box, label } = result;
		const { xmax, xmin, ymax, ymin } = box;

		let rectX = xmin * scaledWidth;
		let rectY = ymin * scaledHeight + 100;
		let rectWidth = (xmax - xmin) * scaledWidth;
		let rectHeight = (ymax - ymin) * scaledHeight;

		stroke(255, 0, 0);
		noFill();
		rect(rectX, rectY, rectWidth, rectHeight);

		fill(255);
		strokeWeight(1);
		textSize(16);
		text(label, rectX, rectY - 10);
	}
}

// The draw function is left empty because updates only occur during detection
function draw() {
	// Leave this function empty if updates only occur during detection
}