Update app.py

app.py

@@ -14,9 +14,19 @@ from models.huggingface import Geolocalizer
 import spacy
 from collections import Counter
 from spacy.cli import download
+from typing import Tuple, List, Optional, Union, Dict
 
 
-def load_spacy_model(model_name="en_core_web_md"):
+def load_spacy_model(model_name: str = "en_core_web_md") -> spacy.Language:
+    """
+    Load the specified spaCy model.
+
+    Args:
+        model_name (str): Name of the spaCy model to load.
+
+    Returns:
+        spacy.Language: Loaded spaCy model.
+    """
     try:
         return spacy.load(model_name)
     except IOError:
@@ -31,9 +41,14 @@ IMAGE_SIZE = (224, 224)
 GEOLOC_MODEL_NAME = "osv5m/baseline"
 
 
-# Load geolocation model
 @st.cache_resource(show_spinner=True)
-def load_geoloc_model() -> Geolocalizer:
+def load_geoloc_model() -> Optional[Geolocalizer]:
+    """
+    Load the geolocation model.
+
+    Returns:
+        Optional[Geolocalizer]: Loaded geolocation model or None if loading fails.
+    """
     with st.spinner('Loading model...'):
         try:
             model = Geolocalizer.from_pretrained(GEOLOC_MODEL_NAME)
@@ -44,31 +59,43 @@ def load_geoloc_model() -> Geolocalizer:
             return None
 
 
-# Function to find the most frequent location
-def most_frequent_locations(text: str):
+def most_frequent_locations(text: str) -> Tuple[str, List[str]]:
+    """
+    Find the most frequent locations mentioned in the text.
+
+    Args:
+        text (str): Input text to analyze.
+
+    Returns:
+        Tuple[str, List[str]]: Description of the most mentioned locations and a list of those locations.
+    """
     doc = nlp(text)
     locations = []
 
-    # Collect all identified location entities
     for ent in doc.ents:
         if ent.label_ in ['LOC', 'GPE']:
             print(f"Entity: {ent.text} | Label: {ent.label_} | Sentence: {ent.sent}")
             locations.append(ent.text)
 
-    # Count occurrences and extract the most common locations
     if locations:
         location_counts = Counter(locations)
-        most_common_locations = location_counts.most_common(2)  # Adjust the number as needed
-        # Format the output to show location names along with their counts
+        most_common_locations = location_counts.most_common(2)
         common_locations_str = ', '.join([f"{loc[0]} ({loc[1]} occurrences)" for loc in most_common_locations])
-
         return f"Most Mentioned Locations: {common_locations_str}", [loc[0] for loc in most_common_locations]
     else:
         return "No locations found", []
 
 
-# Transform image for model prediction
 def transform_image(image: Image) -> torch.Tensor:
