commahack4_chatpid/chatPID.py at main · sincethestudy/commahack4_chatpid · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
from skimage.measure import label, regionprops
from skimage.color import label2rgb
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import gc


from modal import Stub, method
import os
from modal import Image, Secret
import pathlib
import modal
from modal import Mount

def download_models():
    from transformers import AutoModel, AutoTokenizer
    model = AutoModel.from_pretrained("facebook/sam-vit-large")

image = (
    modal.Image.debian_slim()
    .pip_install("torch", "transformers", "Pillow", "torchvision", "matplotlib", "scikit-image", "openai")
    .run_function(download_models)
)

stub = modal.Stub("chatPID", image=image)

local_dir = pathlib.Path(os.getcwd())
remote_dir = "./"

mount = Mount.from_local_dir(local_dir, remote_path=remote_dir)

@stub.cls(gpu="A100", cpu=16, container_idle_timeout=300)
class ImageProcessor:
    def __enter__(self):
        from transformers import AutoModel, pipeline

        import os

        self.model = AutoModel.from_pretrained("facebook/sam-vit-large")
        self.generator = pipeline("mask-generation", model="facebook/sam-vit-large", device="cuda:0")

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    @method()
    def process_mask(self, image: np.array) -> dict:
        from PIL import Image
        import numpy as np
        import matplotlib.pyplot as plt
        from skimage.measure import label, regionprops
        import gc
        import torch

        raw_image = Image.fromarray(image).convert("RGB")
        aspect_ratio = raw_image.size[0] / raw_image.size[1]
        new_height = 100
        new_width = int(new_height * aspect_ratio)
        raw_image = raw_image.resize((new_width, new_height), Image.LANCZOS)
        width, height = raw_image.size
        left = width * 0.2
        right = width * 0.8
        top = height * 0.05
        bottom = height * 0.95
        raw_image = raw_image.crop((left, top, right, bottom))


        outputs = self.generator(raw_image, points_per_batch=256)
        masks = outputs["masks"]

        return masks