Multiple output prediction is possible in cog?

[tzktz]

i want to return two output prediction in below code..

        if tables_crops:
            temp_dir = tempfile.mkdtemp()
            cropped_table_path = Path(temp_dir) / f'cropped_table.jpg'

            # Save the cropped table image in the temporary directory
            cropped_table = tables_crops[0]['image'].convert("RGB")
            cropped_table.save(cropped_table_path)
            print(f"File is successfully saved at {cropped_table_path}")

            bbox_info = [obj['bbox'] for obj in objects if obj['label'] in ['table', 'table rotated']]
            bbox_json = json.dumps(bbox_info, indent=4)

            # Save JSON data to a file
            json_path = Path(temp_dir) / 'bbox_info.json'
            with open(json_path, 'w') as json_file:
                json_file.write(bbox_json)

            # Return the paths of the saved image and JSON file
            return cropped_table_path, json_path

my run comment is : $ cog predict -i file_path=@test4.jpg

but i got error..when i use single prediction it worked smoothly..(ex. return json_path ) my cog version is 0.8.6

can anyone fix this..thx in advance!!! @mattt @zeke

[zeke]

Yep, you wanna do something like this:

from cog import BasePredictor, BaseModel, File, Path

class Output(BaseModel):
    cropped_table: Path
    json: Path

class Predictor(BasePredictor):
    def predict(self) -> Output:
         ...
        return Output(cropped_table=..., json=...)

Bear in mind that cog has a Path that differs from (extends?) the built-in Python path.

See https://github.com/replicate/cog/blob/main/docs/python.md#returning-an-object