Skip to content

Commit

Permalink
Merge pull request #194 from prajwaltulawe/master
Browse files Browse the repository at this point in the history 
Added new export feature - YOLOV5 PyTorch Format
  • Loading branch information
@amitguptagwl
amitguptagwl authored Jan 1, 2024
2 parents 6564316 + 2b91177 commit 7597d2d
Show file tree
Hide file tree
Showing 3 changed files with 201 additions and 16 deletions.
1 change: 1 addition & 0 deletions index.html
View file
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
View file
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
View file

Large diffs are not rendered by default.

0 comments on commit 7597d2d

Please sign in to comment.