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 resnet18
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# Feature extractor using pretrained ResNet18
12class VisualFeatureExtractor:
13    def __init__(self):
14        model = resnet18(pretrained=True)
15        self.model = torch.nn.Sequential(*list(model.children())[:-1]).to(device).eval()
16        self.transform = transforms.Compose([
17            transforms.ToPILImage(),
18            transforms.Resize((224, 224)),
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                features = self.model(tensor).squeeze()
27            return features / features.norm()
28        except:
29            return None
30
31# Memory system for object identity
32class ObjectMemory:
33    def __init__(self):
34        self.memory = {}  # id: feature_vector
35        self.next_id = 1
36
37    def compare(self, feat, threshold=0.9):
38        best_id, best_sim = None, 0.0
39        for obj_id, stored_feat in self.memory.items():
40            sim = cosine_similarity(feat, stored_feat, dim=0).item()
41            if sim > best_sim and sim > threshold:
42                best_id, best_sim = obj_id, sim
43        return best_id, best_sim
44
45    def memorize(self, feat):
46        obj_id = self.next_id
47        self.memory[obj_id] = feat
48        self.next_id += 1
49        return obj_id
50
51# Main application
52def main():
53    cap = cv2.VideoCapture(0)
54    fgbg = cv2.createBackgroundSubtractorMOG2()
55    extractor = VisualFeatureExtractor()
56    memory = ObjectMemory()
57
58    while True:
59        ret, frame = cap.read()
60        if not ret:
61            break
62
63        fgmask = fgbg.apply(frame)
64        _, thresh = cv2.threshold(fgmask, 200, 255, cv2.THRESH_BINARY)
65        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
66
67        for cnt in contours:
68            if cv2.contourArea(cnt) < 1000:
69                continue
70
71            x, y, w, h = cv2.boundingRect(cnt)
72            crop = frame[y:y+h, x:x+w]
73            feat = extractor.extract(crop)
74
75            if feat is None:
76                continue
77
78            matched_id, similarity = memory.compare(feat)
79
80            if matched_id is not None:
81                label = f"Known ID {matched_id} ({similarity*100:.1f}%)"
82                color = (0, 255, 0)
83            else:
84                new_id = memory.memorize(feat)
85                label = f"New Object (ID {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-10), cv2.FONT_HERSHEY_SIMPLEX,
90                        0.6, (255, 255, 255), 2)
91
92        cv2.imshow("AI Object Memory", frame)
93        if cv2.waitKey(1) & 0xFF == 27:  # ESC
94            break
95
96    cap.release()
97    cv2.destroyAllWindows()
98
99if __name__ == "__main__":
100    main()
101