Merge pull request #140 from biigle/patch-1

Update initialize feature vector command to process all volumes
biigle · Feb 2, 2024 · ec1a083 · ec1a083
2 parents 7d1a72a + 66c779c
commit ec1a083
Showing 1 changed file with 37 additions and 33 deletions.
diff --git a/src/Console/Commands/InitializeFeatureVectors.php b/src/Console/Commands/InitializeFeatureVectors.php
@@ -3,6 +3,7 @@
 namespace Biigle\Modules\Largo\Console\Commands;
 
 use Biigle\ImageAnnotation;
+use Biigle\MediaType;
 use Biigle\Modules\Largo\Jobs\InitializeFeatureVectorChunk;
 use Biigle\VideoAnnotation;
 use Biigle\Volume;
@@ -17,8 +18,8 @@ class InitializeFeatureVectors extends Command
      * @var string
      */
     protected $signature = 'largo:initialize-feature-vectors
-        {volume : ID of the volume of which the annotations should be processed}
         {--dry-run : Do not submit jobs to generate feature vectors}
+        {--volume= : Process only annotations of a single volume (ID)}
         {--queue=default : Queue name to push jobs to generate feature vectors to}
         {--chunk-size=1000 : Number of annotations to process in a single job}';
 
@@ -36,65 +37,68 @@ class InitializeFeatureVectors extends Command
      */
     public function handle()
     {
-        $volume = Volume::findOrFail($this->argument('volume'));
-        if ($volume->isImageVolume()) {
+        $query = Volume::select('id')
+            ->where('media_type_id', MediaType::imageId())
+            ->when($this->option('volume'), fn ($q) => $q->where('id', $this->option('volume')));
+
+        $count = $query->clone()->count();
+        $progress = $this->output->createProgressBar($count);
+        $this->info("Processing {$count} image volumes.");
+        $query->lazyById()->each(function ($volume) use ($progress) {
             $this->processImages($volume);
-        } else {
+            $progress->advance();
+        });
+        $progress->finish();
+        $this->line('');
+
+        $query = Volume::select('id')
+            ->where('media_type_id', MediaType::videoId())
+            ->when($this->option('volume'), fn ($q) => $q->where('id', $this->option('volume')));
+
+        $count = $query->clone()->count();
+        $progress = $this->output->createProgressBar($count);
+        $this->info("Processing {$count} video volumes.");
+        $query->lazyById()->each(function ($volume) use ($progress) {
             $this->processVideos($volume);
-        }
+            $progress->advance();
+        });
+        $progress->finish();
+        $this->line('');
     }
 
     protected function processImages(Volume $volume)
     {
         $chunkSize = intval($this->option('chunk-size'));
         $loopChunkSize = max($chunkSize, 10000);
 
-        $query = ImageAnnotation::join('images', 'images.id', '=', 'image_annotations.image_id')
-            ->where('images.volume_id', $volume->id);
-
-        $count = $query->count();
-        $this->info("Processing {$count} image annotations.");
-        $p = $this->output->createProgressBar($count);
-
-        $query->select('image_annotations.id')
-            ->chunkById($loopChunkSize, function ($chunk) use ($p, $chunkSize) {
-                $chunk->chunk($chunkSize)->each(function ($c) use ($p) {
+        ImageAnnotation::join('images', 'images.id', '=', 'image_annotations.image_id')
+            ->where('images.volume_id', $volume->id)
+            ->select('image_annotations.id')
+            ->chunkById($loopChunkSize, function ($chunk) use ($chunkSize) {
+                $chunk->chunk($chunkSize)->each(function ($c) {
                     $job = new InitializeFeatureVectorChunk($c->pluck('id')->all(), []);
                     if (!$this->option('dry-run')) {
                         Queue::pushOn($this->option('queue'), $job);
                     }
-                    $p->advance($c->count());
                 });
             });
-
-        $p->finish();
-        $this->line('');
     }
 
     protected function processVideos(Volume $volume)
     {
         $chunkSize = intval($this->option('chunk-size'));
         $loopChunkSize = max($chunkSize, 10000);
 
-        $query = VideoAnnotation::join('videos', 'videos.id', '=', 'video_annotations.video_id')
-            ->where('videos.volume_id', $volume->id);
-
-        $count = $query->count();
-        $this->info("Processing {$count} video annotations.");
-        $p = $this->output->createProgressBar($count);
-
-        $query->select('video_annotations.id')
-            ->chunkById($loopChunkSize, function ($chunk) use ($p, $chunkSize) {
-                $chunk->chunk($chunkSize)->each(function ($c) use ($p) {
+        VideoAnnotation::join('videos', 'videos.id', '=', 'video_annotations.video_id')
+            ->where('videos.volume_id', $volume->id)
+            ->select('video_annotations.id')
+            ->chunkById($loopChunkSize, function ($chunk) use ($chunkSize) {
+                $chunk->chunk($chunkSize)->each(function ($c) {
                     $job = new InitializeFeatureVectorChunk([], $c->pluck('id')->all());
                     if (!$this->option('dry-run')) {
                         Queue::pushOn($this->option('queue'), $job);
                     }
-                    $p->advance($c->count());
                 });
             });
-
-        $p->finish();
-        $this->line('');
     }
 }