ads / app /services /ai_processor.py
rastof9's picture
Create services/ai_processor.py
07685f0 verified
raw
history blame contribute delete
878 Bytes
import cv2
import pytesseract
from transformers import pipeline
class AIPipeline:
def __init__(self):
self.nlp = pipeline("text-classification", model="roberta-base")
self.detector = cv2.dnn.readNet("yolov4.weights", "yolov4.cfg")
def process_ad(self, ad):
results = {
"sentiment": self.nlp(ad.content),
"ocr": self._extract_ocr(ad.media),
"objects": self._detect_objects(ad.media)
}
return results
def _extract_ocr(self, media):
if media and media.type == "image":
return pytesseract.image_to_string(media.path)
return None
def _detect_objects(self, media):
if media and media.type == "image":
img = cv2.imread(media.path)
self.detector.setInput(img)
return self.detector.forward()
return None