TYPE_NAME = "FEATURE_VECTOR" class Post(Base):
  from rep0st.db.feature import Feature
  from rep0st.db.tag import Tag

  __tablename__ = 'post'
  # Post id.
  id = Column(Integer, primary_key=True, index=True, autoincrement=False)
  # Timestamp this post was created.
  created = Column(DateTime(), nullable=False)
  # Path of the image on pr0gramm servers.
  image = Column(String(256), nullable=False, index=True)
  # Path of the thumbnail on pr0gramm servers.
  thumb = Column(String(256), nullable=False)
  # Path of the fullsize image on pr0gramm servers. Optional.
  fullsize = Column(String(256))
  # Width of the media in pixels.
  width = Column(Integer(), nullable=False)
  # Height of the media in pixels.
  height = Column(Integer(), nullable=False)
  # True if the media has audio.
  audio = Column(Boolean(), nullable=False)
  # URL of the source of the image. Optional.
  source = Column(String(512))
  # Flags for the post encoded as a bitset. If the bit is set,
  # the post is marked with the given tag.
  # Bit 0: SFW
  # Bit 1: NSFW
  # Bit 2: NSFL
  # Bit 3: NSFP
  flags = Column(Integer(), nullable=False)
  # Name of the user that uploaded the post.
  user = Column(String(32), nullable=False)
  # Type of the media in the post.
  # - IMAGE: Static images. (jpg, png)
  # - ANIMATED: Animated images. (gif)
  # - VIDEO: Videos. (mp4, webm)
  type = Column(Enum(Type), nullable=False, index=True)
  # Status of the post.
  status = Column(
      Enum(Status), nullable=False, index=True, default=Status.NO_MEDIA)
  # True if the post is deleted on pr0gramm.
  deleted = Column(Boolean(), nullable=False, default=False)
  # List of features associated with this post.
  features = relationship(Feature)
  # List of tags associated with this post.
  tags = relationship(Tag)

  def __json__(self):
    return {
        'id': self.id,
        'user': self.user,
        'created': self.created.isoformat(),
        'is_sfw': self.is_sfw(),
        'is_nsfw': self.is_nsfw(),
        'is_nsfl': self.is_nsfl(),
        'image': self.image,
        'thumb': self.thumb,
        'fullsize': self.fullsize,
    }

  def as_indexed_doc(self):

    def feauture_key_func(feature: Feature):
      return feature.id

    return {
        'meta': {
            'id': self.id
        },
        'created':
            int(self.created.timestamp() * 1000),
        'flags': [flag.value for flag in self.get_flags()],
        'type':
            self.type.value,
        'tags': [tag.tag for tag in self.tags],
        'frames': [{
            'id': key,
            'features': {v.type: v.data for v in valuesiter}
        } for key, valuesiter in groupby(
            sorted(self.features, key=feauture_key_func), key=feauture_key_func)
                  ],
    }

  def is_sfw(self):
    return self.flags & 1 != 0

  def is_nsfw(self):
    return self.flags & 2 != 0

  def is_nsfl(self):
    return self.flags & 4 != 0

  def is_nsfp(self):
    return self.flags & 8 != 0

  def get_flags(self) -> List[Flag]:
    flags = []
    if self.is_sfw():
      flags.append(Flag.SFW)
    if self.is_nsfw():
      flags.append(Flag.NSFW)
    if self.is_nsfl():
      flags.append(Flag.NSFL)
    if self.is_nsfp():
      flags.append(Flag.NSFP)
    return flags

  def get_flag_by_importance(self) -> Flag:
    if self.is_nsfl():
      return Flag.NSFL
    if self.is_nsfw():
      return Flag.NSFW
    if self.is_nsfp():
      return Flag.NSFP
    return Flag.SFW

  def __str__(self):
    return "Post(id=" + str(self.id) + ")"

  def __repr__(self):
    return "Post(id=" + str(self.id) + ")" class PostRepository(Repository[int, Post]):

  @inject
  def __init__(self, session_provider: ProviderOf[Session]) -> None:
    super().__init__(int, Post, session_provider)

  @transactional()
  def get_latest_post_id(self) -> int:
    session = self._get_session()
    id = session.query(func.max(Post.id)).scalar()
    return 0 if id is None else id

  @transactional()
  def get_posts(self, type: Optional[str] = None):
    session = self._get_session()
    if type is not None:
      return session.query(Post).filter(Post.type == type)
    else:
      return session.query(Post)

  @transactional()
  def get_posts_missing_features(self, type: Optional[Type] = None):
    session = self._get_session()
    q = session.query(Post)
    if type:
      q = q.filter(Post.type == type)
    return q.filter(
        and_(Post.status == Status.NOT_INDEXED, Post.deleted == False))

  @transactional()
  def post_count(self) -> int:
    session = self._get_session()
    return session.query(func.count(Post.id)).scalar() class PostRepositoryModule(Module):

  def configure(self, binder):
    binder.install(Rep0stDatabaseModule)
    binder.bind(PostRepository) class PostIndex(Index[Post]):

  @inject
  def __init__(self, elasticsearch: Elasticsearch):
    super().__init__(Post, elasticsearch=elasticsearch)

  def _index_post_from_post(self, post: DBPost) -> Post:
    index_post = Post()
    index_post.meta.id = post.id
    index_post.created = post.created
    index_post.type = post.type.value
    index_post.flags = [flag.value for flag in post.get_flags()]
    index_post.tags = [tag.tag for tag in post.tags]
    index_post.frames = [
        Frame(
            id=feature.id,
            feature_vector=[float(n / 255.0)
                            for n in feature.data])
        for feature in post.features
        if feature.type == FEATURE_VECTOR_TYPE
    ]
    return index_post

  def add_posts(self, posts: Iterable[Post]):

    def _it():
      for post in posts:
        yield self._index_post_from_post(post)

    self.save_all(_it())

  def find_posts(self, feature_vector):
    response = self.search().update_from_dict({
        'size': 50,
        'fields': [],
        '_source': False,
        'min_score': 0.3,
        'query': {
            'nested': {
                'path': 'frames',
                'query': {
                    'elastiknn_nearest_neighbors': {
                        'field': 'frames.feature_vector',
                        'vec': feature_vector,
                        'model': 'lsh',
                        'similarity': 'l2',
                        'candidates': 500
                    },
                },
            },
        },
    }).execute()

    for post in response:
      yield SearchResult(post.meta.score, post.meta.id) class PostIndexModule(Module):

  def configure(self, binder):
    binder.install(Rep0stElasticsearchModule)
    binder.bind(PostIndex) class AnalyzeService:
  _analyzers: List[Analyzer] = None

  @inject
  def __init__(self, analyzers: Analyzers) -> None:
    self._analyzers = set(analyzers)
    for analyzer in self._analyzers:
      log.info(f'Registered analyzer {analyzer.get_type()}')

  def analyze(self, image: numpy.ndarray):
    input_image = InputImage(image, cv2.cvtColor(image, cv2.COLOR_BGR2GRAY))

    result = {}
    for analyzer in self._analyzers:
      start = time.time()
      result[analyzer.get_type()] = analyzer.analyze(input_image)
      end = time.time()
      time_taken = end - start
      log.debug(
          f'Analyzed image {image.shape} with analyzer {analyzer.get_type()} in {time_taken * 1000:.2f}ms'
      )
    return result class AnalyzeServiceModule(Module):

  def configure(self, binder):
    binder.install(FeatureVectorAnalyzerModule)
    # TODO(rhollander): Enable other analyzers.
    # binder.install(AverageHashAnalyzerModule)
    # binder.install(DifferenceHashAnalyzerModule)
    # binder.install(PerceptualHashAnalyzerModule)
    # binder.install(WaveletHashAnalyzerModule)
    binder.bind(AnalyzeService) class DecodeMediaService:

  def _decode_image(self, data: numpy.ndarray) -> numpy.ndarray:
    try:
      img = imdecode(data, IMREAD_COLOR)
      if img is None:
        raise ImageDecodeException("Could not decode image")
      return img
    except:
      raise ImageDecodeException("Could not decode image")

  def decode_image_from_buffer(self, data: bytes) -> Iterable[numpy.ndarray]:
    try:
      data = numpy.frombuffer(data, dtype=numpy.uint8)
    except (IOError, OSError) as e:
      raise NoMediaFoundException('Could not data from buffer') from e
    yield self._decode_image(data)

  def decode_image_from_file(self, file: BinaryIO) -> Iterable[numpy.ndarray]:
    try:
      data = numpy.fromfile(file, dtype=numpy.uint8)
    except (IOError, OSError) as e:
      raise NoMediaFoundException(
          f'Could not read data from file {file}') from e
    yield self._decode_image(data)

  def decode_video_from_file(self, file: BinaryIO) -> Iterable[numpy.ndarray]:
    cmd = ffmpeg.input(
        'pipe:',
        vsync=0,
        skip_frame='nokey',
        hide_banner=None,
        threads=1,
        loglevel='error').output(
            'pipe:', vcodec='ppm', format='rawvideo')
    proc = subprocess.Popen(
        cmd.compile(),
        stdin=file,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE)

    while True:
      format = _readline(proc.stdout)
      if not format:
        break
      if format != 'P6':
        raise NotImplementedError(
            'frames returned by ffmpeg cannot be decoded due to an unsupported format'
        )
      width, height = [int(x) for x in _readline(proc.stdout).split(' ')]
      max_value = int(_readline(proc.stdout))
      if max_value != 255:
        raise NotImplementedError(f'max_value has to be 255, it is {max_value}')

      in_bytes = proc.stdout.read(width * height * 3)
      if not in_bytes:
        raise BufferError('could not read the full frame')
      in_frame = numpy.frombuffer(in_bytes,
                                  numpy.uint8).reshape([height, width, 3])
      in_frame = cvtColor(in_frame, COLOR_RGB2BGR)
      yield in_frame

    retcode = proc.wait(timeout=1)

    if retcode != 0:
      err = proc.stderr.read().decode('utf-8')
      raise SyntaxError(err) class DecodeMediaServiceModule(Module):

  def configure(self, binder: Binder):
    binder.bind(DecodeMediaService)
import logging
from typing import Collection, NamedTuple
from injector import Binder, Module, inject, singleton
from rep0st.analyze.feature_vector_analyzer import TYPE_NAME as FEATURE_VECTOR_TYPE
from rep0st.db.post import Post, PostRepository, PostRepositoryModule
from rep0st.index.post import PostIndex, PostIndexModule
from rep0st.service.analyze_service import AnalyzeService, AnalyzeServiceModule
from rep0st.service.media_service import DecodeMediaService, DecodeMediaServiceModule



log = logging.getLogger(__name__)


class PostSearchServiceModule(Module):

  def configure(self, binder: Binder):
    binder.install(PostIndexModule)
    binder.install(AnalyzeServiceModule)

    binder.install(PostRepositoryModule)