from pathlib import Path

import numpy as np

from .base_dataset import BaseDataset
from .builder import DATASETS


@DATASETS.register_module()
class GoalClassification(BaseDataset):
    """
    进球判定数据集
    """

    def load_annotations(self):
        """
        0是负样本：没进球
        1是正样本：进球了
        Returns:
        """
        positive_images = []
        negative_images = []
        for p in ["/mnt/nfs-storage/yujiannan/data/20211007", "/mnt/nfs-storage/yujiannan/data/20211103"]:
            positive_images.extend(Path(p).joinpath("positive").iterdir())
            negative_images.extend(Path(p).joinpath("negative").iterdir())
        positive_count = len(positive_images)
        negative_count = len(negative_images)
        annotations = []
        start = 0
        end = 1
        if self.data_prefix == "train":
            start = 0
            end = 8
        elif self.data_prefix == "val":
            start = 8
            end = 9
        elif self.data_prefix == "test":
            start = 9
            end = 10

        for i in range(0, positive_count, 10):
            for image_index in range(i + start, i + end):
                if image_index >= positive_count:
                    break
                image = positive_images[image_index]
                annotations.append(dict(img_prefix=image.parent,
                                        img_info={'filename': image.name},
                                        gt_label=np.array(1, dtype=np.int64)))
        for i in range(0, negative_count, 10):
            for image_index in range(i+start, i+end):
                if image_index >= negative_count:
                    break
                image = negative_images[image_index]
                annotations.append(dict(img_prefix=image.parent,
                                        img_info={'filename': image.name},
                                        gt_label=np.array(0, dtype=np.int64)))
        return annotations

from pathlib import Path

import numpy as np

from .base_dataset import BaseDataset
from .builder import DATASETS


@DATASETS.register_module()
class GoalClassification(BaseDataset):
    """
    进球判定数据集
    """

    def load_annotations(self):
        """
        0是负样本：没进球
        1是正样本：进球了
        Returns:
        """
        positive_images = []
        negative_images = []
        for p in ["/mnt/nfs-storage/yujiannan/data/20211007", "/mnt/nfs-storage/yujiannan/data/20211103"]:
            positive_images.extend(Path(p).joinpath("positive").iterdir())
            negative_images.extend(Path(p).joinpath("negative").iterdir())
        positive_count = len(positive_images)
        negative_count = len(negative_images)
        annotations = []
        start = 0
        end = 1
        if self.data_prefix == "train":
            start = 0
            end = 8
        elif self.data_prefix == "val":
            start = 8
            end = 9
        elif self.data_prefix == "test":
            start = 9
            end = 10

        for i in range(0, positive_count, 10):
            for image_index in range(i + start, i + end):
                if image_index >= positive_count:
                    break
                image = positive_images[image_index]
                annotations.append(dict(img_prefix=image.parent,
                                        img_info={'filename': image.name},
                                        gt_label=np.array(1, dtype=np.int64)))
        for i in range(0, negative_count, 10):
            for image_index in range(i+start, i+end):
                if image_index >= negative_count:
                    break
                image = negative_images[image_index]
                annotations.append(dict(img_prefix=image.parent,
                                        img_info={'filename': image.name},
                                        gt_label=np.array(0, dtype=np.int64)))
        return annotations

进球判定数据集 0是负样本：没进球
         1是正样本：进球了
         Returns: