File size: 907 Bytes
1e008dd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
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")
|