From a04e6e241b77732abcc62dd50ae0d1fe89abcaa3 Mon Sep 17 00:00:00 2001
From: Ryan Shpeherd <rotaryprophet@gmail.com>
Date: Sat, 26 Sep 2026 10:23:55 -0400
Subject: [PATCH] Updates to ai worker

---
 aiworker/run.pl | 32 +++++++++++++++++++++++++++-----
 1 file changed, 27 insertions(+), 5 deletions(-)

diff --git a/aiworker/run.pl b/aiworker/run.pl
index 2c0503f..c105fa9 100755
--- a/aiworker/run.pl
+++ b/aiworker/run.pl
@@ -9,15 +9,24 @@ my $seconds = $ENV{'SECONDS_PER_FRAME'} || 30;
 my $ollama = $ENV{'OLLAMA_URL'} || 'http://10.0.1.7:11434';
 my $api = $ENV{'API_BASE_URL'} || 'http://10.0.1.7:3000';
 my $task_count = $ENV{'TASK_COUNT'} || 1;
+my $image_count = $ENV{'IMAGE_COUNT'} || 10; # Number of images to extract from the video
+my $min_time = $ENV{'MIN_TIME'} || 5; # Minimum time between images
 my $sensitivity = 20;
 my $resolution = "672:672";
-my $model = "hf.co/Jommarn/UNSEEN_Gemma_4_26B_NSFW-GGUF:Q8_0";
+#my $model = "hf.co/Jommarn/UNSEEN_Gemma_4_26B_NSFW-GGUF:Q8_0";
+my $model = "hf.co/Jommarn/UNSEEN_Gemma_4_26B_NSFW-GGUF:Q5_K_M";
 
 chdir("/tmp") || die "Unable to change directory to /tmp: $!\n";
 
 my $ua = LWP::UserAgent->new;
 $ua->timeout(60);
 
+if(@ARGV) {
+    CheckPath($ARGV[0], {'duration' => 0});
+
+    exit;
+}
+
 for(my $i=0; $i<$task_count; $i++) {
     my $ref = GetNext();
     my $task = $ref->{'task'};
@@ -26,7 +35,7 @@ for(my $i=0; $i<$task_count; $i++) {
     my $task_id = $task->{'id'};
     my $video_id = $video->{'id'};
     my $path = $video->{'file_path'};
-    my $result = CheckPath($path);
+    my $result = CheckPath($path, $video);
 
     next unless $result;
     
@@ -37,12 +46,21 @@ print "Done\n";
 exit(0);
 
 sub CheckPath {
-    my ($path) = @_;
+    my ($path, $video) = @_;
     die "File path not set\n" unless $path;
     die "File does not exist: $path\n" unless -e $path;
 
     unlink(<*.png>);
-    system("ffmpeg -i \"$path\" -vf fps=1/$seconds,scale=$resolution output_%04d.png");
+
+    # Pick a frame interval that spreads $image_count frames across the video,
+    # but never more frequently than 1 frame per $min_time seconds.
+    my $interval = $seconds; # fallback if duration is unknown
+    if (my $duration = $video->{'duration'}) {
+        my $desired = $duration / $image_count;
+        $interval = ($desired > $min_time) ? $desired : $min_time;
+    }
+
+    system("ffmpeg -i \"$path\" -vf fps=1/$interval,scale=$resolution output_%04d.png 2>&1");
 
 
     my $total=0;
@@ -66,9 +84,13 @@ sub CheckPath {
         my $url = "$ollama/api/generate";
         my $res = GetRest($url, $ref);
         print "$file\t$res\n";
+        #Remove thought information from the response
+        $res =~ s/<\|channel>.*?<channel\|>//s;
+        $res =~ s/^\s+|\s+$//g;
         $total++;
-        if($res=~/yes$/i) {
+        if($res=~/yes/i) {
             $yes++;
+            print "\tYES\n";
         }
         unlink($file);
     }