SortScrews/collect.py at main · ATATC/SortScrews · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
from csv import writer, reader
from json import dump
from os import makedirs, remove
from typing import override

import cv2
import numpy as np

from sort_screws import Camera

DATASET_DIR: str = "SortScrews"


class Collector(Camera):
    def __init__(self, *, append_from: int | None = None) -> None:
        super().__init__(512, upside_down=True)
        self.images_dir: str = f"{DATASET_DIR}/images"
        self.csv_path: str = f"{DATASET_DIR}/labels.csv"
        if not append_from:
            makedirs(self.images_dir)
            with open(self.csv_path, "w") as f:
                f.write("filename,class\n")
        # runtime
        self.class_id: int = 0
        self.num_cases: int = append_from if append_from else 0

    @override
    def job(self, frame: np.ndarray, roi: np.ndarray, bbox: tuple[int, int, int, int]) -> bool:
        x1, y1, x2, y2 = bbox
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(frame, f"Class: {self.class_id}", (20, 40), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2,
                    cv2.LINE_AA)
        cv2.putText(frame, f"Num cases: {self.num_cases}", (20, 80), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 255, 255),
                    2,
                    cv2.LINE_AA)
        cv2.imshow("Camera Preview", frame)
        key = self.wait_key()
        if key == ord("c"):
            filename = f"case_{self.num_cases:03d}.png"
            cv2.imwrite(f"{self.images_dir}/{filename}", roi)
            with open(self.csv_path, "a", newline="") as f:
                w = writer(f)
                w.writerow([filename, self.class_id])
            self.num_cases += 1
            print(f"Saved {filename} (class {self.class_id})")
        elif key == ord("x") and self.num_cases > 0:
            with open(self.csv_path, "r") as f:
                rows = list(reader(f))
            if len(rows) > 1:
                with open(self.csv_path, "w", newline="") as f:
                    w = writer(f)
                    w.writerows(rows[:-1])
                remove(f"{self.images_dir}/{rows[-1][0]}")
                print(f"Removed last entry")
            self.num_cases -= 1
        elif key == ord("q"):
            return True
        elif ord("0") <= key <= ord("9"):
            self.class_id = key - ord("0")
            print(f"Set class to {self.class_id}")
        return False


if __name__ == "__main__":
    app = Collector()
    app.run()
    with open(f"{DATASET_DIR}/types.json", "w") as _f:
        dump([
            {"class_id": 0, "description": "background / no screw"},
            {"class_id": 1, "description": "flat 1.5cm"},
            {"class_id": 2, "description": "round 2.5cm"},
            {"class_id": 3, "description": "flat 3.0cm"},
            {"class_id": 4, "description": "flat 3.5cm"},
            {"class_id": 5, "description": "flat 6.0cm"},
            {"class_id": 6, "description": "round 7.5cm"}
        ], _f)