from transformers import Pipeline , AutoModel, AutoImageProcessor, pipeline
from PIL import Image
from loadimg import load_img
import torch
REPO_NAME = "p1atdev/MangaLineExtraction-hf"

class MangaLineExtractionPipe(Pipeline):
  def __init__(self,**kwargs):
    self.processor = AutoImageProcessor.from_pretrained("p1atdev/MangaLineExtraction-hf", trust_remote_code=True)
    Pipeline.__init__(self,**kwargs)
    self.model.eval()
  def _sanitize_parameters(kwargs):
    return {}, {} , {}
  def preprocess(self,image):
    # takes any image type and returns a pil image
    image = load_img(image, output_type = "pil").convert("RGB")
    processed =  self.processor(image, return_tensors="pt")
    return processed.pixel_values
  def _forward(self,inputs):
    return self.model(inputs)
  def postprocess(self,outputs):
    return Image.fromarray(outputs.pixel_values[0].numpy().astype("uint8"), mode="L")