datachain 0.2.11__py3-none-any.whl → 0.2.12__py3-none-any.whl

datachain/lib/hf_image_to_text.py

@@ -1,97 +0,0 @@
-import numpy as np
-import torch
-from PIL import Image, ImageOps, UnidentifiedImageError
-from transformers import (
-    AutoProcessor,
-    Blip2ForConditionalGeneration,
-    Blip2Processor,
-    LlavaForConditionalGeneration,
-)
-
-from datachain.query import Object, udf
-from datachain.sql.types import String
-
-DEFAULT_FIT_BOX = (500, 500)
-
-
-def encode_image(raw):
-    try:
-        img = Image.open(raw)
-    except UnidentifiedImageError:
-        return None
-    img.load()
-    img = img.convert("RGB")
-    return ImageOps.fit(img, DEFAULT_FIT_BOX)
-
-
-def infer_dtype(device):
-    if device == "cpu":
-        return torch.float32
-    return torch.float16
-
-
-@udf(
-    params=(Object(encode_image),),  # Columns consumed by the UDF.
-    output={
-        "description": String,
-        "error": String,
-    },  # Signals being returned by the UDF.
-    batch=64,
-    method="describe",
-)
-class BLIP2describe:
-    def __init__(self, device="cpu", model="Salesforce/blip2-opt-2.7b", max_tokens=300):
-        self.torch_dtype = infer_dtype(device)
-        self.processor = Blip2Processor.from_pretrained(model)
-        self.model = Blip2ForConditionalGeneration.from_pretrained(
-            model, torch_dtype=self.torch_dtype
-        )
-        self.device = device
-        self.model.to(device)
-        self.max_tokens = max_tokens
-
-    def describe(self, imgs):
-        images = np.squeeze(np.asarray(imgs))
-        inputs = self.processor(images=images, return_tensors="pt").to(
-            self.device, self.torch_dtype
-        )
-
-        generated_ids = self.model.generate(**inputs, max_new_tokens=self.max_tokens)
-        generated_text = self.processor.batch_decode(
-            generated_ids, skip_special_tokens=True
-        )
-        return [(desc.strip(), "") for desc in generated_text]
-
-
-@udf(
-    params=(Object(encode_image),),  # Columns consumed by the UDF.
-    output={
-        "description": String,
-        "error": String,
-    },  # Signals being returned by the UDF.
-    batch=16,
-    method="describe",
-)
-class LLaVAdescribe:
-    def __init__(self, device="cpu", model="llava-hf/llava-1.5-7b-hf", max_tokens=300):
-        self.device = device
-        self.torch_dtype = infer_dtype(device)
-        self.processor = AutoProcessor.from_pretrained(model)
-        self.model = LlavaForConditionalGeneration.from_pretrained(
-            model, torch_dtype=self.torch_dtype, low_cpu_mem_usage=True
-        )
-        self.model.to(device)
-        self.max_tokens = max_tokens
-        self.prompt = "USER: <image>\nDescribe this picture\nASSISTANT:"
-
-    def describe(self, imgs):
-        images = np.squeeze(np.asarray(imgs))
-        inputs = self.processor(
-            text=[self.prompt] * len(imgs), images=images, return_tensors="pt"
-        ).to(self.device, self.torch_dtype)
-
-        generated_ids = self.model.generate(**inputs, max_new_tokens=self.max_tokens)
-        generated_text = self.processor.batch_decode(
-            generated_ids, skip_special_tokens=True
-        )
-        return [(desc.split("ASSISTANT:")[-1].strip(), "") for desc in generated_text]

datachain/lib/hf_pipeline.py

