Scale usage to many thousands of images, disables previews

- Apparently, tfjs doesn't do memory management for you!

.gitignore

@@ -104,7 +104,7 @@ venv.bak/
 .mypy_cache/
 
 #screenshot data
-*.png
+images/*.png
 *.jpg
 *.tar
 

webui/src/app/components/eyeballer/eyeballer.component.html

@@ -1,12 +1,13 @@
 <mat-toolbar class="app-toolbar" color="primary"  fxLayout="row" fxLayoutAlign="start center">
   <h2>
-    <i class="fas fa-eye fa-fw"></i>
-    Eyeballer
+    <img src="assets/eyeballer_logo.png" width="46" height="60">
   </h2>
-
+  Eyeballer
+  <span fxFlex></span>
+  <span>Bishop Fox R&D</span>
   <span fxFlex></span>
 
-  <button mat-raised-button  (click)="restart()">
+  <button mat-raised-button (click)="restart()">
     <i class="fas fa-fw fa-undo"></i>
     Restart
   </button>
@@ -21,14 +22,8 @@
 <!-- Drop Zone -->
 <div *ngIf="tfFilesCompleted && !eyeballing && !eyeballCompleted">
   <mat-card>
-
-    <button mat-raised-button (click)="startEyeball()">
-      Start
-    </button>
-
-    <br />
-    <br />
-
+    Drop your screenshots below:
+    <br><br>
     <ngx-dropzone (change)="onSelect($event)">
       <ngx-dropzone-label>
         Files dropped here are <span style="font-weight: bold">not uploaded anywhere.</span> No data leaves your computer at all on this tool.
@@ -39,15 +34,15 @@
   <br />
 
   <mat-card><span style="font-weight: bold">Note:</span> For best results,
-    ensure your screenshots are taken natively in a <span style="font-weight: bold">16:10 aspect ratio.</span> (ie: 1440x900 or other)</mat-card>
+    ensure your screenshots are taken natively in a <span style="font-weight: bold">16:10 aspect ratio.</span> (ie: 1440x900 or other)
+  </mat-card>
+  <mat-card>
+    Source code available at <a href="https://github.com/BishopFox/eyeballer" target="_blank" rel="noopener noreferrer">GitHub</a>
+  </mat-card>
 </div>
 
 <!-- Progress Bar -->
 <div *ngIf="tfFilesCompleted && eyeballing && !eyeballCompleted">
-  <mat-card>
-    <h3>Loading Images: {{this.images.size}} of {{this.imageCount}}</h3>
-    <mat-progress-bar mode="determinate" [value]="loadPercent()"></mat-progress-bar>
-  </mat-card>
   <mat-card>
     <h3>Classifying: {{this.eyeballedCount}} of {{this.imageCount}}</h3>
     <mat-progress-bar mode="determinate" [value]="eyeballPercent()"></mat-progress-bar>
@@ -61,65 +56,59 @@
 <div *ngIf="tfFilesCompleted && !eyeballing && eyeballCompleted">
   <div class="results-container">
     <h2>Results</h2>
-    <div>
-      <div>
-        <button class="include-button" mat-raised-button [disabled]="selected.webapp === 0" (click)="selected.webapp = 0; updateSelections()">Must Include</button>
-        <button mat-raised-button [disabled]="selected.webapp === 1" (click)="selected.webapp = 1; updateSelections()">Webapp</button>
-        <button class="reject-button" mat-raised-button [disabled]="selected.webapp === 2" (click)="selected.webapp = 2; updateSelections()">Never Include</button>
+    <mat-toolbar>
+      <mat-toolbar-row style="height:40px;">
+        <span class="button-spacer"></span>
+        <button class="label-button" mat-raised-button color="primary" [disabled]="selected.webapp === 0" (click)="selected.webapp = 0; updateSelections()">Must Include</button>
+        <button class="label-button" mat-raised-button color="basic" [disabled]="selected.webapp === 1" (click)="selected.webapp = 1; updateSelections()">Webapp</button>
+        <button class="label-button" mat-raised-button color="accent" [disabled]="selected.webapp === 2" (click)="selected.webapp = 2; updateSelections()">Never Include</button>
         <label class="label-count">{{ classifications.webapp.length }}</label>
-      </div>
-      <div>
-        <button class="include-button" mat-raised-button [disabled]="selected.oldLooking === 0" (click)="selected.oldLooking = 0; updateSelections()">Must Include</button>
-        <button mat-raised-button [disabled]="selected.oldLooking === 1" (click)="selected.oldLooking = 1; updateSelections()">Old Looking</button>
-        <button class="reject-button" mat-raised-button [disabled]="selected.oldLooking === 2" (click)="selected.oldLooking = 2; updateSelections()">Never Include</button>
+        <span class="button-spacer"></span>
+      </mat-toolbar-row>
+      <mat-toolbar-row style="height:40px;">
+        <span class="button-spacer"></span>
+        <button class="label-button" mat-raised-button color="primary" [disabled]="selected.oldLooking === 0" (click)="selected.oldLooking = 0; updateSelections()">Must Include</button>
+        <button class="label-button" mat-raised-button color="basic" [disabled]="selected.oldLooking === 1" (click)="selected.oldLooking = 1; updateSelections()">Old Looking</button>
+        <button class="label-button" mat-raised-button color="accent" [disabled]="selected.oldLooking === 2" (click)="selected.oldLooking = 2; updateSelections()">Never Include</button>
         <label class="label-count">{{ classifications.oldLooking.length }}</label>
-      </div>
-      <div>
-        <button class="include-button" mat-raised-button [disabled]="selected.loginPage === 0" (click)="selected.loginPage = 0; updateSelections()">Must Include</button>
-        <button mat-raised-button [disabled]="selected.loginPage === 1" (click)="selected.loginPage = 1; updateSelections()">Login Page</button>
-        <button class="reject-button" mat-raised-button [disabled]="selected.loginPage === 2" (click)="selected.loginPage = 2; updateSelections()">Never Include</button>
+        <span class="button-spacer"></span>
+      </mat-toolbar-row>
+      <mat-toolbar-row style="height:40px;">
+        <span class="button-spacer"></span>
+        <button class="label-button" mat-raised-button color="primary" mat-raised-button [disabled]="selected.loginPage === 0" (click)="selected.loginPage = 0; updateSelections()">Must Include</button>
+        <button class="label-button" mat-raised-button color="basic" [disabled]="selected.loginPage === 1" (click)="selected.loginPage = 1; updateSelections()">Login Page</button>
+        <button class="label-button" mat-raised-button color="accent" [disabled]="selected.loginPage === 2" (click)="selected.loginPage = 2; updateSelections()">Never Include</button>
         <label class="label-count">{{ classifications.loginPage.length }}</label>
-      </div>
-      <div>
-        <button class="include-button" mat-raised-button [disabled]="selected.custom404 === 0" (click)="selected.custom404 = 0; updateSelections()">Must Include</button>
-        <button mat-raised-button [disabled]="selected.custom404 === 1" (click)="selected.custom404 = 1; updateSelections()">Custom 404</button>
-        <button class="reject-button" mat-raised-button [disabled]="selected.custom404 === 2" (click)="selected.custom404 = 2; updateSelections()">Never Include</button>
+        <span class="button-spacer"></span>
+      </mat-toolbar-row>
+      <mat-toolbar-row style="height:40px;">
+        <span class="button-spacer"></span>
+        <button class="label-button" mat-raised-button color="primary" mat-raised-button [disabled]="selected.custom404 === 0" (click)="selected.custom404 = 0; updateSelections()">Must Include</button>
+        <button class="label-button" mat-raised-button color="basic" [disabled]="selected.custom404 === 1" (click)="selected.custom404 = 1; updateSelections()">Custom 404</button>
+        <button class="label-button" mat-raised-button color="accent" [disabled]="selected.custom404 === 2" (click)="selected.custom404 = 2; updateSelections()">Never Include</button>
         <label class="label-count">{{ classifications.custom404.length }}</label>
-      </div>
-      <div>
-        <button class="include-button" mat-raised-button [disabled]="selected.parked === 0" (click)="selected.parked = 0; updateSelections()">Must Include</button>
-        <button mat-raised-button [disabled]="selected.parked === 1" (click)="selected.parked = 1; updateSelections()">Parked Domain</button>
-        <button class="reject-button" mat-raised-button [disabled]="selected.parked === 2" (click)="selected.parked = 2; updateSelections()">Never Include</button>
+        <span class="button-spacer"></span>
+      </mat-toolbar-row>
+      <mat-toolbar-row style="height:40px;">
+        <span class="button-spacer"></span>
+        <button class="label-button" mat-raised-button color="primary" mat-raised-button [disabled]="selected.parked === 0" (click)="selected.parked = 0; updateSelections()">Must Include</button>
+        <button class="label-button" mat-raised-button color="basic" [disabled]="selected.parked === 1" (click)="selected.parked = 1; updateSelections()">Parked Domain</button>
+        <button class="label-button" mat-raised-button color="accent" [disabled]="selected.parked === 2" (click)="selected.parked = 2; updateSelections()">Never Include</button>
         <label class="label-count">{{ classifications.parked.length }}</label>
-      </div>
-    </div>
+        <span class="button-spacer"></span>
+      </mat-toolbar-row>
+    </mat-toolbar>
 
-    <button mat-raised-button (click)="exportResults()">
-      Export
-    </button>
-    Total Selected: {{ selectedScreens.length }}
-
-    <div>
-      Preview:
-    </div>
-
-    <div class="eyeballed-class-container">
-      <div *ngFor="let key of selectedScreens.slice(0, 8); let index = index" [attr.data-index]="index">
-        <mat-card class="result-card">
-          <mat-card-header>
-            <mat-card-title>
-                {{ key }}
-            </mat-card-title>
-          </mat-card-header>
-          <img class="preview-image" style="width:100%;" mat-card-image src="{{images.get(key)}}">
-          <span class="webapp-label">{{ isWebapp(key) }}</span>
-          <span class="custom-404-label">{{ isCustom404(key) }}</span>
-          <span class="login-label">{{ isLoginPage(key) }}</span>
-          <span class="old-looking-label">{{ isOldLooking(key) }}</span>
-          <span class="parked-label">{{ isParked(key) }}</span>
-        </mat-card>
-      </div>
-    </div>
+    <mat-card>
+      <span style="font-weight: bold">Total Selected:</span> {{ selectedScreens.length }}
+    </mat-card>
+    <br>
+    <mat-card>
+      <button class="export-button" mat-raised-button color="primary" (click)="exportResults()">
+        Export
+      </button>
+      <a href="https://github.com/BishopFox/eyeballer#what-the-labels-mean" target="_blank" rel="noopener noreferrer">What do these labels mean?</a>
+    </mat-card>
   </div>
 
 </div>

webui/src/app/components/eyeballer/eyeballer.component.scss

@@ -49,10 +49,6 @@
     opacity: .5;
 }
 
-// .result-card:hover div {
-//     display: block;
-// }
-
 .result-card > div {
     display: none;
     position: absolute;
@@ -71,26 +67,6 @@
     flex: auto;
 }
 
-.label-count {
-    margin-left: 10px;
-}
-
-.include-button {
-  background-color: #3fa600;
-}
-
-.include-button:disabled {
-  background-color: #006802;
-}
-
-.reject-button {
-  background-color: red;
-}
-
-.reject-button:disabled {
-  background-color: #700000;
-}
-
 .img-container {
   position: relative;
   text-align: center;
@@ -131,3 +107,22 @@
   margin-left: 1px;
   margin-right: 1px;
 }
+
+.label-count {
+  margin-left: 10px;
+  width: 50px!important;
+  min-width: unset!important;
+}
+
+.label-button{
+  width: 130px!important;
+  min-width: unset!important;
+}
+
+.button-spacer {
+  flex: 1 1 auto;
+}
+
+.export-button {
+  margin-right: 10px;
+}

webui/src/app/components/eyeballer/eyeballer.component.ts

@@ -11,17 +11,16 @@ import * as tf from '@tensorflow/tfjs';
 export class EyeballerComponent implements OnInit {
 
   offset = tf.scalar(127.5);
-  images = new Map<string, string>();
-  confidence = 0.6;
+  confidence = 0.5;
 
   width = 224;
   height = 224;
 
   tfFilesCompleted = false;
   tfFiles: File[] = [];
+  model: any = null;
 
   imageCount = 0;
-  loadedCount = 0;
   finishedLoading = false;
   eyeballing = false;
   eyeballCompleted = false;
@@ -56,18 +55,63 @@ export class EyeballerComponent implements OnInit {
   }
 
   async onSelect(event) {
-    console.log(event);
     this.imageCount = event.addedFiles.length;
-    await Promise.all(event.addedFiles.map(async (file) => {
-      const dataString = await this.dataURI(file);
-      this.images.set(file.name, dataString);
-    }));
+    this.eyeballing = true;
+    await this.classifyImages(event.addedFiles);
+    this.eyeballing = false;
+    this.eyeballCompleted = true;
+    await this.updateSelections();
   }
 
-  onRemove(event) {
-    this.images.delete(event.name);
+  async classifyImages(files) {
+    this.model = await tf.loadLayersModel(tf.io.browserFiles(this.tfFiles));
+    for(const file of files) {
+      await this.classify(file);
+      this.eyeballedCount++;
+    }
   }
 
+  async classify(file) {
+    const dataString = await this.dataURI(file);
+    const img = new Image(this.width, this.height);
+    img.src = dataString;
+    img.onerror = () => {
+      return;
+    }
+    img.onload = () => {
+      tf.tidy(() => {
+        const tensor = tf.browser.fromPixels(img)
+          .resizeNearestNeighbor([224, 224])
+          .toFloat()
+          .sub(this.offset)
+          .div(this.offset)
+          .expandDims();
+        const resultTensor = <tf.Tensor<tf.Rank>> this.model.predict(tensor);
+        const predictions = resultTensor.dataSync();
+        tensor.dispose();
+        resultTensor.dispose();
+        if (predictions[0] > this.confidence) {
+          this.classifications.custom404.push(file.name);
+        }
+        if (predictions[1] > this.confidence) {
+          this.classifications.loginPage.push(file.name);
+        }
+        if (predictions[2] > this.confidence) {
+          this.classifications.webapp.push(file.name);
+        }
+        if (predictions[3] > this.confidence) {
+          this.classifications.oldLooking.push(file.name);
+        }
+        if (predictions[4] > this.confidence) {
+          this.classifications.parked.push(file.name);
+        }
+      });
+    };
+
+  }
+
+
+
   isWebapp(key) {
     if(this.classifications.webapp.includes(key)) {
       return "Webapp"
@@ -122,12 +166,6 @@ export class EyeballerComponent implements OnInit {
     return new File([blob], base);
   }
 
-  async startEyeball() {
-    await this.eyeballScan();
-    await this.updateSelections();
-    console.log(this.classifications);
-  }
-
   async updateSelections() {
     // First thing: load up all the images into the set
     let selectedScreensSet = new Set([]);
@@ -200,67 +238,10 @@ export class EyeballerComponent implements OnInit {
     this.selectedScreens = Array.from(selectedScreensSet);
   }
 
-  async eyeballScan(): Promise<void> {
-    console.log('eyeballing ...');
-    this.eyeballing = true;
-    const model = await tf.loadLayersModel(tf.io.browserFiles(this.tfFiles));
-    const keys = Array.from(this.images.keys());
-    await Promise.all(keys.map(async (key) => {
-      await this.classifyImage(key, model);
-    }));
-  }
-
-  async classifyImage(key: string, model: tf.LayersModel) {
-    const img = new Image(this.width, this.height);
-    img.src = this.images.get(key);
-    console.log("Queued up image: " + key)
-    this.loadedCount++;
-    img.onload = () => {
-      console.log(`classifying: ${key}`);
-      const tensor = tf.browser.fromPixels(img)
-        .resizeNearestNeighbor([224, 224])
-        .toFloat()
-        .sub(this.offset)
-        .div(this.offset)
-        .expandDims();
-      const predictions = (<tf.Tensor<tf.Rank>> model.predict(tensor)).dataSync();
-      console.log(`${predictions}`);
-      if (predictions[0] > this.confidence) {
-        console.log(`Custom 404: ${key}`);
-        this.classifications.custom404.push(key);
-      }
-      if (predictions[1] > this.confidence) {
-        console.log(`Login Page: ${key}`);
-        this.classifications.loginPage.push(key);
-      }
-      if (predictions[2] > this.confidence) {
-        console.log(`webapp: ${key}`);
-        this.classifications.webapp.push(key);
-      }
-      if (predictions[3] > this.confidence) {
-        console.log(`Old Looking: ${key}`);
-        this.classifications.oldLooking.push(key);
-      }
-      if (predictions[4] > this.confidence) {
-        console.log(`Parked: ${key}`);
-        this.classifications.parked.push(key);
-      }
-      this.eyeballedCount++;
-
-      if (this.eyeballedCount >= this.imageCount) {
-        console.log('eyeballed all images');
-        this.eyeballing = false;
-        this.eyeballCompleted = true;
-        this.updateSelections();
-      }
-    };
-  }
-
   async dataURI(file: File): Promise<string> {
     const buf = await file.arrayBuffer();
     let ext = file.name.split('.').reverse()[0]?.toLocaleLowerCase();
     if (!["jpg", "jpeg", "png", "gif", "bmp"].some(allow => allow === ext)) {
-      console.log(`Unknown file type ${ext}, defaulting to jpg`);
       ext = "jpg";
     }
     return `data:image/${encodeURIComponent(ext)};base64,${base64.encode(buf)}`;
@@ -270,10 +251,6 @@ export class EyeballerComponent implements OnInit {
     return (this.eyeballedCount / this.imageCount) * 100;
   }
 
-  loadPercent() {
-    return (this.images.size / this.imageCount) * 100;
-  }
-
   restart() {
     window.location.reload();
   }