Create app.py

app.py

@@ -0,0 +1,83 @@
+from transformers import CLIPProcessor, CLIPModel
+import numpy as np
+from typing import Union
+import PIL
+
+class ZeroShotCLF():
+  def __init__(self, model_id = "openai/clip-vit-base-patch32"):
+    """
+    Clasificador de imágenes zero-shot usando CLIP.
+
+    Métodos:
+      - set_classes: sirve para especificar las posibles clases. Se debe usar antes de predict o evaluate.
+      - predict: sirve para clasificar una o varias imágenes. Se debe usar antes set_classes.
+      - evaluate: sirve para evaluar el modelo. Se debe usar antes set_classes.
+    """
+    self.model_id = model_id
+    self.processor = CLIPProcessor.from_pretrained(model_id)
+    self.model = CLIPModel.from_pretrained(model_id)
+
+  def set_classes(self, clases: list[str], tipos: Union[str, list[str]] = "image"):
+    """
+    Establece las posibles clases en las que deberemos clasificar.
+
+    Inputs:
+      - clases: lista de strings con todas las clases a considerar.
+      - tipos: string o lista de strings con los tipos de imagen. Este tipo sirve para especificar si buscas fotografías, dibujos... Valor por defecto: "image".
+      Si la variable es string, se utilizará el mismo tipo para todas las clases, si es una lista, deberá tener la misma longitud que la variable clases.
+
+    Outputs:
+      - Lista de textos a usar para cada clase. Se usará el siguiente esquema: "a __ of a __", donde en el primer hueco se rellenará con el tipo y el segundo por la clase.
+      Por ejemplo, si el tipo es "photo" y la clase es "dog", el texto será "a photo of a dog".
+    """
+    if isinstance(tipos, str):
+      tipos = [tipos] * len(clases)
+    else:
+      assert len(tipos) == len(clases), "¡clases y tipos deben tener la misma longitud!"
+
+    self.clases = clases
+    self.tipos = tipos
+    self.text = list(map(lambda par: "a %s of a %s" % par, zip(tipos, clases)))
+    return self.text
+
+  def predict(self, images):
+    """
+    Predice las clases de las imágenes. Es necesario usar previamente el método set_classes.
+
+    Inputs:
+      - images: imagen PIL o lista de imágenes PIL a clasificar.
+
+    Outputs:
+      - Nombre de la clase predicha (string) o lista de las clases predichas (lista de strings).
+      - Índice de la clase predicha (int) o lista de los índices de las clases predichas (lista de ints).
+      - Vector de probabilidades de las clases (array) o matriz de probabilidades de las clases (ndarray de dimensiones número de imágenes por número de clases).
+    """
+    assert 'text' in dir(self), "¡No se han especificado las clases! Usa set_classes antes que predict."
+
+    inputs = self.processor(text=self.text, images=images, return_tensors="pt", padding=True)
+    outputs = self.model(**inputs)
+
+    probs = outputs.logits_per_image.softmax(dim=1)
+
+    ind_clases = probs.argmax(1)
+
+    if isinstance(images, list):
+      ind_clases = ind_clases.tolist()
+      pred_clases = list(map(lambda ind: self.clases[ind], ind_clases))
+      return pred_clases, ind_clases, probs
+
+    ind_clases = int(ind_clases)
+    return self.clases[ind_clases], ind_clases, probs.ravel()
+
+clf = ZeroShotCLF()
+
+def clasifica(imagen, clases_, tipos_):
+    clases = clases_.split(", ")
+    tipos = tipos_.split(", ")
+    if len(tipos) == 1:
+        tipos = tipos[0]
+    textos = clf.set_classes(clases, tipos)
+    probs = clf.predict(PIL.Image.fromarray(imagen))[-1][0]
+    return {textos[i]: probs[i] for i in range(len(textos))}
+
+gr.Interface(fn=clasifica, inputs=["image", "text", "text"], outputs="text").launch(share=False)