@@ -1,90 +0,0 @@
-import json
-
-from PIL import (
-    Image,
-    UnidentifiedImageError,
-)
-from transformers import pipeline
-
-from datachain.query import Object, udf
-from datachain.sql.types import JSON, String
-
-
-def read_image(raw):
-    try:
-        img = Image.open(raw)
-    except UnidentifiedImageError:
-        return None
-    img.load()
-    return img.convert("RGB")
-
-
-def read_object(raw):
-    return raw.read()
-
-
-def read_text(raw):
-    return read_object(raw).decode("utf-8")
-
-
-@udf(
-    params=(Object(read_image),),  # Columns consumed by the UDF.
-    output={
-        "model_output": JSON,
-        "error": String,
-    },  # Signals being returned by the UDF.
-    method="image_processor",
-)
-class ImageHelper:
-    def __init__(self, model, device, **kwargs):
-        self.helper = pipeline(model=model, device=device)
-        self.kwargs = kwargs
-
-    def image_processor(self, imgs):
-        result = self.helper(
-            imgs,
-            **self.kwargs,
-        )
-        return (json.dumps(result), "")
-
-
-@udf(
-    params=(Object(read_text),),  # Columns consumed by the UDF.
-    output={
-        "model_output": JSON,
-        "error": String,
-    },  # Signals being returned by the UDF.
-    method="text_processor",
-)
-class TextHelper:
-    def __init__(self, model, device, **kwargs):
-        self.helper = pipeline(model=model, device=device)
-        self.kwargs = kwargs
-
-    def text_processor(self, text):
-        result = self.helper(
-            text,
-            **self.kwargs,
-        )
-        return (json.dumps(result), "")
-
-
-@udf(
-    params=(Object(read_object),),  # Columns consumed by the UDF.
-    output={
-        "model_output": JSON,
-        "error": String,
-    },  # Signals being returned by the UDF.
-    method="raw_processor",
-)
-class RawHelper:
-    def __init__(self, model, device, **kwargs):
-        self.helper = pipeline(model=model, device=device)
-        self.kwargs = kwargs
-
-    def raw_processor(self, obj):
-        result = self.helper(
-            obj,
-            **self.kwargs,
-        )
-        return (json.dumps(result), "")

datachain/lib/image_transform.py

@@ -1,103 +0,0 @@
-import os
-
-import fsspec
-from PIL import Image
-
-from datachain.catalog import get_catalog
-from datachain.query import DatasetRow, Object, udf
-
-
-def load_image(raw):
-    img = Image.open(raw)
-    img.load()
-    return img
-
-
-@udf(
-    output=DatasetRow.schema,
-    params=(Object(load_image), *tuple(DatasetRow.schema.keys())),
-)
-class ImageTransform:
-    def __init__(
-        self,
-        *,
-        image_filter,
-        bucket_name,
-        prefix,
-        output_folder,
-        file_prefix="",
-        vtype="",
-    ):
-        # Once we fix the UDF decorator situation, it would make more sense to put this
-        # into a child class and make apply_filter an abstractmethod.
-        self.image_filter = image_filter
-        self.folder_name = output_folder
-        self.file_prefix = file_prefix
-        self.prefix = prefix
-        self.vtype = vtype
-
-        catalog = get_catalog()
-        self.client, _ = catalog.parse_url(os.path.join(self.prefix, bucket_name))
-
-    def apply_filter(self, image):
-        return image.filter(self.image_filter)
-
-    def save(self, image, source, name, format):
-        # Make name for new image
-        new_name = f"{self.file_prefix}{name}"
-
-        # Do writeback
-        blob_name = os.path.join(self.folder_name, new_name)
-        urlpath = os.path.join(source, blob_name)
-        cloud_file = fsspec.open(urlpath=urlpath, mode="wb")
-        with cloud_file as fp:
-            image.save(fp, format=format)
-
-        # Get the blob info
-        info_ = self.client.fs.info(urlpath)
-        info = self.client.convert_info(info_, self.folder_name)
-        info.name = new_name
-        return info
-
-    def __call__(
-        self,
-        image,
-        *args,
-    ):
-        # Build a dict from row contents
-        record = dict(zip(DatasetRow.schema.keys(), args))
-        record["is_latest"] = record["is_latest"] > 0  # needs to be a bool
-
-        # yield same row back
-        yield DatasetRow.create(**record)
-
-        # Don't apply the filter twice
-        if record["parent"] == self.folder_name:
-            return
-
-        # Apply the filter
-        image_b = self.apply_filter(image)
-
-        # Save the image and get the cloud object info
-        entry = self.save(
-            image_b, record["source"], name=record["name"], format=image.format
-        )
-
-        # Build the new row
-        yield DatasetRow.create(
-            name=entry.name,
-            source=record["source"],
-            parent=self.folder_name,
-            size=entry.size,
-            location=record["name"]
-            if not record["parent"]
-            else f"{record['parent']}/{record['name']}",
-            vtype=self.vtype,
-            dir_type=record["dir_type"],
-            owner_name=entry.owner_name,
-            owner_id=entry.owner_id,
-            is_latest=record["is_latest"],
-            last_modified=entry.last_modified,
-            version=entry.version,
-            etag=entry.etag,
-        )

