FindMyImage.py

import torch
import gc
import sys, os
import getopt
import subprocess
import pathlib
from PIL import Image

#I have no idea what I'm doing
#this is probably a terrible idea
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)),'MetaCLIP'))

from MetaCLIP.src.mini_clip.factory import create_model_and_transforms
from MetaCLIP.src.mini_clip.tokenizer import tokenize
from os import listdir
from os.path import isfile, isdir, join, splitext
from timeit import default_timer as timer

#ginormous images break things
#if you get that "bomb" exception, turn on the print
#that goes file by file and increase this to cover it
Image.MAX_IMAGE_PIXELS = 220498432

validExtensions = {
    #might want .tga and .tif but I had some it didn't like
    #palettized stuff seems to be trouble
    ".jpg", ".jpeg", ".png", ".gif", ".JPG", ".webp"
}

def CheckExtension(fileExt:str):
    for ve in validExtensions:
        if(ve == fileExt):
            return True
    #print(f"Skipping unhandled extension: " + fileExt)
    return False


def SortByConfidence(result):
    #confidence is in index 1
    return result[1][1]


#given a bunch of phrases or search terms or whatsits, see which
#one the image is most like
def ClassifyImage(filePath:str, searchTexts:list, myIndexes:list, text_features:list, results:dict):
    #check the extension
    extension = splitext(filePath)
    
    valid = CheckExtension(extension[1])
    if not valid:
        return
    
    #if you crash out with a bad file turn this on
    #so you can tell which file is causing it
    if bShowFiles:
        print(filePath)
    
    #if a file, check it
    image = preprocess(Image.open(filePath)).unsqueeze(0)
    
    # Calculate similarities between an image and multiple texts
    image_features = model.encode_image(image)
    image_features /= image_features.detach().norm(dim=-1, keepdim=True)
    
    similarities = (100.0 * image_features @ text_features.T).softmax(dim=-1)
    
    top_probs, top_indices = similarities.topk(1, dim=-1)
    
    #is the top score one of our items or one of the generics?
    if(top_indices[0] in myIndexes):
        predicted = searchTexts[top_indices[0].item()]
        confidence = top_probs[0].item()
        
        #keep track of which text and score
        fileNTerm = [ top_indices[0].item(), confidence ]
        results[filePath] = fileNTerm
        if(bShowHits):
            print(f"Predicted: {predicted}, Confidence: {confidence:.2f} {filePath}")


#classify all the images in a directory, with optional recursion
def ClassifyDirectory(dirPath:str, bRecurse:bool, searchTexts:list, myIndexes:list, text_features:list, results:dict):
    #grab a list of files and dirs
    listing = listdir(dirPath)
    
    for f in listing:
        filePath = join(dirPath, f)
        #if a file checkit
        if isfile(filePath):
            ClassifyImage(filePath, searchTexts, myIndexes, text_features, results)
        elif isdir(filePath):
            if(bRecurse):
                ClassifyDirectory(filePath, bRecurse, searchTexts, myIndexes, text_features, results)


#default options
bRecurse:bool = False   #recurse?
searchTerms = []        #user specified search words or phrases
searchArgs:str = ''     #searchstuff from command line as a string
topX:int = 5            #top X results
bDisplay:bool = False   #use OS image viewer?
bShowFiles:bool = False #spam file paths
bShowHits:bool = False  #spam all hits
bColors:bool = True     #show color on print result

#default to home dir pictures
#todo test this on windows?
dirPath:str = pathlib.Path.home().as_posix() + "/Pictures"

#check command line args
numArgs = len(sys.argv)

bValidArgs:bool = True
for i in range(numArgs):
    arg = sys.argv[i]
    if arg == '--recurse':
        bRecurse = True
    elif arg == '--dirpath':
        dirPath = sys.argv[i + 1]
    elif arg == '--searchterms':
        searchArgs = sys.argv[i + 1]
    elif arg == '--top':
        topX = int(sys.argv[i + 1])
    elif arg == '--display':
        bDisplay = True
    elif arg == '--showfiles':
        bShowFiles = True
    elif arg == '--showhits':
        bShowHits = True
    elif arg == '--nocolors':
        bColors = False
    elif arg == '--nocolours':
        bColors = False
    elif arg.startswith('-'):
        bValidArgs = False
        

if not bValidArgs:
    print("Usage:")
    print("--dirpath      The directory to search for images within")
    print("--recurse      Delve into subdirectories of dirpath")
    print("--top X        Find the top X matches if any (default is 5)")
    print("--searchterms  A comma seperated list of search words or phrases")
    print("--display      Kick off os image viewer for top 5")
    print("--showfiles    List each file as it is classified (spammy)")
    print("--showhits     Show every time an image scores (spammy)")
    print("--nocolors     Print in plain terminal color")
    sys.exit()

#grab search terms from command line if any
searchTerms = searchArgs.split(',')
searchTerms = list(map(str.strip, searchTerms))


# Load model
print("Loading model...")
model, _, preprocess = create_model_and_transforms('ViT-B-32-quickgelu', pretrained='metaclip_400m')

print(f"Searching {dirPath} with recursion {bRecurse}")
startTime = timer()

#a bunch of common generic stuff to measure against
#this is a bit of a hack, and probably won't cover
#every case.  Should probably stick this in a csv or json or something
commonTexts = ["animal", "building", "bread", "pizza", "mountain", "beach",
         "road", "tree", "flower", "car", "text", "tweet", "furniture", "fire",
         "bowl", "chopsticks", "captions", "weapon"]

#combine command line texts with generic
combinedTexts = list(set(commonTexts + searchTerms))

#keep track of the indices of the command line stuff
myIndexes = []

for txt in combinedTexts:
    if txt in searchTerms:
        myIndexes.append(combinedTexts.index(txt))

text_tokens = tokenize(combinedTexts)
text_features = model.encode_text(text_tokens)
text_features /= text_features.norm(dim=-1, keepdim=True)

#keep track of hits
results = {}

ClassifyDirectory(dirPath, bRecurse, combinedTexts, myIndexes, text_features, results)

#show the top X results
top = dict(sorted(results.items(), key = SortByConfidence, reverse=True)[:topX])

normalCode = "\033[0m"

if bColors:
    for k, v in top.items():
        colorCode = "\033[91m"
        if v[1] > 0.9:
            colorCode = "\033[92m"
        elif v[1] > 0.8:
            colorCode = "\033[94m"
        elif v[1] > 0.7:
            colorCode = "\033[93m"
        print(f"{colorCode}{combinedTexts[v[0]]} : {v[1] * 100:.0f}%{normalCode} : {k}")
else:
    for k, v in top.items():
        print(f"{combinedTexts[v[0]]} : {v[1] * 100:.0f}% : {k}")

endTime = timer()
dt = endTime - startTime
print(f"Finished in {dt:.2f} seconds.")

#reverse order so the #1 appears on top
rev = reversed(top.items())

if bDisplay:
    imageViewer = {'linux':'xdg-open',
                   'win32':'explorer',
                   'darwin':'open'}[sys.platform]
    for k, v in rev:
        #handle spaces and such
        ppath = pathlib.PosixPath(k)    
        subprocess.run([imageViewer, ppath.as_uri()])