Update data models module.

demo/models.py

@@ -1,13 +1,23 @@
-from sqlmodel import SQLModel, Session, Field, create_engine, select
-from typing import Optional
+import os
 from datetime import datetime
-from PIL import Image, ExifTags
-from uuid import uuid4
 from pathlib import Path
-import os
+from typing import Optional
+from uuid import uuid4
+
+import pandas as pd
+from PIL import Image, ExifTags
+from sqlalchemy.orm import load_only
+from sqlmodel import SQLModel, Session, Field, create_engine, select
+
+SQLITE_FILENAME = 'database.db'
+SQLITE_URL = f'sqlite:///{SQLITE_FILENAME}'
+ENGINE = create_engine(SQLITE_URL, echo=False)
+IMAGES_STORAGE_FOLDER = './images'
 
 
 class Photo(SQLModel, table=True):
+    """Data model for the photo entity."""
+
     __table_args__ = {"extend_existing": True}
     id: int = Field(default=None, primary_key=True)
     filepath: str
@@ -18,26 +28,72 @@ class Photo(SQLModel, table=True):
 
 
 class WalrusCoord(SQLModel, table=True):
+    """Data model for the walrus'es coordinates instance."""
+
     __table_args__ = {"extend_existing": True}
     id: int = Field(default=None, primary_key=True)
     photo_id: int = Field(default=None, foreign_key='photo.id')
-    x: float
-    y: float
+    x: int
+    y: int
     is_young: bool = Field(default=False)
 
 
-sqlite_filename = 'database.db'
-sqlite_url = f'sqlite:///{sqlite_filename}'
-engine = create_engine(sqlite_url, echo=False)
-IMAGES_STORAGE_FOLDER = './images'
+def add_walrus_coords(
+        coords: list, classes: list, photo: Photo, session: Session
+):
+    """Add and new record with walruses coordinates."""
+    for (x, y), cls in zip(coords, classes):
+        session.add(WalrusCoord(photo_id=photo.id, x=x, y=y, is_young=cls))
+    session.commit()
+
+
+def add_photo(image: Image, fname: str, session: Session):
+    """Add and new record in Photo table."""
+    uid = uuid4()
+    ext = os.path.splitext(fname)[1]
+    folder = Path(IMAGES_STORAGE_FOLDER)
+    if not folder.exists():
+        folder.mkdir(parents=True)
+        print(f'Created directory: {folder}')
+    fpath = (folder / f'{uid}{ext}').absolute()
+    exif_data = get_exif_data(image)
+    lat, lon = get_lat_lon(exif_data)
+
+    image.save(fpath)
+    photo = Photo(
+        filepath=str(fpath),
+        fname=fname,
+        latitude=lat,
+        longitude=lon
+    )
+    session.add(photo)
+    session.commit()
+    return photo
+
+
+def add_record_in_db(image: Image, fname:str, coords, classes):
+    """Add new data into WalrusCoords and Photo tables."""
+    with Session(ENGINE) as session:
+        ph = add_photo(image, fname, session)
+        add_walrus_coords(coords, classes, ph, session)
 
 
 def create_db_and_tables():
-    SQLModel.metadata.create_all(engine)
+    """Create table."""
+    SQLModel.metadata.create_all(ENGINE)
+
+
+def check_db_exists():
+    """Check if table is exist."""
+    return os.path.exists(SQLITE_FILENAME)
 
 
 def get_exif_data(image):
-    """Returns a dictionary from the exif data of an PIL Image item. Also converts the GPS Tags"""
+    """
+    Returns a dictionary from the exif data of an PIL Image item.
+
+    Also converts the GPS Tags.
+    """
     exif_data = {}
     info = image._getexif()
     if info:
@@ -74,13 +130,13 @@ def get_lat_lon(exif_data):
     lat = None
     lon = None
 
-    if "GPSInfo" in exif_data:
-        gps_info = exif_data["GPSInfo"]
+    if 'GPSInfo' in exif_data:
+        gps_info = exif_data['GPSInfo']
 
