IoU evaluator

.idea/detector.iml

@@ -4,7 +4,7 @@
     <content url="file://$MODULE_DIR$">
       <excludeFolder url="file://$MODULE_DIR$/.venv" />
     </content>
-    <orderEntry type="inheritedJdk" />
+    <orderEntry type="jdk" jdkName="detector" jdkType="Python SDK" />
     <orderEntry type="sourceFolder" forTests="false" />
   </component>
 </module>

.idea/misc.xml

@@ -3,5 +3,5 @@
   <component name="Black">
     <option name="sdkName" value="Python 3.9 (detector)" />
   </component>
-  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.9 (detector)" project-jdk-type="Python SDK" />
+  <component name="ProjectRootManager" version="2" project-jdk-name="detector" project-jdk-type="Python SDK" />
 </project>

.idea/workspace.xml

@@ -4,10 +4,16 @@
     <option name="autoReloadType" value="SELECTIVE" />
   </component>
   <component name="ChangeListManager">
-    <list default="true" id="d7806539-b6d6-42e7-bb45-1565f5d54891" name="Changes" comment="clean requirements.txt">
+    <list default="true" id="d7806539-b6d6-42e7-bb45-1565f5d54891" name="Changes" comment="fix confidence">
+      <change afterPath="$PROJECT_DIR$/evaluator.py" afterDir="false" />
+      <change afterPath="$PROJECT_DIR$/tests/coco8.zip" afterDir="false" />
+      <change afterPath="$PROJECT_DIR$/tests/project-2-at-2024-01-20-11-43-00ad4911.zip" afterDir="false" />
+      <change afterPath="$PROJECT_DIR$/yolo_dataset.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/.idea/detector.iml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/detector.iml" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/.idea/misc.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/misc.xml" afterDir="false" />
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/Dockerfile" beforeDir="false" afterPath="$PROJECT_DIR$/Dockerfile" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/requirements.txt" beforeDir="false" afterPath="$PROJECT_DIR$/requirements.txt" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/demo.py" beforeDir="false" afterPath="$PROJECT_DIR$/demo.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/model.py" beforeDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
@@ -41,28 +47,36 @@
     <option name="hideEmptyMiddlePackages" value="true" />
     <option name="showLibraryContents" value="true" />
   </component>
