Chanlefe commited on
Commit
131383f
·
verified ·
1 Parent(s): 07560ee

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +17 -0
app.py CHANGED
@@ -2,6 +2,23 @@ import torch
2
  from PIL import Image
3
  from transformers import AutoProcessor, AutoModelForImageClassification
4
  import gradio as gr
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5
 
6
  # Load model and processor from Hugging Face
7
  model = AutoModelForImageClassification.from_pretrained("google/siglip2-base-patch16-naflex")
 
2
  from PIL import Image
3
  from transformers import AutoProcessor, AutoModelForImageClassification
4
  import gradio as gr
5
+ import pytesseract
6
+
7
+ def classify_meme(image: Image.Image):
8
+     # OCR: extract text from image
9
+     extracted_text = pytesseract.image_to_string(image)
10
+
11
+     # Process image with SigLIP2 model
12
+     inputs = processor(images=image, return_tensors="pt").to(model.device)
13
+     with torch.no_grad():
14
+         outputs = model(**inputs)
15
+         probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
16
+     predictions = {labels[i]: float(probs[0][i]) for i in range(len(labels))}
17
+
18
+     return {
19
+         "Predictions": predictions,
20
+         "Extracted Text": extracted_text.strip()
21
+     }
22
 
23
  # Load model and processor from Hugging Face
24
  model = AutoModelForImageClassification.from_pretrained("google/siglip2-base-patch16-naflex")