PyCS/pycs/management/result.py

import click
import flask
import simplejson as json

from flask.cli import AppGroup
from tqdm import tqdm

from pycs import app
from pycs import database as db

result_cli = AppGroup("result", short_help="Result operations")



@result_cli.command("export")
@click.argument("project_id")
@click.argument("indent", required=False)
@click.argument("output", required=False)
def export(project_id, output, indent):
    """ Export results for a specific project or for all projects """
    if project_id == "all":
        projects = db.Project.query.all()
        app.logger.info(f"Exporting results for all projects ({len(projects)})!")
        if output is None:
            output = "output.json"

    else:
        project = db.Project.query.get(project_id)
        if project is None:
            app.logger.error(f"Could not find project with ID {project_id}!")
            return
        app.logger.info(f"Exporting results for project {project}!")
        projects = [project]
        if output is None:
            output = f"output_project_{int(project_id):04d}.json"

    app.logger.info(f"Exporting to {output}")

    results = []

    for project in projects:
        project_files = [
            dict(**f.serialize(),
                results=[
                    dict(**r.serialize(), label=r.label.serialize() if r.label is not None else None)
                        for r in f.results.all()
                ])
                for f in tqdm(project.files.all()) if f.results.count() != 0
            ]

        results.append(dict(
            project_id=project.id,
            files=project_files,
            labels=[lab.serialize() for lab in project.labels.all()],
        ))


    if indent is not None:
        indent = int(indent)

    with open(output, "w", encoding="utf-8") as out_f:
        flask.json.dump(results, out_f, app=app, indent=indent)




@result_cli.command("restore")
@click.argument("infile")
@click.option("--dry-run", is_flag=True)
def restore(infile, dry_run):

    with open(infile) as f:
        results = json.load(f)

    for project_results in results:
        project = db.Project.get_or_404(project_results["project_id"])
        for file_results in project_results["files"]:
            file = db.File.get_or_404(file_results["id"])

            assert file.path == file_results["path"]

            # first check for new and changed results
            for _result in file_results["results"]:

                if not _is_data_valid(**_result):
                    continue

                result = get_result_or_none(file, **_result)

                user1 = _result["origin_user"]
                data1 = _result["data"]
                ref1 = (_result["label"] or {}).get("reference")
                # lab1 = (_result["label"] or {}).get("id")


                if result is None:
                    # we have a new result entry
                    if not dry_run:
                        file.create_result(
                            result_type="bounding-box",
                            origin="user",
                            origin_user=user1,
                            label=ref1,
                            data=data1,
                            commit=True
                        )
                    print(" | ".join([
                        f"Project #{project.id:< 6d}"
                        f"File #{file.id:< 6d} [{file.name:^30s}]",
                        "[New Result]",
                        f"User: {user1 or '':<10s}",
                        f"Data: {data1}, Label-Ref: {ref1}",
                        ])
                    )

                    continue

                assert result.file_id == _result["file_id"]
                user0 = result.origin_user
                data0 = result.data
                ref0 = getattr(result.label, "reference", None)
                # lab0 = getattr(result.label, "id", None)

                is_same_data = _check_data(data0, data1)

                if is_same_data and (ref0 == ref1 or ref1 is None):
                    # nothing to change
                    continue

                print(" | ".join([
                    f"Project #{project.id:< 6d}"
                    f"File #{file.id:< 6d} [{file.name:^30s}]",
                    ]), end=" | "
                )
                if not is_same_data:
                    # data was updated
                    print(" | ".join([
                        "[Data updated]",
                        f"User: {user1 or '':<10s}",
                        f"Data: {data0} -> {data1}"
                        ]), end=" | "
                    )
                    assert user1 is not None
                    if not dry_run:
                        result.origin_user = user1
                        result.data = data1

                if ref0 != ref1:
                    assert user1 is not None
                    if not dry_run:
                        result.origin_user = user1
                    if ref1 is None:
                        # label was deleted
                        print("[Label Deleted]")
                        if not dry_run:
                            result.label_id = None
                    else:
                        # label was updated
                        print(" | ".join([
                            "[Label updated]",
                            f"User: {user0 or '':<10s} -> {user1 or '':<10s}",
                            f"{ref0 or 'UNK':<6s} -> {ref1 or 'UNK':<6s}"
                            ])
                        )
                        label = project.label_by_reference(ref1)
                        if not dry_run:
                            result.label_id = label.id
                else:
                    print()

                if not dry_run:
                    result.commit()

            # then check for deleted results
            for result in file.results.all():
                if result.origin != "user" or result.type != "bounding-box":
                    continue

                found = False
                for _result in file_results["results"]:
                    if not _is_data_valid(**_result):
                        continue

                    if _check_data(result.data, _result["data"]):
                        found = True
                        break

                if not found:
                    print(" | ".join([
                        f"Project #{project.id:< 6d}"
                        f"File #{file.id:< 6d} [{file.name:^30s}]",
                        "[Result deleted]",
                        f"{result.data}",
                        f"{result.label}",
                        ])
                    )

                    if not dry_run:
                        result.delete()

def _is_data_valid(*, data, type, origin, **kwargs):

    wh = (None, None) if data is None else (data["w"], data["h"])

    return (type != "labeled-image" and
        origin == "user" and
        0 not in wh)

def _check_data(data0, data1):

    if None in (data0, data1):
        return data0 == data1 == None

    for key in data0:
        if data1.get(key) != data0.get(key):
            return False

    return True

def get_result_or_none(file: db.File, id: int, data: dict, **kwargs):

    result = db.Result.query.filter(
        db.Result.id==id, db.Result.file_id==file.id).one_or_none()

    if result is not None:
        return result

    for other_results in file.results.all():
        if _check_data(data, other_results.data):
            # import pdb; pdb.set_trace()
            return other_results