-        gps_latitude = _get_if_exist(gps_info, "GPSLatitude")
-        gps_latitude_ref = _get_if_exist(gps_info, 'GPSLatitudeRef')
-        gps_longitude = _get_if_exist(gps_info, 'GPSLongitude')
-        gps_longitude_ref = _get_if_exist(gps_info, 'GPSLongitudeRef')
+        gps_latitude = gps_info.get('GPSLatitude', None)
+        gps_latitude_ref = gps_info.get('GPSLatitudeRef', None)
+        gps_longitude = gps_info.get('GPSLongitude', None)
+        gps_longitude_ref = gps_info.get('GPSLongitudeRef', None)
 
         if gps_latitude and gps_latitude_ref and gps_longitude and gps_longitude_ref:
             lat = _convert_to_degress(gps_latitude)
@@ -93,62 +149,30 @@ def get_lat_lon(exif_data):
 
     return lat, lon
 
-def add_photo(image: Image, fname: str, session: Session):
-    uid = uuid4()
-    ext = os.path.splitext(fname)[1]
-    fpath = (Path(IMAGES_STORAGE_FOLDER) / f'{uid}{ext}').absolute()
-    exif_data = get_exif_data(image)
-    lat, lon = get_lat_lon(exif_data)
-
-    image.save(fpath)
-    photo = Photo(
-        filepath=str(fpath),
-        fname=fname,
-        latitude=lat,
-        longitude=lon
-    )
-    session.add(photo)
-    session.commit()
-    return photo
-
-
-def add_walrus_coords(coords: list, photo: Photo, session: Session):
-    for x, y in coords:
-        session.add(WalrusCoord(photo_id=photo.id, x=x, y=y))
-    session.commit()
-
-
-coords = [
-    [0.2, 0.2],
-    [0.6, 0.6],
-]
-
-import pandas as pd
-from sqlalchemy.orm import load_only
 
 def create_df():
-    with Session(engine) as session:
+    """Create a Pandas DataFrame from the database."""
+    with Session(ENGINE) as session:
         photos = session.exec(
-            select(Photo).options(load_only('id', 'fname', 'created', 'latitude', 'longitude'))
-            # select(
-            #     Photo.fname,
-            #     Photo.created,
-            #     Photo.latitude,
-            #     Photo.longitude,
-            # )
+            select(Photo).options(
+                load_only('id', 'fname', 'created', 'latitude', 'longitude')
+            )
         ).all()
-    # print(photos)
-    df = pd.DataFrame(list(map(lambda x: x.dict(), photos)))
-    # df.set_index('id', inplace=True)
-    df = df[['id', 'fname', 'created', 'latitude', 'longitude']]
+    data = list(map(lambda x: x.dict(), photos))
+    for i in range(len(data)):
+        walruses_coords = get_points(data[i]['id'])
+        data[i]['walruses_count'] = len(walruses_coords)
+    df = pd.DataFrame(data)
+    df = df[
+        ['id', 'fname', 'created', 'walruses_count', 'latitude', 'longitude']
+    ]
     df.sort_values(by=['created'], inplace=True)
     return df
-        # for ph in photos:
-        #     print(ph.json())
 
 
 def get_path(photo_id):
-    with Session(engine) as session:
+    """Get a path to the source image of photo id."""
+    with Session(ENGINE) as session:
         photo = session.exec(
             select(Photo).where(Photo.id == photo_id)
         ).one()
@@ -159,30 +183,9 @@ def get_path(photo_id):
 
 
 def get_points(photo_id: int):
-    with Session(engine) as session:
+    """Get all walruses coordinates by given photo."""
+    with Session(ENGINE) as session:
         coords = session.exec(
             select(WalrusCoord).where(WalrusCoord.photo_id == photo_id)
         ).all()
     return coords
-
-    # if photo is not None:
-
-def main():
-    # df = create_df()
-    # print(df)
-    # create_db_and_tables()
-
-    image = Image.open('/home/fateev/dev/walruses/DJI_0005 (2).jpg')
-
-    # with Session(engine) as session:
-    #     ph = add_photo(image, 'walruses.jpg', session)
-    #     add_walrus_coords(coords, ph, session)
-
-
-    # create_photos()
-    # select_photos()
-
-
-if __name__ == '__main__':
-    # print(datetime.utcnow())
-    main()