From a9fe73a4839a03ad9675ec54166e252d2bb5a4f9 Mon Sep 17 00:00:00 2001
From: evazion <noizave@gmail.com>
Date: Tue, 28 Jun 2022 02:07:40 -0500
Subject: [PATCH] ai tags: save ai tags on upload.

Save the AI tags when a media asset is uploaded.
---
 Gemfile                   |  1 +
 Gemfile.lock              |  3 +++
 app/logical/media_file.rb | 13 +++++++++++++
 app/models/ai_tag.rb      |  4 +++-
 app/models/media_asset.rb |  9 +++++++++
 5 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/Gemfile b/Gemfile
index b431fb297ba..5eba396d7fe 100644
--- a/Gemfile
+++ b/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
diff --git a/Gemfile.lock b/Gemfile.lock
index c6085151c50..2c49c6860f0 100644
--- a/Gemfile.lock
+++ b/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
diff --git a/app/logical/media_file.rb b/app/logical/media_file.rb
index 5ac19db150b..9d49f058b52 100644
--- a/app/logical/media_file.rb
+++ b/app/logical/media_file.rb
@@ -189,6 +189,19 @@ def preview(width, height, **options)
     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,
diff --git a/app/models/ai_tag.rb b/app/models/ai_tag.rb
index f88562648c3..448c701d3f1 100644
--- a/app/models/ai_tag.rb
+++ b/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
diff --git a/app/models/media_asset.rb b/app/models/media_asset.rb
index d27ea81e942..a683d061b20 100644
--- a/app/models/media_asset.rb
+++ b/app/models/media_asset.rb
@@ -27,6 +27,7 @@ class Error < StandardError; end
 
   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 @@ def file=(file_or_path)
       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!