+    """
+    Transform the input image for model prediction.
+
+    Args:
+        image (Image): Input image.
+
+    Returns:
+        torch.Tensor: Transformed image tensor.
+    """
     transform = transforms.Compose([
         transforms.Resize(IMAGE_SIZE),
         transforms.ToTensor(),
@@ -77,7 +104,17 @@ def transform_image(image: Image) -> torch.Tensor:
     return transform(image).unsqueeze(0)
 
 
-def check_location_match(location_query, most_common_locations):
+def check_location_match(location_query: dict, most_common_locations: List[str]) -> bool:
+    """
+    Check if the predicted location matches any of the most common locations.
+
+    Args:
+        location_query (dict): Predicted location details.
+        most_common_locations (List[str]): List of most common locations.
+
+    Returns:
+        bool: True if a match is found, False otherwise.
+    """
     name = location_query['name']
     admin1 = location_query['admin1']
     cc = location_query['cc']
@@ -88,8 +125,16 @@ def check_location_match(location_query, most_common_locations):
     return False
 
 
-# Fetch city GeoJSON data
-def get_city_geojson(location_name: str) -> dict:
+def get_city_geojson(location_name: str) -> Optional[dict]:
+    """
+    Fetch the GeoJSON data for the specified city.
+
+    Args:
+        location_name (str): Name of the city.
+
+    Returns:
+        Optional[dict]: GeoJSON data of the city or None if fetching fails.
+    """
     geolocator = Nominatim(user_agent="predictGeolocforImage")
     try:
         location = geolocator.geocode(location_name, geometry='geojson')
@@ -99,8 +144,16 @@ def get_city_geojson(location_name: str) -> dict:
         return None
 
 
-# Fetch media from URL
-def get_media(url: str) -> list:
+def get_media(url: str) -> Optional[List[Tuple[str, str]]]:
+    """
+    Fetch media URLs and associated text from the specified URL.
+
+    Args:
+        url (str): URL to fetch media from.
+
+    Returns:
+        Optional[List[Tuple[str, str]]]: List of tuples containing media URLs and associated text or None if fetching fails.
+    """
     try:
         response = requests.get(url)
         response.raise_for_status()
@@ -112,8 +165,17 @@ def get_media(url: str) -> list:
         return None
 
 
-# Predict location from image
-def predict_location(image: Image, model: Geolocalizer) -> tuple:
+def predict_location(image: Image, model: Geolocalizer) -> Optional[Tuple[List[float], dict, Optional[dict], float]]:
+    """
+    Predict the location from the input image using the specified model.
+
+    Args:
+        image (Image): Input image.
+        model (Geolocalizer): Geolocation model.
+
+    Returns:
+        Optional[Tuple[List[float], dict, Optional[dict], float]]: Predicted GPS coordinates, location query, city GeoJSON data, and processing time or None if prediction fails.
+    """
     with st.spinner('Processing image and predicting location...'):
         start_time = time.time()
         try:
@@ -130,8 +192,14 @@ def predict_location(image: Image, model: Geolocalizer) -> tuple:
             return None
 
 
-# Display map in Streamlit
-def display_map(city_geojson: dict, gps_degrees: list) -> None:
+def display_map(city_geojson: dict, gps_degrees: List[float]) -> None:
+    """
+    Display a map with the specified city GeoJSON data and GPS coordinates.
+
+    Args:
+        city_geojson (dict): GeoJSON data of the city.
+        gps_degrees (List[float]): GPS coordinates.
+    """
     map_view = pdk.Deck(
         map_style='mapbox://styles/mapbox/light-v9',
         initial_view_state=pdk.ViewState(
@@ -156,8 +224,13 @@ def display_map(city_geojson: dict, gps_degrees: list) -> None:
     st.pydeck_chart(map_view)
 
 
-# Display image
 def display_image(image_url: str) -> None:
+    """
+    Display an image from the specified URL.
+
+    Args:
+        image_url (str): URL of the image.
+    """
     try:
         response = requests.get(image_url)
         response.raise_for_status()
@@ -169,8 +242,16 @@ def display_image(image_url: str) -> None:
         st.error(f"An error occurred: {e}")
 
 
-# Scrape webpage for text and images
-def scrape_webpage(url: str) -> tuple:
+def scrape_webpage(url: str) -> Union[Tuple[Optional[str], Optional[List[str]]], Tuple[None, None]]:
+    """
+    Scrape the specified webpage for text and images.
+
+    Args:
+        url (str): URL of the webpage to scrape.
+
+    Returns:
+        Union[Tuple[Optional[str], Optional[List[str]]], Tuple[None, None]]: Extracted text and list of image URLs or None if scraping fails.
+    """
     with st.spinner('Scraping web page...'):
         try:
             response = requests.get(url)
@@ -185,27 +266,31 @@ def scrape_webpage(url: str) -> tuple:
             return None, None
 
 
-def main():
+def main() -> None:
+    """
+    Main function to run the Streamlit app.
+    """
     st.title('Welcome to Geolocation Guesstimation Demo 👋')
 
-    # Define page navigation using the sidebar
     page = st.sidebar.selectbox(
         "Choose your action:",
         ("Home", "Images", "Social Media", "Web Pages"),
-        index=0  # Default to Home
+        index=0
     )
 
     st.sidebar.success("Select a demo above.")
     st.sidebar.info(
         """
         - Web App URL: <https://yunusserhat-guesstimatelocation.hf.space/>
-        """)
+        """
+    )
 
     st.sidebar.title("Contact")
     st.sidebar.info(
         """
         Yunus Serhat Bıçakçı at [yunusserhat.com](https://yunusserhat.com) | [GitHub](https://github.com/yunusserhat) | [Twitter](https://twitter.com/yunusserhat) | [LinkedIn](https://www.linkedin.com/in/yunusserhat)
-        """)
+        """
+    )
 
     if page == "Home":
         st.write("Welcome to the Geolocation Predictor. Please select an action from the sidebar dropdown.")
@@ -220,7 +305,10 @@ def main():
         web_page_url_page()
 
 
-def upload_images_page():
+def upload_images_page() -> None:
+    """
+    Display the image upload page for geolocation prediction.
+    """
     st.header("Image Upload for Geolocation Prediction")
     uploaded_files = st.file_uploader("Choose images...", type=["jpg", "jpeg", "png"], accept_multiple_files=True)
     if uploaded_files:
@@ -230,7 +318,7 @@ def upload_images_page():
                 st.image(image, caption=f'Uploaded Image: {file.name}', use_column_width=True)
                 model = load_geoloc_model()
                 if model:
-                    result = predict_location(image, model)  # Assume this function is defined elsewhere
+                    result = predict_location(image, model)
                     if result:
                         gps_degrees, location_query, city_geojson, processing_time = result
                         st.write(
@@ -240,7 +328,10 @@ def upload_images_page():
                             st.write(f"Processing Time (seconds): {processing_time}")
 
 
-def social_media_page():
+def social_media_page() -> None:
+    """
+    Display the social media analysis page.
+    """
     st.header("Social Media Analyser")
     social_media_url = st.text_input("Enter a social media URL to analyse:", key='social_media_url_input')
     if social_media_url:
@@ -270,7 +361,6 @@ def social_media_page():
                             if city_geojson:
                                 display_map(city_geojson, gps_degrees)
                                 st.write(f"Processing Time (seconds): {processing_time}")
-                            # Check for match and notify
                             if check_location_match(location_query, most_common_locations):
                                 st.success(
                                     f"The predicted location {location_name} matches one of the most frequently mentioned locations!")
@@ -278,7 +368,10 @@ def social_media_page():
                     st.error(f"Failed to fetch image at URL {media_url}: HTTP {response.status_code}")
 
 
-def web_page_url_page():
+def web_page_url_page() -> None:
+    """
+    Display the web page URL analysis page.
+    """
     st.header("Web Page Analyser")
     web_page_url = st.text_input("Enter a web page URL to scrape:", key='web_page_url_input')
     if web_page_url:
@@ -307,7 +400,6 @@ def web_page_url_page():
                                 if city_geojson:
                                     display_map(city_geojson, gps_degrees)
                                     st.write(f"Processing Time (seconds): {processing_time}")
-                                # Check for match and notify
                                 if check_location_match(location_query, most_common_locations):
                                     st.success(
                                         f"The predicted location {location_name} matches one of the most frequently mentioned locations!")