diff --git a/src/Console/Commands/InitializeFeatureVectors.php b/src/Console/Commands/InitializeFeatureVectors.php index 4da2c3d..7d7e1a7 100644 --- 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,12 +37,33 @@ 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) @@ -49,26 +71,17 @@ 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) @@ -76,25 +89,16 @@ 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(''); } }