You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
18 lines
560 B
18 lines
560 B
import sys |
|
from PIL import Image |
|
from cog import BasePredictor, Input, Path |
|
|
|
sys.path.extend(["src/clip", "src/blip"]) |
|
|
|
from clip_interrogator import Interrogator, Config |
|
|
|
|
|
class Predictor(BasePredictor): |
|
def setup(self): |
|
config = Config(device="cuda:0", clip_model_name='ViT-L/14') |
|
self.ci = Interrogator(config) |
|
|
|
def predict(self, image: Path = Input(description="Input image")) -> str: |
|
"""Run a single prediction on the model""" |
|
image = Image.open(str(image)).convert("RGB") |
|
return self.ci.interrogate(image)
|
|
|