ysn-rfd/text-dataset-tiny-code-script-py-format

1import cv2
2import numpy as np
3import torch
4import torchvision.transforms as transforms
5from torchvision.models import mobilenet_v3_small
6from torch.nn.functional import cosine_similarity
7
8# Use GPU if available
9device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
11# Lightweight feature extractor using MobileNetV3
12class FastFeatureExtractor:
13    def __init__(self):
14        model = mobilenet_v3_small(pretrained=True).features
15        self.model = torch.nn.Sequential(*list(model.children())[:-1]).to(device).eval()
16        self.transform = transforms.Compose([
17            transforms.ToPILImage(),
18            transforms.Resize((96, 96)),
19            transforms.ToTensor()
20        ])
21
22    def extract(self, image):
23        try:
24            tensor = self.transform(image).unsqueeze(0).to(device)
25            with torch.no_grad():
26                feat = self.model(tensor).mean([2, 3]).squeeze()
27            return feat / feat.norm()
28        except:
29            return None
30
31# Simple memory with similarity threshold
32class ObjectMemory:
33    def __init__(self, threshold=0.88):
34        self.memory = {}
35        self.next_id = 1
36        self.threshold = threshold
37
38    def match(self, feat):
39        best_id, best_sim = None, 0.0
40        for obj_id, ref_feat in self.memory.items():
41            sim = cosine_similarity(feat, ref_feat, dim=0).item()
42            if sim > best_sim and sim > self.threshold:
43                best_id, best_sim = obj_id, sim
44        return best_id, best_sim
45
46    def add(self, feat):
47        obj_id = self.next_id
48        self.memory[obj_id] = feat
49        self.next_id += 1
50        return obj_id
51
52# Main app
53def main():
54    cap = cv2.VideoCapture(0)
55    fgbg = cv2.createBackgroundSubtractorMOG2()
56    extractor = FastFeatureExtractor()
57    memory = ObjectMemory()
58
59    while True:
60        ret, frame = cap.read()
61        if not ret:
62            break
63
64        fg = fgbg.apply(frame)
65        _, thresh = cv2.threshold(fg, 200, 255, cv2.THRESH_BINARY)
66        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
67
68        for cnt in contours:
69            if cv2.contourArea(cnt) < 1200:
70                continue
71
72            x, y, w, h = cv2.boundingRect(cnt)
73            roi = frame[y:y+h, x:x+w]
74            feat = extractor.extract(roi)
75
76            if feat is None:
77                continue
78
79            matched_id, similarity = memory.match(feat)
80            if matched_id:
81                label = f"Known #{matched_id} ({similarity*100:.1f}%)"
82                color = (0, 255, 0)
83            else:
84                new_id = memory.add(feat)
85                label = f"New Object #{new_id}"
86                color = (0, 0, 255)
87
88            cv2.rectangle(frame, (x, y), (x+w, y+h), color, 2)
89            cv2.putText(frame, label, (x, y-8), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
90
91        cv2.imshow("Fast Object Understanding", frame)
92        if cv2.waitKey(1) & 0xFF == 27:  # ESC to exit
93            break
94
95    cap.release()
96    cv2.destroyAllWindows()
97
98if __name__ == "__main__":
99    main()
100