ai tags: save ai tags on upload.

Save the AI tags when a media asset is uploaded.

Gemfile

@@ -56,6 +56,7 @@ gem "public_suffix"
 gem "elastic-apm"
 gem "debug"
 gem "ffaker"
+gem "composite_primary_keys"
 
 group :development do
   gem 'rubocop', require: false

Gemfile.lock

@@ -133,6 +133,8 @@ GEM
     codecov (0.6.0)
       simplecov (>= 0.15, < 0.22)
     coderay (1.1.3)
+    composite_primary_keys (14.0.4)
+      activerecord (~> 7.0.2)
     concurrent-ruby (1.1.10)
     crass (1.0.6)
     daemons (1.4.1)
@@ -528,6 +530,7 @@ DEPENDENCIES
   capybara
   clockwork
   codecov
+  composite_primary_keys
   crass
   daemons
   debug

app/logical/media_file.rb

@@ -189,6 +189,19 @@ class MediaFile
     nil
   end
 
+  # Return a set of AI-inferred tags for this image. Performs an API call to
+  # the Autotagger service. The Autotagger service must be running, otherwise
+  # it will return an empty list of tags.
+  #
+  # @return [Array<AITag>] The list of AI tags.
+  def ai_tags(autotagger: AutotaggerClient.new)
+    tags = autotagger.evaluate(self)
+
+    tags.map do |tag, score|
+      AITag.new(tag: tag, score: (100*score).round)
+    end
+  end
+
   def attributes
     {
       path: path,

app/models/ai_tag.rb

@@ -1,11 +1,13 @@
 # frozen_string_literal: true
 
 class AITag < ApplicationRecord
+  self.primary_keys = :media_asset_id, :tag_id
+
   belongs_to :tag
   belongs_to :media_asset
   has_one :post, through: :media_asset
 
-  validates :score, inclusion: { in: (0.0..1.0) }
+  validates :score, inclusion: { in: (0..100) }
 
   def self.named(name)
     name = $1.downcase if name =~ /\A(rating:.)/i

app/models/media_asset.rb

@@ -27,6 +27,7 @@ class MediaAsset < ApplicationRecord
 
   scope :public_only, -> { where(is_public: true) }
   scope :private_only, -> { where(is_public: false) }
+  scope :without_ai_tags, -> { where.not(AITag.where("ai_tags.media_asset_id = media_assets.id").select(1).arel.exists) }
 
   # Processing: The asset's files are currently being resized and distributed to the backend servers.
   # Active: The asset has been successfully uploaded and is ready to use.
@@ -279,6 +280,14 @@ class MediaAsset < ApplicationRecord
       self.duration = media_file.duration
       self.media_metadata = MediaMetadata.new(file: media_file)
       self.pixiv_ugoira_frame_data = PixivUgoiraFrameData.new(data: media_file.frame_data, content_type: "image/jpeg") if is_ugoira?
+      self.ai_tags = media_file.preview(360, 360).ai_tags # XXX should do this in parallel with thumbnail generation.
+    end
+
+    def regenerate_ai_tags!
+      with_lock do
+        ai_tags.each(&:destroy!)
+        update!(ai_tags: variant(:"360x360").open_file.ai_tags)
+      end
     end
 
     def expunge!