PocketFinder/PocketFinder.py at main · ReadyMouse/PocketFinder · 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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
import requests
from typing import List, Dict
import time
from dataclasses import dataclass
from GooglePlacesPhotoReviews import GooglePlacesPhotoReviews
from PoolTableInference import PoolTableInference
from GooglePlacesVenueFinder import GooglePlacesVenueFinder
import os
import pdb

# Note this is unfinished, but a holding spot as I finished modules
# M. Bennett December 2024

def main():
    # Example usage
    cred_json_path ='pocketfinder-a0ced62ce802.json'
    api_key='api_key.txt'
    output_dir = 'nearby_BOS'
    os.makedirs(output_dir, exist_ok=True)

    # Example coordinates for Burlington, MA
    # TODO: Make place type an input field. Resturant, bar, etc
    venue_Finder = GooglePlacesVenueFinder()
    venues = venue_Finder.getPlaceID(latitude=42.4883,
                           longitude=-71.2338,
                           radius=1000)


    # TODO: How many venues did it find?
    print(f'Found {len(venues.items())} venues.')

    # Initialize photo finder
    photo_finder = GooglePlacesPhotoReviews(
        cred_json_path=cred_json_path,
        api_key=api_key,
        output_dir=output_dir
    )

    # Intialuze the inference engine
    pool_table_inference = PoolTableInference(
        model_path = './First-working-copy/weights/best.pt',
        conf_threshold=0.5,
        output_dir=output_dir
    )

    # Process each venue category
    for category, venue_list in venues.items():
        for venue in venue_list:
            print('-----------------')
            print(f"Processing venue: {venue.name}")
            place_id = venue.place_id  # Fixed attribute name
            # Get photos and reviews
            try:
                photos = photo_finder.get_place_photos(place_id)
                reviews = photo_finder.get_place_reviews(place_id)

                # Make sure venue object has these attributes
                if not hasattr(venue, 'photos_urls'):
                    venue.photos_urls = []
                if not hasattr(venue, 'reviews'):
                    venue.reviews = []

                venue.photos_urls.extend(photos)
                venue.reviews.extend(reviews)

                print(f"Number of photos found: {len(photos)}")
                print(f"Number of reviews found: {len(reviews)}")

                # Rate limiting
                time.sleep(0.2)

                # Process images for pool tables if photos exist
                table_prob = []
                if venue.photos_urls:
                    try:
                        for item in venue.photos_urls:
                            results = pool_table_inference.run_inference(image_path=item,
                                                                        save_negative=True # to save photos that don't have pool tables
                                                                        )
                            #pdb.set_trace()
                            if results[os.path.basename(item)]['class_name'] == 'pool_table':
                                confidence = results[os.path.basename(item)]['confidence']
                            else:
                                confidence = 1 - results[os.path.basename(item)]['confidence']

                            table_prob.append(confidence)

                        # TODO: math on the set of probabilities of a pool table
                        max_conf = max(table_prob)
                        print(f'There is a {max_conf:.1%} chance that {venue.name} has a pool table.')

                    except Exception as e:
                        print(f"Error in pool table inference for {venue.name}: {str(e)}")

                # TODO: Add text review processing

            except Exception as e:
                print(f"Error processing venue {venue.name}: {str(e)}")
            continue

if __name__ == "__main__":
    main()