-  <component name="PropertiesComponent">{
-  &quot;keyToString&quot;: {
-    &quot;Python.demo.executor&quot;: &quot;Run&quot;,
-    &quot;RunOnceActivity.OpenProjectViewOnStart&quot;: &quot;true&quot;,
-    &quot;RunOnceActivity.ShowReadmeOnStart&quot;: &quot;true&quot;,
-    &quot;git-widget-placeholder&quot;: &quot;main&quot;,
-    &quot;last_opened_file_path&quot;: &quot;/Users/famer.me/PycharmProjects/detector&quot;,
-    &quot;node.js.detected.package.eslint&quot;: &quot;true&quot;,
-    &quot;node.js.detected.package.tslint&quot;: &quot;true&quot;,
-    &quot;node.js.selected.package.eslint&quot;: &quot;(autodetect)&quot;,
-    &quot;node.js.selected.package.tslint&quot;: &quot;(autodetect)&quot;,
-    &quot;nodejs_package_manager_path&quot;: &quot;npm&quot;,
-    &quot;settings.editor.selected.configurable&quot;: &quot;settings.qodana&quot;,
-    &quot;vue.rearranger.settings.migration&quot;: &quot;true&quot;
+  <component name="PropertiesComponent"><![CDATA[{
+  "keyToString": {
+    "Python.demo.executor": "Run",
+    "Python.evaluator.executor": "Debug",
+    "Python.extract.executor": "Run",
+    "Python.yolo_dataset.executor": "Run",
+    "RunOnceActivity.OpenProjectViewOnStart": "true",
+    "RunOnceActivity.ShowReadmeOnStart": "true",
+    "git-widget-placeholder": "main",
+    "last_opened_file_path": "/Users/famer.me/PycharmProjects/detector/tests",
+    "node.js.detected.package.eslint": "true",
+    "node.js.detected.package.tslint": "true",
+    "node.js.selected.package.eslint": "(autodetect)",
+    "node.js.selected.package.tslint": "(autodetect)",
+    "nodejs_package_manager_path": "npm",
+    "settings.editor.selected.configurable": "settings.qodana",
+    "vue.rearranger.settings.migration": "true"
   }
-}</component>
+}]]></component>
   <component name="RecentsManager">
     <key name="CopyFile.RECENT_KEYS">
+      <recent name="$PROJECT_DIR$/tests" />
+      <recent name="$PROJECT_DIR$" />
+    </key>
+    <key name="MoveFile.RECENT_KEYS">
+      <recent name="$PROJECT_DIR$/tests" />
       <recent name="$PROJECT_DIR$" />
     </key>
   </component>
-  <component name="RunManager" selected="Python.demo">
+  <component name="RunManager" selected="Python.evaluator">
     <configuration name="demo" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
       <module name="detector" />
       <option name="ENV_FILES" value="" />
@@ -86,6 +100,52 @@
       <option name="INPUT_FILE" value="" />
       <method v="2" />
     </configuration>
+    <configuration name="evaluator" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
+      <module name="detector" />
+      <option name="ENV_FILES" value="" />
+      <option name="INTERPRETER_OPTIONS" value="" />
+      <option name="PARENT_ENVS" value="true" />
+      <envs>
+        <env name="PYTHONUNBUFFERED" value="1" />
+      </envs>
+      <option name="SDK_HOME" value="" />
+      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$" />
+      <option name="IS_MODULE_SDK" value="true" />
+      <option name="ADD_CONTENT_ROOTS" value="true" />
+      <option name="ADD_SOURCE_ROOTS" value="true" />
+      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
+      <option name="SCRIPT_NAME" value="$PROJECT_DIR$/evaluator.py" />
+      <option name="PARAMETERS" value="" />
+      <option name="SHOW_COMMAND_LINE" value="false" />
+      <option name="EMULATE_TERMINAL" value="false" />
+      <option name="MODULE_MODE" value="false" />
+      <option name="REDIRECT_INPUT" value="false" />
+      <option name="INPUT_FILE" value="" />
+      <method v="2" />
+    </configuration>
+    <configuration name="yolo_dataset" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
+      <module name="detector" />
+      <option name="ENV_FILES" value="" />
+      <option name="INTERPRETER_OPTIONS" value="" />
+      <option name="PARENT_ENVS" value="true" />
+      <envs>
+        <env name="PYTHONUNBUFFERED" value="1" />
+      </envs>
+      <option name="SDK_HOME" value="" />
+      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$" />
+      <option name="IS_MODULE_SDK" value="true" />
+      <option name="ADD_CONTENT_ROOTS" value="true" />
+      <option name="ADD_SOURCE_ROOTS" value="true" />
+      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
+      <option name="SCRIPT_NAME" value="$PROJECT_DIR$/yolo_dataset.py" />
+      <option name="PARAMETERS" value="" />
+      <option name="SHOW_COMMAND_LINE" value="false" />
+      <option name="EMULATE_TERMINAL" value="false" />
+      <option name="MODULE_MODE" value="false" />
+      <option name="REDIRECT_INPUT" value="false" />
+      <option name="INPUT_FILE" value="" />
+      <method v="2" />
+    </configuration>
     <configuration name="detector" type="Python.FastAPI">
       <option name="file" value="$PROJECT_DIR$/main.py" />
       <module name="detector" />
@@ -103,6 +163,8 @@
     </configuration>
     <recent_temporary>
       <list>
+        <item itemvalue="Python.evaluator" />
+        <item itemvalue="Python.yolo_dataset" />
         <item itemvalue="Python.demo" />
       </list>
     </recent_temporary>
@@ -110,7 +172,7 @@
   <component name="SharedIndexes">
     <attachedChunks>
       <set>
-        <option value="bundled-python-sdk-7a29c1521ef0-c986f194a52a-com.jetbrains.pycharm.pro.sharedIndexes.bundled-PY-233.11799.298" />
+        <option value="bundled-python-sdk-50da183f06c8-2887949eec09-com.jetbrains.pycharm.pro.sharedIndexes.bundled-PY-233.13135.95" />
       </set>
     </attachedChunks>
   </component>
@@ -124,7 +186,8 @@
       <updated>1702301546100</updated>
       <workItem from="1702301554378" duration="4025000" />
       <workItem from="1702450051021" duration="10065000" />
-      <workItem from="1702476484538" duration="7523000" />
+      <workItem from="1702476484538" duration="9613000" />
+      <workItem from="1705757850315" duration="10355000" />
     </task>
     <task id="LOCAL-00001" summary="init commit">
       <option name="closed" value="true" />
@@ -222,7 +285,47 @@
       <option name="project" value="LOCAL" />
       <updated>1702483217316</updated>
     </task>
-    <option name="localTasksCounter" value="13" />
+    <task id="LOCAL-00013" summary="clean requirements.txt">
+      <option name="closed" value="true" />
+      <created>1702484185625</created>
+      <option name="number" value="00013" />
+      <option name="presentableId" value="LOCAL-00013" />
+      <option name="project" value="LOCAL" />
+      <updated>1702484185625</updated>
+    </task>
+    <task id="LOCAL-00014" summary="use env">
+      <option name="closed" value="true" />
+      <created>1702484512308</created>
+      <option name="number" value="00014" />
+      <option name="presentableId" value="LOCAL-00014" />
+      <option name="project" value="LOCAL" />
+      <updated>1702484512308</updated>
+    </task>
+    <task id="LOCAL-00015" summary="use /run/secrets">
+      <option name="closed" value="true" />
+      <created>1702484869890</created>
+      <option name="number" value="00015" />
+      <option name="presentableId" value="LOCAL-00015" />
+      <option name="project" value="LOCAL" />
+      <updated>1702484869890</updated>
+    </task>
+    <task id="LOCAL-00016" summary="use /run/secrets">
+      <option name="closed" value="true" />
+      <created>1702484897551</created>
+      <option name="number" value="00016" />
+      <option name="presentableId" value="LOCAL-00016" />
+      <option name="project" value="LOCAL" />
+      <updated>1702484897551</updated>
+    </task>
+    <task id="LOCAL-00017" summary="fix confidence">
+      <option name="closed" value="true" />
+      <created>1702485964999</created>
+      <option name="number" value="00017" />
+      <option name="presentableId" value="LOCAL-00017" />
+      <option name="project" value="LOCAL" />
+      <updated>1702485964999</updated>
+    </task>
+    <option name="localTasksCounter" value="18" />
     <servers />
   </component>
   <component name="TypeScriptGeneratedFilesManager">
@@ -261,9 +364,14 @@
     <MESSAGE value="add test image" />
     <MESSAGE value="remove opencv" />
     <MESSAGE value="clean requirements.txt" />
-    <option name="LAST_COMMIT_MESSAGE" value="clean requirements.txt" />
+    <MESSAGE value="use env" />
+    <MESSAGE value="use /run/secrets" />
+    <MESSAGE value="fix confidence" />
+    <option name="LAST_COMMIT_MESSAGE" value="fix confidence" />
   </component>
   <component name="com.intellij.coverage.CoverageDataManagerImpl">
     <SUITE FILE_PATH="coverage/detector$demo.coverage" NAME="demo Coverage Results" MODIFIED="1702302380580" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$" />
+    <SUITE FILE_PATH="coverage/detector$evaluator.coverage" NAME="evaluator Coverage Results" MODIFIED="1705768159641" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$" />
+    <SUITE FILE_PATH="coverage/detector$extract.coverage" NAME="yolo_dataset Coverage Results" MODIFIED="1705764465837" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$" />
   </component>
 </project>

demo.py

@@ -1,4 +1,3 @@
-
 import time
 from PIL import Image
 import streamlit as st
@@ -38,4 +37,4 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    main()

evaluator.py

@@ -0,0 +1,53 @@
+from typing import Callable
+
+from PIL.Image import Image
+from coco_eval import CocoEvaluator
+from pycocotools.coco import COCO
+from tqdm import tqdm
+
+from yolo_dataset import YoloDataset
+from yolo_fire import model
+
+image_loader = Callable[[str], Image]
+
+
+def evaluate(coco_gt: COCO, loader: image_loader, confidence_threshold=0.6):
+    # initialize evaluator with ground truth (gt)
+    evaluator = CocoEvaluator(coco_gt=coco_gt, iou_types=["bbox"])
+
+    print("Running evaluation...")
+    for image_id, annotations in tqdm(coco_gt.imgToAnns.items()):
+        # get the inputs
+        image = coco_gt.imgs[image_id]
+        results = model(source=loader(image["file_name"]))
+        for result in results:
+            coco_anns = yolo_boxes_to_coco_annotations(image_id, result.boxes, confidence_threshold=confidence_threshold)
+            if len(coco_anns) == 0:
+                continue
+            evaluator.update(coco_anns)
+
+    if len(evaluator.eval_imgs["bbox"]) == 0:
+        print("No detections!")
+        return
+    evaluator.synchronize_between_processes()
+    evaluator.accumulate()
+    evaluator.summarize()
+
+
+def yolo_boxes_to_coco_annotations(image_id: int, yolo_boxes, confidence_threshold=0.6):
+    return [
+        {
+            "image_id": image_id,
+            "category_id": box.cls.tolist()[0],
+            "area": box.xywh.tolist()[0][2] * box.xywh.tolist()[0][3],
+            "bbox": box.xywh.tolist()[0],
+            "score": box.conf.tolist()[0],
+        }
+        for box in yolo_boxes if box.conf.tolist()[0] > confidence_threshold
+    ]
+
+
+if __name__ == '__main__':
+    yolo_dataset = YoloDataset.from_zip_file('tests/coco8.zip')
+    coco_gt = yolo_dataset.to_coco()
+    evaluate(coco_gt=coco_gt, loader=yolo_dataset.load_image, confidence_threshold=0.1)

model.py

@@ -1,93 +0,0 @@
-import torch
-import random
-from PIL import ImageDraw
-import torchvision.transforms as T
-
-# COCO Classes
-CLASSES = [
-    'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
-    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A',
-    'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
-    'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack',
-    'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
-    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
-    'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass',
-    'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
-    'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
-    'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A',
-    'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
-    'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A',
-    'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
-    'toothbrush'
-]
-
-# Standard PyTorch mean-std Input Image Normalization
-transform = T.Compose([
-    T.Resize(500),
-    T.ToTensor(),
-    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-])
-
-
-# For Output Bounding Box Post-processing
-def box_cxcywh_to_xyxy(x):
-    x_c, y_c, w, h = x.unbind(1)
-    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
-         (x_c + 0.5 * w), (y_c + 0.5 * h)]
-    return torch.stack(b, dim=1)
-
-
-def rescale_bboxes(out_bbox, size):
-    img_w, img_h = size
-    b = box_cxcywh_to_xyxy(out_bbox)
-    b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
-    return b
-
-
-# Pre-processing on Image
-def image_processing(im, model, transform, confidence=0.9):
-    # im = Image.open(image_path)
-    img = transform(im).unsqueeze(0)
-
-    outputs = model(img)
-    probas = outputs['pred_logits'].softmax(-1)[0, :, :-1]
-    keep = probas.max(-1).values > confidence
-    bboxes_scaled = rescale_bboxes(outputs['pred_boxes'][0, keep], im.size)
-
-    return probas[keep], bboxes_scaled
-
-
-# Helper Functions for Plotting BBoxes
-def plot_one_box(x, img, color=None, label=None, line_thickness=None):
-    width, height = img.size
-    tl = line_thickness or round(0.002 * (width + height) / 2) + 1  # line/font thickness
-    color = color or (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
-    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
-    img_draw = ImageDraw.Draw(img)
-    img_draw.rectangle((c1[0], c1[1], c2[0], c2[1]), outline=color, width=tl)
-    if label:
-        tf = max(tl - 1, 1)  # font thickness
-        x1, y1, x2, y2 = img_draw.textbbox(c1, label, stroke_width=tf)
-        img_draw.rectangle((x1, y1, x2, y2), fill=color)
-        img_draw.text((x1, y1), label, fill=(255, 255, 255))
-
-
-# Ploting Bounding Box on img
-def add_bboxes(pil_img, prob, bboxes):
-    for p, coord in zip(prob, bboxes.tolist()):
-        cl = p.argmax()
-        text = f'{CLASSES[cl]}: {p[cl]: 0.2f}'
-        plot_one_box(x=coord, img=pil_img, label=text)
-
-    return pil_img
-
-
-def detect(im, confidence):
-    # Load model
-    model = torch.hub.load('facebookresearch/detr', 'detr_resnet101', pretrained=True)
-    model.eval()
-
-    scores, boxes = image_processing(im, model, transform, confidence / 100)
-    im = add_bboxes(im, scores, boxes)
-
-    return im

yolo_dataset.py

@@ -0,0 +1,199 @@
+# Author: Ming Yang
+# Date: 2023/01/20
+# Description: Traverse the zip file but not decompress it.
+# Suppose the zip file contains yolo format annotation files.
+# /
+# ├── classes.txt
+# ├── images
+# │   ├── 1.jpg
+# │   ├── 2.jpg
+# │   └── ...
+# └── labels
+#     ├── 1.txt
+#     ├── 2.txt
+#     └── ...
+import io
+import json
+from datetime import datetime
+from typing import Optional
+from zipfile import ZipFile
+from PIL import Image
+from pycocotools import coco
+from pycocotools.coco import COCO
+
+yolo_label = {
+    'class': str,
+    'class_id': int,
+    'x_center': float,
+    'y_center': float,
+    'width': float,
+    'height': float
+}
+
+yolo_image = {
+    'image_name': str,
+    'image': Image,
+    'labels': list[yolo_label]
+}
+
+coco_annotation = {
+    "id": int,
+    "image_id": int,  # the id of the image that the annotation belongs to
+    "category_id": int,  # the id of the category that the annotation belongs to
+    # "segmentation": RLE or [polygon],
+    "area": float,
+    "bbox": [float, float, float, float],  # [x,y,width,height]
+    "iscrowd": bool,  # 0 or 1,
+}
+
+coco_category = {
+    "id": int,
+    "name": str,
+    "supercategory": Optional[str],
+}
+
+coco_image = {
+    "id": int,
+    "width": int,
+    "height": int,
+    "file_name": str,
+    "date_captured": Optional[datetime],
+}
+
+coco_dataset = {
+    "images": list[coco_image],  # list of all images in the dataset
+    "annotations": list[coco_annotation],  # list of all annotations in the dataset
+    "categories": list[coco_category]  # list of all categories
+}
+
+
+class YoloImage:
+    def __init__(self, image_name: str, image: Image, labels: list[yolo_label]):
+        self.image_name = image_name
+        self.image = image
+        self.labels = labels
+
+    def __repr__(self):
+        return f'YoloImage(image_name={self.image_name}, image={self.image}, labels={self.labels})'
+
+    def to_coco_image(self, id: int) -> coco_image:
+        return {
+            "id": id,
+            "width": self.image.width,
+            "height": self.image.height,
+            "file_name": self.image_name,
+        }
+
+    def to_coco_annotations(self, image_id: int, ann_id_start: int) -> list[coco_annotation]:
+        ann_id = ann_id_start
+        annotations: list[coco_annotation] = []
+        for label in self.labels:
+            ann_id = ann_id + 1
+            annotations.append({
+                "id": ann_id,
+                "image_id": image_id,
+                "category_id": label['class_id'],
+                "area": label['width'] * label['height'],
+                "bbox": [label['x_center'] - label['width'] / 2, label['y_center'] - label['height'] / 2,
+                         label['width'], label['height']],
+                "iscrowd": False,
+            })
+        return annotations
+
+
+class YoloDataset:
+    _zip_file: ZipFile
+    _classes: list[str]
+    _images: list[str]
+    _labels: list[str]
+
+    def __init__(self, zip_file: ZipFile, classes=None, images=None, labels=None):
+        if labels is None:
+            labels = []
+        if images is None:
+            images = []
+        if classes is None:
+            classes = []
+        self._zip_file = zip_file
+        self._classes = classes
+        self._images = images
+        self._labels = labels
+
+    @staticmethod
+    def from_zip_file(zip_file: str) -> 'YoloDataset':
+        zip_file = ZipFile(zip_file, 'r')
+        namelist = zip_file.namelist()
+        root_name = namelist[0]
+        namelist = list(filter(lambda x: not zip_file.getinfo(x).is_dir(), namelist))
+        if 'classes.txt' in namelist:
+            classes = zip_file.read('classes.txt').decode('utf-8').split('\n')
+        else:
+            classes = []
+        images = list(filter(lambda x: x.startswith(root_name + 'images'), namelist))
+        labels = list(filter(lambda x: x.startswith(root_name + 'labels'), namelist))
+        assert len(images) == len(labels) and len(images) > 0
+        images.sort()
+        labels.sort()
+        for image, label in zip(images, labels):
+            image_name = image.split('/')[-1]
+            label_name = label.split('/')[-1]
+            assert image_name.split('.')[0] == label_name.split('.')[0]
+        return YoloDataset(zip_file, classes, images, labels)
+
+    def __len__(self):
+        return len(self._images)
+
+    def __getitem__(self, index: int) -> YoloImage:
+        image_name = self._images[index]
+        labels = self._zip_file.read(self._labels[index]).decode('utf-8').split('\n')
+        labels = list(filter(lambda x: len(x) > 0, labels))
+        labels = list(map(lambda x: x.split(' '), labels))
+        labels = list(map(lambda x: {
+            'class': self._classes[int(x[0])] if len(self._classes) > int(x[0]) else 'unknown',
+            'class_id': int(x[0]),
+            'x_center': float(x[1]),
+            'y_center': float(x[2]),
+            'width': float(x[3]),
+            'height': float(x[4])
+        }, labels))
+
+        return YoloImage(image_name, Image.open(self._zip_file.open(self._images[index])), labels)
+
+    def __iter__(self):
+        for i in range(len(self)):
+            yield self[i]
+
+    def load_image(self, image_name: str) -> Image:
+        return Image.open(self._zip_file.open(image_name))
+
+    def to_coco(self) -> COCO:
+        images: list[coco_image] = []
+        annotations: list[coco_annotation] = []
+        categories: list[coco_category] = []
+        ann_id = 0
+        for i in range(len(self)):
+            image = self[i]
+            images.append(image.to_coco_image(i))
+            annotations.extend(image.to_coco_annotations(i, ann_id))
+            ann_id = ann_id + len(image.labels)
+        for i in range(len(self._classes)):
+            categories.append({
+                "id": i,
+                "name": self._classes[i],
+                "supercategory": None,
+            })
+
+        coco_ds = coco.COCO()
+        coco_ds.dataset = {
+            "images": images,
+            "annotations": annotations,
+            "categories": categories,
+        }
+        coco_ds.createIndex()
+        return coco_ds
+
+
+if __name__ == '__main__':
+    dataset = YoloDataset.from_zip_file('tests/coco8.zip')
+    coco = dataset.to_coco()
+    print(coco)