|
1 import shelve |
|
2 |
|
3 class TagDB (object) : |
|
4 _imgs_cache = {} |
|
5 |
|
6 def __init__ (self, read_only=True) : |
|
7 self.img_tags = shelve.open("img_tags", read_only and 'r' or 'c') |
|
8 self.tag_imgs = shelve.open("tag_imgs", read_only and 'r' or 'c') |
|
9 |
|
10 def tag (self, img, tag) : |
|
11 """ |
|
12 Associate the given image with the given tag |
|
13 """ |
|
14 |
|
15 if img not in self.img_tags : |
|
16 self.img_tags[img] = set([tag]) |
|
17 else : |
|
18 s = self.img_tags[img] |
|
19 s.add(tag) |
|
20 self.img_tags[img] = s |
|
21 |
|
22 |
|
23 if tag not in self.tag_imgs : |
|
24 self.tag_imgs[tag] = set([img]) |
|
25 else : |
|
26 s = self.tag_imgs[tag] |
|
27 s.add(img) |
|
28 self.tag_imgs[tag] = s |
|
29 |
|
30 print "%s <-> %s" % (img, tag) |
|
31 |
|
32 def imgs (self, tags) : |
|
33 """ |
|
34 Get the set of images that have the given set of tags |
|
35 """ |
|
36 |
|
37 cache_key = "/".join(tags) |
|
38 |
|
39 if cache_key in self._imgs_cache : |
|
40 return self._imgs_cache[cache_key] |
|
41 |
|
42 if not tags : |
|
43 return set(self.img_tags.keys()) |
|
44 |
|
45 img_sets = [self.tag_imgs[tag] for tag in tags] |
|
46 |
|
47 res = None |
|
48 |
|
49 for img_set in img_sets : |
|
50 if res : |
|
51 res = res & img_set |
|
52 else : |
|
53 res = img_set |
|
54 |
|
55 self._imgs_cache[cache_key] = res |
|
56 |
|
57 return res |
|
58 |
|
59 def tags (self, tags) : |
|
60 """ |
|
61 Get the set of tags that are present in the set of images specified by these tags, sorted by count |
|
62 |
|
63 This is currently implemented quite inefficiently... giev SQL db? |
|
64 """ |
|
65 |
|
66 imgs = self.imgs(tags) |
|
67 |
|
68 ret = [] |
|
69 |
|
70 for tag in self.tag_imgs.keys() : |
|
71 if tag in tags : |
|
72 continue |
|
73 |
|
74 count = len(self.tag_imgs[tag] & imgs) |
|
75 |
|
76 if count : |
|
77 ret.append((tag, count)) |
|
78 |
|
79 def my_cmp ((at, ac), (bt, bc)) : |
|
80 return cmp((ac, at), (bc, bt)) |
|
81 |
|
82 ret.sort(reverse=True) |
|
83 |
|
84 return ret |
|
85 |
|
86 def close (self) : |
|
87 self.img_tags.close() |
|
88 self.tag_imgs.close() |
|
89 |