datachain/lib/iptc_exif_xmp.py

@@ -1,76 +0,0 @@
-import json
-
-from PIL import (
-    ExifTags,
-    Image,
-    IptcImagePlugin,
-    TiffImagePlugin,
-    UnidentifiedImageError,
-)
-
-from datachain.query import Object, udf
-from datachain.sql.types import JSON, String
-
-
-def encode_image(raw):
-    try:
-        img = Image.open(raw)
-    except UnidentifiedImageError:
-        return None
-    return img
-
-
-@udf(
-    params=(Object(encode_image),),  # Columns consumed by the UDF.
-    output={
-        "xmp": JSON,
-        "exif": JSON,
-        "iptc": JSON,
-        "error": String,
-    },  # Signals being returned by the UDF.
-    method="image_description",
-)
-class GetMetadata:
-    def cast(self, v):  # to JSON serializable types
-        if isinstance(v, TiffImagePlugin.IFDRational):
-            return float(v)
-        if isinstance(v, tuple):
-            return tuple(self.cast(t) for t in v)
-        if isinstance(v, bytes):
-            return v.decode(encoding="utf-8", errors="ignore")
-        if isinstance(v, dict):
-            for kk, vv in v.items():
-                v[kk] = self.cast(vv)
-            return v
-        if isinstance(v, list):
-            return [self.cast(kk) for kk in v]
-        return v
-
-    def image_description(self, img):
-        (xmp, exif, iptc) = ({}, {}, {})
-        if img is None:
-            error = "Image format not understood"
-            return ({}, {}, {}, error)
-        error = ""
-        xmp = img.getxmp()
-        img_exif = img.getexif()
-        img_iptc = IptcImagePlugin.getiptcinfo(img)
-
-        if img_iptc:
-            for k, v in img_iptc.items():
-                iptc[str(k)] = self.cast(v)
-
-        if img_exif:
-            for k, v in img_exif.items():
-                v = self.cast(v)
-                if k in ExifTags.TAGS:
-                    exif[ExifTags.TAGS[k]] = v
-                if k in ExifTags.GPSTAGS:
-                    exif[ExifTags.GPSTAGS[k]] = v
-
-        return (
-            json.dumps(xmp),
-            json.dumps(exif),
-            json.dumps(iptc),
-            error,
-        )

datachain/lib/unstructured.py

@@ -1,41 +0,0 @@
-import shutil
-import tempfile
-
-from unstructured.partition.auto import partition
-from unstructured.staging.base import convert_to_dataframe
-
-from datachain.lib.udf import Mapper
-from datachain.query import Stream
-from datachain.sql.types import JSON, String
-
-
-class PartitionObject(Mapper):
-    def __init__(self):
-        super().__init__(
-            [
-                Stream(),
-            ],
-            {
-                "elements": JSON,
-                "title": String,
-                "text": String,
-                "error": String,
-            },
-        )
-
-    def encode_object(self, raw):
-        fname = str(raw).replace(">", "").replace("<", "")
-        output = tempfile.TemporaryFile()
-        shutil.copyfileobj(raw, output)
-        elements = partition(file=output, metadata_filename=fname)
-        output.close()
-        return elements
-
-    def __call__(self, stream):
-        with stream:
-            elements = self.encode_object(stream)
-
-        title = str(elements[0])
-        text = "\n\n".join([str(el) for el in elements])
-        df = convert_to_dataframe(elements)
-        return (df.to_json(), title, text, "")

datachain/text/__init__.py

@@ -1,3 +0,0 @@
-from datachain.lib.text import convert_text
-
-__all__ = ["convert_text"]