Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added new export feature - YOLOV5 PyTorch Format #194

Merged
merged 4 commits into from
Jan 1, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions index.html
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,7 @@
</script>
<script src="js/settings.js"></script>
<script src="js/thirdparty/svg.min.js"></script>
<script src="js/thirdparty/jszip.min.js"></script>
<script src="js/thirdparty/svg.draw.min.js"></script>
<script src="js/thirdparty/svg.select.min.js"></script>
<script src="js/thirdparty/svg.resize.min.js"></script>
Expand Down
203 changes: 187 additions & 16 deletions js/savefile.js
Original file line number Diff line number Diff line change
Expand Up @@ -2,22 +2,25 @@ function selectFileTypeToSave(){
$.dialog({
title: 'Save/Export as',
content: `<div style="text-align:center;">
<div>
<button class="btn btn-primary savebtn" onclick="javascript:saveAsNimn()" id="saveAsNimn">Project file</button>
</div>
<div>
<button class="btn btn-primary savebtn" onclick="javascript:saveAsDlibXML()" id="saveAsNimn">Dlib XML</button>
</div>
<div>
<button class="btn btn-primary savebtn" onclick="javascript:saveAsDlibPts()" id="saveAsNimn">Dlib pts</button>
</div>
<div>
<button class="btn btn-primary savebtn" onclick="javascript:saveAsCOCO()" id="saveAsCOCO">COCO JSON</button>
</div>
<div>
<button class="btn btn-primary savebtn" onclick="javascript:saveAsPascalVOC()" id="saveAsPascalVOC">Pascal VOC XML</button>
</div>
<div>`,
<div>
<button class="btn btn-primary savebtn" onclick="javascript:saveAsNimn()" id="saveAsNimn">Project file</button>
</div>
<div>
<button class="btn btn-primary savebtn" onclick="javascript:saveAsDlibXML()" id="saveAsNimn">Dlib XML</button>
</div>
<div>
<button class="btn btn-primary savebtn" onclick="javascript:saveAsDlibPts()" id="saveAsNimn">Dlib pts</button>
</div>
<div>
<button class="btn btn-primary savebtn" onclick="javascript:saveAsCOCO()" id="saveAsCOCO">COCO JSON</button>
</div>
<div>
<button class="btn btn-primary savebtn" onclick="javascript:saveAsPascalVOC()" id="saveAsPascalVOC">Pascal VOC XML</button>
</div>
<div>
<button class="btn btn-primary savebtn" onclick="javascript:saveAsYoloV5Pytorch()" id="saveAsYoloV5Pytorch">YOLO V5 Pytorch</button>
</div>
<div>`,
escapeKey: true,
backgroundDismiss: true,
});
Expand Down Expand Up @@ -131,6 +134,174 @@ function saveAsPascalVOC(){

}

/**
* Save labelled data as YOLO supported TXT Files file.
* It will export files in a zip format
*/

function saveAsYoloV5Pytorch() {
// COUNT TOTAL MASKED IMAGES
var totalLabbeledImages = 0;
for (const key in labellingData) {
if(labellingData[key].shapes[0] && labellingData[key].shapes[0].points.length > 0){
totalLabbeledImages++;
}
}

// MODAL FOR RECEIVING DATA SPLIT INPUTS
$.dialog({
title: `Split ${totalLabbeledImages} Labeled Images into :`,
content: `<div class="col w-75 m-auto">
<div class="row">
<div ref="label-data" class="col"> Train </div>
<div ref="label-data" class="col"> Test </div>
<div ref="label-data" class="col"> Valid </div>
</div>
<div class="row justify-content-between">
<input type="text" class="col" value="${Math.floor(totalLabbeledImages/100*70)}" onchange="javascript:validateImageSplit()" style="width: 20px;" placeholder="">
<input type="text" class="col" value="${Math.floor(totalLabbeledImages/100*30)}" onchange="javascript:validateImageSplit()" style="width: 20px;" placeholder="">
<input type="text" class="col" value="${Math.floor(totalLabbeledImages/100*10)}" onchange="javascript:validateImageSplit()" style="width: 20px;" placeholder="">
</div>
</div>
<div class="d-flex justify-content-center">
<span class="mt-2 w-75" id="splitMsg" > </span>
</div>
<div class="d-flex" style="text-align:center;">
<button class="btn btn-primary savebtn" onclick="javascript:yoloDataRendering()">Generate</button>
</div>`,
escapeKey: true,
backgroundDismiss: true,
});
}

function validateImageSplit(){
// COLLECTING USER INPUT VALUES
const valueElements = document.querySelectorAll('[style="width: 20px;"]');
const trainImages = parseInt(valueElements[0].value);
const testImages = parseInt(valueElements[1].value);
const validImages = parseInt(valueElements[2].value);

var totalLabbeledImages = 0;
for (const key in labellingData) {
if(labellingData[key].shapes[0] && labellingData[key].shapes[0].points.length > 0){
totalLabbeledImages++;
}
}

var remainingImages = totalLabbeledImages -trainImages -testImages -validImages;
if(remainingImages > 0){
document.getElementById('splitMsg').innerText = `${remainingImages} Labelled Images remaining to split.`;
}

if(trainImages + testImages + validImages == totalLabbeledImages){
document.getElementById('splitMsg').innerText = "";
return [trainImages, testImages, validImages, true];
}else{
showSnackBar(`Total of ( Train+Test+Valid ) must be ${totalLabbeledImages}.`);
return [0, 0, 0, false];
}
}

function yoloDataRendering() {
var [train, test, valid, isValuesValid] = validateImageSplit();

// CHECK FOR VALID INPUTS
if (!isValuesValid) return;

// COLLECT AND SAVE UNIQUE LABELS
const labels = new Set();
for (const image in labellingData) {
const img = labellingData[image];
img.shapes.forEach((shape) => {
labels.add(shape.label);
});
}

const zip = new JSZip();
const finalLabels = [...labels];

var trainImgLimit = train;
var testImgLimit = train + test;
var validImgLimit = testImgLimit + valid;
var currImage = 0;

// CREATE DATA.YAML FILE
const dataYamlFile = `train: ../train/images\nval: ../valid/images\ntest: ../test/images\n\nnc: ${finalLabels.length}\nnames: ['${finalLabels.join("','")}']`;
zip.file("data.yaml", dataYamlFile);

// GENERATE YOLO REQUIRED FORMAT DATA
for (const image in labellingData) {
let outputArr = [];
let fileName = image.substring(0, image.lastIndexOf(".")) + ".txt";

// RECONSTRUCT IMAGES
let base64Image = document.querySelector(`[label="${image}"]`).src;
var byteCharacters = atob(base64Image.split(",")[1]);
var byteNumbers = new Array(byteCharacters.length);
for (var i = 0; i < byteCharacters.length; i++) {
byteNumbers[i] = byteCharacters.charCodeAt(i);
}
var byteArray = new Uint8Array(byteNumbers);
var blob = new Blob([byteArray], { type: "image/jpg" }); // Specify the appropriate MIME type

// ITRATE THROUGH EACH SHAPE
labellingData[image].shapes.forEach((shape) => {
let currArr = [];
currArr.push(finalLabels.indexOf(shape.label));

// CALCULATE POINTS
if (shape.points[0][0]) {
let coordinates = shape.points;
coordinates.forEach((point) => {
let xNorm = point[0] / labellingData[image].size.width;
let yNorm = point[1] / labellingData[image].size.height;
currArr.push(xNorm, yNorm);
});
} else {
let [xMin, yMin, w, h] = shape.points;
const yoloXCenter = (xMin + w / 2) / labellingData[image].size.width;
const yoloYCenter = (yMin + h / 2) / labellingData[image].size.height;
const yoloBoxWidth = w / labellingData[image].size.width;
const yoloBoxHeight = h / labellingData[image].size.height;
currArr.push(yoloXCenter, yoloYCenter, yoloBoxWidth, yoloBoxHeight);
}
outputArr.push(currArr.join(" "));
});

// ADD FILES TO ZIP IF LABELLED
if (outputArr.length > 0) {
if (currImage < trainImgLimit) {
saveDirectory = "train";
} else if (currImage < testImgLimit) {
saveDirectory = "test";
} else if (currImage <= validImgLimit) {
saveDirectory = "valid";
}
currImage++;

zip.file(`${saveDirectory}/images/${image}`, blob);
zip.file(`${saveDirectory}/labels/${fileName}`, outputArr.join("\n"));
}
}

// SET FILE NAME
var curTimeStamp = new Date();
var timeStamp = `${curTimeStamp.getDate()}/ ${curTimeStamp.getMonth() + 1}/ ${curTimeStamp.getFullYear()}/
${curTimeStamp.getHours()}: ${curTimeStamp.getMinutes()}: ${curTimeStamp.getSeconds()}`;

// CREATE ZIP AND SAVE
zip.generateAsync({ type: "blob" })
.then(function (content) {
saveAs(content, timeStamp + ".zip");
showSnackBar(`File will be downloaded automatically`);
analytics_reportExportType("yoloV5PyTorch");
})
.catch(function (error) {
showSnackBar("Error occoured while creating ZIP file");
console.error("Error creating ZIP file:", error);
});
}

/**
* Save given data to a file
* @param {*} data
Expand Down
13 changes: 13 additions & 0 deletions js/thirdparty/jszip.min.js

Large diffs are not rendered by default.